Дата публикации
science

Apple научила классические RNN масштабироваться до 7B параметров и работать не хуже трансформеров

Что открыли

Apple показала, что классические рекуррентные нейросети (RNN) можно обучать параллельно и масштабировать до миллиардов параметров — без отказа от нелинейностей.

Ключевые результаты из работы ParaRNN: Unlocking Parallel Training of Nonlinear RNNs for Large Language Models (ICLR 2026, Oral):

  • Параллельное обучение RNN по последовательности ускорили до 665× по сравнению с классическим пошаговым (см. эксперимент на Figure 1).
  • Apple обучила первые классические RNN на 7 млрд параметров (ParaGRU и ParaLSTM), которые по perplexity и бенчмаркам конкурируют с трансформером и Mamba2 (Figure 2, Table 1).
  • На задачах языкового моделирования модели от 400M до 7B параметров показывают качество на уровне современных архитектур.
  • На синтетических задачах, где важна работа с внутренним состоянием, нелинейные RNN (ParaGRU, ParaLSTM) показывают 100% точность там, где Mamba2 и трансформер заметно проседают (Table 2).
  • На инференсе RNN сохраняют главное преимущество: время генерации токена не растёт с длиной контекста, в отличие от трансформеров (Figure 8).
  • Apple выложила open-source фреймворк ParaRNN для автоматического параллельного обучения нелинейных RNN, с тремя уровнями производительности: чистый PyTorch, CUDA-ускорение и полностью fused CUDA-ядро.

Суть идеи: Apple берёт нелинейную RNN, переписывает её как систему уравнений по всей последовательности и решает эту систему итерациями метода Ньютона. Каждая итерация сводится к линейной state space модели (как в Mamba), которую уже можно эффективно распараллелить.

Как исследовали

Кто и где

Исследование провела команда Apple Machine Learning Research. Работа принята на ICLR 2026 как Oral-доклад — это верхний уровень отбора конференции по глубокому обучению.

Что именно сделали технически

  1. Переформулировали RNN как систему уравнений.

    • Вместо пошагового обновления скрытого состояния (h_l = \sigma_l(h_{l-1}, x_l)) Apple рассматривает все скрытые состояния всей последовательности сразу как набор неизвестных.
  2. Применили метод Ньютона.

    • Метод Ньютона итеративно решает нелинейную систему, на каждом шаге линейризуя её через якобианы.
    • Локальные производные (якобианы) по скрытому состоянию играют роль матриц состояния (A_l) в линейной SSM.
    • Получается: каждая итерация Ньютона — это линейная SSM, которую можно посчитать параллельно через parallel scan.
  3. Ограничили структуру якобианов.

    • Классические GRU и LSTM дают плотные якобианы — хранить и умножать такие матрицы при больших размерностях скрытого состояния слишком дорого.
    • Apple разработала ParaGRU и ParaLSTM — адаптации классических GRU и LSTM, в которых матрицы в определении ячейки имеют ненулевые элементы только на главной диагонали.
    • Это даёт:
      • Для ParaGRU — диагональные якобианы.
      • Для ParaLSTM — 2×2 блочно-диагональные якобианы.
    • Хранение и умножение таких матриц масштабируется линейно по размерности скрытого состояния, а не квадратично/кубически.
  4. Реализовали три варианта вычислений.

    • Чистый PyTorch — для прототипирования, использует автодиф, работает с любыми RNN, но без оптимизации под масштаб.
    • PyTorch + CUDA-ускоренный parallel reduction — для ячеек с (блочно-)диагональными якобианами.
    • Полностью fused CUDA-ядро — одна CUDA-функция, которая делает сразу:
      • итерации Ньютона,
      • сборку системы,
      • параллельное сканирование (reduction) по последовательности.
    • В Figure 7 показано, что fully-fused реализация ParaGRU и ParaLSTM по времени сопоставима с Mamba на всех длинах последовательности.

Эксперименты и масштаб

  • Диапазон размеров моделей: от 400M до 7B параметров.
  • Крупнейшие модели:
    • Mamba2 — 6,96B параметров.
    • ParaLSTM — 6,76B параметров.
    • ParaGRU — 6,76B параметров.
    • Трансформер — 6,89B параметров.
  • На языковом моделировании измеряли perplexity (PPL, меньше — лучше).
  • Для оценки «как LLM» использовали lm-eval-harness с задачами:
    • Arc-C,
    • HellaSwag (HSwag),
    • OpenBookQA (OBQA),
    • WinoGrande (WinoG),
    • PiQA,
    • MMLU.
  • На синтетических задачах тестировали способность модели:
    • отслеживать внутреннее состояние,
    • хранить и извлекать информацию (MQAR, k-hop, Parity).

Сколько итераций Ньютона

