Проблема: Разная природа фаз инференса
По мере роста числа запросов к большим языковым моделям (LLM), совокупные затраты на инференс становятся критичнее разовых затрат на обучение. Традиционные подходы масштабирования (увеличение ширины или глубины сети) увеличивают стоимость обеих фаз одновременно, так как каждый дополнительный слой обрабатывается и при префилле, и при декодировании. Однако эти фазы предъявляют разные требования к оборудованию:
- Prefill (префилл): Параллельная обработка промпта, обычно ограничена вычислительной мощностью (compute-bound).
- Decode (декодирование): Последовательная генерация токенов, часто ограничена пропускной способностью памяти (memory-bandwidth-bound).
Решение: Dual-Flow Transformer
Авторы (Liming Liu, Mingze Wang, Tuo Zhao) предлагают архитектуру, которая отделяет основную вычислительную цепочку от дополнительных вычислений для предсказания продолжения. Ключевые особенности:
- Primary Flow (Основной поток): Полноценная причинно-следственная языковая модель, обрабатывающая промпт и записывающая ключ-значение (KV) кэш.
- Auxiliary Flow (Вспомогательный поток): Отсутствует на этапе префилла. Активируется только начиная с последней позиции промпта для добавления вычислений предсказания продолжения. Он не записывает постоянное состояние и не влияет на основной поток.
Архитектурные детали и эффективность
Оба потока разделяют основные матрицы внимания (attention), многослойные перцептроны (MLP) и выходные матрицы, но используют отдельные эмбеддинги токенов и легкие механизмы связи. Это позволяет:
- Переиспользовать загруженные веса и кэшированные ключи/значения при групповом выполнении запросов.
- В моделях с экспертами (MoE) сделать независимыми контроль за стоимость префилла, стоимость декодирования и качество предсказания.
Результаты и сравнение
Эксперименты показали, что Dual-Flow достигает меньшей ошибки валидации (validation loss) по сравнению с базовыми архитектурами при сопоставимых параметрах. Исследователи изучили два режима:
- Увеличение вычислений на этапе decode при фиксированных вычислениях на этапе prefill.
- Перераспределение фиксированного бюджета экспертов decode между двумя потоками.
Результаты демонстрируют компромисс между качеством на этапах prefill и decode, подтверждая потенциал специфичного для фаз распределения экспертов.
| Характеристика | Traditional Transformer | Dual-Flow Transformer |
|---|---|---|
| Обработка Prefill | Основной поток (Compute-bound) | Основной поток (Compute-bound) |
| Обработка Decode | Основной поток (Memory-bound) | Основной + Вспомогательный поток |
| Запись KV Cache | Да (в обоих потоках) | Только в основном потоке |
| Масштабируемость | Рост стоимости обеих фаз | Независимый контроль затрат |
| Качество (Validation Loss) | Базовый уровень | Ниже (лучше) при matched-token сравнении |
Источник: arXiv cs.AI ↗