- Приближённый градиент

Ключевое нововведение HRM — это способность эффективно обучать такие глубокие рекуррентные процессы. Модель обходит стороной ресурсоёмкое по памяти и биологически неправдоподобное обратное распространение ошибки во времени (Backpropagation Through Time, BPTT), которое требует O(T) памяти. Вместо этого она использует одношаговую аппроксимацию градиента, теоретически обоснованную в Deep Equilibrium Models (DEQ, https://arxiv.org/abs/1909.01377).

Этот подход использует теорему о неявной функции (Implicit Function Theorem, IFT), которая позволяет вычислить градиент неподвижной точки, не разворачивая вычисления. Аппроксимируя обратную матрицу Якоби (I - J_F)?? единичной матрицей I, модель может вычислять градиенты с постоянным расходом памяти O(1). По сути, это упрощение предполагает, что каждый рекуррентный шаг является стабильным уточнением, позволяя модели вычислять градиент путём обратного распространения ошибки только через самый последний вычислительный шаг, а не разворачивая всю историю.

В итоге градиент от выхода ко входу течёт через финальное состояние H-модуля в финальное состояние L-модуля и затем уже на вход. Сходу кажется, что мы много теряем, не пропуская градиент через все финальные состояния L и соответствующие им H, но может в следующей версии.

- Deep supervision

Для дальнейшей стабилизации обучения в HRM используется глубокий супервижн (deep supervision), вдохновлённый принципом, что нейронные осцилляции регулируют обучение в мозге. Может я не до конца понял идею, но мне кажется само наличие H и L модулей уже прямая отсылка к динамике, разворачивающейся на разных частотах, со всеми этими альфа, бета, тета ритмами. Но с другой стороны этот супервижн можно рассматривать как неявный модуль ещё более высокого уровня, чем H, я бы назвал его S.

Для каждого сэмпла (x,y) делается множество forward pass HRM, каждый из которых называется сегментом. В конце каждого из них вычисляется лосс и обновляются параметры. Важно, что скрытое состояние отсоединяется от вычислительного графа перед передачей в следующий сегмент, что действует одновременно как регуляризатор и эффективный обучающий сигнал. То есть градиенты из сегмента m+1 не влияют на сегмент m. Это сильно напоминает подход с recycling в AlphaFold 2 (@gonzo_ML649), где 3D структура белка с выхода системы отправлялась снова на вход на последующее улучшение.

Количество сегментов определяется динамически через ACT.

- Adaptive computational time (ACT)

Это вообще моя любимая тема, про ACT я писал давно и много (@gonzo_ML71).

Вдохновляясь способностью мозга переключаться между быстрым, автоматическим мышлением («Система 1») и медленным, обдуманным («Система 2»), HRM интегрирует ACT. Алгоритм Q-learning решает, «остановиться» или «продолжить» вычисления после каждого сегмента, основываясь на состоянии H-модуля.

Количество сегментов с вероятностью ? ограничивается сверху M_max (гиперпараметр) и снизу M_min (случайная величина полученная равномерным сэмплингом из множества {2, …, M_max}, и с вероятностью 1?? устанавливается в 1. Остановка происходит когда превысили M_max или когда Q-значение для остановки превысило Q для продолжения и мы уже превысили M_min.

Общая функция потерь объединяет стандартный лосс для задачи sequence-to-sequence с лоссом Q-обучения, позволяя модели динамически распределять вычислительные ресурсы. В целом ACT процесс позволяет достигать качества сравнимого с фиксированным прогоном M_max сегментов, тратя заметно меньше вычислительных ресурсов.

Для inference-time scaling достаточно установить повыше M_max без какого-либо обучения

Другой недавний пример использования ACT в рекурсиях -- работа про Mixture-of-Recursions (MoR, @gonzo_ML_podcasts489)

- Имплементация