В экспериментах на ParaGRU и ParaLSTM хватило трёх итераций Ньютона, чтобы сойтись к тому же скрытому состоянию, что и при классическом последовательном применении нелинейной RNN.

Это критично: если бы нужно было десятки итераций, выигрыш от параллельности пропал бы. Здесь же три итерации дают итоговое ускорение до 665× по сравнению с наивным последовательным обучением.

Что это меняет на практике

1. RNN снова становятся реальным вариантом для LLM

Долгое время у RNN было очевидное преимущество на инференсе: O(1) время на токен, независимо от длины контекста. Но обучение по последовательности шло строго по шагам и плохо масштабировалось.

Теперь Apple показывает:

  • Нелинейные RNN можно обучать с параллелизацией по длине последовательности.
  • Classical RNN на 7B параметров по perplexity и downstream-задачам конкурируют с трансформером и Mamba2.

Цифры из Table 1 (чем ниже PPL, тем лучше):

  • Mamba2 (6,96B): PPL = 8,62.
  • ParaLSTM (6,76B): PPL = 9,16.
  • ParaGRU (6,76B): PPL = 9,19.
  • Трансформер (6,89B): PPL = 9,55.

RNN не «выигрывают всухую», но находятся в одном классе качества. Для продакшн-сценариев, где критичны задержки и энергопотребление, это важный сигнал: можно выбирать не только между трансформером и SSM, но и между нелинейной RNN с параллельным обучением.

2. Инференс быстрее и стабильнее по длине контекста

Figure 8 показывает ключевую практическую разницу:

  • У трансформера время генерации токена растёт с длиной контекста (O(L)).
  • У RNN и Mamba2 — практически константа: один шаг обновления скрытого состояния.

Что это даёт:

  • Для чат-ботов и ассистентов с длинной историей диалога — ровную, предсказуемую задержку.
  • Для генерации кода, текста или аудио «на лету» — стабильный FPS/скорость.
  • Для он-дивайс сценариев (смартфоны, ноутбуки, гарнитуры) — более экономное использование памяти и вычислений.

3. Нелинейности в RNN — не роскошь, а реальный бонус

Mamba и другие SSM используют линейное обновление состояния. Это позволяет параллелизацию, но ограничивает типы динамики, которые модель может выразить.

Apple показывает, что если добавить нелинейность в шаг рекурсии и при этом научиться обучать такую модель параллельно, то на задачах, где нужно хранить и обновлять состояние, нелинейные RNN выигрывают у линейных SSM.

Table 2 (точность на синтетических задачах MQAR, k-hop, Parity):

  • Transformer: MQAR 100%, k-hop 78%, Parity 53%.
  • Mamba2: MQAR 100%, k-hop 98%, Parity 51%.
  • ParaGRU: MQAR 100%, k-hop 100%, Parity 100%.
  • ParaLSTM: MQAR 100%, k-hop 100%, Parity 100%.

Вывод: когда вам важно, чтобы модель надёжно запоминала и извлекала внутреннее состояние, нелинейные RNN дают ощутимое преимущество.

4. Для кого это особенно интересно

  • Разработчики он-дивайс LLM:

    • смартфоны, ноутбуки, гарнитуры, автомобили,
    • где память и энергопотребление ограничены,
    • где задержка на токен критична.
  • Команды, которым важен throughput:

    • генерация длинных текстов или кода,
    • потоковая обработка событий и логов,
    • задачи, где длина контекста может быть очень большой, а скорость должна оставаться высокой.
  • Исследователи архитектур:

    • теперь можно экспериментировать с нелинейными RNN на уровне миллиардов параметров,
    • не застревая в узком горлышке последовательного обучения.

Минусы и ограничения:

  • Параллельное обучение нелинейных RNN — нетривиальная инженерная задача. Чтобы получить максимальный выигрыш, нужен CUDA-код и работа с памятью GPU на низком уровне.
  • ParaGRU и ParaLSTM — упрощённые версии классических GRU/LSTM с ограниченным смешиванием компонент скрытого состояния. Это цена за структурированные якобианы.
  • Метод Ньютона требует аккуратного дизайна ячеек. Если нелинейность «жёсткая», можно получить медленную или нестабильную сходимость.

Что это значит для вас

Если вы ML-инженер или исследователь

  • Появился готовый фреймворк для параллельного обучения нелинейных RNN — ParaRNN.
  • Можно:
    • взять свою RNN-ячейку,
    • реализовать один метод recurrence_step,
    • и сразу получить параллельное обучение по последовательности.

Это снижает порог входа, если вы хотите исследовать:

  • альтернативы трансформерам для LLM,
  • новые типы рекуррентных блоков,
  • гибридные архитектуры (например, трансформер + нелинейная RNN как память).

