💬 NLP

Training Large Reasoning Models Efficiently via Progressive Thought Encoding

This paper introduces Progressive Thought Encoding, a parameter-efficient fine-tuning method that enables Large Reasoning Models to achieve significantly higher accuracy on complex mathematical benchmarks while overcoming memory bottlenecks by encoding intermediate reasoning into fixed-size vectors, thereby eliminating the need for full-cache backpropagation during reinforcement learning.

Original authors: Zeliang Zhang, Xiaodong Liu, Hao Cheng, Hao Sun, Chenliang Xu, Jianfeng Gao

Published 2026-02-20

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zeliang Zhang, Xiaodong Liu, Hao Cheng, Hao Sun, Chenliang Xu, Jianfeng Gao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a very difficult math problem, like a complex puzzle. You have a brilliant assistant (the AI) who can think through the steps, but this assistant has a very strange limitation: their working memory is tiny.

They can only hold about 10 pages of notes in front of them at once. If the solution requires 100 pages of thinking, the moment they write page 11, page 1 falls off the desk and is lost forever.

The Problem: The "Falling Notes" Dilemma

In the world of Artificial Intelligence, these "notes" are called tokens, and the "desk" is the KV Cache (a memory buffer).

The Old Way (Full Memory): To solve hard problems, the AI used to be given a massive desk that could hold all 100 pages. But this desk is so huge and heavy that it costs a fortune to rent (expensive GPUs) and takes forever to organize.
The "Sliding Window" Way (The Cheap Desk): To save money, researchers tried giving the AI a small, fixed-size desk. As new notes came in, old ones were pushed off the edge.
- The Flaw: If the AI is solving a problem where the answer on page 90 depends on a clue from page 5, and page 5 fell off the desk, the AI gets confused and fails. It's like trying to finish a story when you've forgotten the beginning.

The Solution: "Progressive Thought Encoding"

The authors of this paper came up with a clever trick called Progressive Thought Encoding.

Instead of letting the old notes fall off the desk and disappear into the trash, they teach the AI a new skill: Summarizing on the fly.

Here is how it works, using a creative analogy:

The "Traveling Journalist" Analogy

Imagine the AI is a journalist writing a long investigative report.

The Constraint: The editor says, "You can only keep the last 10 pages of your notebook on your desk. The rest must be cleared."
The Old Method: The journalist just throws away the old pages. When they need to reference a fact from page 3, they can't find it. The story falls apart.
The New Method (Progressive Thought Encoding):
- Every time a page is about to be thrown off the desk, the journalist doesn't just toss it.
- Instead, they quickly write a one-sentence summary of that page's most important clues into a special Traveling Journal (this is the "vector representation" or "LoRA adapter").
- They then throw the full page away, but they keep the Journal.
- As they write more, they keep adding new summaries to the Journal.
- When they need to remember something from the beginning, they don't look at the lost pages; they read the Journal.

Why This is a Game-Changer

1. It's Like a "Magic Backpack"
The AI learns to compress its entire history of thinking into a tiny, lightweight "backpack" (the model weights). Even though the desk (cache) is small, the backpack contains the essence of everything the AI has thought about so far.

2. Cheaper and Faster Training
Because the AI doesn't need a massive desk to hold all its thoughts, it can train on much cheaper computers. The paper shows this cuts the memory usage by nearly 50%. It's like switching from renting a warehouse to using a backpack.

3. Smarter Thinking
Surprisingly, this method didn't just save money; it made the AI smarter.

When the AI is forced to summarize its thoughts, it learns to focus on what's actually important.
In tests on hard math competitions (like AIME), this method beat the standard "big desk" methods by a huge margin (up to 23% better on some tests).

The Bottom Line

The paper solves a major bottleneck in AI: How do we make super-smart AI that can think for a long time without needing a super-expensive computer?

By teaching the AI to "summarize its past thoughts" before forgetting them, the researchers created a system that is:

Efficient: Uses half the memory.
Scalable: Can run on smaller, more accessible hardware.
Accurate: Actually solves harder problems than before because it doesn't lose the "thread" of its reasoning.

It's the difference between a student who forgets the first half of the exam instructions because their notebook is too small, and a student who writes a perfect cheat-sheet summary of every page they've read, allowing them to ace the test even with a tiny notebook.

1. Problem Statement

Large Reasoning Models (LRMs) excel at complex tasks like mathematics and coding by generating long chains of thought (CoT). However, training these models via Reinforcement Learning (RL) faces a critical efficiency bottleneck:

