С нуля
EN Открыть

Программа курса

Неделя 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; её спектральный радиус (наибольший модуль собственного значения) описывает асимптотическое поведение её степеней. В нелинейной RNN добавляется зависящая от состояния производная tanh: спектрального радиуса одного W_h недостаточно.

  • 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 нужно учитывать весь набор зависящих от состояния Якобианов и их сингулярные числа. Даже при ρ(W_h) > 1 насыщение 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). Для causal decoder состояние на позиции t содержит только префикс: S_t = Σ_{j≤t} k_j v_jᵀ, выход q_t S_t. При равных размерностях ключей и значений состояние имеет размер H×H; глобальное KᵀV открыло бы будущие токены. Цена обоих подходов: хуже точное извлечение из далёкого контекста. Отсюда гибриды, чередующие слои внимания и SSM

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

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

  • Шаги 1–3 шли от состояния к наблюдениям: RNN и SSM сжимают историю в h_t. Теорема Такенса (1981) даёт обратный результат для гладкой динамики на компактном многообразии размерности d: для generic пары динамики и наблюдаемой окно X_t = (x_t, x_{t−τ}, …, x_{t−(E−1)τ}) с E = 2d+1 задаёт вложение. Произвольные наблюдаемая и задержка такой гарантии не имеют. Версии для фрактальных аттракторов используют другие условия; d исходной теоремы не произвольная оценка размерности аттрактора
  • У идеального маятника при фиксированной энергии регулярная замкнутая траектория одномерна. Два измерения положения могут различить фазу при подходящей задержке, но задержка в целый период повторит то же число и ничего не добавит. Проверяй задержку, а не только число координат
  • Мост к 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⁹. Число FFN-параметров и основные FFN-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 среди всех T · k назначений) получается из выбора top-k и градиента не имеет, поэтому градиент течёт через Pᵢ (средняя по батчу вероятность роутера для i): ∂L_aux/∂Pᵢ = α·E·fᵢ, сильнее всего давит вероятность перегруженного. Множитель E делает значение при равномерной нагрузке равным α при любом E; для Switch k = 1. α маленький, порядка 10⁻²: сильный лосс спорит с основной задачей

Ёмкость эксперта. На пальцах. В батче 96 токенов, E = 8, top-2, capacity factor (запас ёмкости) CF = 1.25: у эксперта 1.25·96·2/8 = 30 мест. Эксперту выпало 37 назначений: 7 лишних назначений он не обрабатывает. При top-2 другой эксперт ещё может обработать тот же токен; только при потере всех назначений остаётся один 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 и сравни траектории градиента при слабом и сильном насыщении tanh. Наклон log ρ используй как контроль для скалярной линейной рекуррентности, а не как точное предсказание для tanh.

Математика (трек 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 уменьшает затухание при forget-гейтах около 1, почему не гарантирует его отсутствие, и связать это с residual stream.
  • Могу реализовать RNN и LSTM с нуля и показать разницу в затухании графиком.
  • Могу за 3 минуты объяснить, почему линейный SSM обучается свёрткой, а Mamba обучается scan'ом.
  • Могу посчитать для MoE с E экспертами и top-k рост параметров и FLOPs на токен.
  • Могу посчитать вспомогательный лосс и ёмкость эксперта и объяснить, почему смещение роутера не входит в веса смешивания.
  • Могу посчитать лосс с MTP на маленьком примере и объяснить, как модуль MTP становится черновиком на инференсе.

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

  1. Почему клиппинг градиента спасает от взрыва, но не от затухания?
  2. Что в Mamba «селективно» и почему из-за этого нельзя обучать свёрткой?
  3. Что произойдёт с MoE без балансировки нагрузки и как это чинят?
  4. 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. Партнёра нет: запиши ответы на диктофон и оцени по рубрике через сутки.