MBA-RAG: Подход на основе многорукого бандита для адаптивной генерации с дополнением поиском через сложность вопроса

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

• ID статьи: 2412.01572
• Название: MBA-RAG: a Bandit Approach for Adaptive Retrieval-Augmented Generation through Question Complexity
• Авторы: Xiaqiang Tang, Qiang Gao, Jian Li, Nan Du, Qi Li, Sihong Xie
• Принадлежность: Гонконгский научно-технический университет (Гуанчжоу), Tencent Hunyuan, Уханьский университет, Государственный университет Айовы
• Категория: cs.AI
• Дата публикации: 1 января 2025 г. (arXiv v4)
• Ссылка на статью: https://arxiv.org/abs/2412.01572
• Ссылка на код: https://github.com/FUTUREEEEEE/MBA

Аннотация

Генерация с дополнением поиском (RAG) значительно улучшает производительность языковых моделей в задачах, требующих знаний. Однако существующие системы RAG либо выполняют поиск без дифференциации, либо полагаются на жесткие однокласовые классификаторы для выбора методов поиска, что приводит к неэффективности и субоптимальной производительности при запросах различной сложности. Для решения этих проблем в статье предлагается фреймворк на основе обучения с подкреплением, который динамически выбирает наиболее подходящую стратегию поиска в зависимости от сложности запроса. Метод использует алгоритм многорукого бандита, рассматривая каждый метод поиска как различное "плечо", и балансирует исследование и эксплуатацию для адаптации процесса выбора. Кроме того, вводится динамическая функция вознаграждения, которая балансирует точность и эффективность, штрафуя методы, требующие большего количества шагов поиска, даже при получении правильного результата. Метод достигает новых результатов SOTA на нескольких наборах данных с одним и несколькими переходами, одновременно снижая стоимость поиска.

Предпосылки и мотивация исследования

Определение проблемы

Существующие системы RAG имеют следующие основные проблемы:

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

Мотивация исследования

Основная мотивация этой работы заключается в разработке фреймворка, который может:

1. Динамически адаптироваться к сложности запроса: Интеллектуально выбирать стратегию поиска в зависимости от сложности проблемы
2. Балансировать точность и эффективность: Минимизировать вычислительные затраты при сохранении качества ответов
3. Поддерживать исследование нескольких стратегий: Позволить нескольким стратегиям производить правильные ответы вместо принудительного выбора единого "оптимального" пути

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

1. Предложение фреймворка MBA-RAG: Первое применение алгоритма многорукого бандита к выбору стратегии поиска в системах RAG, реализующее динамическую адаптивную выборку
2. Разработка динамической функции вознаграждения: Инновационное объединение точности и вычислительной эффективности путем штрафования методов с высокой стоимостью для оптимизации использования ресурсов
3. Достижение производительности SOTA: Получение лучших результатов на 6 наборах данных при одновременном снижении стоимости поиска на 20%
4. Предоставление гибкого механизма контроля: Использование частичного информационного контроля вместо строгого однозначного контроля, позволяя модели исследовать несколько эффективных стратегий

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

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

Для заданного запроса x система RAG должна:

1. Этап поиска: Модуль R извлекает релевантные документы D для запроса x
2. Этап генерации: LLM использует x и D для генерации ответа ā = LLM(yt|x,D)

В данной работе это переформулируется как задача многорукого бандита, где каждый метод поиска (без поиска, одиночный поиск, множественный поиск) рассматривается как "плечо".

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

1. Кодирование запроса и выбор плеча

• Кодировщик: Использует DistilBERT для кодирования пользовательского запроса, генерируя распределение действий z = fθ(x)
• Стратегия выбора: Применяет ε-жадную стратегию для балансирования исследования и эксплуатации:
  • С вероятностью (1-ε) выбирает a = argmax(z)
  • С вероятностью ε случайно выбирает метод генерации

2. Алгоритм обучения

Целевая функция минимизирует квадратичную ошибку между фактическим вознаграждением ra и предсказанным вознаграждением fθ(x)a:

min_θ (ra - fθ(x)a)²

Правило обновления параметров:

θt+1 = θt - α∇θ((ra - fθ(x)a)²)

3. Динамическая функция вознаграждения

ra = A(y, ŷa) - λC(a)

где:

• A(y, ŷa): метрика качества генерации (например, точное совпадение)
• C(a): вычислительная стоимость метода a (например, количество шагов поиска)
• λ: масштабирующий коэффициент для балансирования точности и эффективности

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

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

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

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

Наборы данных QA с одним переходом:

