Неделя 19. Другие архитектуры: RNN, SSM, MoE

Фаза 5. Расширение · неделя 19 из 24

Учиться в приложении: тьютор, задачи с кодом →

Ядро: RNN и затухание градиента, идея SSM и Mamba, MoE: роутинг и балансировка · Глубина: LSTM в деталях, linear attention, гауссовские процессы (05-ГЛУБИНА) · ≈ 11 ч ядро / 23 ч всё

Вся неделя отвечает на один вопрос: чем платить за контекст. Трансформер хранит историю целиком (KV-кэш O(S)), зато обучается параллельно. RNN сжимает историю в состояние фиксированного размера: O(1) на токен при генерации, но обучение последовательное, а далёкое забывается. SSM и гибриды пытаются получить и то, и другое.

Шаг 1. Vanilla RNN и затухание по времени

На пальцах. Градиент идёт назад через 40 шагов и на каждом умножается на одно и то же число. Множитель 0.9: 0.9⁴⁰ ≈ 0.015, от сигнала остаётся полтора процента. Множитель 1.1: 1.1⁴⁰ ≈ 45, сигнал взорвался. В RNN вместо числа стоит матрица W_h, и роль множителя играет её спектральный радиус (наибольший модуль собственного значения). Удержать его ровно на 1 не выходит, а производная tanh (не больше 1) только уменьшает множитель.

h_t = tanh(W_h h_{t−1} + W_x x_t). Развёртка по времени (unroll: копия ячейки на каждый шаг) даёт сеть глубины T с общими весами
Вывести: ∂h_T/∂h_k = Π_{t=k+1..T} diag(1 − h_t²) W_h. Это произведение T − k Якобианов, которое ведёт себя как ρ(W_h)^{T−k}. Меньше единицы означает затухание, больше единицы означает взрыв
Вывод: проблема не в tanh, а в многократном умножении на одну матрицу. Взрыв лечится клиппингом (неделя 3), затухание лечится только сменой архитектуры

Шаг 2. LSTM: аддитивная память

На пальцах. В LSTM по пути памяти c градиент умножается не на W, а на forget-гейт, число от 0 до 1, которое сеть выставляет сама на каждом шаге. Гейт 0.99 за 40 шагов даёт 0.99⁴⁰ ≈ 0.67: две трети сигнала доходят. Гейт 0.9 дал бы те же полтора процента, что у RNN. Разница в том, что этот множитель сеть может держать около единицы, не ломая всё остальное.

c_t = f_t ⊙ c_{t−1} + i_t ⊙ g_t. По прямому пути ∂c_t/∂c_{t−1} = diag(f_t): умножения на W нет, при f_t ≈ 1 градиент проходит сотни шагов. Поэтому bias forget-гейта инициализируют единицей
Это та же идея, что residual stream (неделя 6): аддитивное обновление даёт градиенту путь в обход нелинейностей

Шаг 3. SSM и Mamba

На пальцах. Одномерная линейная рекуррентность h_t = 0.5·h_{t−1} + x_t, выход y_t = h_t, вход (2, 4, 6). Рекуррентно: h₁ = 2, h₂ = 1 + 4 = 5, h₃ = 2.5 + 6 = 8.5. Свёрткой с ядром (1, 0.5, 0.25): y₃ = 1·6 + 0.5·4 + 0.25·2 = 8.5. Число то же, но теперь все y_t можно считать одновременно. Если коэффициент 0.5 зависит от входа, как в Mamba, ядро у каждой позиции своё, и единой свёртки больше нет.

SSM (state space model, модель со скрытым состоянием и линейной рекуррентностью). Линейная рекуррентность h_t = A h_{t−1} + B x_t, y_t = C h_t разворачивается в свёртку с ядром (CB, CAB, CA²B, …). Главный вывод недели: линейность покупает оба режима: обучение свёрткой параллельно, генерацию рекуррентно за O(1)
Mamba: B, C и шаг Δ зависят от входа, то есть модель решает, что записать и что забыть. Ядро перестаёт быть постоянным, свёртка не работает. Её заменяет параллельный scan (расчёт всех префиксов ассоциативной операции за O(log T) параллельных шагов), возможный потому, что композиция линейных обновлений ассоциативна
Linear attention: без softmax (QKᵀ)V = Q(KᵀV), состояние KᵀV размера H×H. Цена обоих подходов: хуже точное извлечение из далёкого контекста. Отсюда гибриды, чередующие слои внимания и SSM