Если вы продуктовый менеджер или архитектор систем

  • В стеке архитектур для LLM теперь есть ещё один реальный кандидат:
    • нелинейная RNN на миллиардах параметров,
    • с качеством, сравнимым с трансформером и Mamba2,
    • с лучшим профилем инференса по задержке и памяти.

Это особенно интересно, если вы:

  • планируете запускать большие модели на устройстве, а не только в облаке;
  • хотите уменьшить стоимость генерации токенов при длинных контекстах;
  • строите системы, где важна работа с внутренним состоянием (сложные диалоги, агенты, долгие сценарии взаимодействия).

Если вы уже используете модели вроде GPT-5 или Claude 4

ParaRNN не пытается заменить трансформеры прямо сейчас. Но для вас важны несколько моментов:

  • В долгосрочной перспективе часть задач может уйти с трансформеров на RNN/SSM, когда нужно:
    • постоянное время на токен,
    • экономия памяти,
    • хорошее состояние внутри модели.
  • Внутри ассистентов и агентов, которых вы используете, могут появиться гибридные архитектуры, где:
    • трансформер отвечает за богатое контекстное понимание,
    • RNN — за быстрый, дешёвый инференс и долговременное состояние.

Пользовательский опыт от этого, скорее всего, проявится в виде:

  • более быстрых ответов при длинных диалогах;
  • меньшего количества обрывов контекста;
  • появления он-дивайс функций, которые раньше требовали облака.

Как запустить

Apple выложила ParaRNN как open-source фреймворк. Ключевая идея — вы описываете только логику одного шага рекуррентной ячейки, а всё остальное (метод Ньютона, сборка якобианов, параллельное сканирование, оптимизации) берёт на себя библиотека.

Определение своей RNN-ячейки

Базовый пример из ParaRNN:

class MyRNNCell( BaseRNNCell ):
 def recurrence_step( self, h, x, system_parameters ):
 h_new = ... # \sigma( h, x; params )
 return h_new

Дальше ParaRNN автоматически:

  • применит метод Ньютона к вашей рекуррентной системе;
  • соберёт якобианы;
  • выберет нужную схему параллельного reduction;
  • при необходимости задействует CUDA-ускорение или fully-fused ядро.

Три уровня производительности

  1. Pure PyTorch

    • Для быстрых экспериментов и прототипов.
    • Использует автодиф PyTorch.
    • Работает с любыми RNN, без требований к структуре якобианов.
    • Не оптимизирован под большие модели и очень длинные последовательности.
  2. CUDA-Accelerated

    • Для ячеек с диагональными или блочно-диагональными якобианами.
    • Использует кастомные CUDA-ядра для параллельного reduction.
    • Подходит для ParaGRU, ParaLSTM и похожих структурированных ячеек.
  3. Fully-Fused

    • Одна CUDA-функция, которая объединяет:
      • итерации Ньютона,
      • сборку системы,
      • параллельный scan.
    • Требует CUDA-реализации recurrence_step.
    • Даёт максимальную скорость, сопоставимую с Mamba по времени прямого прохода на разных длинах последовательности (Figure 7).

Архитектура ParaRNN модульная: вы можете экспериментировать с разными типами ячеек, структурами якобианов и конфигурациями солвера, унаследовавшись от базовых классов и реализовав свои операции reduction.

Кому это нужно прямо сейчас, а кому можно подождать

Подходит:

  • исследовательским группам, которые работают с архитектурами LLM и SSM;
  • командам, строящим он-дивайс или low-latency LLM;
  • разработчикам фреймворков и библиотек вокруг PyTorch и CUDA.

Пока рано:

  • если вы просто используете готовые API крупной модели через HTTP и не контролируете архитектуру;
  • если ваша инфраструктура заточена под трансформеры и менять стек в ближайший год не планируется;
  • если у вас нет ресурсов на разработку и поддержку CUDA-кода.

Итог

Apple показывает, что ограничение «RNN плохо масштабируются по длине последовательности» — не фундаментальный закон природы, а инженерная проблема, которую можно решить.

  • Нелинейные RNN теперь можно обучать параллельно по последовательности.
  • На масштабе до 7B параметров они дают качество, сопоставимое с трансформерами и Mamba2.
  • На инференсе они сохраняют свои сильные стороны: константное время на токен и экономию памяти.
  • Нелинейности в рекуррентном шаге дают лучшее управление внутренним состоянием, чем линейные SSM.

ParaRNN делает эксперименты с такими моделями доступными исследователям и инженерам: вы описываете RNN-ячейку, а фреймворк берёт на себя тяжёлую часть — метод Ньютона, якобианы и параллельное обучение.

RNN не просто «вернулись», а получили шанс занять своё место рядом с трансформерами и state space моделями в стеке архитектур для больших языковых моделей.


Читайте также