Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Языковую модель научили сохранять между генерируемыми токенами больше контекста, чем помещается в один выбранный вариант. В препринте Dor Tirosh, Ido Amos и Mor Geva, который не прошёл рецензирование и где все числа получили сами авторы, LIFT решил последовательности из 12 шагов со 100% точностью, тогда как Transformer, который служил учителем, показал 3%. Метод касается архитектуры и предобучения базовых моделей, поэтому его нельзя добавить к готовой LLM через обычную донастройку.
Модель передаёт следующему шагу несколько вариантов
В обычном авторегрессионном Transformer результат глубоких слоёв не возвращается в нижние слои на следующем шаге. Модель выбирает токен, снова подаёт его на вход и теряет остальные варианты вместе с информацией о том, насколько они были вероятны. Такой канал переносит около 17 бит при типичном размере словаря.
LIFT передаёт вместе с токеном разреженное распределение вероятностей. Оно хранит наиболее вероятные продолжения и их веса. Отдельный слой преобразует это распределение через таблицу векторных представлений токенов, объединяет результат с новым токеном и передаёт их стандартным блокам Transformer.
Состояние строится из выходного распределения, а не из внутреннего представления другой модели. Благодаря этому учитель и ученик работают в общем пространстве: каждой координате соответствует один и тот же токен. Внутренние представления разных моделей так не совпадают, поэтому ученику было бы сложнее воспроизводить их и затем использовать собственные состояния при генерации.
Во время подготовки данных замороженная модель-учитель обрабатывает текст и выдаёт распределение для каждой позиции. LIFT получает эти состояния на вход и одновременно учится предсказывать следующий токен и следующее состояние. Поскольку состояния рассчитаны заранее, все позиции последовательности можно обучать параллельно, как в обычном предобучении.
Ближе к концу обучения учителя убирают и подают модели её собственные состояния, чтобы уменьшить расхождение между обучением и генерацией. Дополнительный слой содержит 3% параметров миллиардной модели. При обработке запроса LIFT умеет читать весь текст параллельно без состояний, а затем включает обратную связь при последовательной генерации ответа.
Обратная связь лучше всего проявилась в пошаговых задачах
Авторы проверили LIFT на языковых моделях от 135 млн до 1 млрд параметров. Крупнейшую модель обучали на 100 млрд токенов. Во всех масштабах её сравнивали со стандартным Transformer при одинаковом числе обучающих токенов, а также при сопоставимых вычислительных затратах.
При одинаковом объёме текста LIFT лучше моделировал язык и набирал более высокие результаты в задачах на рассуждение и выполнение процедур. При равных вычислениях он оставался на уровне стандартного Transformer или опережал его. На арифметике с несколькими операндами и продолжении закономерностей LIFT соответствовал модели, обученной на вдвое большем корпусе, либо превосходил её.
Синтетическая задача S5 отделяет запоминание состояния от знания языка: модель должна последовательно применять перестановки и назвать итог. Обычному Transformer фиксированной глубины нужно заново восстанавливать промежуточный результат, а LIFT переносит его дальше через состояние. Модель также смогла выйти за длины, встречавшиеся при обучении, когда получала состояния от рекуррентного учителя.
Проверка охватывает небольшие по современным меркам базовые модели на архитектуре OLMo 2, языковое моделирование, наборы задач на рассуждение и процедурные тесты. Учитель использовал ту же архитектуру и совместимый словарь. Работа поэтому показывает преимущество механизма в заданном диапазоне, но не позволяет переносить размер выигрыша на более крупные модели или учителей с другим словарём.
Когда LIFT стоит закладывать в план обучения
Для команды, которая обучает базовую модель с нуля, LIFT предлагает компромисс между рекуррентными вычислениями и параллельным предобучением. Учитель требует отдельного прохода по корпусу, но ученик не разворачивает состояния последовательно для каждого токена. В сравнении при равных вычислениях эта дополнительная работа не устранила выигрыш на проверенных масштабах.
Метод особенно уместен, если продукт требует хранить промежуточное состояние: считать цепочки операций, продолжать формальные процедуры или отслеживать изменения объекта по длинной последовательности. Здесь один выбранный токен служит слишком узким каналом, а распределение сохраняет несколько возможных продолжений.
Для продукта, который использует внешнюю LLM через API, планы не меняются: заменить архитектуру поставщика или повторить её предобучение нельзя. Обычная донастройка также не добавит канал состояния, потому что LIFT меняет вход модели, функцию обучения и режим генерации.
При собственном предобучении придётся заложить совместимого учителя, подготовку его распределений и отдельный этап адаптации к состояниям самого ученика. На стороне применения добавятся вычисления для слоя объединения и передачи состояния между токенами. При этом последовательность не создаёт нового принципиального барьера для генерации: авторегрессионная модель и без неё выпускает токены по одному.
Практический результат работы — не готовая замена Transformer, а архитектурная гипотеза для следующего цикла обучения. Проверять её имеет смысл сначала на задачах продукта, где ответ зависит от переноса промежуточного результата, и сравнивать при одинаковых вычислительных затратах, а не только при одинаковом корпусе.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