Врезка. Окно наблюдений как состояние

На пальцах. Маятник, видим только отклонение, округлённо: 0, 7, 10, 7, 0, −7, −10, −7, 0. Значение 7 встречается дважды: по пути от центра и по пути обратно. По одному числу не понять, куда маятник летит. Пара «предыдущее, текущее» уже различает: (0, 7) значит «удаляется», (10, 7) значит «возвращается». Разность соседних чисел даёт скорость, а положение и скорость вместе и есть полное состояние маятника.

Шаги 1–3 шли от состояния к наблюдениям: RNN и SSM сжимают историю в h_t. Для детерминированной системы верно и обратное (теорема Такенса, 1981): окно X_t = (x_t, x_{t−τ}, …, x_{t−(E−1)τ}) из значений одной наблюдаемой величины с шагом задержки τ в типичном случае взаимно однозначно задаёт состояние, если E > 2d (d: размерность аттрактора, множества, по которому движется система)
У маятника без трения траектория замкнута, d = 1, и теорема гарантирует успех при E = 3. В примере хватило двух: гарантия берёт с запасом
Мост к LLM: окно контекста тоже может служить состоянием, и трансформеру с окном S не обязательно что-то сжимать. Но текст не маломерная детерминированная система, поэтому это аналогия, а не теорема
Не путать омонимы. State space в SSM пришло из теории управления: состояние h_t линейное, скрытое и задано уравнением модели. State space reconstruction в нелинейной динамике означает восстановление неизвестного состояния чужой системы по окну наблюдений. Корень общий, смысл разный

Шаг 4. Mixture of Experts

На пальцах. FFN блока весит 100 млн параметров. Заменим его на 8 экспертов такого же размера и роутер (маленький линейный слой, который выбирает экспертов для токена) с top-2. Параметров в FFN стало 800 млн, а каждый токен проходит только через двух экспертов, то есть через 200 млн. Вычислений стало вдвое больше, а не в 8 раз. Коллапс: при top-2 из 8 каждому эксперту положено 25% токенов. Если на старте эксперт 3 случайно получил 40%, он учится быстрее, становится лучше, и роутер шлёт ему ещё больше.

FFN → E экспертов, роутер выбирает top-k. Параметры растут в E раз, FLOPs на токен примерно в k. Число параметров и объём вычислений развязаны
Без балансировки роутинг коллапсирует: популярные эксперты учатся быстрее и становятся популярнее. Ниже два способа устроить экспертов иначе и три способа держать нагрузку ровной

Shared и routed experts. Часть экспертов включена всегда (shared, общие), остальных выбирает роутер (routed, маршрутизируемые): y = x + Σ_{j∈shared} FFN_j(x) + Σ_{i∈top-k} gᵢ·FFNᵢ(x). Общие берут то, что нужно любому токену (грамматика, частые слова), и маршрутизируемым не приходится дублировать это знание в каждом из E экспертов (DeepSeekMoE).

Fine-grained experts. На пальцах. 16 экспертов ширины F, top-2: токен проходит через 2F нейронов, а различных пар экспертов C(16, 2) = 120. Разрежем каждого на 4 узких эксперта ширины F/4: их 64, берём top-8. Через токен проходят те же 8·F/4 = 2F нейронов, параметров те же 64·F/4 = 16F, а сочетаний C(64, 8) ≈ 4.4·10⁹. FLOPs и память не изменились, а способов собрать под токен «своего» эксперта стало в 37 млн раз больше.

Вспомогательный лосс (Switch Transformer). На пальцах. E = 4, доли назначений в батче f = (0.5, 0.25, 0.125, 0.125), средние вероятности роутера P = (0.4, 0.3, 0.15, 0.15). α·E·Σ fᵢ·Pᵢ = α·4·(0.2 + 0.075 + 0.01875 + 0.01875) = 1.25·α. При ровной нагрузке f = P = (¼, ¼, ¼, ¼) вышло бы 1.0·α.

