Как вы делаете длинный контекст для RAG (100k+ токенов в контексте)?

Краткий тезис

Обработка контекста длиной 100k+ токенов в RAG требует комбинации аппаратных и алгоритмических техник. Ключевые подходы: KV cache quantization для снижения потребления памяти, sliding window attention для фокуса на ближайших токенах, summarization для сжатия старых частей диалога, selective pruning для удаления неважных токенов и retrieval|contextual retrieval для извлечения только релевантных фрагментов из длинных документов. Ни один метод не работает идеально в одиночку — нужна гибридная архитектура, балансирующая между полнотой контекста, скоростью и точностью.

1. Термин: Длинный контекст (Long context) в RAG

Длинный контекст — это ситуация, когда суммарное количество токенов, передаваемых в LLM (промпт + извлечённые документы + история диалога), превышает 100k токенов. Стандартные модели (например, GPT-4 с 128k контекстом) могут его вместить, но возникают проблемы:

• Потребление памяти: KV cache (кэш ключей и значений attention) растёт линейно с длиной контекста. Для 100k токенов он может занимать десятки гигабайт в FP16.
• Время инференса: внимание (attention) имеет квадратичную сложность O(n²) по длине контекста, что делает каждый шаг генерации медленным.
• Lost in the middle: LLM хуже запоминает информацию из середины длинного контекста, что снижает качество ответов.

Поэтому для Context RAG с длинным контекстом нужны специальные техники, описанные ниже.

2. KV cache quantization (INT4/INT8)

KV cache — это матрицы ключей (K) и значений (V) для каждого слоя и каждого токена, которые сохраняются во время генерации, чтобы не пересчитывать attention для уже обработанных токенов. При длинном контексте KV cache становится узким местом по памяти.

Квантизация (quantization) — преобразование чисел из FP16 (16 бит) в INT8 (8 бит) или INT4 (4 бита) с минимальной потерей точности. Это уменьшает размер KV cache в 2–4 раза.

Как реализовать: библиотеки вроде vLLM, TensorRT-LLM, Hugging Face Transformers (через quantization_config). Пример кода для загрузки модели с KV cache quantization в INT8:

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

quant_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_enable_fp32_cpu_offload=False
)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-chat-hf",
    quantization_config=quant_config,
    device_map="auto"
)

Ограничения: квантизация может незначительно снизить качество (perplexity), но для RAG это обычно приемлемо.

3. Sliding window attention

Sliding window attention — механизм, при котором каждый токен «видит» только фиксированное окно предыдущих токенов (например, 4096), а не весь контекст. Это снижает сложность attention с O(n²) до O(n × w), где w — размер окна.

Как применяется в RAG с длинным контекстом:

• Ближайшие токены (последние 32k) обрабатываются через sliding window — они наиболее важны для текущего ответа.
• Для distant токенов (старые части диалога, далёкие документы) используется retrieval поверх контекста: из длинной истории извлекаются только релевантные фрагменты и вставляются в окно.

Пример архитектуры:

[System prompt] + [Retrieved chunks] + [Sliding window of last 32k tokens]

Модели с поддержкой sliding window: Mistral (window=4096), Qwen2.5 (window=131072, но с внутренним sliding).

Плюсы: значительное ускорение и экономия памяти. Минусы: потеря информации за пределами окна, если retrieval не сработает.

4. Summarization (сжатие контекста)

Summarization — периодическое сжатие старых частей диалога или длинных документов в краткие саммари. Это позволяет хранить суть без хранения всех токенов.

Стратегия:

• Каждые N токенов (например, 10k) запускается LLM, которая генерирует саммари предыдущего блока.
• Саммари сохраняется в отдельном буфере и включается в контекст вместо полного текста.
• Для иерархического сжатия можно делать саммари саммари (древовидное summarization).

Пример кода (псевдо) для агента с саммари:

class LongContextAgent:
    def __init__(self, llm, max_tokens=100000, summary_interval=10000):
        self.llm = llm
        self.max_tokens = max_tokens
        self.summary_interval = summary_interval
        self.history = []
        self.summaries = []

    def add_message(self, message):
        self.history.append(message)
        if len(self.history) >= self.summary_interval:
            # сжимаем старые сообщения в саммари
            text_to_summarize = "\n".join(self.history[:self.summary_interval])
            summary = self.llm.generate(f"Summarize the following conversation:\n{text_to_summarize}")
            self.summaries.append(summary)
            self.history = self.history[self.summary_interval:]  # удаляем сжатые сообщения

    def get_context(self):
        # собираем контекст: саммари + последние сообщения
        return "\n".join(self.summaries) + "\n" + "\n".join(self.history)

Проблемы: потеря деталей, возможное искажение смысла, дополнительный latency на генерацию саммари.

5. Selective pruning (на основе attention)

Selective pruning — удаление из контекста токенов, которые модель считает неважными на основе весов attention. Идея: если на какой-то токен почти не обращают внимания последующие слои, его можно выбросить.

Как работает:

1. После первого forward pass (prefill) получаем attention scores для каждого токена.
2. Сортируем токены по суммарному вниманию (или по attention к последнему токену).
3. Оставляем только top-k% токенов (например, 50%), остальные удаляем из KV cache.

Инструменты: библиотека FastChat (техника StreamingLLM), LLMLingua (сжатие промпта через small LM).

Пример с LLMLingua:

from llmlingua import PromptCompressor

compressor = PromptCompressor()
compressed_prompt = compressor.compress(
    prompt,  # длинный контекст
    rate=0.5,  # сжимаем до 50%
    force_tokens=['\n', '?'],  # сохраняем важные токены
)

Ограничения: pruning может удалить ключевую информацию, если attention распределён равномерно. Требует калибровки на конкретной задаче.

6. Contextual retrieval (RAG поверх контекста)

Contextual retrieval — это не просто поиск по внешней базе знаний, а извлечение релевантных фрагментов из самого длинного контекста (истории диалога, документа). То есть мы применяем RAG внутри контекста.

Как реализовать:

• Разбиваем длинный контекст на чанки (например, по 512 токенов).
• Для каждого чанка вычисляем эмбеддинг (через embedding model).
• При запросе пользователя ищем среди этих чанков наиболее релевантные (косинусная близость).
• Добавляем найденные чанки в начало контекста (или в sliding window).

Пример с LlamaIndex:

from llama_index.core import Document, VectorStoreIndex
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# длинный документ разбиваем на чанки
documents = [Document(text=long_text)]
index = VectorStoreIndex.from_documents(documents, embed_model=HuggingFaceEmbedding())
retriever = index.as_retriever(similarity_top_k=5)

# при запросе извлекаем чанки
query = "What is the capital of France?"
nodes = retriever.retrieve(query)
context = "\n".join([node.text for node in nodes])

Плюсы: позволяет работать с контекстом, который не влезает в окно модели. Минусы: дополнительный latency на поиск, необходимость хранить эмбеддинги для чанков контекста.

7. Комбинирование подходов: гибридная архитектура

На практике для 100k+ контекста используют комбинацию всех техник. Пример архитектуры:

1. Prefill: весь контекст загружается, но KV cache квантизуется до INT4.
2. Sliding window: во время генерации используется окно 32k токенов.
3. Summarization: каждые 10k токенов истории сжимаются в саммари (асинхронно, в фоне).
4. Selective pruning: после prefill удаляются 30% токенов с наименьшим attention.
5. Contextual retrieval: если модель запрашивает детали из удалённой части контекста, запускается retrieval по чанкам.

Результат: контекст 200k токенов умещается в память GPU 80GB, latency увеличивается всего в 2-3 раза по сравнению с 4k контекстом.

8. Инструменты и библиотеки

9. Оценка качества длинного контекста

Метрики для проверки, что длинный контекст работает корректно:

• Perplexity на длинных последовательностях — не должна резко расти после определённой длины.
• Needle-in-a-haystack (иголка в стоге сена): вставляем факт в середину длинного контекста и проверяем, может ли модель его извлечь.
• Faithfulness (фактологическая точность) ответов при использовании summarization/pruning — не теряются ли важные детали.
• Latency и потребление памяти — должны оставаться в допустимых пределах.

10. Проблемы и ограничения

• Компромисс скорость/качество: агрессивное сжатие (pruning, summarization) может уничтожить нужную информацию.
• Дополнительный latency: каждый вызов summarization или retrieval добавляет время.
• Сложность отладки: трудно понять, почему ответ неверен — из-за потери контекста или плохого retrieval.
• Зависимость от модели: не все LLM одинаково хорошо работают с длинным контекстом; некоторые (например, GPT-4 Turbo) имеют нативное преимущество.

Пет-проект для закрепления

Задача: Реализовать чат-бота с памятью на 200k токенов, который может поддерживать диалог на основе длинной истории.

Инструменты: Python, LangChain, Hugging Face Transformers (модель Mistral-7B с sliding window), FAISS для retrieval, библиотека llmlingua.

Шаги:

1. Загрузите модель Mistral-7B-Instruct с sliding_window=4096 и KV cache quantization INT8.
2. Реализуйте класс LongMemoryChat:
   • Храните историю диалога в списке.
   • При добавлении нового сообщения проверяйте общую длину токенов.
   • Если превышает 100k, запустите summarization старых сообщений (через LLM) и замените их саммари.
   • Также используйте LLMLingua для pruning неважных токенов в каждом новом сообщении (сжимайте до 70%).
3. Для retrieval по истории: при каждом запросе вычисляйте эмбеддинг запроса и ищите топ-3 чанка из истории (через FAISS), добавляйте их в начало контекста.
4. Протестируйте на синтетическом диалоге из 500 сообщений (сгенерируйте через LLM). Измерьте время ответа и точность извлечения фактов (needle test).

Ожидаемый результат: Бот отвечает на вопросы, ссылаясь на события из начала диалога, latency не превышает 5 секунд на токен, память GPU < 40GB.

Связь с другими вопросами

Навигация

• Предыдущий: 630
• Следующий: 632
• Индекс: 00. Индекс разборов