Интересно, что это отличается от латентного ризонинга в стиле Coconut (https://t.me/gonzo_ML/3567), там он был на уровне токенов при авторегрессионной генерации, тут же он скорее на уровне глубины вызова модели, разворачивание идёт в другом измерении. Авторы пробовали другое количество фич, как в большую, так и в меньшую стороны. Один из вариантов разбивал `z` на множество фич `z_i` и каждая итерация рекурсии обновляла одну конкретную фичу из них, перенося все остальные как есть. Качество упало относительно варианта с двумя фичами (но не удивлюсь, если следующая работа найдёт этому причины, типа проблем с градиентным сигналом, и предложит хаки, чтобы это заработало). Авторы объясняют это тем, что нет какой-то причины для разбиения `z` на множество частей. Другая крайность, с одной фичей z_H, привела к ещё более серьёзному падению качества. Авторы объясняют это необходимостью для модели хранить решение `y` внутри `z` (но может надо просто отскейлить скрытый эмбеддинг?). В предложенной схеме отпадает и необходимость в двух отдельных сетях H и L (и соответственно в 2x параметров) -- решаемая задача `z ← f_L(x + y + z)` или `y ← f_H(y + z)` определяется по наличию или отсутствию `x` на входе (но что это значит на уровне обучения модели и структуры эмбеддингов -- для меня вопрос). В итоге имеем одну сеть вместо двух, и абляции показали, что это даёт улучшение на Sudoku-Extreme с 82.4% до 87.4% (но тоже не удивлюсь, если при скейлинге окажется, что это таки играет, просто недообучили и не вбухали компьюта сколько надо). 🗼 Другие архитектурные модификации Авторы пробовали докинуть больше слоёв в модели, но получили оверфиттинг. Как они пишут __“Surprisingly, we found that adding layers decreased generalization due to overfitting.”__ Не знаю, честно говоря, что тут такого прямо surprising, ML 101 какой-то, ну да ладно. В современный век видимо все привыкли, что надо скейлить. В другую сторону, уменьшая слои, но увеличивая количество рекурсий, чтобы эффективная глубина и количество компьюта оставались теми же, обнаружили, что 2 слоя это оптимум, дают на Sudoku-Extreme улучшение с 79.5% до 87.4%, а число параметров в два раза меньше (5M вместо 10M). Авторы ссылаются на работу “Fixed Point Diffusion Models” (https://arxiv.org/abs/2401.08741), где вроде тоже два слоя оказались оптимальны в контексте deep equilibrium diffusion models, но там перформанс был аналогичен более тяжёлым моделям, а здесь он прямо выше. Less is more. Маленькая сеть в сочетании с глубокой рекурсией и deep supervision позволяет обойти оверфиттинг при малом количестве данных. Интересно, как было бы при скейлинге датасета. Был эксперимент под влиянием MLP-Mixer (https://t.me/gonzo_ML/776) с заменой self-attention на MLP, работающем на всей длине последовательности, поскольку он требует меньше параметров для случая, когда длина контекста (L) меньше скрытой размерности (D). Это улучшило результат на Sudoku-Extreme с 74.7% до 87.4%, но ухудшило для Maze-Hard и ARC-AGI, требующих большего контекста. TRM упрощает механизм адаптивного времени вычислений (ACT). Можно отказаться от отдельного вычисления для continue loss, достаточно иметь halting probability, тогда минус один forward pass модели, но по-прежнему достаточно точное определение, когда модели нужно остановиться, что значительно ускоряет процесс. Это дало слабое улучшение с 86.1% до 87.4% (непонятно какие здесь доверительные интервалы). Также имплементировали экспоненциальное скользящее среднее (EMA, 0.999) для весов, поскольку на малом количестве данных модель быстро оверфитится и начинает расходиться, тоже улучшает качество с 79.9% до 87.4% (ну или скорее ухудшает с 87.4% до 79.9%, когда от полной модели это отнимают). Как я понимаю, берётся предыдущее сглаженное значение параметров модели с весом 0.999 и добавляется новое с весом 0.001.