Memory & Compute Costs: RL requires long rollouts to obtain outcome-based rewards. Autoregressive decoding during these rollouts dominates memory (KV cache) and compute usage.
The Trade-off: To mitigate memory usage, standard approaches use sliding-window caches or token pruning. However, discarding intermediate reasoning tokens disrupts long-context understanding, leading to a significant degradation in reasoning quality and training effectiveness.
The Core Question: Can LRMs be trained efficiently under strict memory budgets (fixed-size caches) without sacrificing reasoning accuracy?

2. Methodology: Progressive Thought Encoding

The authors propose Progressive Thought Encoding, a parameter-efficient fine-tuning (PEFT) method that allows models to "remember" evicted tokens by encoding them into the model's weights rather than discarding them.

Key Components:

Cache-Constrained Optimization: The method formalizes RL training under a cache policy $D$ that prunes the KV cache online. The goal is to make the policy under partial-cache rollouts ( $\pi^D_\theta$ ) approximate the full-cache policy ( $\pi_\theta$ ).
Learning from Evicted Tokens: Instead of simply dropping tokens when the cache fills, the model learns from them before eviction.
- Global Query Vector ( $q_g$ ): A learnable vector that aggregates information from all evicted context encountered so far.
- Context State ( $S_e$ ): A compact latent representation of the evicted tokens, computed using global query, key, and value projections.
- LoRA Adaptation: The information from $S_e$ is used to update Low-Rank Adaptation (LoRA) weights ( $\Delta W$ ). This effectively "distills" the reasoning history of the evicted tokens into the model's parameters.
Iterative Update: As the cache fills and new tokens are evicted, the context state $S_e$ is updated and normalized, and the LoRA weights are recomputed. This allows the model to maintain a constant memory footprint while retaining a "summary" of the entire reasoning trajectory.
Training Strategy:
- Question Tokens: Permanently retained in the cache (acting as "sink tokens" to anchor the prompt).
- Thinking Tokens: Subject to a sliding-window eviction strategy.
- Objective: The model is trained using Grouped Reinforcement Policy Optimization (GRPO) with the cache-aware objective, minimizing the KL divergence between the cache-constrained policy and a reference full-cache policy.

3. Key Contributions

Problem Formalization: Identified and formalized the inefficiency of RL training for LRMs under long rollouts as a cache-constrained optimization problem.
Novel Algorithm: Introduced Progressive Thought Encoding, a technique that learns from evicted tokens to update lightweight LoRA adapters, preserving long-context reasoning capacity without increasing cache size.
Efficiency & Scalability: Demonstrated that this method enables scalable RL training by significantly reducing peak memory usage while maintaining or improving reasoning accuracy.
Comprehensive Evaluation: Validated the approach across three models (Qwen2.5-3B/7B, DeepSeek-R1-Distill-Llama-8B) and six challenging mathematical benchmarks.

4. Experimental Results

The method was evaluated on Math500, Olympiad-Bench, Minerva Math, AMC, AIME2024, and AIME2025.

Accuracy Gains:
- Achieved an average improvement of +19.3% over LoRA-based fine-tuning and +29.9% over non-fine-tuned LRMs.
- On the challenging AIME2024/2025 benchmarks, the method achieved up to +23.4% absolute accuracy improvement over baselines under tight cache budgets.
- Specifically, on DeepSeek-R1-Distill-Llama-8B, it improved average accuracy to 45.6% (a +10.7 gain over LoRA).
Memory & Compute Efficiency:
- Peak GPU Memory: Reduced by nearly 50% compared to full-cache LoRA training (e.g., dropping from ~89% to ~60% on the 8B model).
- FLOPs: Reduced attention computation by nearly half (e.g., 3.6 TFLOPs vs. 5.7 TFLOPs for Qwen2.5-7B).
Robustness:
- Unlike sliding-window baselines (LoRAc) which degrade rapidly as cache size shrinks, Progressive Thought Encoding maintains high accuracy even with very small cache windows (e.g., 768 tokens).
- Scaling: The method scales favorably with longer generation lengths (up to 64K tokens), whereas other methods plateau.

5. Significance

This work addresses a fundamental barrier in the deployment of Large Reasoning Models: the conflict between the need for long-context reasoning and the hardware limitations of memory.

Practical Impact: It makes RL training of LRMs feasible on consumer-grade or limited enterprise hardware by drastically lowering memory requirements.
Theoretical Insight: It demonstrates that reasoning information does not need to be stored in the KV cache to be effective; it can be compressed and stored in the model's weights (via LoRA) during the training process.
Future Direction: The paper opens new avenues for adaptive eviction strategies and efficient inference-time optimization, suggesting that "thinking" can be encoded into parameters rather than just stored in memory.

In summary, Progressive Thought Encoding transforms the problem of memory-constrained reasoning from a limitation into a manageable optimization task, enabling more efficient, scalable, and accurate training of next-generation reasoning models.