REFRAG: Переосмысление декодирования на основе RAG

Основная информация

• ID статьи: 2509.01092
• Название: REFRAG: Rethinking RAG based Decoding
• Авторы: Xiaoqiang Lin, Aritra Ghosh, Bryan Kian Hsiang Low, Anshumali Shrivastava, Vijai Mohan
• Учреждения: Meta Superintelligence Labs, National University of Singapore, Rice University
• Классификация: cs.CL cs.AI cs.LG
• Дата публикации: 14 октября 2025 г. (препринт arXiv)
• Ссылка на статью: https://arxiv.org/abs/2509.01092

Аннотация

Большие языковые модели (LLM) демонстрируют превосходные способности в использовании внешних знаний для улучшения ответов в многораундовых диалогах и приложениях с агентами, таких как поиск с увеличением контекста (RAG). Однако обработка длинных входных контекстов вызывает значительную системную задержку и требует большого объема памяти для кэша ключ-значение, что приводит к снижению пропускной способности и фундаментальному компромиссу между богатством знаний и системной эффективностью. В данной работе предлагается REFRAG — эффективная структура декодирования, которая улучшает задержку в приложениях RAG посредством сжатия, осознания и расширения. Благодаря использованию структуры разреженности внимания достигнуто ускорение задержки до первого токена в 30,85 раза (улучшение в 3,75 раза по сравнению с предыдущими работами) без потери перплексии. Кроме того, оптимизированная структура позволяет REFRAG расширить размер контекста LLM в 16 раз.

Исследовательский контекст и мотивация

Основные проблемы

1. Узкие места эффективности при обработке длинного контекста: RAG-системы сталкиваются со значительными вычислительными и памятными издержками при обработке длинного контекста, при этом задержка до первого токена (TTFT) растет квадратично, что серьезно влияет на пользовательский опыт.
2. Особенности сценариев RAG: Контекст в RAG состоит в основном из объединенных извлеченных фрагментов, при этом только небольшая часть напрямую связана с запросом. Из-за разнообразия и операций дедупликации эти фрагменты имеют низкую семантическую схожесть, что приводит к блочно-диагональному паттерну внимания.
3. Вычислительная избыточность: Существующие методы рассматривают RAG как общую проблему длинного контекста, игнорируя специфичную для RAG структуру разреженности внимания, что приводит к большому количеству ненужных вычислений.

Исследовательская мотивация

• Требования к эффективности: Острая необходимость веб-масштабных приложений в высокой пропускной способности и низкой задержке
• Оптимизация ресурсов: Снижение потребления памяти и вычислительных издержек, повышение масштабируемости системы
• Сохранение производительности: Значительное повышение эффективности при сохранении производительности модели без деградации

Основные вклады

1. Предложение структуры REFRAG: Первая специализированная эффективная структура декодирования для приложений RAG, поддерживающая сжатие и расширение контекста в произвольных позициях
2. Техника сжатия блочных эмбеддингов: Использование предварительно вычисленных сжатых блочных эмбеддингов вместо исходных токенов для значительной оптимизации задержки и памяти
3. Стратегия выборочного сжатия: Сетевая политика на основе обучения с подкреплением для динамического определения того, какие блоки должны сохранять исходную форму
4. Значительное улучшение производительности: Достижение ускорения TTFT в 30,85 раза, расширение окна контекста в 16 раз без потери производительности
5. Широкая валидация: Проверка эффективности на различных задачах: RAG, многораундовые диалоги, суммаризация длинных документов

Подробное описание методики

Определение задачи

Дана входная последовательность из T токенов x₁, x₂, ..., xₜ, где первые q токенов — это основной ввод (например, вопрос), а последующие s токенов — контекст (например, извлеченные фрагменты), при условии q + s = T. Цель — эффективно генерировать ответ, одновременно минимизируя задержку TTFT и использование памяти.

Архитектура модели

Общий дизайн

REFRAG использует архитектуру кодировщик-декодировщик:

• Декодировщик: Модель только-декодер на основе LLaMA
• Кодировщик: Легкая модель RoBERTa для обработки блоков контекста
• Слой проекции: Отображение блочных эмбеддингов в пространство токенов декодировщика

Основные компоненты

