"Я тут в свободное время продолжаю интересные мне эксперименты и решил вот развить давние темы Миши Бурцева про Memory Transformer (https://arxiv.org/abs/2006.11527). Как вы могли заметить, мне очень нравится архитектура Universal Transformer (UT), я много про неё пишу и даже порывался собрать на ней модель-замену TRM, но не успел (https://t.me/gonzo_ML/4437). Мысли всё равно не оставляли и раз не получилось собрать URM, то почему бы не сделать тогда следующий логичный шаг, добавив память. Что я и сделал. Идея Universal Transformer + Memory мне нравится особенно, потому что это ещё ближе к полноценному компьютеру (хотя и обычный UT по идее уже Turing-complete). Подход сработал. Если взять однослойный UT, добавить динамическую рекурсию этого слоя с механизмом Adaptive Computation Time (ACT), который динамически решает, когда останавливать обработку конкретного токена, и начать учить его решать сложные судоку, то без памяти это не удаётся. Зато если добавить хотя бы 8 токенов памяти, то дело идёт существенно лучше! У голов внимания при этом явно наблюдается разная специализация по части работы с памятью. Интересно. Параллельно удалось разобраться с проблемой большой зависимости от случайного сида. На некоторых сидах обучение происходило (среди чемпионов, что забавно, был сид 42), на некоторых -- нет. Оказалось, что сильная зависимость от сида была проблемой архитектуры, а вернее инициализации роутера в механизме ACT. При правильной инициализации обучение случается на всех протестированных сидах. Код для экспериментов я оформил в виде JAX кода, используя свежее NNX API, гонял на TPU v6e. Код выложил на гитхаб. Постарался сделать его минималистичным, чтобы любой мог начать экспериментировать. Также мне лично нравится практика записывать и сохранять ADR (Architecture Decision Records), в которых отражены важные архитектурные решения с объяснением, почему они были приняты и какие были альтернативы. Это и в классическом SWE очень полезная вещь, и в ML-research мне так же всегда её не хватало -- успеваешь перепробовать много разных подходов, откинув множество разных веток, и если это не логгировать, то легко потеряться. А кроме того, для всех, кто впоследствии будет работать с этим кодом, код превращается из замороженного артефакта в живую летопись, объясняющую почему код именно такой какой есть. Наверное, бывают репозитории с экспериментами, где это понятно и очевидно, но мне в целом такие не попадались. Надо задать тренд! А ещё это хорошее подспорье для агентов, помогающих в рисёче. Такие дела. Читайте, критикуйте, пробуйте, дополняйте. Universal Transformers Need Memory: Depth-State Trade-offs in Adaptive Recursive Reasoning __Grigory Sapunov__ Статья: https://arxiv.org/abs/2604.21999v2 Код: https://github.com/che-shr-cat/utm-jax Ревью: https://arxiviq.substack.com/p/universal-transformers-need-memory Пост про подход: https://gonzoml.substack.com/p/why-i-keep-coming-back-to-universal # TL;DR ЧТО сделали: Исследователи (1 шт.) представили одноблочный Universal Transformer, дополненный явными токенами памяти и модифицированным механизмом Adaptive Computation Time (ACT). Они показали, что устранение неочевидной ловушки при инициализации роутера позволяет этой компактной модели решать сложные комбинаторные задачи на рассуждение (наподобие судоку из датасета Sudoku-Extreme (https://huggingface.co/datasets/sapientinc/sudoku-extreme)), при условии наличия достаточного объёма внутренней памяти. ПОЧЕМУ это важно: Работа предоставляет эмпирические доказательства того, что одной лишь глубины архитектуры — даже если она динамически адаптируется — недостаточно для сложных рассуждений без постоянного, выделенного пространства состояний. Авторы диагностировали и исправили давнюю проблему с инициализацией ACT. Это открывает более стабильный путь к созданию эффективных по числу параметров рекурсивных моделей ""Системы 2"", способных масштабировать вычисления на инференсе."