Head mixing convolution позволяет перемешивать внимание между разными головами в пределах одного временного шага. Все головы внимания разбиваются на группы заданного размера и перемешивание происходит внутри группы (его также можно рассматривать и как небольшой полносвязный слой). Это делается после софтмакса, но при желании можно делать и до, на логитах, тоже получается pre и post (по дефолту).

Итого, возможны четыре варианта блока MTA с разными комбинациями pre/post свёрток. Тут есть простор для оптимизации, так если оба варианта pre или post, то можно объединить это в одну трёхмерную свёртку.

Group normalization with depth scaling использует GroupNorm и независимый скейлинг для каждой головы по рецепту от Differential Transformer (https://arxiv.org/abs/2410.05258, может кстати тоже его разобрать?).

Эксперименты начинают с игрушечной задачи: модели дают последовательность блоков, каждый из N случайных букв. Далее следует L<N букв вопроса. Задача -- найти блок, содержащий все буквы из вопроса в любом порядке. Модель должна вывести все буквы целевого блока, или только его первый либо последний токен (три разные варианта задачи). Для стандартного трансформера задача сложная, так как требует L кусочков информации для определения целевого блока, и их надо закодировать в один query вектор. С MTA должно быть проще, так как он может найти позицию каждой буквы, а потом свёрткой увеличить вес внимания, если все L букв найдены вместе.

Проверили на N=5 и 8, L=2. Пример задачи (надо найти блок с pb):

hjnvt.qfjgt.whftb.bjtpq. ...(many blocks)... .pxjvf.ulhik.qoiax#pb

Обучали на 1M таких блоков, тестировали на отложенных 1K. Трансформер 4 слоя, 2 головы, размерность 256.

У MTA ошибка почти везде ноль или рядом, у обычного трансформера почти везде двузначные числа процентов. Размеры свёрток были c_q=2 (как L), c_k=2N-1, чтобы можно было покрыть весь блок. Свёртка для голов не использовалась.

Следующий эксперимент с LLM. Предобучили 880M модели с архитектурой LLaMa и сравнили обычный трансформер, Differential Transformer и MTA. Обучали на SlimPajama на 105B токенов. В MTA key-query convolution использовали в каждом четвёртом слое, а head convolution в каждом. Свёртки c_q=6, c_k=11, размер группы 2.

По перплексии MTA лучше (GroupNorm при этом важен). На наборе бенчмарков тоже обычно бьёт остальных, но не везде и разница часто в последней цифре (и непонятно какой доверительный интервал -- обучали дважды). В среднем лучше.

Проверили на отдельном пуле long-range dependency задач: LAMBADA, NeedleIn-A-Haystack и BabiLong. На ламбаде однозначно бьёт, на multi-needle (2,4,6) retrieval точность MTA обычно выше, причём без GroupNorm часто лучше. На BabiLong и QA1-5 у MTA тоже всё хорошо.

Приложили сколько-то визуализаций свёрточных ядер, заметное число близко к identity, но есть и более хитрые. Например, один с диагональной структурой, удобен чтобы находить точное совпадение с паттерном. Или есть аналог edge detection, усиливающий первый или последний из последовательных ключей с высоким вниманием. В свёртках по головам частый паттерн это контраст, вычитание одной головы из другой.

Абляции показали, что даже пары MTA слоёв достаточно для превосходства над бейзлайнами. Все предложенные компоненты что-то улучшают по перплексии.

В целом забавно. Кажется, свёртки по q/k это ещё не предел. Для каких-то задач и языков не удивлюсь, если более забористые и менее локальные интеракции рулят. Главное чтоб параметров много не добавляли. Здесь в примере с LLM разница была на уровне 0.001% (+10K параметров на фоне 880M).

По памяти и FLOPS текущая неоптимизированная имплементация сильно проигрывает у использующих обычное scaled dot product attention: памяти раза в три больше надо, флопсов меньше раз в пять. Но это скорее проблема отсутствия оптимизированного ядра для CUDA. Интересно, компиляция через XLA что бы дала.

Brand new 2025 AI Index Report is released!

https://hai.stanford.edu/ai-index/2025-ai-index-report

Вчера в разборе Multi-Token Attention упоминалась статья про Differential Transformer. До её разбора я всё-таки не доберусь, так что продолжаю эксперименты с автоматизированным разбором.

Результат тут.

Пообщались с Иваром Максутовым и Постнаукой о хренах и пряниках разном

Что ждёт профессии, связанные с переводами, в ближайшие годы? Каким станет мышление и коммуникация в эпоху AI? Как LLM справляются с пониманием языка?

https://youtu.be/jWVbaCiN0Tc

Об этом — в подкасте с Григорием Сапуновым, соучредителем и техническим директором компании Intento, специалистом в области ИИ и анализа данных.

Теперь признанная классика. Статьи про Adam и механизм внимания, ещё на базе RNN

https://blog.iclr.cc/2025/04/14/announcing-the-test-of-time-award-winners-from-iclr-2015/

30 мая в Royal Institution в Лондоне будет лекция Хинтона. Есть онлайн.

Discourse: Digital intelligence vs biological intelligence

https://www.rigb.org/whats-on/discourse-digital-intelligence-vs-biological-intelligence

Tech headlines in the last couple of years have been dominated by Artificial Intelligence. But what do we mean by intelligence? What has AI learned from biological intelligence, and how do they still differ?

Acclaimed computer scientist, and winner of the 2024 Nobel Prize in Physics, Geoffrey Hinton will examine the similarities and differences between artificial and biological intelligence, following his decades of ground-breaking work which has enabled the neural networks of today.


Спасибо Мише Бурцеву за ссылку!

Training Large Language Models to Reason in a Continuous Latent Space
Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston, Yuandong Tian
Статья: https://arxiv.org/abs/2412.06769
Код: https://github.com/facebookresearch/coconut
Песня: https://www.youtube.com/watch?v=PKQPey6L42M

Статья, которую надо разобрать, про Coconut и ризонинг в латентном пространстве. Тем более даже Quanta уже написала, а мы всё нет.

📃TL;DR.

Идея проста: языковое пространство может быть не лучшим выбором для ризонинга через CoT (про CoT и ToT тут https://news.1rj.ru/str/gonzo_ML/1885), и ризонинг можно делать не выходя в пространство токенов. Отсюда рождается Coconut (Chain Of CONtinUous Thought). Последнее скрытое состояние можно не декодировать в токен, а сразу подавать его на вход декодера в авторегрессионном процессе генерации как эмбеддинг для следующего шага.


💡Идея

Это интересный подход по нескольким причинам.

Во-первых, прогон всего ризонинга через токены создаёт узкое место. Из одного эмбеддинга могут быть сгенерированы разные токены, эквивалентные и не очень, и всё богатство оригинальной “мысли” может потеряться. Что если его не терять?

Дополнительный минус в том, что на каждый токен уходит одинаковый вычислительный бюджет, хотя токены неравнозначны. BLT (https://news.1rj.ru/str/gonzo_ML/3109) отчасти решал эту проблему, но там и часть авторов та же, что у кокоса.

В общем, почему бы не попробовать ризонинг без языковых ограничений? Тем более, что, во-вторых, это соответствует некоторым данным по neuroimaging, когда в процессе ризонинга языковые области мозга не задействуются. В частности, одна из свежих работ (https://news.1rj.ru/str/gonzo_ML/2797) говорит, что язык оптимизирован для коммуникации, а не для думания.


🛠Реализация

Она, как уже упоминалось в TL;DR, чрезвычайно проста. В процессе ризонинга мы не задействуем выходную голову модели и её слой эмбеддинга, а в качестве входного эмбеддинга на следующем шаге используем выходной эмбеддинг предыдущего шага.

LLM переключается между двумя режимами: языковым (language mode) и латентным (latent mode). Языковой режим — это стандартный режим для LLM с генерацией токенов. Латентный режим -- новый режим с переиспользованием эмбеддинга. Начало и окончание латентного режима обрамляется токенами <bot> и <eot>. Соответственно для всех токенов с индексами между индексами этих двух токенов включается новый режим. Весь процесс полностью дифференцируемый и позволяет обучать модель обычным бэкпропом.

Вопрос, где взять данные для такого обучения. Для этого берутся языковые данные для обычного CoT и реализуется мультиэтапный curriculum. На начальном этапе (stage) модель обучается обычному языковому CoT. На последующих этапах, для шага номер k, первые k шагов языкового размышления убираются, а внутри <bot>/<eot> тэгов появляются k позиций, в каждую из которых записываются эмбеддинги предыдущего шага. Состояние оптимизатора сбрасывается между отдельными этапами.

Во время обучения оптимизируется обычный negative log-likelihood loss, но маскируется и не учитывается лосс для вопроса и латентных мыслей (не буду кавычки каждый раз ставить вокруг слова мысль). Эта objective никак не стимулирует модель сжимать убираемые текстовые мысли, так что есть потенциал выучивания более эффективных репрезентаций ризонинга.

Для примера с N шагами рассуждения надо сделать N+1 forward pass. KV-кеш может помочь не пересчитывать одно и то же, но всё равно последовательная природа вычисления этих этапов не даёт эффективно распараллеливать весь процесс. Оптимизация является отдельным интересным направлением исследования, и я уверен, что за ней дело не станет.

При инференсе главный челлендж это решить, когда нужно уходить в латентный режим и выходить из него. С токеном <bot> всё просто, ставим сразу после вопроса, а для <eot> рассматриваются две стратегии: 1) обучить бинарный классификатор, решающий по эмбеддингу, когда надо переключаться, и 2) добивать паддингом латентные размышления до фиксированной длины. Оба подхода работают, так что по дефолту используют более простой второй.


🧪Оценка

Тестируют подход на трёх датасетах с математическим (GSM8k) и логическим ризонингом (ProntoQA, и новый ProsQA).

Проверяют на предобученной GPT-2. Для математики используют две латентных мысли (c=2) на каждый шаг ризонинга, три этапа (stages, k=3) и обучают шесть эпох на первом этапе и три на остальных. На логике одна латентная мысль на шаг, шесть этапов, обучают пять эпох на этап.

В сравнении используются следующие бейзлайны:
1) обычный CoT с файнтюнингом модели на примерах
2) No-CoT, с обучением модели сразу выдавать ответ
3) iCoT, implicit CoT из https://arxiv.org/abs/2405.14838, который постепенно интернализировал промежуточные шаги рассуждений через последовательный файнтюнинг, там шаги по одному выбрасывались, а в Coconut вместо них появляются латентные шаги, что по идее даёт модели “подумать” побольше.
4) Pause token, когда между вопросом и ответом вставляются специальные токены <pause> (столько же, сколько continuous thoughts у кокоса) -- никакой цепочки размышлений тут нет, но дополнительные токены могут дать модели дополнительные вычислительные возможности.

