Динамическое рассуждение в длинном контексте со сжатой памятью с помощью сквозного обучения с подкреплением

Dynamic Long Context Reasoning over Compressed Memory via End-to-End Reinforcement Learning

February 9, 2026

Авторы: Zhuoen Chen, Dongfang Li, Meishan Zhang, Baotian Hu, Min Zhang

cs.AI

Аннотация

Крупные языковые модели (LLM) сталкиваются со значительными трудностями при обработке длинных контекстов, включая квадратичные вычислительные затраты, забывание информации и фрагментацию контекста, присущую retrieval-augmented generation (RAG). Мы предлагаем когнитивно-инспирированную структуру для эффективного вывода в условиях длинного контекста, основанную на поблочном сжатии и избирательном припоминании из памяти, вместо обработки всех исходных токенов. Данная структура сегментирует длинные входные данные на блоки и кодирует каждый блок в сжатые представления памяти с помощью обученного компрессора. Модуль управления динамически выбирает релевантные блоки памяти, которые затем итеративно обрабатываются рассуждающим модулем с развивающейся рабочей памятью для решения последующих задач. Компрессор и модуль рассуждений совместно оптимизируются посредством сквозного обучения с подкреплением, в то время как модуль управления обучается отдельно как классификатор. Результаты экспериментов показывают, что предложенный метод демонстрирует конкурентоспособную точность на бенчмарках многозвенного логического вывода, таких как RULER-HQA, экстраполирует длину контекста с 7K до 1.75M токенов и предлагает благоприятный баланс между точностью и эффективностью по сравнению с сильными базовыми методами для длинных контекстов. В частности, он позволяет достичь до 2-кратного снижения пикового использования памяти GPU и 6-кратного ускорения вывода по сравнению с MemAgent.

English

Large Language Models (LLMs) face significant challenges in long-context processing, including quadratic computational costs, information forgetting, and the context fragmentation inherent in retrieval-augmented generation (RAG). We propose a cognitively inspired framework for efficient long-context inference based on chunk-wise compression and selective memory recall, rather than processing all raw tokens. The framework segments long inputs into chunks and encodes each chunk into compressed memory representations using a learned compressor. A gating module dynamically selects relevant memory blocks, which are then iteratively processed by a reasoning module with an evolving working memory to solve downstream tasks. The compressor and reasoner are jointly optimized via end-to-end reinforcement learning, while the gating module is trained separately as a classifier. Experimental results show that the proposed method achieves competitive accuracy on multi-hop reasoning benchmarks such as RULER-HQA, extrapolates context length from 7K to 1.75M tokens, and offers a favorable accuracy-efficiency trade-off compared to strong long-context baselines. In particular, it achieves up to a 2 times reduction in peak GPU memory usage and a 6 times inference speedup over MemAgent.