1. Генерация блочных эмбеддингов

   Разбиение контекста: {C₁, C₂, ..., Cₗ}, где L = s/k
   Блочные эмбеддинги: cᵢ = Mₑₙc(Cᵢ)
   Проецированные эмбеддинги: eᶜⁿᵏᵢ = φ(cᵢ)
   

2. Обработка смешанного ввода Ввод декодировщика: {e₁, ..., eᵩ, eᶜⁿᵏ₁, ..., eᶜⁿᵏₗ} Коэффициент сжатия: ≈ k-кратное уменьшение
3. Механизм выборочного сжатия
   • Сетевая политика RL πθ определяет, какие блоки остаются несжатыми
   • Последовательный выбор на основе блочных эмбеддингов и масок
   • Функция вознаграждения: отрицательная логарифмическая перплексия

Технические инновации

1. Сжатие в произвольных позициях: Преодоление ограничения существующих методов, поддерживающих только префиксное сжатие, с поддержкой сжатия и расширения контекста в произвольных позициях
2. Переиспользование предварительных вычислений: Блочные эмбеддинги могут быть предварительно вычислены и закэшированы, избегая повторных вычислительных издержек
3. Адаптивный коэффициент сжатия: Динамическая регулировка коэффициента сжатия через политику RL без необходимости пересчета блочных эмбеддингов
4. Сохранение авторегрессивного характера: Поддержание причинной структуры декодировщика, обеспечивающей поддержку многораундовых диалогов и задач суммаризации

Экспериментальная установка

Наборы данных

• Предварительное обучение: Набор данных SlimPajama (20B токенов), включающий 50% ArXiv + 50% Book
• Оценка: Наборы данных Book, ArXiv, PG19, Proof-pile
• Нижестоящие задачи:
  • RAG: 1,1M образцов, охватывающих наборы данных QA из 5 доменов
  • Многораундовые диалоги: TopiOCQA, ORConvQA, QReCC
  • Суммаризация: Суммаризация длинных документов ArXiv и PubMed

Метрики оценки

• Метрики эффективности: TTFT, TTIT (задержка между токенами), пропускная способность
• Метрики производительности: Перплексия, точность, F1-оценка, ROUGE-оценка
• Метрики памяти: Использование памяти кэша KV

Методы сравнения

• Варианты LLaMA: LLaMA-Full Context, LLaMA-No Context, LLaMA-32K
• Существующие методы: CEPE, REPLUG
• Различные коэффициенты сжатия: REFRAG8, REFRAG16, REFRAG32

Детали реализации

• Базовая модель: LLaMA-2-7B
• Кодировщик: RoBERTa-Large (355M параметров)
• Стратегия обучения: Обучение по программе + предварительное прогревание задачи реконструкции
• Оптимизатор: AdamW, пиковая скорость обучения 5e-5
• Оборудование: 8 узлов × 8 GPU H100

Результаты экспериментов

Основные результаты

Производительность задержки

При длине контекста 16K:

• Ускорение TTFT: 16,53 раза (с кэшем), 8,59 раза (без кэша)
• По сравнению с CEPE: Улучшение TTFT в 2,01 раза (с кэшем), 1,04 раза (без кэша)
• При k=32: Достижение ускорения TTFT в 30,85 раза, в 3,75 раза быстрее CEPE

Производительность модели

REFRAG16 показывает среднее улучшение перплексии на 9,3% по сравнению с CEPE при одновременном достижении значительного ускорения.

Абляционные исследования

Необходимость обучения по программе

Обучение по программе критически важно для успеха задачи реконструкции.

Роль задачи реконструкции

Предварительное обучение задаче реконструкции значительно улучшает эффект непрерывного предварительного обучения.

Выборочное сжатие RL

При одинаковом коэффициенте сжатия 8, REFRAG16+RL постоянно превосходит REFRAG8, что подтверждает эффективность стратегии динамического сжатия.

Производительность нижестоящих задач

Задачи RAG

При сильном извлекателе, при одинаковых ограничениях задержки:

• REFRAG с 8 фрагментами vs LLaMA с 1 фрагментом: среднее улучшение 1,22%
• При слабом извлекателе улучшение более выраженное: 1,93%

Многораундовые диалоги

При установке с 10 фрагментами REFRAG превосходит LLaMAFT на всех трех наборах данных, особенно в сценариях с длинной историей диалога.

Анализ примеров