L_aux = α·E·Σᵢ fᵢ·Pᵢ. fᵢ (доля токенов, отправленных эксперту i) получается из выбора top-k и градиента не имеет, поэтому градиент течёт через Pᵢ (средняя по батчу вероятность роутера для i): ∂L_aux/∂Pᵢ = α·E·fᵢ, сильнее всего давит вероятность перегруженного. Множитель E делает значение при балансе равным 1 при любом E. α маленький, порядка 10⁻²: сильный лосс спорит с основной задачей

Ёмкость эксперта. На пальцах. В батче 96 токенов, E = 8, top-2, capacity factor (запас ёмкости) CF = 1.25: у эксперта 1.25·96·2/8 = 30 мест. Эксперту выпало 37 назначений: 7 лишних он не обрабатывает, и они идут дальше только по residual.

capacity = CF · B·T · k / E. Ёмкость нужна, потому что формы тензоров на ускорителе фиксированы. Больше CF значит меньше потерянных токенов, но больше пустых мест и памяти

Балансировка без вспомогательного лосса (смещение роутера, DeepSeek-V3). На пальцах. Скоры токена s = (0.70, 0.55, 0.52, 0.30), top-2 выбрал бы экспертов 1 и 2. Но 1 и 2 в последнее время перегружены, 3 недогружен, и накопились смещения b = (−0.02, −0.02, +0.04, 0). Выбор идёт по s + b = (0.68, 0.53, 0.56, 0.30): эксперты 1 и 3. Веса смешивания берутся из исходных скоров: 0.70 / (0.70 + 0.52) ≈ 0.57 и 0.52 / 1.22 ≈ 0.43.

Смещение bᵢ решает только, кого выбрать, и не входит в веса gᵢ, поэтому выход не искажается добавкой без смысла. После шага bᵢ уменьшают на γ у перегруженных и увеличивают на γ у недогруженных (γ: маленькая скорость обновления). Градиента у b нет, основной лосс его не видит
Уточнения к DeepSeek-V3. Гейт сигмоидный: sᵢ = sigmoid(xᵀeᵢ) (eᵢ вектор эксперта), каждый эксперт оценивается отдельно, без softmax по всем; веса нормируют по выбранным, gᵢ = sᵢ / Σ_{j∈top-k} sⱼ. И «совсем без лосса» неверно: остаётся балансирующий лосс на уровне последовательности с очень маленьким α, чтобы одна последовательность не легла на пару экспертов
Expert parallelism: два all-to-all (каждый GPU отправляет каждому свою часть токенов) на слой (неделя 14). Объём на слой пропорционален токены × D × top-k; сколько это в миллисекундах внутри узла и между узлами, посчитано в задаче 5 на салфетке недели 14. В тренажёре есть задача про top-k роутинг: с неё удобно начать

MTP как вспомогательная цель обучения (DeepSeek-V3). Устройство модулей глубины разобрано в неделе 11 (шаг 3); здесь взгляд со стороны обучения. На пальцах. Последовательность из 6 токенов. Основная голова получает 5 целей (каждая позиция предсказывает следующий токен), модуль глубины 1 ещё 4 (каждая позиция предсказывает токен через один). Лосс основной головы 2.0, модуля 3.0 (на два шага вперёд угадывать труднее), вес λ = 0.3: L = 2.0 + 0.3·3.0 = 2.9. В общий ствол приходит градиент от обеих целей, вспомогательная входит с весом 0.3.

Зачем: больше сигнала на каждую позицию, и представлению приходится «планировать» на шаг вперёд. Выигрыш заметнее у больших моделей и на коде; маленькой модели доп. цель может мешать
Цена: один блок трансформера на каждую глубину и его FLOPs при обучении. Эмбеддинг и выходная голова общие с моделью, так что параметров добавляется немного. λ держат маленьким и снижают к концу обучения
На инференсе модуль можно выбросить, и модель работает как обычная. Или сделать из него черновик для speculative decoding: если второй токен принимается в 85% случаев, за проход выходит в среднем 1 + 0.85 = 1.85 токена
Частая ошибка в коде: цель модуля глубины j в позиции t равна tokens[t + j + 1], а не tokens[t + j]. Хвост, где цели нет, маскируют; при одной глубине основная цель обязана совпасть с обычной CE со сдвигом

