Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Языковую модель научились увеличивать по ходу обучения так, чтобы обработка длинного запроса не дорожала вместе с новой ёмкостью. В препринте, который не прошёл рецензирование и приводит числа, полученные самими авторами, SST достигла меньшей функции потерь, чем оба обычных MoE-трансформера при сопоставимом вычислительном бюджете. Подход рассчитан прежде всего на собственные модели и системы, где запросы длиннее ответов.
Новые параметры не участвуют в подготовке KV-кэша
KITE разделяет модель на область, которая строит ключи и значения для механизма внимания, и область, которая только читает их. При обработке запроса модель сохраняет эти ключи и значения в KV-кэше, чтобы не пересчитывать весь контекст для каждого следующего токена.
Обычное увеличение трансформера расширяет обе части вычислений. Каждый добавленный слой участвует и в обработке запроса, и в генерации ответа. KITE добавляет параметры только после участка, который формирует KV-кэш, поэтому большой блок не приходится запускать для всех токенов уже известного контекста.
Конкретную архитектуру авторы назвали Step Scale Transformer, или SST. Сначала они обучили исходную MoE-модель на 33,8 млрд параметров, затем расширили её примерно до 67 млрд. Получилась конструкция из двух последовательных башен по 18 слоёв: первая создаёт KV-кэш, вторая читает его и предсказывает следующий токен.
При расширении веса исходной модели копируют в обе башни, после чего их обучают совместно. Градиенты проходят через всю конструкцию, поэтому первая башня не заморожена и продолжает приспосабливаться к новой.
Во время массовой обработки запроса работает только первая башня. Вторая обрабатывает последнюю позицию, чтобы получить первый токен ответа. Затем каждый новый токен проходит через обе башни: SST экономит на длинном входе, но генерация одного токена требует больше вычислений, чем у исходной модели.
SST выиграла при равном бюджете обучения
Сравнение провели внутри одного семейства MoE-моделей. Для SST и двух классических трансформеров сохранили одинаковые токенизатор, данные, схему обучения и набор проверок. В бюджет SST включили оба этапа, а вычислительную нагрузку оценивали по числу операций с плавающей точкой.
При сопоставимых суммарных затратах итоговая функция потерь SST составила 1,5900. Ближайший результат классической модели — 1,5921; второй классический вариант также уступил SST. Чем ниже функция потерь, тем точнее модель предсказывает следующий токен на обучающих данных.
Для сценария, где предварительная обработка запроса даёт 75% вычислительной нагрузки, расчётная стоимость инференса SST оказалась на 6,7% ниже меньшей классической модели и на 31,6% ниже близкой по общему размеру. Это оценка количества вычислений, а не замер пропускной способности на GPU.
SST также получила лучший результат на всех выбранных задачах: OpenBookQA, MMLU, GSM8K, MATH, HumanEval, MBPP и BBH. Проверка охватывает одну конфигурацию SST и две модели того же семейства; суммарно SST обучали примерно на 391 млрд токенов. Такой эксперимент показывает работоспособность конкретной схемы, но не устанавливает общее правило масштабирования для других архитектур.
Планы стоит менять только для моделей с дорогим входом
KITE влияет на планы команд, которые обучают и разворачивают собственные базовые модели. Вместо обучения крупной модели с нуля можно сохранить исходный блок, добавить читающую KV-кэш часть и продолжить совместное обучение. Архитектура при этом не требует отказываться от MoE или смешанного механизма внимания.
Главный кандидат — агентная система, которая много раз отправляет модели накопленную историю, результаты инструментов и новые инструкции. В такой нагрузке обработка входа занимает заметную долю вычислений, а KITE оставляет её на уровне меньшей исходной модели.
Для сервисов с короткими запросами и длинной генерацией преимущество может исчезнуть: при создании ответа работают обе башни. Перед сменой архитектуры нужно разделить реальную нагрузку на обработку незакэшированного входа и декодирование, а затем проверить задержку и пропускную способность на целевом оборудовании. Расчётная экономия сама по себе не гарантирует такого же сокращения серверных расходов.
Командам, которые вызывают стороннюю модель через API, работа не предлагает немедленной замены. Она меняет выбор архитектуры на уровне разработчика модели, а не способ интеграции готового сервиса.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