В статье представлены результаты визуализации внимания, подтверждающие предположение о блочно-диагональной разреженности: значения внимания между различными фрагментами значительно ниже, чем внимание внутри фрагментов.

Связанные работы

Поиск с увеличением контекста для языкового моделирования

• REALM: Первое предложение предварительного обучения маскированной языковой модели с увеличением контекста
• RETRO: Использование перекрестного внимания и сквозного предварительного обучения
• FiD: Параллельная обработка фрагментов и объединение скрытых состояний

Эффективные LLM с длинным контекстом

• Сжатое внимание: Снижение сложности внимания, но не решение проблемы памяти
• StreamingLLM: Использование агрегирования внимания для уменьшения кэша KV
• CEPE: Метод перекрестного внимания, но ограничен применением к префиксам

Сжатие Трансформеров

• Compressive Transformer: Сжатие кэша KV, но без улучшения TTFT
• Рекурсивное сжатие: Невозможность предварительного вычисления и переиспользования эмбеддингов

Выводы и обсуждение

Основные выводы

1. Разреженность, специфичная для RAG: Блочно-диагональный паттерн внимания в сценариях RAG предоставляет возможности для специализированной оптимизации
2. Значительное повышение эффективности: Ускорение TTFT в 30,85 раза без потери производительности доказывает эффективность метода
3. Широкая применимость: Превосходная производительность на различных задачах с длинным контекстом

Ограничения

1. Ограничения коэффициента сжатия: Эксперименты показывают значительное снижение производительности при k=64, существует верхний предел сжатия
2. Издержки кодировщика: Хотя и легкий, все еще требует дополнительных вычислений кодирования
3. Сложность обучения: Требует обучения по программе и многоэтапной стратегии обучения

Будущие направления

1. Более высокие коэффициенты сжатия: Исследование более эффективных методов сжатия для преодоления текущих ограничений
2. Сквозная оптимизация: Интеграция стратегии сжатия на этапе предварительного обучения
3. Расширение на мультимодальность: Расширение метода на мультимодальные сценарии, такие как зрительно-языковые

Глубокая оценка

Преимущества

1. Точное выявление проблемы: Точное определение специфичности сценариев RAG и возможностей оптимизации
2. Разумный дизайн метода: Умное проектирование сжатия блочных эмбеддингов и выборочной стратегии
3. Полная экспериментальная валидация: Охват различных задач и подробные абляционные исследования
4. Высокая практическая ценность: Значительное улучшение производительности обеспечивает важное прикладное значение
5. Сильная техническая инновативность: Выделяются инновационные точки, такие как сжатие в произвольных позициях и переиспользование предварительных вычислений

Недостатки

1. Недостаток теоретического анализа: Отсутствие теоретического анализа верхних пределов коэффициента сжатия
2. Выбор кодировщика: Недостаточное исследование влияния различных архитектур кодировщика
3. Долгосрочные зависимости: Способность обработки экстремально длинного контекста требует дальнейшей проверки
4. Вычислительная сложность: Обучение RL увеличивает сложность системы

Влияние

1. Академический вклад: Открывает новое направление исследований в оптимизации RAG-систем
2. Промышленная ценность: Может быть непосредственно применена в крупномасштабных развертываниях RAG
3. Воспроизводимость: Авторы обещают открыть исходный код, что способствует распространению метода

Применимые сценарии

1. Веб-поиск: Оптимизация задержки в крупномасштабных сценариях извлечения
2. Вопросно-ответные системы на основе знаний: Сложные вопросы, требующие интеграции нескольких фрагментов документов
3. Интеллектуальные помощники: Управление контекстом в многораундовых диалогах
4. Анализ документов: Суммаризация и анализ длинных документов

Библиография

В статье цитируется большое количество связанных работ, включая:

• Guu et al. (2020) - REALM предварительное обучение с увеличением контекста
• Borgeaud et al. (2022) - RETRO крупномасштабное поиск-увеличенное генерирование
• Yen et al. (2024) - CEPE параллельное кодирование контекста
• Touvron et al. (2023) - базовая модель LLaMA

Общая оценка: Это высококачественная исследовательская статья, предлагающая инновационное решение для узких мест эффективности в RAG-системах. Метод разумно спроектирован, экспериментально хорошо валидирован, обладает выдающейся практической ценностью и вносит важный вклад в развитие этой области.