Код → nanolm/rnn.py: VanillaRNN, LSTM (со сверкой против nn.RNN и nn.LSTM) и gradient_norm_over_time, норма градиента лосса последнего шага по каждому h_t. Задание в exercises/rnn.py, проверка: NANOLM_IMPL=exercises pytest tests/test_rnn.py -v. Код → nanolm/moe.py: роутер route (softmax или сигмоидный гейт, смещение влияет только на выбор), load_balancing_loss, bias_update, ёмкость expert_capacity, подсчёт moe_param_counts и слой MoE с общими и маршрутизируемыми экспертами. Задание в exercises/moe.py: NANOLM_IMPL=exercises pytest tests/test_moe.py -v. Проверь руками, что при перекосе загрузки лосс балансировки больше α, а смещение выравнивает загрузку, не трогая веса смешивания. Демонстрация: python scripts/rnn_vanishing.py. У RNN норма падает экспоненциально, у LSTM с bias forget-гейта 1 падает медленнее, с bias 3 почти не падает. Дополнительно повтори замер для RNN со спектральным радиусом W_h 0.9, 1.0 и 1.1 и сверь наклон log‖∂L/∂h_t‖ с log ρ. Затухание перестаёт быть словом и становится наклоном прямой.

Математика (трек D): D29: затухание и взрыв градиентов в RNN; D30: стационарное распределение и время возврата.

Интервью-вопрос недели: «Mamba рекуррентна. Почему она обучается параллельно и чем за это платит?» Структура на 3 минуты: линейная рекуррентность = свёртка → селективность ломает свёртку → её заменяет scan благодаря ассоциативности → цена: состояние фиксированного размера, слабее точное извлечение → отсюда гибриды.

Источники: Hochreiter & Schmidhuber, LSTM (1997); Pascanu et al., *On the difficulty of training recurrent neural networks* (2013); Gu et al., S4 (2021); Gu & Dao, Mamba (2023); Fedus et al., Switch Transformers (2021); Lepikhin et al., GShard (2020); Dai et al., DeepSeekMoE (2024); Wang et al., Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts (2024); DeepSeek-AI, DeepSeek-V3 Technical Report (2024); Takens, Detecting strange attractors in turbulence (1981).

Глубже: 05-ГЛУБИНА, разделы «Недели 19 и 22. Регрессия на гауссовских процессах» и «Что ещё добавить по мелочи».

Результаты недели

Могу вывести ∂h_T/∂h_k для RNN и через спектральный радиус объяснить затухание и взрыв.
Могу показать через ∂c_t/∂c_{t−1}, почему LSTM не затухает, и связать это с residual stream.
Могу реализовать RNN и LSTM с нуля и показать разницу в затухании графиком.
Могу за 3 минуты объяснить, почему линейный SSM обучается свёрткой, а Mamba обучается scan'ом.
Могу посчитать для MoE с E экспертами и top-k рост параметров и FLOPs на токен.
Могу посчитать вспомогательный лосс и ёмкость эксперта и объяснить, почему смещение роутера не входит в веса смешивания.
Могу посчитать лосс с MTP на маленьком примере и объяснить, как модуль MTP становится черновиком на инференсе.

Самопроверка

Почему клиппинг градиента спасает от взрыва, но не от затухания?
Что в Mamba «селективно» и почему из-за этого нельзя обучать свёрткой?
Что произойдёт с MoE без балансировки нагрузки и как это чинят?
Fine-grained experts не меняют FLOPs и число параметров. Что тогда они выигрывают?

Mock-интервью недели (1 и 2 из 12). С этой недели mock идут по два в неделю; протоколы сессий A–E собраны в неделе 23, оценка по рубрикам. (1) ML coding, 45 минут по протоколу сессии B: attention с маской и GQA с пустого файла. (2) Rapid-fire, сессия A: 30 вопросов из самопроверок недель 1–18. Партнёра нет: запиши ответы на диктофон и оцени по рубрике через сутки.