Shapley Context Pruning: новый метод для эффективной реранжировки контекста в RAG-системах

Исследователи предложили новый подход к реранжировке контекста в системах Retrieval-Augmented Generation (RAG). Работа, опубликованная на arXiv, представляет Shapley Context Pruning (SCP) — фреймворк, использующий теорию кооперативных игр для оценки важности элементов контекста.

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

SCP моделирует контекст как кооперативную игру, где каждый элемент получает «выигрыш» (значение Шепли), отражающий его вклад в итоговый ответ. Для вычисления используется архитектура Deep Sets с предобученными языковыми моделями в качестве энкодеров предложений, что позволяет аппроксимировать пермутационно-инвариантную функцию ценности.

Чтобы обеспечить масштабируемость, авторы применяют сэмплирование Монте-Карло для обучения и вывода, доказывая теоретические границы ошибок и гарантии сохранения топ-K подмножеств. Это делает метод пригодным для реальных сценариев с большими объёмами данных.

Эксперименты включали задачи recall поддерживающих предложений, тесты Needle-in-the-Haystack (NIAH), длинноконтекстный QA и многошаговое рассуждение. SCP показал конкурентоспособные результаты по сравнению с сильными базовыми линиями, уступая лишь кросс-энкодерам по точности, но выигрывая по скорости.

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

Разработчики отмечают, что SCP — это шаг к более прозрачным и эффективным системам дополненной генерации, где решения об отборе контекста обоснованы математически. Дальнейшие работы могут расширить применение на мультимодальные задачи и потоковые сценарии.