Простое изменение в обучении, которое открывает возможности предвидения, более быстрых выводов и более эффективного рассуждения.
Делиться
Введение
Стандартные большие языковые модели (LLM) обучаются на основе простой цели: прогнозирования следующего токена (NTP). Максимизируя вероятность появления непосредственно следующего токена xt+1 с учётом предыдущего контекста, модели достигают выдающейся беглости и способности к рассуждениям.
Однако этот подход действительно неэффективен, поскольку модели приходится тратить одинаковое количество вычислений на предсказание слов-заполнителей (например, «the», «and», «have») и слов, несущих информацию (например, «red», «apple», «lazy»). Ситуация усугубляется тем фактом, что более 50% слов в английском языке являются словами-заполнителями (Nordquist, 2024)3. В связи с этим возникает практический вопрос: требуется ли для предсказания всех слов полный цикл вывода, или модели уже содержат слова-заполнители в скрытом состоянии задолго до их предсказания?
Мотивация для MTP
Идея о том, что трансформеры способны обрабатывать больше, чем просто следующий шаг, подтверждается недавними эмпирическими исследованиями. (Pal et al., 2023)1 продемонстрировали, что внутренние представления моделей трансформеров часто кодируют траектории будущего текста задолго до их генерации.
Для иллюстрации исследователи провели эксперимент по «трансплантации». Они извлекли скрытые состояния из модели, обрабатывающей предложение «Madison Square Garden is located in…» — непосредственно перед тем, как она собиралась предсказать следующее слово как «New». Затем они поместили этот вектор в модель, обрабатывающую совершенно не связанный с предложением контекст, например, «Расскажите мне что-нибудь о…». Несмотря на не связанную с предложением подсказку, модель авторегрессивно завершила предложение, записав его как «Расскажите мне что-нибудь о Нью-Йорке». Это подтвердило, что модель кодировала не только следующий токен, но и всю будущую последовательность.
Чтобы использовать этот скрытый потенциал LLM, исследователи из Meta FAIR (Gloeckle et al., 2024)2 предлагают новый подход. Вместо того, чтобы рассматривать это предвидение как побочный продукт, они явно используют его в качестве цели обучения. Поставив перед моделью задачу одновременно предсказывать «n» будущих токенов в каждой позиции вместо одного, они фактически смогли заставить модель прогнозировать. Авторы демонстрируют, что парадигма многотокенного прогнозирования (MTP) обеспечивает значительно более высокую производительность в различных бенчмарках, одновременно увеличивая скорость вывода до 3 раз по сравнению с базовым уровнем.
Архитектура MTP: распараллеливание прогнозирования
Если информация для следующих нескольких токенов уже встроена в текущие скрытые состояния LLM, то возникает архитектурный вопрос: как нам извлечь эту информацию заранее, не увеличивая требования к вычислительным ресурсам по сравнению со стандартным NTP?
Предлагаемая авторами архитектура направлена на модификацию существующей магистральной сети трансформатора для одновременного прогнозирования n будущих токенов. В отличие от стандартной парадигмы NTP, где потери кросс-энтропии минимизируются только для ближайшего следующего токена (xt+1), многотокенное прогнозирование (MTP) минимизирует средние потери по n различным выходным заголовкам:
Для реализации этого авторы разделяют модель на два компонента:
- Общая магистраль (fs): основная часть модели представляет собой стандартную магистраль преобразователя, задачей которой является обработка предложенного контекста x1:t в информационно-плотное глобальное представление zt, которое будет использоваться для всех последующих прогнозов.
- Независимые головки (fh_i): выходной сигнал магистрали поступает на n независимых головок. Каждая головка имеет свой слой преобразователя и отвечает за прогнозирование будущего смещения токена (например, головка 1 предсказывает t+1, головка 2 предсказывает t+2 и т. д.).
В конечном итоге выходные данные каждой отдельной головы передаются в общий слой извлечения, который реализован как простая линейная проекция скрытого измерения модели на длину словаря. Диаграмма ниже суммирует наиболее важные аспекты архитектуры MTP:
Преодоление узкого места памяти
Описанная выше архитектура представляет собой существенную инженерную проблему: использование памяти графического процессора.
Размер словаря (V) больших языковых моделей обычно находится в диапазоне от 32 до 256 тысяч слов, что астрономически много. Это делает исходные оценки предсказаний для каждого слова в словаре, также называемые выходными логитами, очень большими. В стандартной конфигурации NTP модели требуется материализовать эти логиты только один раз за шаг, что делает её легко контролируемой. Однако в конфигурации MTP одновременно генерируются n различных наборов этих огромных логитов, что может легко перегрузить память графического процессора. Это делает метод MTP непрактичным для исследователей, если только они не будут радикально уменьшать размеры пакетов, что замедлит весь процесс обучения.
Авторы обходят это узкое место с помощью стратегии последовательного прямого и обратного прохода. Вместо того, чтобы вычислять потери для всех n голов одновременно, обучающий цикл последовательно их перебирает:
- Общая магистраль вычисляет скрытое состояние zt.
- Модель вычисляет логиты для головы 1, вычисляет потери, распространяет градиенты обратно по всей модели и немедленно удаляет логиты из памяти.
- Затем этот процесс повторяется для головки 2, головки 3 и т. д.
Благодаря удалению этих массивных логит-векторов из памяти после каждого вычисления в голове пиковое использование памяти в процессе обучения остаётся равным O(V), а не O(nV). Это позволяет обучать модели MTP пакетами того же размера, что и стандартные модели.
Критические выборы дизайна
Помимо оптимизации памяти, авторы также приняли два конкретных проектных решения, которые важны для понимания показателей производительности и научной обоснованности MTP.
1. Ограничение четности параметров
В модели MTP с n=4 головками четыре дополнительных слоя головок с трансформаторными магистралями приводят к увеличению параметров. Чтобы компенсировать это увеличение, авторы удалили эквивалентное количество слоёв из ствола модели, сделав его более плоским. Это сделано для того, чтобы любые изменения производительности MTP относительно базовой модели можно было отнести исключительно к самой архитектуре MTP, а не к увеличению параметров модели.
Тот факт, что MTP по-прежнему превосходит стандартные модели на основе NTP, несмотря на меньшую глубину ствола, лишь подтверждает достоинства этой архитектуры.
2. Топология головы: параллельная и причинно-следственная
Авторы также экспериментировали с расположением самих головок, в частности сравнивая два подхода:
- Параллельные головки: Это стандартная конструкция MTP, описанная выше. В этой конструкции каждая головка предсказывает свой будущий токен, основываясь только на общем состоянии zt, не учитывая прогнозы других головок.
- Каузальные вершины: В этой конфигурации вершина 2 (предсказывающая t+2) будет получать выходные данные вершины 1 в качестве входных данных. Это создаёт «мини-авторегрессионную» цепочку в конце модели, которая позволяет каждой вершине анализировать состояние предыдущей вершины. Архитектура MTP с n=4 каузальными вершинами представлена ниже:
Удивительно, но параллельный дизайн показал лучшие результаты. Авторы предполагают, что в дизайне с каузальными головками общий ствол «ленился», полагаясь на головки для понимания последовательной информации. Но, заставляя головки действовать независимо, ствол фактически был вынужден выучить глобальное представление, которое могло удовлетворить все головки одновременно. Именно это свойство также проявляется в способности модели планировать будущее, что крайне важно для задач рассуждения.
Экспериментальные результаты: масштаб улучшения
Авторы провели обширные оценки, сравнивая модели MTP со стандартными базовыми показателями прогнозирования следующего токена (NTP) для моделей размером от 300 млн до 13 млрд параметров.
1. «Закон масштабирования» многотокенного прогнозирования
Пожалуй, самым интересным открытием является то, что производительность модели масштабируется с её размером. Для моделей меньшего размера (от 300 млн до 1,3 млрд параметров) разница между MTP и NTP незначительна (часто MTP работает хуже). Но с увеличением размера MTP начинает работать значительно лучше базового уровня. Как показано ниже, MTP превосходит NTP на 17% в тесте MBPP и на 12% в тесте HumanEval.
Возможная причина этого неравенства может заключаться в том, что более крупные модели с большим количеством параметров могут позволить себе выделить больше ресурсов для будущего планирования, чем более мелкие модели. Это позволяет более крупным моделям использовать многотокенную цель для разработки более эффективных рассуждений.
2. Тройное ускорение вывода посредством самоспекуляции
Помимо показателей производительности, MTP также решает одну из самых распространенных проблем в операциях LLM: задержку вывода.
Чтобы в полной мере оценить этот вклад, необходимо сначала понять, что такое спекулятивное декодирование. В стандартном выводе модель должна итеративно генерировать токены. Она должна дождаться генерации xt, прежде чем вычислять xt+1. Спекулятивное декодирование ускоряет этот процесс, используя меньшую и более быструю черновую модель (обычно того же семейства, что и основная модель, но с гораздо меньшим количеством параметров), которая берёт скрытое состояние из основной модели и предсказывает следующие несколько токенов. Затем основной модели поручается проверить все эти токены за один прямой проход, гарантируя их соответствие прогнозам меньшей модели. Поскольку один прямой проход быстрее, чем генерация токенов посредством множества итераций, это приводит к общему ускорению. (Подробнее о спекулятивном декодировании)
Спекулятивное декодирование обычно требует загрузки в память модели меньшего размера, что может потребовать значительных ресурсов памяти. Однако авторы предлагают использовать дополнительные головки MTP, обычно отбрасываемые после обучения, в качестве встроенной модели черновика. Поскольку эти головки имеют общий ствол, они являются высокоточными инструментами для черновика. Используя до четырёх головок для черновика подпоследовательности и затем параллельно проверяя её, MTP обеспечивает трёхкратное ускорение вывода без потери точности.
4. Более быстрое формирование «индукционных головок»
Авторы также анализируют возникновение индукционных возможностей в MTP. Индукционные головки — это цепи в трансформаторах, которые в основном отвечают за способность к сопоставлению с образцом (например, распознавание того, что за [A]…[B]…[A], вероятно, следует [B]). График ниже показывает, что при меньших размерах моделей MTP демонстрирует большую индукционную способность, чем модели NTP аналогичного размера. Это указывает на то, что, заставляя модель предсказывать последствия следующего токена, она создаёт градиентный сигнал, способствующий возникновению распознавания образов и контекстного обучения.
5. Разблокировка обучения на уровне байтов
В более радикальном эксперименте авторы применили MTP к моделям на уровне байтов, которые предсказывают последовательность байтов вместо представлений токенов. Исторически модели на уровне байтов всегда показывали плохие результаты из-за слабого контекстного содержания между байтами, а последовательности байтов, как правило, становятся очень большими. Однако, как показано в таблице ниже, при n=8 заголовках (предсказывая 8 байтов одновременно) модель MTP значительно превосходит базовый NTP с n=1 заголовком, стабильно во всех трёх тестах. Это говорит о том, что модель MTP способна эффективно ориентироваться в байтовом пространстве, позволяя моделям обрабатывать необработанные данные без ущерба для производительности.
Цена предвидения: недостатки и компромиссы
Хотя многотокенное прогнозирование предлагает убедительную альтернативу стандартной парадигме, результаты исследования показывают, что оно не является универсальным решением. Архитектура вносит определённые компромиссы, которые инженерам необходимо учитывать.
1. Регрессия для задач, требующих больших знаний
Хотя MTP улучшает рассуждение (как структурировать ответ), судя по всему, он ухудшает извлечение информации (знание конкретного факта).
Как показано ниже, модели MTP доминируют в тестах генерации кода и рассуждений, но фактически уступают базовому уровню в стандартных задачах обработки естественного языка, включая такие тесты, как MMLU, TriviaQA и ARC Challenge (которые проверяют поиск фактов и знание мира).
Возможным объяснением может быть то, что ответы на вопросы, основанные на припоминании, например, «Какая столица Франции?», требуют точной фокусировки на слове «Париж». Заставляя модель предсказывать несколько токенов одновременно, например, «Париж — город в…», можно ослабить общий сигнал от наиболее важного токена, что снизит эффективность модели на общем бенчмарке. Если ваша цель — создать систему RAG (Retrieval Augmented Generation) или бота для викторин, MTP может оказаться вредным.
2. Чувствительность «Златовласки» n
Здесь нет правила «больше — значит лучше». Авторы обнаружили, что производительность сильно зависит от количества голов (n).
Авторы также пришли к выводу, что количество головок (n) не линейно масштабируется с производительностью MTP. Существует «золотая середина», в которой модель может наиболее эффективно использовать парадигму MTP:
- Слишком мало (n=2): Незначительный выигрыш, поскольку модель не получает достаточного стимула для развития какого-либо предвидения.
- Слишком много (n=8): Производительность быстро падает, поскольку информация для всех 8 головок начинает переполнять скрытое состояние общего ствола.
- В самый раз (n=4): Лучшая производительность
Это вводит новый гиперпараметр, который необходимо настроить. В отличие от прогнозирования следующего токена, которое просто «работает», MTP требует поиска конкретного горизонта, соответствующего сложности ваших данных.
Заключение
Учитывая продемонстрированную способность MTP повышать производительность кодирования и ускорять вывод, остается один очевидный вопрос: если MTP настолько революционен, почему его до сих пор не используют ни одна крупная лаборатория ИИ?
На самом деле ответом на этот вопрос является DeepSeek-V3.
В своём техническом отчёте (Liu et al., 2024)4 команда DeepSeek показала, что MTP был ключевым компонентом в процессе обучения модели. Подобно Meta, они провели интенсивные исследования абляции, сравнивая стандартные модели NTP с MTP как по шкале параметров 15,7B, так и по шкале 228,7B. Используя конфигурацию n=2 во время обучения (предсказывая один дополнительный токен будущего), они обнаружили, что модели, обученные с помощью MTP, стабильно превосходили свои аналоги NTP во всех наборах данных, таких как MMLU, PILE-test, HumanEval, MBPP и т. д. Более того, сохраняя эту вторую предсказательную головку во время вывода для спекулятивного декодирования, как описано ранее, DeepSeek достигла ускорение вывода до 1,8x.
Это успешное развертывание DeepSeek служит практическим подтверждением того, что MTP можно широко использовать в качестве цели обучения в больших языковых моделях, поскольку оно демонстрирует четкий путь к улучшению возможностей рассуждения и эффективности вывода модели с минимальными сопутствующими недостатками.
Если вам нравятся подобные разборы, я поделюсь дополнительной информацией, заметками и пояснениями здесь: https://steadysurfdom.substack.com/
Ссылки
[1] Пал, Койена и др. «Линза будущего: прогнозирование последующих токенов из одного скрытого состояния». Препринт arXiv arXiv:2311.04897 (2023).
[2] Глокль, Фабиан и др. «Лучшие и быстрые большие языковые модели с помощью многотокенного прогнозирования». Препринт arXiv arXiv:2404.19737 (2024).
[3] Нордквист, Р. (20 июля 2024 г.). Определение и примеры функциональных слов в английском языке. ThoughtCo.
[4] Лю, Айсинь и др. «Технический отчет Deepseek-v3». Препринт arXiv arXiv:2412.19437 (2024).
Источник: towardsdatascience.com
