Back to blog

When a Tiny Mamba Learns Arithmetic

The experiment looked small on purpose: train a 380k parameter, 6-layer Mamba-style model on 1,000 examples of two-digit addition without carry. The point was not to solve arithmetic in general. The point was to ask a sharper engineering question: can the model overfit a tiny deterministic task without losing numerical stability?

The answer was mixed in the most useful way. The fourth run stayed stable for 1,000 epochs and produced the target sanity-check answer 21+48=69, but exact-match accuracy stayed low: 1.9% on train and 2.5% on validation. That gap is the story.

Four runs, one stable path

The first attempts showed two separate failure modes. Increasing batch size improved GPU utilization but exposed unstable state-space dynamics. Adding parameter clamping stabilized the recurrent matrix, but a high learning rate still pushed the optimizer into overflow. The successful configuration was more conservative: batch size 256, learning rate 3e-4, and post-step clamping.

RunBatchLRClampOutcome
11281e-3nomanually stopped
22561e-3noFPU overflow at epoch 400
35121e-3yesNaN / FPU overflow
42563e-4yesstable to epoch 1000

What stability actually meant

The key stability principle was simple: the diagonal state transition must remain contracting. If the recurrent state parameter drifts above zero, the scan turns from a decaying system into an expanding one. On GPU, that showed up as invalid floating point operations.

Clamping the state parameter to a negative range kept the recurrence stable. It did not solve optimization by itself, but it separated architectural instability from optimizer instability.

The plateau

The loss curve tells the second half of the story:

EpochLoss
13.0245
1000.8238
2000.2522
6000.2299
10000.2261
Training loss by epoch
13.02451000.82382000.25226000.229910000.2261

Most of the learning happened early. From epoch 200 to 1000, the model only improved by 0.0261 loss points. That plateau suggests the fixed learning rate was too coarse for the final stage, and the large batch size removed the stochastic noise that can help escape shallow traps.

Why exact match stayed low

The model was not random. It often learned the magnitude and partial structure of the answer, but it failed in ways that matter for exact-match evaluation.

Two error classes dominated:

  • off-by-one or digit-level approximations, such as 41+41 -> 84 instead of 82;
  • generation truncation, where the two-token decode budget was too strict after an unwanted token.
Bar chart
Train EM
1.9
Val EM
2.5

This is why the model could pass the sanity check 21+48=69 while still failing the broader exact-match metric.

What this changes next

The next experiment should not only train longer. It should change the task interface and the optimizer loop:

  • add explicit scratchpad tokens so intermediate arithmetic can be represented in the sequence;
  • add learning-rate decay when the loss window flattens;
  • evaluate with a generation budget that can tolerate a stray separator token;
  • track off-by-one and truncation errors separately from total exact match.

The lesson is useful precisely because the model did not reach 100%. The run proved that the state-space core can be made stable, and it exposed the next bottleneck: not recurrence explosion, but the training protocol around a fragile symbolic task.

Когда маленькая Mamba учится арифметике

Эксперимент специально был маленьким: обучить 380k-параметровую 6-слойную Mamba-подобную модель на 1000 примерах сложения двухзначных чисел без переноса. Цель была не в том, чтобы “решить арифметику вообще”. Вопрос был инженерный: сможет ли модель переобучиться на маленькой детерминированной задаче и при этом не потерять численную стабильность?

Ответ получился смешанным, и именно поэтому полезным. Четвёртый запуск стабильно дошёл до 1000 эпох и правильно выдал sanity-check 21+48=69, но exact match остался низким: 1.9% на train и 2.5% на validation. В этом разрыве и находится главная история.

Четыре запуска и один стабильный путь

Первые попытки показали два разных режима отказа. Увеличение batch size лучше нагружало GPU, но проявило нестабильность state-space динамики. Clamping параметров стабилизировал рекуррентную матрицу, но слишком высокий learning rate всё ещё уводил оптимизатор в overflow. Рабочая конфигурация оказалась более осторожной: batch size 256, learning rate 3e-4 и clamping после шага оптимизатора.

RunBatchLRClampOutcome
11281e-3noостановлен вручную
22561e-3noFPU overflow на эпохе 400
35121e-3yesNaN / FPU overflow
42563e-4yesстабильно до эпохи 1000

Что на самом деле значила стабильность

Ключевой принцип был простым: диагональный state transition должен оставаться сжимающим. Если рекуррентный параметр уходит выше нуля, scan превращается из затухающей системы в расширяющуюся. На GPU это проявилось как invalid floating point operation.

Clamping удержал параметр в отрицательной области и стабилизировал рекуррентную часть. Он не решил оптимизацию сам по себе, но отделил архитектурную нестабильность от нестабильности оптимизатора.

Плато

Кривая loss показывает вторую половину истории:

EpochLoss
13.0245
1000.8238
2000.2522
6000.2299
10000.2261
Training loss by epoch
13.02451000.82382000.25226000.229910000.2261

Основное обучение произошло рано. С 200 по 1000 эпоху модель улучшилась всего на 0.0261 loss-пункта. Это плато намекает, что фиксированный learning rate стал слишком грубым для финальной стадии, а большой batch size убрал стохастический шум, который иногда помогает выходить из мелких ловушек.

Почему exact match остался низким

Модель не отвечала случайно. Часто она схватывала порядок величины и часть структуры ответа, но ошибалась так, что exact-match метрика считала это полным промахом.

Доминировали два класса ошибок:

  • off-by-one или digit-level approximation, например 41+41 -> 84 вместо 82;
  • truncation при генерации, когда двух токенов не хватало после лишнего разделителя или шумового токена.
Bar chart
Train EM
1.9
Val EM
2.5

Поэтому модель могла проходить sanity-check 21+48=69, но всё ещё проваливать широкий exact-match.

Что это меняет дальше

Следующий эксперимент должен не просто обучаться дольше. Нужно менять интерфейс задачи и контур оптимизации:

  • добавить scratchpad-токены, чтобы промежуточная арифметика жила в последовательности;
  • включать decay learning rate, когда окно loss становится плоским;
  • оценивать генерацию с запасом по длине, чтобы переживать случайный separator-токен;
  • отдельно считать off-by-one и truncation ошибки, а не только общий exact match.

Ценность эксперимента именно в том, что он не дошёл до 100%. Запуск показал, что state-space ядро можно стабилизировать, и вывел наружу следующий bottleneck: уже не explosion рекурренса, а протокол обучения вокруг хрупкой символической задачи.