Towards robust long-context understanding of large language model via active recap learning

TL;DR

The paper addresses the challenge of robust long-context understanding in LLMs by introducing active recap learning (ARL), a framework that builds recursive memory across long texts. It identifies key tokens using the long-short gap metric $I_{\mathrm{LSG}}(x_i, \theta) = \frac{P_\theta(x_i \mid \boldsymbol{l}_i)}{P_\theta(x_i \mid \boldsymbol{s}_i)}$, retrieves informative preceding segments via loss-based retrieval, and refines them into concise recaps that are reinjected during continued pretraining and used by an inference-time recap agent. The approach yields substantial gains on long-context benchmarks (RULER and LongBench), with notable improvements at longer contexts and across different base models, demonstrating scalable memory augmentation without increasing model capacity. This method provides a practical, data-efficient path to enhance long-range reasoning and summarization in LLMs, compatible with both Transformer-based and RWKV architectures.

Abstract

In this paper, we propose active recap learning (ARL), a framework for enhancing large language model (LLM) in understanding long contexts. ARL enables models to revisit and summarize earlier content through targeted sequence construction during contined pretraining and retrospective summarization at inference. First, we identify key tokens in prepared long context based on loss gaps between long and short forward contexts and find most revant preceding paragraphs, then summarize them using an LLM. Second, ARL equips models with the ability to autonomously generate and utilize these retrospective summaries during inference, thereby establishing a recursive memory mechanism across paragraphs. Experimental results show substantial gains, with ARL achieving a 26.8% improvement on RULER and a 9.44% improvement on LongBench. Overall, ARL offers a simple yet effective continued pretraining-based approach to strengthen long-context understanding, advancing scalable memory augmentation in LLM

Figures (2)

• Figure 1: Overview framework of our proposed ARL with a chunk generating recap example.
• Figure 2: Performance comparison of our proposed recap with origin and scaling on RULER bench.