Training Large Language Models to Reason in a Continuous Latent Space
https://arxiv.org/abs/2412.06769

Очень понравилась недавняя статья от Meta с, как это обычно и бывает, простой и гениальной задумкой – давайте заменим шаги в CoT на какие-нибудь тренируемые репрезентации и посмотрим, эффективнее ли это оказывается для обучения всякому ризонингу (конечно оказывается, что да). Как сейчас все любят отмечать, некоторые токены сложнее остальных, а ресурсов тратится при генерации одинаково, что особенно отражается на длинных CoT-цепочках. Большая часть сгенерированных там слов нужна скорее для поддержания логики языка и структуры предложений, и это мало помогает модели как-то ризонить.

Технически это реализуется так:
- Модель тренируется в несколько этапов. На первом все шаги CoT сохрянются на естественном языке.
- На следующих этапах первые несколько шагов рассуждения заменяются на то, что авторы называют continuous thoughts – по факту это hidden state предыдущего токена. Получается, что вместо обычных входных эмбеддингов у нас получается входная последовательность [e(x1), e(x2)....h(x_i),....,h(x_j-1), e(x_j)....], где e - это эмбеддинги, а h - hiddent states соответственно. По бокам от continuous thoughts ставятся еще специальные служебные токены.
- Каждый следующий этап все больше шагов заменяются на continuous thoughts. В статье авторы на каждом шаге заменяли следующие 2 шага решения и обучали в 3 этапа

Больше всего такой подход повышает качество на QA задачах, например, на ProntoQA этот метод выбивает 99% accuracy, генерируя при этом в 10 раз меньше токенов, чем обычный CoT. На бенче ProsQA, где модели нужно ответить на 5-hop questions (то есть "в уме" подставить пять пропущенных концептов в рассуждении), accuracy повышается на 20%, и используется при этом в 3.5 раза меньше токенов

В статье получились очень интересные аблейшены. Например, даже если заставить модель, обученную таким образом, генерить обычный CoT, то ее качество все равно будет выше. Авторы это атрибутируют к тому, что благодаря способу тренировки, модель учится очень хорошо планировать наперед – если ей доступен только вопрос и последние шаги рассуждения, то ей нужно хорошо понимать, как в итоге к финальным выводам прийти

Еще более интересно – модель учится приоретизировать потенциально более интересные / правильные направления мыслей. Если заставить модель свичнуться на естественный язык после первой continuous thought и посмотреть на вероятности токенов для ответа на следующий шаг, то модель сразу "занижает" вероятности очевидно неправильных вариантов, но при этом примерно с одинаковой уверенностью предсказывает кандидатов, между которому она еще не может решить. После следующих размышлений модель итеративно апдейтит эти вероятности. Это говорит о том, что модель не просто "сжимает" цепочки рассуждений в какой-то более компактный эмбеддинг, а что она оказывается способна держать в голове несколько веток решений параллельно, между которыми она может довольно долго выбирать и только потом сгенерировать первый токен ответа.