Сам Coconut тоже проверяют в трёх режимах:
1) w/o curriculum без мультиэтапного обучения, сразу берутся данные с последнего этапа, где уже нет языковых мыслей, есть только латентные.
2) w/o thought с мультиэтапным обучением и постепенным убиранием языковых шагов рассуждения, но без использования непрерывных латентных мыслей -- по сути похоже на iCoT, но процедура обучения из Coconut
3) Pause as thought c заменой непрерывных мыслей на токены <pause> и сохранением мультиэтапной процедуры


🏁 Результаты

Coconut стабильно лучше LLM без CoT, и лучше CoT на логических задачах. На GSM8k обычный CoT лучше, но у Coconut растёт качество с увеличением числа мыслей на шаг (насыщения там не видно до двух мыслей на шаг, но в приложении попробовали три и там стало хуже, видимо из-за проблем в обучении, нужно что-то менять в процедуре).

Токенов у Coconut при этом меньше в разы.

Coconut с pause as thought в целом тоже неплохо работает, но обычный Coconut лучше.

На авторском датасете ProsQA, где предсказание следующего шага не очень помогает и надо планировать и искать по более сложному графу, обычный CoT не лучше No-CoT, зато Coconut или iCoT существенно улучшают ситуацию.

Без мультиэтапной процедуры с curriculum модель плохо выучивает ризонинг в латентном пространстве. В идеальном мире она бы сама выучила наиболее эффективные непрерывные мысли через бэкпроп, но чего-то не хватает.

Хоть Coconut и призван обходиться без перехода в пространство токенов, это по-прежнему можно делать, если хочется. В одном примере авторы декодировали первую непрерывную мысль и увидели там токены, которые ожидались в промежуточном рассуждении. Это прикольно в плане интерпретируемости.


🔍Анализ

У модели есть интересная возможность переключаться между ризонингом языковым и в латентном пространстве. Авторы дополнительно покопались в латентном ризонинге.

Проверили работу модели на датасете ProsQA с разным числом латентных мыслей, от нуля до шести. Разница здесь только в инференс тайме, модель одна и та же. Также дополнили мультиэтапную процедуру обучения, чтобы модель не забывала более ранние этапы, так что с вероятностью 0.3 замешиваются данные от других этапов.

Сделали более гранулярную классификацию качества ответа модели, теперь есть не просто корректность конечного ответа, а более детальные типы:
1) Correct Path — верный кратчайший путь в графе
2) Longer Path — верный, но не кратчайший
3) Hallucination — путь содержит несуществующие рёбра или несвязный
4) Wrong Target — путь валиден, но не к тому узлу
5) Correct Label и (6) Incorrect Label — для методов, где можно получить только финальный ответ

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