• SQuAD v1.1: задача понимания текста
• Natural Questions: открытый вопрос-ответ
• TriviaQA: вопрос-ответ на основе знаний

Наборы данных QA с несколькими переходами:

• MuSiQue: вопрос-ответ с многошаговым рассуждением
• HotpotQA: вопрос-ответ с многошаговым рассуждением
• 2WikiMultiHopQA: вопрос-ответ на основе Википедии с несколькими переходами

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

Метрики производительности:

• EM (Exact Match): полное совпадение предсказанного результата с истинным ответом
• F1: перекрытие словарного запаса между предсказанным и истинным ответом
• Acc (Accuracy): содержит ли предсказанный ответ истинный ответ

Метрики эффективности:

• Step: количество шагов поиска, требуемых выбранной стратегией поиска

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

1. No-Retrieval: Прямая генерация ответа без поиска
2. Adaptive-Retrieval: Динамическое определение необходимости поиска
3. Self-RAG: Динамическое определение потребности в поиске через самоотражение
4. DRAGIN: Активация поиска на основе неопределенности токена
5. SEAKR: Определение поиска на основе самоосознающейся неопределенности
6. Adaptive-RAG: Использование классификатора для выбора стратегии поиска на основе сложности запроса

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

• Модель кодирования запроса: DistilBERT
• Модель поиска: BM25
• Модель генерации: FLAN-T5-XL (3B)
• Скорость обучения: 5e-5
• Стратегия исследования: ε-жадный алгоритм

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

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

Ключевые выводы

1. Улучшение производительности: MBA-RAG превосходит все базовые методы по всем метрикам производительности
2. Оптимизация эффективности: По сравнению с Adaptive-RAG количество шагов поиска сокращается примерно на 17% (с 2.17 до 1.80)
3. Производительность на наборах данных с одним переходом: Значительное улучшение на SQuAD и TriviaQA с значительным снижением стоимости поиска
4. Производительность на наборах данных с несколькими переходами: Выдающееся улучшение на 2WikiMultiHopQA с сокращением стоимости поиска более чем на 20%

Анализ точности классификации

Точность классификации MBA-RAG составляет 56.1%, что значительно выше:

• Adaptive Retrieval: 42.0%
• Self-RAG: 41.5%
• Adaptive-RAG: 54.0%

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

Сравнение с результатами многоклассового классификатора показывает, что хотя традиционные многоклассовые методы демонстрируют хорошую производительность, стоимость поиска слишком высока (Step достигает 4.514), в то время как MBA-RAG достигает оптимального баланса между производительностью и эффективностью.

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

Развитие систем RAG

1. Традиционный RAG: Фреймворк поиск-генерация, предложенный Lewis et al. (2020)
2. Адаптивный поиск: Методы SEAKR, FLARE и другие реализуют поиск по требованию
3. Осведомленность о сложности: AdaptiveRAG выбирает стратегию на основе сложности запроса

Применение многорукого бандита

В данной работе впервые применяется алгоритм многорукого бандита к системам RAG, предоставляя новую теоретическую основу для выбора стратегии поиска.

Заключение и обсуждение

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

1. Подтверждение эффективности: MBA-RAG достигает производительности SOTA на нескольких наборах данных
2. Улучшение эффективности: Значительное снижение стоимости поиска в среднем на 20%
3. Сильная адаптивность: Способность динамически корректировать стратегию в зависимости от сложности запроса

Ограничения

1. Зависимость от алгоритма: Фреймворк зависит от конкретной структуры алгоритма многорукого бандита
2. Проблемы масштабируемости: Возможные проблемы с адаптивностью при столкновении с новыми невиданными типами запросов
3. Вычислительные требования: Методы обучения с подкреплением могут вносить дополнительные вычислительные затраты

Направления будущих исследований

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

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

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

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

Недостатки

1. Увеличение сложности: Введение дополнительной алгоритмической сложности по сравнению с простыми методами классификации
2. Чувствительность параметров: Коэффициент балансирования λ в функции вознаграждения требует настройки для различных наборов данных
3. Недостаточный теоретический анализ: Отсутствие теоретических гарантий сходимости и оптимальности

Влияние

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

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

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

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

1. Lewis, P., et al. (2020). Retrieval-augmented generation for knowledge-intensive nlp tasks. NeurIPS.
2. Jeong, S., et al. (2024). Adaptive-rag: Learning to adapt retrieval-augmented large language models through question complexity. arXiv preprint.
3. Katehakis, M. N., & Veinott Jr, A. F. (1987). The multi-armed bandit problem: decomposition and computation. Mathematics of Operations Research.

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