Table of Content

• Introduction

• Chapter 1: Broadcasting the Messages with Society of Minds

• Chapter 2: Giving Agents Identities — The Judge and the Tit-for-Tat

• Chapter 3: Agents That Know What They Don’t Know

• Chapter 4: The Memory Problem and MAD with Memory Masking

• Chapter 5: The Signal You Were Ignoring: Diversity-Aware Retention

Introduction

One research thread has been quietly and carefully examining the communication protocol aspect of multi-agent LLMs since 2023. Rather than building elaborate tool-use pipelines or role-playing simulations, it zooms in on the simplest possible form of multi-agent interaction: debate. Multiple LLMs (same task, no tool, just agent talking), arguing with each other until they converge on an answer. The research direction aims to answer these questions:

🧠 Does structured disagreement between LLMs improve reasoning quality? And if so, how should agents communicate to make that happen?

These questions turn out to be much harder than they look. And like most good ideas in deep learning, the story isn’t one clean paper. Instead, it’s a progression of insights, each one exposing a new failure mode of the previous approach. This blog traces that progression from the ground up.

Chapter 1: Broadcasting the Messages with Society of Minds

The idea begins with a philosophical intuition borrowed from Marvin Minsky’s 1988 book The Society of Mind: intelligence emerges not from a single powerful process, but from the interaction of many simpler ones. Du et al. [1] translated this directly into a 👉Society of Minds prompting framework for LLM agents.

The setup is remarkably simple. Given a query q:

• N instances of an LLM each independently answer a question

• Each agent reads all other agents’ responses and updates its own answer

• This repeats for several rounds until a consensus is reached

The debate prompt engineering is simple, something like: “Based on the opinions of other agents, can you give an updated response?” In practice, the debate runs for a fixed number of rounds R. Regardless of whether full consensus is reached, the final answer is obtained by either taking the first agent's response at the last round, majority voting across all agents, or using an LLM as a judge to select among the final responses.

Formally, at round r, agent i generates:

where Gᵣ₋₁ = {gᵣ₋₁,₁, ..., gᵣ₋₁,ₙ} is the full set of previous responses from all agents. The communication topology looks like this:

Despite its simplicity, the empirical results were striking for their time. On arithmetic, factual QA, chess move validity, and other benchmarks, multi-agent debate consistently outperformed single-agent baselines, chain-of-thought, and reflection-based methods.

More interestingly, debate helped even when all agents started wrong. Agents would sometimes arrive at the correct answer through mutual critique, even if no individual agent held it initially.

👀 Society of Mind’s key finding: broadcasting peer responses and iterating creates a self-correction pressure that single-agent reflection simply cannot replicate.

However, there can be many flaws in the approach. When agents were confidently wrong and agreed with each other, the debate failed to correct the error. Models would affirm the incorrect consensus rather than challenge it. The paper even noted that LMs do not express their uncertainty when generating responses, and that correcting this would likely substantially improve performance. That observation planted the seed for several methods that followed.

❌ Society of Mind’s core problem: it broadcasts everything to everyone with no structure, no role differentiation, and no signal about which message to trust.

A simple mitigation is to simplify the communication topology by using a sparse connection graph (👉Sparse MAD [6]). This not only reduces communication costs but also potentially improves performance when the dense connection creates too much noise.

❌ Yet, a critical limitation remains: static topology treats all messages equally, failing to distinguish critical information from redundant noise, risking both useful context being discarded and irrelevant signals being propagated.

This raises an important research question: 🧠 How to meaningfully cut the communication cost? We will come back to this question in Chapters 4 and 5.

Chapter 2: Giving Agents Identities — The Judge and the Tit-for-Tat

Du et al. showed that debate helps. But their agents were identical clones, i.e., they are from the same LLM, same prompt, no differentiation. Every agent played the same role in the same way. 🧠 What if debate works better when agents are structurally forced to be different?

Liang et al. [2] introduced a structurally different MAD formulation motivated by a specific failure mode they called the Degeneration-of-Thought (DoT) problem. The authors’ observation is that when a model is asked to self-reflect, it tends to converge almost immediately to whatever it initially said. Once an LLM has established confidence in its solution, it becomes effectively incapable of generating genuinely novel thinking in subsequent iterations, even when the initial answer is wrong. Self-reflection, despite its appeal, has a stubborn ceiling.

👀 Originally, it is the problem of single agent that self-reflects. It can also happens in MAD if all of the agents are treated equally.

The fix is structural: rather than allowing agents to freely update their opinions, Liang et al. enforce a “tit for tat” state in which agents are specifically instructed to challenge each other’s positions. This is explicitly stated in a meta prompt for all agents:

Then, two debater roles are constructed to take affirmative and negative sides of the argument. To create these roles for the agent, the paper uses prompting techniques:

Finally, a separate judge agent observes the full debate and determines when sufficient disagreement has been expressed to call a conclusion. The judge serves two functions. First, it acts as a moderator, preventing premature convergence by enforcing continued argumentation when the debate hasn’t been resolved enough. Second, it functions as an adaptive stopping criterion: once the debate has sufficiently explored the disagreement space, the judge extracts the final answer rather than waiting for a fixed number of rounds

This is a significant architectural departure from Society of Mind:

❌ Society of Mind: symmetric agents, same prompt, same role, same update rule

✔️ 👉MAD with judge-and-debater: asymmetric roles, adversarial framing, adaptive termination

The results on counter-intuitive arithmetic and commonsense machine translation showed that the tit-for-tat structure genuinely helped prevent premature convergence. In one illustrative example, agents initially made the classic mistake of averaging Alice’s uphill and downhill speeds to compute her average speed. The adversarial framing forced one agent to challenge this, eventually surfacing the correct method of dividing total distance by total time.

👀 Role assignment changes the pressure structure of debate. When agents are forced to maintain opposing stances, they produce reasoning paths they wouldn’t have generated under free updating.

One potential issue is that the judge’s quality matters enormously. When a weaker LLM is used as the judge while stronger models debate, performance degrades because the judge cannot reliably adjudicate between sophisticated arguments it can’t fully evaluate.

❌ The debate’s ceiling is bounded by the judge’s capacity.

Chapter 3: Agents That Know What They Don’t Know

By this point, the field had two strong intuitions. First, debate helps because it creates multiple non-overlapping reasoning paths. Second, the debate’s outcome depends heavily on which agents get the most influence. In Society of Minds, all agents are treated equally. In the judge-and-debater framework, influence is determined by assigned role. But neither approach directly addresses the natural question: 🧠 What if some agents genuinely know more about this particular question than others?

Lin and Hooi [3] built their method, 👉ConfMAD, around this question. The core insight is that LLMs do have internal uncertainty signals [4], and they just don’t communicate them by default. Confidence expression, the explicit verbal or numerical communication of how certain an agent is about its answer, can reshape the influence dynamics of debate in a principled way.

ConfMAD adds confidence expression throughout the debate process. At round r = 0, each debater independently generates a structured triple: (Reason, Answer, Confidence Score). In subsequent rounds, each debater reads the full debate history, which now includes the confidence scores of all prior responses, and conditions its update accordingly. A high-confidence agent saying “42, confidence 95” should exert more pull on the next round’s responses than a hedging agent saying “42, confidence 40.”

The design reveals a subtle tension: miscalibrated confidence scores can actively damage debate quality. If an agent is confidently wrong, it dominates the update and pulls other agents toward the error. If an agent is genuinely uncertain but labels itself as confident, the same harm occurs. The authors, therefore, pair confidence expression with a calibration step, applying temperature scaling to align the model’s expressed confidence with its empirical accuracy.

The confidence dynamics interact with debate in two pathological ways that ConfMAD must guard against:

• Stubbornness: a high-confidence agent refuses to update even when peer responses present compelling evidence for a different answer

• Premature convergence: agents with low confidence collapse immediately to the majority answer before exploring alternative reasoning paths

The calibrated confidence scores help mitigate both: uncertain agents are encouraged to explore longer; confident agents are given appropriate weight without completely dominating.

However, two key limitations remain:

❌ Confidence calibration is never perfect, risking bias in the debate process

❌ There is no mechanism to preserve the diversity of agent reasoning throughout the debate

More importantly, by adding uncertainty to the prompt, the context is getting longer and potentially adding noise and computation complexity. One promising direction to reduce communication cost is to filter messages from low-confidence agents. However, using a fixed threshold (e.g., uncertainty < 0.5) is fragile since it risks discarding valuable information due to the inherent unreliability of uncertainty estimation. This remains an open problem without a satisfying solution.

Chapter 4: The Memory Problem and MAD with Memory Masking

Society of Mind broadcasts everything. ConfMAD broadcasts everything plus confidence scores. But as the number of agents and rounds grows, “everything” becomes an increasingly noisy and redundant context. 🧠 If past responses are treated as memories that agents condition on, what happens when those memories are wrong?

Tian et al. [5] provide both a diagnosis and a theoretical grounding for this problem. Their key observation: MAD agents are vulnerable to erroneous memories. When an incorrect response from round r−1 persists in the context, it obviously corrupts round r generation. And because context accumulates across rounds, early errors have compounding effects. The poison spreads forward through the debate history.

The theoretical insight is crucial: the performance of MAD is bounded above by the quality of the memories it conditions on. Erroneous memories don’t just fail to help; they actively degrade the agent that reads them.

👉MAD-M2 (MAD with Memory Masking) addresses this by inserting an evaluation-and-masking step between each debate round. After round r generates responses Gᵣ, agents evaluate each response in Gᵣ before allowing them to become memories for round r+1. Two masking strategies are proposed:

Subjective Masking. Each agent uses the LLM itself as an evaluator, labeling each peer response as “YES” (keep), “NO” (discard), or “NOT SURE” (treated as YES or NO depending on a strictness hyperparameter). This is reasoning-aware: the agent reads the content and makes a judgment about its correctness.

Objective Masking. Use perplexity as a proxy for reliability (similar to confidence). High perplexity implies the model was uncertain when generating that response, making it more likely to contain errors. Only the response with the lowest perplexity is retained. This avoids the computational cost of LLM-based evaluation but relies on the assumption that perplexity tracks correctness, which holds reasonably for strong models but breaks down for weaker ones.

👀 An interesting interaction effect emerges: the best masking strategy depends on the base model’s capability. For weaker models (Qwen2.5-7B), subjective masking outperforms objective. In this case, the model can reason about peer correctness even if its own generations are imperfect. For stronger models (QwQ-32B) on hard benchmarks like AIME, objective perplexity-based masking wins decisively, because stronger models’ internal probability estimates are better calibrated.

The three-step pipeline can be summarized as:

1. Generate responses Gᵣ from all N agents

2. Evaluate and mask: produce M ∈ {0,1}ᴺ and retain 𝒮ᵣ = {gᵣ,ᵢ | Mᵢ = 1}

3. Condition round r+1 generation on 𝒮ᵣ rather than the full Gᵣ

Despite looking promising, both masking strategies suffer from limitations:

❌ Subjective masking adds significant token overhead; each agent must evaluate all peer responses before generating their own. Subjective evaluation is not always correct, risking losing important messages (if the agent already knows which answer is incorrect to mask, then it does not need debate).

❌ Objective masking is computationally cheap but assumes perplexity ≈ correctness, and it has the same problem as ConfMAD.

Chapter 5: The Signal You Were Ignoring: Diversity-Aware Retention

Each method so far has addressed a different dimension of the debate communication problem. Society of Minds establishes the baseline: peer communication helps. The judge-and-debater structure enforces role-based disagreement. ConfMAD adds confidence as a weight on peer influence. MAD-M2 filters by correctness signal (subjective or objective). But all of them share a blind spot.

🧠 What if the most important property of a retained response isn’t its confidence or its estimated correctness but how much new information it adds relative to what other agents already believe?

👉Diversity-Aware Retention (DAR [7]) proposes this reframing. DAR argues that the bottleneck in debate isn’t the quality of individual responses. It’s the information content of what gets broadcast. Think about it:

• A high-confidence response that echoes what everyone already thinks? Adds nothing.

• A lower-confidence response that takes a completely different reasoning path? Adds everything.

👀 So instead of asking “how confident is this response?”, DAR asks “how different is this response from what everyone else already said?”

This principle can be implemented using a Filtering Module F, an LLM, that is prompted to retain disagreeing messages.

Here’s how DAR works in practice:

• 🎯 Diversity-first selection: at each round, DAR picks the responses that disagree the most with each other and with the current majority vote. Not the most confident. Not the most fluent. The most informationally distinct.

• 📌 Index-based retention: retained messages are preserved exactly as written, no summarization, no paraphrasing. What the agent said is what other agents read. Authentic disagreement stays authentic.

• 📊 Soft uncertainty signals: each peer response comes with an uncertainty score (average negative log-likelihood over answer tokens), passed directly into the prompt. Agents decide how much to trust each peer themselves — no hard threshold, no costly tuning.

Some advantages of DAR compared to MAD-M2:

This also reflects in the empirical results:

❌ DAR's diversity-first selection is not without risk. Retaining the most disagreeing responses assumes divergence is always informative, but a minority view can simply be wrong. This exposes a fundamental tension: too much consensus and debate collapses into echo chambers; too much divergence and agents get misled by confidently incorrect outliers. The optimal balance between convergence and divergent thinking remains an open and unsettled question in MAD research.

References

[1] Du, Y., Li, S., Torralba, A., Tenenbaum, J. B., and Mordatch, I. Improving Factuality and Reasoning in Language Models through Multiagent Debate. ICML 2024.

[2] Liang, T., He, Z., Jiao, W. et al. Encouraging Divergent Thinking in Large Language Models through Multi-Agent Debate. EMNLP 2024.

[3] Lin, Z. and Hooi, B. Enhancing Multi-Agent Debate System Performance via Confidence Expression. EMNLP 2025.

[4] Nguyen, M., Gupta, S., & Le, H. (2025/2026). Probabilities Are All You Need: A Probability-Only Approach to Uncertainty Estimation in Large Language Models. AAAI 2026.

[5] Tian, H., Feng, X., Zhao, Z. et al. Multi-Agent Debate with Memory Masking. ICLR 2026.

[6] Li, Yunxuan, Yibing Du, Jiageng Zhang, Le Hou, Peter Grabowski, Yeqing Li, and Eugene Ie. "Improving multi-agent debate with sparse communication topology." In Findings of the Association for Computational Linguistics: EMNLP 2024, pp. 7281-7294. 2024.

[7] Nguyen, M., Nguyen, A., Nguyen, D., Venkatesh, S., and Le, H. Hear Both Sides: Efficient Multi-Agent Debate via Diversity-Aware Message Retention. arXiv:2603.20640 (2026).

Table of Contents

• Introduction

  • Why RL Post-Training Matters for Reasoning?

  • From Test-Time to Post-Training: What’s the Shift?

• DeepSeek-R1: The Tipping Point

  • Formalizing the Post-training Framework

  • The RL Algorithm: Group-Based Policy Optimization (GRPO)

  • On the Reward Choice

  • Other Interesting Insights

• Beyond GRPO: For A Better RL Post-Training

  • Alternative RL Algorithms

  • Reward Shaping

  • Optimizing the Training Pipeline

Introduction

Reasoning remains one of the few capabilities where LLMs still fall short. When the conversation turns to boosting reasoning in LLMs, the default reaction is to immediately reach for test-time tricks: self-consistency, searching, multi-step decoding, or heavy ensembles. These methods are undeniably useful, but they all share a fundamental limitation: they are external scaffolds layered on top of an unchanged model.

Why RL Post-Training Matters for Reasoning?

Reasoning is ultimately about shaping the model’s internal search process, including how it attends, decomposes, checks, and revises its predictions. Supervised finetuning alone gives you patterns but rarely instills the principles needed for multi-step, error-sensitive reasoning. RL provides a way to push the model to explore diverse behaviors until it figures out the principles of reasoning on its own. When done right, the model starts solving tasks cleanly without huge decode budgets.

Over time, the model learns to prioritize reasoning strategies that generalize across tasks, rather than memorizing superficial patterns from the training data. This is why even relatively small models can outperform larger counterparts on multi-step reasoning benchmarks: the improvement comes from better internal reasoning processes, not from bigger parameter counts or expensive test-time tricks.

From Test-Time to Post-Training: What’s the Shift?

Test-time approaches rely on external scaffolds, generating multiple trajectories, performing searching and voting on the best answer, or sampling repeatedly to increase reliability. The model itself doesn’t change, so inference is expensive, and reasoning traces remain fragile depending on the samples.

Post-training RL takes a fundamentally different approach. By using reward-driven model updates, it reshapes the model’s internal heuristics, allowing it to plan, check, and revise its own reasoning based on the training rewards. As summarized in the figure below, post-training RL trades upfront training cost for cheaper, more reliable inference, cleaner reasoning traces, and stronger generalization. This shift marks the difference between forcing a model to reason through brute-force search and actually teaching it to reason internally.

It’s also important to understand how post-training RL differs from the other two major forms of post-training: supervised fine-tuning (SFT) and preference training.
SFT teaches the model to imitate reasoning patterns, but it requires ground-truth reasoning traces and can only learn what’s shown in the ground-truth data. There’s no pressure to explore, self-correct, or discover better reasoning paths. Preference training (like DPO or MRPO) adds a ranking signal, nudging the model toward preferred outputs, but it still operates passively: the model isn’t encouraged to search, plan, or discover new strategies, only to reshape its distribution around examples humans liked.

RL is different. RL gives the model room to experiment, take multi-step trajectories, encounter failure, and adjust its internal heuristics based on cumulative reward. That exploration loop is exactly what makes reasoning skills intrinsic rather than decorative. Put simply:

❌ Test-time scaling searches for the best outputs

❌ SFT teaches outputs

❌ Preference training teaches preferences over outputs

✔️ RL teaches the process that produces those outputs

And that’s why post-training RL is emerging as the most reliable way to move beyond surface-level reasoning gains. It turns reasoning from something we approximate at inference into something the model genuinely knows how to do. With that foundation in place, we can now study representative works on RL post-training, also called Reinforcement Finetuning (RFT).

DeepSeek-R1: The Tipping Point

DeepSeek began as an ambitious open-source initiative pushing the boundaries of efficient LLM training. Unlike many labs chasing scale at any cost, DeepSeek focused on pushing reasoning quality without gigantic models or massive inference budgets. That philosophy culminated in 👉DeepSeek-R1 [1], their breakthrough reasoning model.

R1 is built on a deceptively simple idea: outcome reward is all you need for RFT. It raises a research question:🧠 If a task has a verifiable final answer, can we train the model directly with policy gradients and a verifiable reward? For domains like math, logical inference, and structured problem solving, where correctness is binary, this setup is ideal. Instead of collecting huge preference datasets, relying on imitation or complicated process reward models, the model explores, fails, receives a clean reward, and gradually internalizes the principles of stepwise reasoning.

The DeepSeek team not only proved that this is doable, but they also open-sourced almost the entire training pipeline, showed how to make verifiable-reward RL work at scale, and did so at a shockingly low cost. Training was cheap, the resulting API was 27× cheaper than OpenAI’s reasoning models, and in many benchmarks, the performance was effectively equivalent.

The moment constituted a powerful slap in the face to the prevailing view that expensive GPUs (promoted by NVIDIA) and high-budget models (like those from OpenAI) were indispensable for advanced AI:

👀 R1 demonstrated something the field had suspected but never proven at this scale: Great reasoning doesn’t require trillion-parameter models. Rather, it requires the right RL formulation.

Formalizing the Post-training Framework

DeepSeek-R1 is an RL post-training framework that helps LLMs acquire real reasoning ability purely through reinforcement learning without needing human-annotated reasoning traces.

In the paper [1], the authors present 2 DeepSeek-R1 versions. In the original “R1-Zero” variant, the authors start from a base LLM (DeepSeek-V3-Base) and apply a pure RL pipeline with a verifiable final-answer reward (for math, code, STEM tasks) and without any supervised fine-tuning (SFT). Over thousands of training updates, the model begins to exhibit emergent reasoning behaviors: self-verification, chain-of-thought generation, reflection, and dynamic strategy adaptation.

The final, polished DeepSeek-R1 extends this with a multi-stage pipeline: a cold-start using a small curated CoT dataset to improve readability, followed by RL, then rejection-sampling to build fresh SFT data, then another RL pass (including domain-general prompts).

A key is what happens after the first RL run. DeepSeek takes the improved RL policy and uses it to generate many candidate reasoning traces for new problems. Instead of directly using the raw samples, they apply rejection sampling:

• Keep only trajectories with correct final answers

• Optionally rank by readability, coherence, or internal consistency

• Discard the rest completely

This gives you a large dataset of high-quality, self-generated reasoning demonstrations. It’s like mining the model’s own best moments and discarding the noise.

👀 Crucially, this dataset is much cleaner than the original SFT data with no human variation, a fully standardized structure, and perfectly aligned with the target format introduced during warm-start.

Then, DeepSeek performs an SFT pass on the new data. This step accomplishes two things:

1. It distills the stable reasoning behaviors discovered during RL into a single, deterministic forward pass.

2. It smooths out the variance of RL sampling, making the model’s outputs cleaner and more consistent at inference.

At the end, the key component is the Pure-RL block. To understand it, we begin by framing LLM reasoning as an RL problem. An LLM acts as a conditional policy π_θ over text, and each reasoning trajectory can be viewed within the standard RL tuple (S,A,R):

• State s_t​: the prompt q plus the partial generation up to step t:

• Action a_t​: the next token x_t​ drawn from the policy:

• Reward r: a final sequence-level reward given after the full output is completed:

where a* is the ground-truth answer, and R is the reward function that can measure how accurate the output x_1:T is compared to the true answer a*. Usually, a simple exact match function with the ability to extract the final answer from the output is used. Because the reward is computed based on the ground-truth answer given by the dataset, it is:

✔️ Verifiable

✔️ Not requiring human annotations

This simple reward mechanism is crucial to DeepSeek-R1’s success, establishing it as an affordable alternative model. Yet, one question remains: 🧠 If the method is so straightforward, why has no one successfully implemented it before?

The answer lies in the RL, which demands specialized techniques to work effectively in the LLM context:

• RL is Inherently Hard to Train: Standard RL algorithms are notoriously sample-inefficient and often exhibit high variance. Training requires meticulous tuning and robust techniques to stabilize the learning process.

• Need for a Strong Base Model: RL algorithms typically cannot “bootstrap” a successful model from scratch. They require a “strong enough” base model to provide a solid foundation of language understanding and coherence. The base model used for DeepSeek-R1 (DeepSeek-V3), which is actually quite substantial and not small at all, provides this necessary strength to enable RL fine-tuning to work effectively.

• Sparse Reward Problem: The task of aligning an LLM often results in a sparse reward signal. The model only receives feedback on the quality of its final response, making it difficult for the RL agent to determine which specific intermediate steps or tokens contributed to the success (or failure) of the output.

• Necessity for a Novel RL Algorithm: Due to the issues of high variance and sparse rewards, the authors could not simply apply a naive or off-the-shelf RL technique. They had to adopt a new RL algorithm, named GRPO, specifically designed to be stable, scalable, and effective when applied to the vast parameter space of a powerful base LLM in a practical setting.

The RL Algorithm: Group-Based Policy Optimization (GRPO)

DeepSeek-R1’s RL algorithm is based on a well-known trust-region algorithm, called PPO [2]. To remind you, the original PPO objective is like this:

where w_t is the important sampling ratio for action a_t.

Here, DeepSeek-R1 modifies PPO with a group-based policy optimization scheme suitable for text space, which results in a new group-based advantage function. Specifically, for each query q, the model samples a group of G candidate outputs:

Each output o_i represents the whole sequence of tokens and receives a scalar reward r_i. Instead of fitting a value function and using the value function estimation in calculating the advantage, DeepSeek uses relative ranking inside the group to compute advantages:

This yields a stable, variance-reduced advantage signal without a critic. The optimization objective is a PPO-style clipped surrogate, coupled with a KL constraint:

Here is the summary of GRPO’s features:

• It operates on a whole sentence level, not a step/token level.

• It avoids learning value networks and lets the reward signal directly shape reasoning behaviors through group-based advantage estimations.

• The objective function is averaged over G outputs, reducing variance in the gradients

• The DPO clip forces the update to be small, increasing stability.

• The KL further forces the new LLM to be close to the base model, maintaining general pre-trained performance.

On the Reward Choice

As discussed earlier, the reward is simple and verifiable based on the ground-truth. However, to support reward calculation, the output of LLM must follow certain formats to allow later extraction of the final answers. If the model rambles, changes styles mid-generation, or forgets to delimit its conclusion, the verifier can’t score it, even if the reasoning is mostly correct. Therefore, DeepSeek-R1 proposes using the following prompt to ensure output format:

Once the output is parseable, the reward becomes a clean deterministic function:

where a is the extracted answer between <answer></answer> and a* is the true answer.

Other Interesting Insights

One of the most surprising findings is that you don’t need any chain-of-thought supervision to unlock multi-step reasoning. The model invents a reasoning style because it is the only reliable way to consistently maximize outcome reward. This is a major shift in thinking:

Reliable Answer⇒Reliable Internal Search

The model rediscovers analysis → checking → correction, not because we teach it, but because outcome reward makes messy one-shot answers unstable, failing to generalize to many input questions seen during training. As proved by the results, pure RL post-training is already competitive against other baselines:

Second, the paper also indicates the benefit of later SFT or distillation that compresses RL behaviors. After the first RL run, DeepSeek rejects low-quality samples and distills the remaining high-quality traces back into the model. You might expect a noticeable drop in performance because SFT is usually seen as “blurry” compared to RL. Instead, the distilled model actually becomes more stable and often more accurate.

Why? RL discovers a diverse set of good reasoning strategies, but SFT selects only the clean, high-confidence trajectories and forces the policy to reproduce them deterministically. The result is a smoother, more reliable reasoning model without the sampling volatility of pure RL.

👀 RL finds the behavior; SFT perfects it.

This approach of distillation from RL-trained big models to smaller ones is very effective:

Finally, one of the most striking aspects of DeepSeek-R1 is how much emergent structure arises even though the reward is extremely simple. There is no preference model, no human-written quality annotations, no chain-of-thought supervision baked into the reward. Yet the model discovers behaviors that look almost like they were explicitly trained for: the “aha moment”.

👀 Although the behavior is indeed interesting and never seen in training data, subsequent papers (e.g., in [3]) show that the base LLM intrinsically possessed this capability, though it manifested with a lower probability before fine-tuning.

Beyond GRPO: For A Better RL Post-Training

In the previous section, we saw how GRPO laid the foundation for post-training RL: by estimating advantages across sampled reasoning trajectories, GRPO reduces variance compared to naïve policy gradients and reliably nudges the model toward higher-reward outputs. However, it remains rudimentary and leaves plenty of room for improvement. GRPO is still rudimentary in several ways:

• Residual variance in long reasoning chains: As tasks grow in complexity, the gradient signal from sampled trajectories can remain noisy, making learning slow or unstable.

• Vulnerability to rare but correct trajectories: GRPO updates can inadvertently reduce the probability of correct but infrequent outputs, especially under sparse rewards.

• Scalability bottlenecks: When applied to very large models or long, multi-step reasoning, GRPO can become computationally expensive and less stable.

Alternative RL Algorithms

To improve GRPO, many researchers have come up with alternative algorithms for RL post-training. For example, researchers introduced Decoupled Clip Dynamic Sampling Policy Optimization (👉DAPO [4]), which is a set of algorithmic and engineering fixes to make RL practical and reproducible at scale.

Trick 0: Remove KL term

Standard GRPO uses a KL penalty to constrain policy divergence. The authors exclude this term because long-CoT reasoning requires significant distributional shifts from the initial model, rendering the restriction unnecessary.

Trick 1: Clip-Higher

Recall that the clip in GRPO is to constrain the policy update, stabilizing the training. However, because of the constraints, standard clipping restricts probability updates: high-probability “exploitation” tokens can grow freely, but low-probability “exploration” tokens struggle. For example: with ε = 0.2 and positive advantage:

• Token A: π_θold=0.9 → upper bound 0.9⋅1.2=1.08

• Token B: π_θold=0.01 → upper bound 0.01⋅1.2=0.012

Here, high-probability tokens (A) easily increase, but rare tokens (B) barely move. This limits exploration and slows scaling. DAPO fixes this with the Clip-Higher strategy. Clip-Higher relaxes the upper bound for low-probability tokens, enabling more diverse sampling and helping the model discover better reasoning trajectories. In practice, the idea is implemented by simply modifying the clip part:

By using different hyperparameters for low and high clipping thresholds, the authors increase the value of ε_high to create space for the increase of low-probability tokens. while keeping ε_low small because increasing it will drive these token probabilities toward zero, collapsing the effective sampling space.

Trick 2: Token-Level Policy Gradient Loss

The original GRPO setup reduces the loss at the sample level: it calculates the probability of the whole response by averaging token probabilities within a sample, then averages across samples, giving every sample the same weight. In implementation, it looks like this:

Because the token contribution is normalized by the length of the response, tokens inside long responses effectively count less. This creates two issues:

1. Good long answers are undertrained, as the model can’t fully absorb the reasoning patterns in key tokens in the long answers.

2. Bad long answers are under-penalized as gibberish repetition isn’t punished strongly enough, causing entropy and response length to drift upward

To fix this, DAPO switches to a token-level policy gradient loss, giving each token its fair contribution to the update by not normalizing by response length but by the total number of tokens. The loss is:

Under this formulation, longer sequences naturally exert more influence on the gradient than shorter ones. More importantly, useful token patterns, whether they appear early or deep inside the reasoning trace, are updated proportionally, and harmful patterns get penalized no matter where they occur.

Trick 3: Dynamic Sampling

A subtle issue in GRPO-style algorithms is that when a prompt’s outputs all achieve 100% accuracy, every sample in that group receives the same reward. This collapses the advantage to zero, leading to zero gradients, which means weaker updates, higher noise sensitivity, and degraded sample efficiency.

Worse, as training progresses, the number of “all-correct” prompts steadily increases, so the effective number of learning prompts per batch shrinks. DAPO addresses this by oversampling and filtering prompts whose accuracy is exactly 1 or 0. Only prompts with meaningful gradient signals are retained, ensuring every batch contains a stable number of informative prompts. Before each update, the system dynamically samples until the batch is fully populated with prompts that provide non-zero gradient contribution, keeping training stable even as accuracy improves.

Trick 4: Overlong Reward Shaping

Most RL pipelines cap the maximum generation length and simply truncate responses that run past it. The default practice is to assign a punitive reward to these truncated samples. But in long-CoT reasoning, this turns out to be harmful: a perfectly good chain of thought may get cut off just because it’s long, and the model is punished for the wrong reason.

DAPO proposes to mask the loss of truncated outputs (Overlong Filtering), making training dramatically more stable and performance on benchmarks like AIME jumps immediately. But filtering alone isn’t enough. Long responses can still waste tokens and reduce efficiency. So DAPO introduces a Soft Overlong Punishment, a gentle, length-aware shaping term:

This function grades the penalty smoothly as the model approaches the max-length zone, which is L_cache from L_max, and only applies a full −1 reward when it exceeds L_max.

The result is incrementally better as we apply these tricks one by one:

At the same time, concurrent researchers have started investigating other changes to the GRPO algorithm. Paper 👉Dr. GRPO [5] reveals 2 limitations of GRPO as a result of dividing the advantage by response length and per-question return variance.

Both seem harmless at first glance, but together they create two systematic biases:

❌ Length bias: If a response is correct, shorter outputs get larger gradients because you’re dividing a positive advantage by fewer tokens. If a response is incorrect, longer outputs get less penalty because they’re cushioned by a bigger length divisor. Over time, this pushes the policy toward long incorrect chains and short correct ones. Simple fix:

👀 This is similar to the token-level policy loss trick in DAPO paper.

❌ Difficulty bias: Normalizing by the per-question standard deviation means that questions where all responses are almost correct (or almost wrong) get disproportionate weight. Easy or impossible examples dominate updates, while medium-difficulty ones, where reasoning actually matters, get suppressed. This can be simply fixed by removing the std denominator:

These tricks accelerate the progress of RL post-training:

In addition to the above simple tricks, other papers look for fundamental changes to the RL algorithm. One of the big limitations of pure on-policy RL for reasoning (GRPO/PPO-style) is that the model can only explore what it can generate. If the base model can’t produce high-quality chains of thought, then RL mostly just amplifies its existing bias patterns instead of teaching it genuinely new reasoning skills.

The paper, Learning to Reason under Off-Policy Guidance (👉LUFFY [6]), introduces a simple but surprisingly effective idea: use off-policy guidance such as reasoning traces from a much stronger model (e.g., DeepSeek R1), and fold them directly into the RL pipeline so the learner can “see” good reasoning even before it can generate any. But unlike naïve imitation, the goal here is not to copy the teacher; it’s to blend demonstrations into GRPO in a way that preserves exploration and lets the student ultimately outgrow the off-policy samples.

Formally, let G_on and G_off denote on-policy and off-policy trajectory groups, sampled from π_θold and π_ϕ, respectively. LUFFY computes advantages using the union of both:

This ensures that high-quality off-policy trajectories receive larger advantages early in training (as on-policy rollout is worse), accelerating learning without overriding on-policy exploration as the policy improves. Then, two PPO-styled objectives for on- and off-policy training are used to update the current policy. Here, the off-policy term uses its own importance sampling ratio:

Together, we have a mixed objective:

where Z is the total number of tokens, serving as a normalization term.

👀 Because π_θ​ is much closer to πθold​ than to π_ϕ​, the off-policy ratio gradually​ becomes smaller, naturally tempering gradients from the off-policy data.

Now, one problem arises. Naively mixing off-policy data leads to rapid entropy collapse: the model overcommits to high-probability actions that coincide with off-policy tokens, eliminating exploratory behaviour required for multi-step reasoning. To counter this, LUFFY applies a shaping transformation to the off-policy importance ratio f(r) and removes the clip function, yielding the shaped off-policy gradient:

If we derive the gradient contribution for each candidate off-policy token τ′_t at time t:

If we don’t use the shaping function, f(r) = r, f’(r)=1, then the gradient magnitude is bounded by π_θ​(1−π_θ​). This term is tiny when:

• π_θ is near 0 → model thinks the token is impossible

• π_θ is near 1 → model is already confident

This is bad because the model struggles to get gradient signals from new reasoning moves (often starting with low-probability tokens). Thus, they propose shaping a function like this:

where γ is set as 0.1. This has the key effect:

• When π_θ is small, f′ is large⇒boosts gradients

• When π_θ is large, f′ shrinks⇒dampens near-certain tokens

This improves the exploration and avoids entropy collapse:

and better performance overall:

Reward Shaping

Improving the RL algorithm is only part of the story. The other part is how we tell the model what success looks like. Reward design shapes the entire search landscape that the model explores. A good reward doesn’t just score outputs; it pulls the model toward the kinds of internal reasoning moves we want it to adopt. And unlike test-time scaling, reward shaping changes the model’s internal dynamics, not just its decoding strategy. The right reward can turn a passive pattern-matcher into an active problem-solver.

One important element in determining reasoning quality is the CoT length: Give models more thinking time, and they often reason better. But when you fine-tune a model on long-form CoT traces and then optimize it with RL, the model doesn’t merely preserve the long patterns; it tends to amplify them. Both Llama-3.1-8B and Qwen-2.5-Math-7B quickly push their CoTs longer and longer until they hit the context limit. Once trajectories exceed the max length, the model gets a negative or zero reward, not because the reasoning is bad, but because the sequence literally doesn’t fit.

👀 Similar problem like Trick 4 in the DAPO paper.

To fix this, the authors propose a length-based reward shaping, named 👉Cosine Reward [7]. It introduces a very clean idea: make the reward sensitive to the CoT length.

Instead of the classic “1 for correct, 0 for incorrect” sparse reward, they use a piecewise cosine reward that encodes three simple principles:

1. Correct > wrong: Correct answers always get a higher reward than wrong ones.

2. Shorter > longer (if correct): Among correct solutions, shorter CoTs are better.

3. Longer > shorter (if wrong): If the model is wrong, encourage it to think longer next time.

To implement this, they propose a reward formula:

where:

This function makes sure that as L_gen increases to L_max, the reward smoothly changes between the two reward hyperparameters. Depending on the correctness, the reward change can be decreased (correct) or increased (incorrect):

The addition of Cosine Reward to RL training helps stabilize training as the number of training iterations increases:

👀 As we will see later, the drop of performance due to repoonse length suddenly jumps to max is one of the “training collapses“ often occured in RL post-training.

However, introducing a length-based reward like Cosine Reward can be problematic. As for incorrect answers, a longer response is encouraged. The model exhibited reward-hacking behavior, inflating CoT length on difficult problems through repetition rather than genuine reasoning.

To address this, the authors introduce an N-gram repetition penalty: apply token-level penalties directly at the positions where repetition occurs. By recording N-grams at every step, they can detect an N-gram repetition to add a penalty (as a negative reward, for instance).

These reward shaping techniques collectively help improve the reasoning performance on math datasets significantly:

As we can see, length-aware reward shaping helps stabilize emergent CoT scaling by nudging models toward efficient reasoning trajectories. However, even sophisticated length-based rewards face a key limitation: they assume the model is already capable of producing coherent, partially correct reasoning. This assumption breaks down for weaker models such as tiny LLMs(≤1B parameters), where reasoning is fragile, and outcome rewards are extremely sparse. This leads to a natural question: 🧠 How do we guide RL training when the model fails to produce any correct trajectories for long periods of time?

This is where the 👉Memory-R+ [8] paper contributes a major new idea:
Memory-augmented intrinsic rewards. The Memory-R+ method proposes a shift in perspective: instead of relying solely on external rewards (correct/incorrect), the model should also learn from its own past reasoning, much like humans rely on episodic memory.

The approach introduces two episodic memory banks:

• Success Memory (Mₛ): stores reasoning traces that led to the correct answer

• Failure Memory (M_f): stores reasoning traces that produced incorrect results

These memories are stored as embeddings in a shared representation space using a Sentence Transformer encoder. The key idea is simple but powerful:

• Reward the model when its reasoning resembles past successes

• Reward the model when it avoids past failures

This creates a dense, performance-driven intrinsic reward signal that tiny models desperately need. Given a new query, the memory system:

1. Looks up similar past queries in both success and failure memories using k-NN in embedding space.

2. Retrieves the associated reasoning traces (not just final answers).

3. Evaluates the new response based on:

   • similarity to successful reasoning (exploit)

   • dissimilarity to failed reasoning (explore)

Exploit Reward

The success memory provides a set of retrieved successful reasoning traces B. Their embeddings are averaged to form a centroid:

A generated response a earns a higher reward when it is closer to this centroid:

This encourages generalizable reasoning patterns, not rote memorization.
It captures structure like “identify quantities → set up equation → solve”, even if the surface text differs.

Explore Reward

Failure memory provides a set of incorrect reasoning traces. The authors measure novelty as the inverse similarity to the closest failure:

If the model repeats a failed reasoning pattern, its reward goes down. If it proposes a novel direction, the reward goes up. This creates a natural curriculum:

• Early: encourages broad exploration (most reasoning is wrong).

• Later: focuses on fine-grained distinctions between near-correct and correct reasoning.

Finally, each intrinsic reward is min-max normalized over a sliding window, then combined:

The results are improvements across many tiny LLMs, including one with Thinking Mode (Qwen3-0.6B):

Training Instability and Collapse in LLMs

The paper also analyzes the training collapse issues with RL post-training. LLMs often over-optimize simpler rewards, such as format rewards, at the expense of correctness, a phenomenon known as reward mode collapse. Models without intrinsic rewards focus on easy metrics, while Memory-R+ balances exploitation (aligning with successful reasoning) and exploration (avoiding past failures), preventing collapse.

Response length collapse is another issue: models either under-generate (too short) or over-generate (too long) sequences, producing meaningless outputs. Length-based rewards like Cosine can worsen this. Memory-R+ stabilizes training by providing dense memory-based feedback, ensuring reasonable response lengths while improving correctness.

While Memory-R+ effectively stabilizes tiny LLM training through episodic memory and intrinsic rewards, it primarily operates at the response level, assigning rewards based on entire reasoning chains or key intermediate patterns. However, recent research has highlighted the benefits of dense process rewards that provide feedback at each reasoning step, rather than only at the outcome.

👀 Token- or step-level feedback allows the model to understand which intermediate decisions are correct, improving training efficiency and reasoning fidelity.

Despite their advantages, dense rewards are rarely deployed at scale in RL for LLMs. The challenges are threefold:

❌ Defining process rewards: Assigning meaningful credit to intermediate steps is non-trivial, as some seemingly “incorrect” steps may contribute indirectly to a correct final answer.

❌ Scalability of online updates: Updating dense reward models (process reward models, PRMs) online requires frequent retraining on step-level labels, which is costly and often infeasible.

❌ Extra modeling cost: Conventional dense reward methods require separate reward models trained with expensive annotations, adding significant overhead.

The 👉PRIME [9] framework (Process Reinforcement through Implicit Rewards) offers an elegant solution. Instead of requiring step-level labels, PRIME leverages implicit process reward modeling to generate token-level dense rewards derived from standard outcome labels. Essentially, a single reward model can infer dense rewards for each step, which are updated online with policy rollouts. This approach:

✔️ Provides fine-grained, step-level feedback for improved credit assignment.

✔️ Reduces reward sparsity without additional annotation cost.

✔️ Mitigates reward hacking by updating the reward model along with the policy, maintaining alignment between the model and its reward signal.

The Implicit PRM is trained only with outcome-level labels but can produce token-level dense rewards. To be specific, the reward model π_ϕ​ is trained with the outcome reward:

We can train it using cross-entropy loss (after applying the sigmoid function to normalize the reward) based on whether the output is correct or not, such that sequences with higher outcome rewards should be assigned a higher reward.

Then, during inference to generate dense rewards for RL post-training, the reward at step t (i.e., for token y_t​) is:

where:

• π_ϕ​ = the Implicit PRM (reward model)

• π_ref = reference model (often the initial SFT or base LM)

• β = scaling factor

👀 The reward measures how much more likely the PRM is to generate this token compared to the reference model, which implicitly reflects intermediate correctness or quality at that token.

Because we now have a reward for every step, we need to calculate the advantage function used in the RL algorithm (not necessarily GPO). Here, they propose to use a leave-one-out (LOO) baseline, resulting in the advantage function for each step t:

The advantage consists of the implicit process reward and outcome reward components. With a step-based advantage, they can use the standard step-level PPO as the RL objective:

❌ However, training an additional reward model and using step-level PPO can be computationally expensive.

Optimizing the Training Pipeline

So far, most advances in post‑training RL for reasoning optimize which algorithm to use or how to shape the reward. But there’s a third: how we structure the training itself, e.g., via curriculum learning. Curriculum learning is inspired by human education: instead of exposing a model to the hardest tasks from the start, we begin with easier tasks and progressively increase difficulty. So the key task is to estimate the difficulty of the data sample.

In the 👉AdaRFT [10] paper, the authors propose to use precomputed difficulty scores for each problem, which can come from human annotations, empirical success rates, or a separate difficulty-estimation mode. For example, they use an external LLM as the difficulty estimator. The problem now is choosing the right estimator model because not all models are suitable for difficulty estimation:

• Too strong (e.g., OpenAI o1, DeepSeek-R1): These models solve most problems on the first attempt, leaving little variance across problems. As a result, easy and hard problems cannot be distinguished effectively.

• Too weak (e.g., LLaMA 3 1B): These models fail on nearly every problem, producing insufficient signal to guide curriculum adaptation.

Thus, they select Qwen 2.5 MATH 7B as the difficulty estimator because it exhibits balanced problem-solving ability: it succeeds on moderately difficult problems but struggles with the most challenging ones. Then, for each problem i, the difficulty score d_i​ is defined as:

This represents the empirical average accuracy of the estimator on problem i. With the difficulty score, the curriculum goal is to assign the samples with difficulty most suitable to the current model (target difficulty), not too easy, not too hard. As such, starting with an initial level of difficulty target, as the LLM improves over time, the target difficulty increases; if performance drops, it decreases. At each step, the model is finetuned on problems closest to the current target difficulty, ensuring steady, aligned progression. This is how it works in RL post-training:

1. Dynamic Curriculum Sampling

   • Compute the absolute difference between each problem’s difficulty and the current target difficulty.

   • Select a batch of problems that are closest to the target, keeping tasks neither too easy nor too hard.

2. Policy Update

   • The policy model generates responses for the selected batch.

   • Compute rewards based on correctness and update the policy using a reinforcement learning algorithm (e.g., PPO, GRPO, REINFORCE++).

3. Target Difficulty Update

   • The average reward over the batch determines whether to increase or decrease the target difficulty.

   • Smooth updates are ensured using a tanh function, while the difficulty is clipped within valid bounds.

👀 This algorithm introduces many hyperparameters, which can be problematic for tuning.

That said, the results have improved significantly:

In addition to difficulty estimation via external models, there can be a simpler way to define sample difficulty: rely on the length of the question. In the 👉FASTCURL [11] paper, the researchers analyze the DEEPSEEK-R1-DISTILL-QWEN-1.5B model and discover that longer input prompts generally produce longer output responses:

This observation motivated a simple but effective hypothesis: For complex reasoning tasks, the complexity of a problem correlates with the length of the solution the model must generate. Using this principle, the paper divides the training dataset into three subsets based on average input prompt length:

• L1: Short CoT reasoning problems (simpler tasks, easy)

• L2: The original dataset (baseline tasks, medium)

• L3: Long CoT reasoning problems (more complex tasks, hard)

Below are the statistics of the datasets:

Therefore, the paper proposes to divide training into stages that correspond to different context lengths and data set lengths (difficulties). They tested with many configurations and found many good strategies. For example, they use 5 training stages corresponding to:

• Context length: 8K, 16K, 24K, 16K, 16K

• Datasets: L1, L2, L3, L2, L2

The intuition is first to progress from easy to hard (L1 to L3), then come back to the medium difficulty (L2) until convergence. No theory is given, but empirically, it works best:

Despite these good results, two major obstacles limit LLM performance in the normal curriculum setting:

❌ Difficulty Shift: In education science, there is a Difficulty Shift phenomenon, a.k.a., the model’s perception of a problem’s difficulty changes dynamically as it learns. A problem considered “hard” initially may become “easy” later, making a static difficulty score, as in AdaRFT or FastCuRL, obsolete.

Inspired by this principle, the paper 👉ADCL+EGSR [12] focuses on optimizing the data progression via a dynamic data assignment. There are two contributions in the paper: ADCL and EGSR, where the former focuses on curriculum improvement and the latter helps with training guidance. First, instead of using a static task ordering, ADCL periodically re-estimates the difficulty of upcoming batches based on the model’s current state:

1. Initially, difficulty scores δ₀ are assigned to all samples in the dataset D using the base model parameters θ₀​.

2. The dataset is sorted by these scores and partitioned into sequential batches B₁,B₂,...,B_K​.

3. After training on batch B_k​, the model parameters are updated to θ_k.

4. ADCL re-evaluates difficulty scores δ_k+1 for the next batch’s samples and re-sorts the batch internally according to the new difficulty estimation before proceeding to the next iteration

Besides curriculum training, the paper also aims to address one issue of current RL post-training:

❌ On-Policy Exploration Limits: Standard post-training, rely solely on self-generated trajectories. While this can improve problem-solving for tasks the model can occasionally solve, it struggles when the model’s current knowledge is insufficient to produce non-zero reward outputs. This “zero-reward” scenario halts learning because gradient updates vanish, i.e., no signal for training.

To address this issue, they propose to incorporate training trajectories generated from an expert policy. This idea is very similar to the aforementioned LUFFY paper. However, instead of following the correct way of defining important sampling for on and off-policy data like this:

The paper indicates that directly using an expert policy π_ϕ​ is often impractical. The expert policy may be inaccessible or rely on incompatible tokenization, making probability ratio calculations infeasible. Even when available, the distributional mismatch between π_ϕ and the current policy π_θold​ can result in extreme importance weights, leading to unstable updates. So they instead use expert demonstrations to guide the current policy in generating improved trajectories:

Here, the expert demonstration g is used in the prompt. This leads to a different training objective:

Here, the offline trajectories are generated by π_θold, and the important sampling ratio for offline trajectories is now not significantly different from that for online trajectories. The advantages are also mixed between the two online and offline returns, as expected. The performance of the final model, ADCL+EGSR, is better than the standard RL post-training:

❌ A big limitation of the approach is that ADCL requires recomputing the difficulty scores for each sample frequently, which slows down training.

🧠 The question of how to build an efficient adaptive RL curriculum during LLM post-training remains open and calls for continued research [13].

References

[1] Guo, D., Yang, D., Zhang, H. et al. DeepSeek-R1 incentivizes reasoning in LLMs through reinforcement learning. Nature 645, 633–638 (2025). https://doi.org/10.1038/s41586-025-09422-z

[2] Schulman, John, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. “Proximal policy optimization algorithms.” arXiv preprint arXiv:1707.06347 (2017).

[3] There May Not be Aha Moment in R1-Zero-like Training — A Pilot Study. https://sail.sea.com/blog/articles/62

[4] Yu, Q., Zhang, Z., Zhu, R. et al. (2025). DAPO: An open-source LLM reinforcement learning system at scale. NeurIPS 2025.

[5] Liu, Zichen, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin. “Understanding r1-zero-like training: A critical perspective.” COLM 2025.

[6] Yan et al. (2025). Learning to Reason under Off-Policy Guidance. NeurIPS 2025.

[7] Yang, S., Tong, Y., Niu, X., Neubig, G., & Yue, X. (2025). Demystifying Long Chain-of-Thought Reasoning. In Proceedings of the 42nd International Conference on Machine Learning (ICML 2025).

[8] Hung Le, Van Dai Do, Dung Nguyen, and Svetha Venkatesh. Reasoning Under 1 Billion: Memory-Augmented Reinforcement Learning for Large Language Models. Transactions on Machine Learning Research (TMLR), 2025.

[9] Cui, Ganqu, Lifan Yuan, Zefan Wang, Hanbin Wang, Yuchen Zhang, Jiacheng Chen, Wendi Li et al. “Process reinforcement through implicit rewards.” arXiv preprint arXiv:2502.01456 (2025).

[10] Shi, Taiwei, Yiyang Wu, Linxin Song, Tianyi Zhou, and Jieyu Zhao. “Efficient reinforcement finetuning via adaptive curriculum learning.” arXiv preprint arXiv:2504.05520 (2025).

[11] Song, M., Zheng, M., Li, Z., Yang, W., & Luo, X. FastCuRL: Curriculum Reinforcement Learning with Stage-wise Context Scaling for Efficient Training R1-like Reasoning Models. Findings of the Association for Computational Linguistics: EMNLP 2025.

[12] Zhang, E., Yan, X., Lin, W., Zhang, T., & Lu, Q. Learning Like Humans: Advancing LLM Reasoning Capabilities via Adaptive Difficulty Curriculum Learning and Expert-Guided Self-Reformulation. EMNLP 2025.

[13] Do, Dai, Manh Nguyen, Svetha Venkatesh, and Hung Le. “SPaRFT: Self-Paced Reinforcement Fine-Tuning for Large Language Models.” arXiv preprint arXiv:2508.05015 (2025).

Table of Contents

• Why Test-Time Matters

• The First Breakthrough: Test-Time Compute

• Scaling Up Test-Time Efforts

• Process Reward Models: Defining What’s "Good Reasoning"

• Search and Planning with RL

• Conclusion & What’s Next

Why Test-Time Matters

A long-standing principle in machine learning is the clean separation between training and inference. Traditional pipelines front-load all the optimization into training, assuming that, once deployed, models will perform inference in a fixed, feed-forward manner. The model learns once and acts passively forever after.

In the context of LLMs, a dominant strategy to improve performance has been simple: scale-up training ⬆️. Larger models, more data, and longer context windows have powered much of the progress to date. And this strategy works until it doesn’t. As we push toward ever-larger scales, we begin to see diminishing returns:

More parameters don’t necessarily lead to better reasoning. And more data doesn’t guarantee generalization, especially to tasks that are unfamiliar or structurally complex. Why? Because we’re running up against a fundamental bottleneck: the data itself.

As Ilya Sutskever famously put it, “We have but one internet.” That single, finite corpus has already been scraped, filtered, and optimized for training. There is no untapped “second internet” to unlock new capabilities. This saturation of training-time gains is prompting a natural question: 🧠 What else can we optimize?

A new direction is taking shape: test-time compute. Rather than relying solely on what was learned during pretraining, this paradigm focuses on what can be actively computed during inference. The insight is simple but powerful: not all reasoning can (or should?) happen in a single forward pass. Some problems demand more thought. It pays to think longer, search deeper, and adapt on the fly while solving the task.

This idea isn’t new in the broader AI landscape. Fields like robotics, planning, and control have long leveraged test-time optimization to dynamically adjust actions, internal states, or policies based on the current situation. In these domains, the agent doesn’t freeze after training. It continues to optimize in real-time, guided by the specifics of the environment or task at hand. 🧠 So why hasn’t this philosophy fully carried over to LLMs?

Because LLMs are typically treated as static predictors. Once trained, they produce outputs via greedy decoding or sampling, which is efficient, but fixed. This works fine for fluent language generation. Without a mechanism to adapt at test time, LLMs tend to default to surface-level heuristics, even when the task demands deeper computation. This is where test-time reinforcement learning shines as a promising bridge between passive prediction and active, deliberative reasoning.

The First Breakthrough: Test-Time Compute

As pretraining-driven gains begin to plateau, the field has turned to post-training methods to push LLM capabilities further. Techniques like alignment through RLHF and instruction tuning have become standard tools to refine behavior after pretraining. But another powerful and increasingly relevant direction is emerging: test-time compute (TTC).

Instead of relying solely on what was baked in during pretraining, TTC explores what a model can compute on the fly, at the moment of inference. What if reasoning didn’t stop when training ended? What if, during deployment, models could think longer, adapt dynamically, and search more strategically to solve hard problems?

This shift in mindset is already underway. OpenAI’s o1 models mark a pivot toward deeper inference-time reasoning:

👀 While the exact techniques remain undisclosed, OpenAI has stated that o1 is “trained with reinforcement learning to perform complex reasoning” [2]. The role of RL is indisputable.

The result is amazing:

• On Codeforces, o1 ranks in the 89th percentile—an unprecedented level of programming ability for a language model.

• On AIME, a prestigious U.S. math competition, o1 ranks among the top 500 students.

• On GPQA, a benchmark of PhD-level science questions, o1 surpasses human experts in physics, chemistry, and biology.

The results reveal a compelling insight: allocating more computational resources at test time by allowing the model to perform multiple reasoning passes, search over alternatives, or optimize responses can yield improvements comparable to, or even surpassing, those from scaling training-time compute:

👀 While the compute costs may appear similar in magnitude, they reflect fundamentally different trade-offs. Training-time compute is a one-time, centralized investment that doesn’t directly affect the end-user experience. In contrast, test-time compute occurs during inference and consumes resources every time the model is used, placing a computational burden on the user or deployment system.

These achievements aren’t just about language. They reflect System 2 reasoning: deliberate, multi-step, and strategic. Unlike System 1, which governs fast, automatic responses like language completion or basic recall, System 2 engages when tasks demand controlled, sequential thinking, such as solving a math proof, writing code, or evaluating competing hypotheses. A summary of System 1 and System 2 properties is given below:

As we can see, System 2 is more powerful with deliberate reasoning. However, this deeper reasoning comes at a cost. High TTC is expensive, making such models suitable only for a narrow slice of high-stakes or compute-rich applications.

This raises two pressing research questions:

1. 🧠 How can we equip LLMs with the ability to scale their test-time reasoning?
   More forward passes? Search over reasoning chains? External tools? Many strategies exist, but Reinforcement Learning (RL) offers a particularly principled framework.

2. 🧠 More importantly, how can we make TTC efficient? Inference-time RL can guide the search process efficiently. Rather than brute-force reasoning, models can learn when and how to allocate extra thinking time for prioritizing hard problems, rethinking weak outputs, or dynamically adjusting their decoding strategy.

We can organize the approaches to scaling test-time compute by identifying which stages of the LLM inference pipeline offer opportunities for enhancement:

One way to increase test-time compute is through prompting, encouraging the model to “think longer” by crafting prompts that trigger deeper reasoning. Beyond prompting, more structured approaches involve fine-tuning the model’s internal weights or modifying its outputs dynamically. In the following sections, we’ll investigate the growing ecosystem of test-time reinforcement learning methods, with a focus on these last strategies where RL plays a central role in interfering with the output generation process.

Scaling Up Test-Time Efforts

Scaling test-time compute can be as simple as prompting the model to reflect more deeply, or as sophisticated as dynamically fine-tuning internal representations and outputs. The more sophisticated the approach, the more it engages System 2-style reasoning, which is deliberate, strategic, and multi-step:

One early attempt to enhance LLM reasoning through structured test-time compute is 👉Tree of Thoughts (ToT). Building on Chain-of-Thought prompting, ToT treats reasoning as an explicit search process through a space of intermediate “thoughts.”

Rather than committing to a single reasoning path, ToT asks LLMs to generate multiple candidate thoughts at each step, evaluates them using scoring or voting prompts, and searches through the resulting tree (e.g., via BFS or DFS) to find high-quality solutions.

The core steps are:

1. Define what a “thought” means for the task using prompting (e.g., via few-shot examples).

2. Generate candidate thoughts using LLMs and prompts.

3. Evaluate their quality using LLMs and prompts.

4. Search for the most promising reasoning path using tree-based search algorithms like BFS or DFS.

Example: The 24 Game

Consider the classic 24 Game: you're given a set of numbers, and the challenge is to apply arithmetic operations to reach the target number 24. The task requires reasoning and deciding which operations to apply, in what order, and how intermediate steps interact.

1. Thought corresponds to one or more steps of applying operations to the input numbers 4,9,10,13. For example, use + for the first two numbers: 4+9=13, resulting in 10,13,13 left.

2. Generate several “thoughts“ using the Propose Prompt (see Figure below, box (a)). This results in several thought candidates: e.g., 4+9=13 or 10-4=6

3. Evaluate the candidate by prompting another LLM with Value Prompt (see Figure below, box (b)). The LLM will decide whether it is possible to continue the thought

4. Do the tree search based on LLM evaluations with the node as the candidate thought. If LLMs say it is impossible to continue. In this case, “4+9=13” node is impossible to continue, so the search algorithm will not go further. Instead, it chooses “10-4=6” to continue the generation and search process.

From this example, we can construct a general framework for test-time scaling of LLMs as follows:

There are 3 main components of the framework:

• Verification: Evaluate the quality or correctness of generated outputs.

• Generation / Search: Produces reasoning candidates through sampling or exploration.

• Improvement with Feedback: Refines the model or outputs using signals from verification or external supervision.

As we shall see, the most recent methods aim to enhance one or more of these components. The last one will be discussed in my next blog. Now, let’s explore the first component in the next section.

Process Reward Models: Defining What’s "Good Reasoning"

For many reasoning-intensive tasks, such as math word problems, code generation, or symbolic logic, the final answer is verifiable. That is, given a question and a proposed solution, we can automatically check whether the answer is correct. This verifiability makes it possible to supervise models during training by rewarding correct outputs, even if we don't label every intermediate step.

However, this introduces a subtle challenge: not all paths to the correct answer are equally good. A model might stumble into the right solution by chance, through memorization, or via logically inconsistent steps. If we only reward the final answer, we risk reinforcing spurious or incoherent reasoning.

This realization has spurred interest in Process Reward Models (PRMs), which aim to evaluate and guide the reasoning trajectories of LLMs, rather than just their final outputs.

From Outcome-Based to Process-Based Evaluation

Traditional reinforcement learning approaches often focus on outcome-based rewards, evaluating the correctness of the final answer. However, this method overlooks the variety of the reasoning path taken to arrive at that answer. PRMs address this gap by assigning rewards to intermediate reasoning steps, promoting coherent and logical progression throughout the problem-solving process.

While verifying a final answer can often be automated (e.g., is the math answer correct?), supervising the reasoning process itself is a much harder problem. There’s no straightforward way to automatically label whether an intermediate step is valid or logically sound. An early paper named 👉“Let’s Verify Step by Step” tackles this head-on by manually labeling the reasoning process using human annotators [6]. The researchers asked labelers to evaluate the correctness of each step in the solution paths generated by LLMs.

This resulted in PRM800K, which contains 800K step-level labels, annotated across 75K model-generated solutions spanning 12K problems. Each step is labeled for correctness, enabling fine-grained reward modeling for reasoning.

The researchers then aim to train reward models in two regimes:

• Large-scale: They use the labeled dataset to fine-tune GPT-4, aiming to build high-quality outcome and process reward models (ORMs and PRMs). While this setup pushed performance to the state-of-the-art, differences in training data made direct comparisons between ORMs and PRMs tricky.

• Small-scale: To isolate the effect of supervision and enable controlled ablations, they trained smaller models from scratch. This allows using large models to create labeled data for training small models. This setup allowed them to probe how process supervision affects reasoning quality under more constrained conditions.

Here, the paper focuses on training reliable reward models rather than running a full RL loop. They evaluate the reward models by using the best-of-N approach:

1. Generate N candidate solutions for each problem.

2. Use the reward model for assigning scores to the solutions to rank them.

👀 For PRM, when determining a solution-level score, they either use the minimum or the product of step-level scores as a reduction.

3. Pick the top-ranked one.

4. Check if the final answer is correct.

The evaluation metric is the fraction of problems where the top-ranked solution is correct. The results reveal that PRM can detect the wrong steps in reasoning solutions. Compared to other methods, such as ORM or Majority Voting, PRM scales much better as N increases:

❌ One major bottleneck for training PRMs is the need for high-quality, step-level annotations, especially in domains like math, where reasoning is complex and annotators need domain expertise.

❌ Human labeling at this level of granularity is expensive and time-consuming, making it difficult to scale PRM training for practical use.

Therefore, recent papers propose an automatic process supervision framework to reduce reliance on manual labeling. For example, in the 👉MATH-SHEPHERD paper [7], the researchers define a step's quality by how likely it is to lead to the correct final answer. Intuitively, here's how it works:

• Start with a math problem and its ground-truth final answer.

• From an intermediate step in a candidate solution, decode multiple possible future reasoning paths using a fine-tuned LLM.

• Check how often these paths arrive at the correct final answer.

• The more reliably a step leads to correct answers, the higher its correctness score.

In particular, given a reasoning step s_i​, they generate a set of N completions (i.e., possible continuations of reasoning) using a fine-tuned LLM:

where a_j​ is the final answer reached in the j-th continuation. The set of all decoded answers is:

To score the quality y_si​​ of step s_i​, we consider two estimation strategies:

• Hard Estimation (HE): A step is labeled as "good" (reward =1) if any of its completions reaches the correct answer a∗:

• Soft Estimation (SE): Calculate the step-reward score as the frequency with which the completions arrive at the correct answer:

How to Train a Process Reward Model (PRM)

Formally, a PRM maps a problem P and a reasoning step S to a (positive) reward R:

The training objective treats step-level scoring as a binary classification problem. For a solution with K reasoning steps {s₁,s₂,…,s_K}, the PRM learns to assign scores r_si∈[0,1] indicating the quality of each step, i.e., >0.5 good, otherwise bad. The cross-entropy loss is applied over all steps:

where:

• y_si​​∈{0,1} is the target label indicating whether step s_i is good,

• r_si is the PRM's predicted score (typically after applying a sigmoid).

The training results show that Soft Estimation is better than Hard Estimation:

Once we’ve trained a PRM to score each reasoning step, we can take the minimum score across all steps in a solution as the ranking score in the best-of-N approach. The authors even go further by integrating self-consistency with reward models:

1. First, generate multiple candidate solutions with

2. Group these solutions by their predicted final answer.

3. Within each group, aggregate their PRM scores.

4. Finally, select the answer that is not just the most frequent, but the one supported by high reward scores.

Formally, the final answer chosen under this combined scoring strategy is:

More importantly, combining fine-tuning LLMs with the test-time strategy above (SHEPHERD) yields awe-inspiring results:

However, treating the PRM training task as a classification problem, i.e., labeling each step s_i as either "correct" or "incorrect", is not always perfect:

❌ It overlooks the structure of the trajectory, including the dependencies among steps, the cascading effect of earlier errors, and the varying importance of each step in solving the overall task.

❌ Each state is treated in isolation, leading to reward assignments that are often misaligned with the actual contribution of a step toward reaching a correct final answer.

To address this, a recent work introduces a new perspective: 👉Process Q-value Modeling (PQM) [8]. Instead of predicting correctness as a binary label, PQM frames the reasoning process as a Markov Decision Process (MDP) and aims to estimate the Q-value—that is, the expected probability that a particular action (a reasoning step) from a given state will lead to a correct final answer.

In this framework:

• The state s_h=(x,a_1:h−1) represents the prompt x and prior steps a_1:h−1.

• The action a_h is the generated reasoning at step h.

• The policy π(a_h∣s_h) represents the LLM.

• The Q-value Q(s_h, a_h) estimates how promising this step is in achieving the correct final answer.

Here, the Q-value implicitly defines a process reward: good steps raise the probability of success, and bad ones diminish it. This leads to a comparative training objective, where the PRM learns to rank better trajectories (or sub-trajectories) higher than weaker ones. Crucially, this also recovers classification-based PRMs as a special case when only terminal correctness is available and steps are independent.

The paper offers several theoretical results:

• The Q-value for a reasoning step can be defined as the inverse sigmoid of the probability that a trajectory leads to a correct final answer, under a given policy.

• Then, the Q-value represents the likelihood that a partial action sequence will result in a correct answer, making it a natural fit for evaluating intermediate reasoning steps.

• Under deterministic settings, the advantage function (a form of temporal differences of values) can be treated as a reward function due to potential-based reward shaping.

• Among correct steps, later ones have higher Q-values; among incorrect steps, later ones have lower Q-values.

• The Q-value of the first correct step is greater than that of the initial state, which is in turn greater than that of the first incorrect step.

• Final findings: Q-values of incorrect steps decrease over time, followed by the initial state's Q-value, then increasing Q-values of correct steps, establishing a clear separation between correct and incorrect reasoning trajectories.

A central goal in training PQM with rankings is not just to distinguish between correct and incorrect actions but to reflect the magnitude of the difference between them. While the classical Plackett-Luce (PL) model provides a way to capture ranking relationships, it falls short in our setting where Q-value gaps between correct and incorrect steps can be highly uneven and consequential. Therefore, the authors propose a Q-margin-based comparative loss that is sensitive to Q-value differences, capturing not only which step is better but also how much better.

Based on the theoretical findings above, the optimal Q-values induce a strict ordering:

where W and C denote the set of wrong and correct step indices, respectively. To enforce this structure during training, the authors define the loss:

Here, ζ is a margin hyperparameter that magnifies the penalty for incorrect steps by adding a positive bias to their exponentiated Q-values. Let’s explain the equation:

• The first term models the relative ranking among incorrect steps W. The PL-style ranking is used to penalize incorrect steps that receive higher Q-values than worse ones. Starting the sum at t=2 ensures we compare each step against all higher-ranked (i.e., better, smaller indices) incorrect steps. For example, if the model is estimating wrongly, like Q_w2>Q_w1, the loss is greater, and thus minimizing the loss helps.

• The second term compares correct steps C with all incorrect ones and “less correct“ steps. The margin ζ encourages a clear separation between correct and incorrect actions, not just in rank but in actual Q-value magnitude.

While the full loss L_theorem captures nuanced structure, it assumes accurate annotation of incorrect steps, which is rarely guaranteed in practice. Most datasets label all steps after the first mistake as "wrong," even if some are less severe. This motivates a simplified yet robust loss:

This loss ignores the internal ranking among wrong steps (the first term in L_theorem) and concentrates on learning a clear separation margin between the correct and incorrect actions.

The paper demonstrates better performance than other PRM training losses, such as using binary classification’s cross-entropy or MSE loss:

From a different perspective, researchers argue that intermediate, or process, rewards should act as fine-grained supervision, guiding the model at each step in a way that supports the long-term objective of arriving at the right outcome. This perspective differs from the conventional approach, which often assumes process rewards should strictly reflect the mathematical correctness or semantic relevance of each step. In contrast, large language models may need to take detours—generating seemingly trivial or redundant steps—to ultimately find a successful reasoning trajectory. Penalizing such steps too early or too harshly can prematurely collapse promising search paths. Therefore, the paper proposes focusing on rewarding the progress, which can be estimated by advantage functions rather than the Q-value [9]. They name the method 👉Process Advantage Verifiers (PAV).

To see why Q-value may not be an ideal estimation of progress, imagine we're running a beam search to explore different reasoning paths. Each path (or trace) is a sequence of steps, and we want to keep the most promising ones in the beam. A natural approach might be to retain those with the highest Q-values. But Q-values inherently entangle two things—the value of the action and the value of the state it came from. That creates a problem. If we compare actions from different states purely based on their absolute Q-values, we risk keeping steps that decrease our likelihood of success, simply because they come from "stronger" states:

Rather than looking at raw Q-values, we should focus on the change they induce—that is, the progress a step makes toward success. This is captured by the advantage function:

Advantage tells us whether a step improved or hurt our chances of solving the problem. A positive advantage means the step made meaningful progress; a negative one means it sets us back. By supervising the model with these advantage values, we reward real progress, not just proximity to good outcomes.

Furthermore, the advantage doesn’t need to be computed under the same policy that we’re training. As the paper argues, we may want to compute these progress signals under a stronger prover policy μ, different from our current base policy π. This prover acts like an expert verifier, assessing whether a step truly advances reasoning under its own more accurate value estimates.

The authors then propose components of the loss to train the policy:

• Outcome reward (of course) of the base policy:

• The advantage of the prover policy:

Using policy gradient methods, we can have the gradient to update the policy as:

Now, the remaining question is: 🧠 How Should We Choose the Prover Policy μ?

• Don’t set μ=π (base policy):

  • This reduces to standard outcome-based training.

  • The advantage A_π adds no extra supervision, especially harmful if π is weak, since both Q_π and A_π will be close to zero.

• Don’t use a very weak μ:

  • Just like a weak base policy, a weak prover produces flat or noisy advantage signals.

  • Beam search and training signals become unreliable.

• Don’t use an overly strong prover μ:

  • A strong prover can succeed regardless of trivial or irrelevant steps (e.g., restating the question).

  • As a result, Q_μ stays constant across such steps, leading to A_μ​≈0.

  • Training then rewards superficial behavior and fails to improve reasoning.

To prove the points, the author conducts a synthetic experiment (didactic setup):

• Goal: Train a policy π to generate a sequence that contains a hidden sub-sequence y∗ from a vocabulary V={1,2,…,15}.

• Reward: Sparse and terminal — the agent only receives a reward of 1 if y∗ appears in the output y; otherwise, 0.

• Prover policy μ: A procedural policy controlled by a scalar γ>0; as γ increases, μ becomes more effective, reaching perfect accuracy as γ→∞.

• Assumption: The setup uses Oracle access to true Q_π and A_μ​ to avoid approximation errors from learned models.

The results prove the point that A is better than Q, and middle γ is best.

The results on the MATH dataset confirm the points:

Search and Planning with RL

The second component of test-time scaling is the generation and search process. RL allows us to help LLMs actively explore the space of reasoning steps to discover new, high-reward trajectories. When paired with search, RL becomes a powerful tool, not just for improving final answers, but for shaping how the model thinks.

Imagine solving a math problem. You might try several intermediate calculations, discard a few, refine your strategy, and finally reach a solution. This iterative process of proposing, testing, and revising steps is precisely what we want LLMs to emulate. Pure next-token prediction is too myopic for this, while search gives the model the "lookahead" it needs.

TTC Scaling Using Classic Search Algorithms

Early approaches, such as ToT [4], involve using straightforward brute-force search algorithms, such as Breadth-First Search (BFS) and Depth-First Search (DFS). For example, a simple search framework will look like this

• Step Sampling: At each step, use the question and prior reasoning to sample B next steps via an LLM.

• Reasoning Tree:

  • Organize reasoning as a tree:

    • Root: question

    • Node: current generated text

    • Edge: reasoning step

    • Leaf: final answer

  • A search algorithm selects a node to expand.

• Expansion: Add the selected node's step to the prompt and sample new steps.

• Stopping: Halt when an answer is found or a step/computation limit is reached.

The key here is the search algorithm. 👉MindStar (M*) [10] proposes an LLM-inference-time search procedure using Levin Tree Search, which is a best-first tree search algorithm. This kind of search avoids trying to select all nodes to expand. Thus, a criterion should be employed to choose which node or reasoning step to select. Here, the authors assess each reasoning step with the aforementioned PRM.

Formally, given a current node n_d​ and a candidate's next step e_d​, the PRM outputs a reward:

A high r_d​ means the step is likely valid and consistent with prior reasoning. Once we score the child nodes with PRM, we must decide which node to expand next. Two strategies are examined:

• Beam search selects the (top) next steps with the highest rewards:

• Levin Tree Search (LevinTS): LevinTS improves on beam search by balancing cost and probability:

  • Cost: the cost of the path (the smaller the better)

  • Probability: reflects how likely a node leads to a solution (the higher the better)

The paper proposes the cost function as:

where:

• i_tok = number of tokens in n

• τ = temperature (controls search aggressiveness)

And the probability π(n) is recursively defined as:

where n has parent n′ connected by an edge e′.

LevinTS expands nodes in order of increasing f(n)/π(n), offering a guaranteed token budget:

where N_g is a set of target nodes.

We can see these search algorithms improve LLMs significantly:

Sometimes, reasoning involves complicated planning, and to test LLM reasoning, we can use planning benchmarks such as maze navigation. Inspired by this, 👉LLM-A* paper aims to evaluate and improve LLMs on path planning [11].

The gist of LLM-A*: a hybrid algorithm that combines the optimality of A* with the high-level reasoning capabilities of an LLM. Together with A*, the LLM acts as a planning guide, suggesting waypoints (intermediate targets) based on its understanding of the start, goal, and obstacles.

• The algorithm begins like A*: define the start and goal, initialize the OPEN = {s₀} and CLOSED ={∅} lists, and use a cost function f(s) = g(s) + h(s). Here, A* employs a heuristic function h(s) to estimate the cost from a node s to the goal, and a cost function g(s) to track the exact cost from the start to s.

• But before the search begins, the LLM is prompted with the environment information and asked to generate a target list T: a series of waypoints that form a plausible path to the goal.

• As A* proceeds, it modifies its f(s) score to bias toward the current LLM-suggested waypoint t:

👀 The cost(t,s) term encourages A* to prioritze solutions that contain LLMs’ waypoints.

• Once the current target t is reached, it updates to the next in the sequence, continuously steering the local search using global LLM knowledge.

• Crucially, the algorithm also verifies that all waypoints in T are outside obstacle zones and contain the correct start and goal nodes, correcting any LLM hallucination.

LM-A* uses several prompting techniques to extract better target paths from the language model:

• Few-Shot Learning: Demonstrates the LLM 5 by showing 5 demo examples of correct paths before prompting for a new one, thereby increasing reliability without requiring training.

• Chain-of-Thought (CoT): Encourages the LLM to reason step-by-step rather than jumping to the answer. This is particularly helpful in more complex environments where multi-step logic is essential.

• Recursive Path Evaluation (RePE): Breaks the planning into three LLM-driven stages:

  1. Analyze the environment,

  2. Generate a step,

  3. Evaluate that step.

  This recursion mimics approaches like ReAct and Self-Reflection but focuses entirely on internal reasoning, without environmental feedback.

Finally, the results demonstrate that LLM-A* outperforms LLM alone significantly in the benchmark.

❌ One big issue with relying on classic search algorithms is the choice and design of the cost function f(). Prior approaches often hinge on carefully engineered utility functions (e.g., number of tokens for text and distance for path planning), which introduce practical limitations.

An elegant alternative is to frame multi-step reasoning as an MDP. Here, each reasoning trajectory maps to a path through a decision space, where the state is the input prompt plus previous steps, the action is the next proposed token, and the Q-value can be used as the cost function. The paper 👉Q* leverages this structure by adopting a best-first decoding strategy, inspired by A* search [12].

👀 A bit different from the aforementioned papers like PQM that model action as a reasoning step consisting of multiple tokens, the action in Q* is single token.

Specifically, Q* casts reasoning as a heuristic search problem. Each partial reasoning path s_t​ is scored using an A*-style utility function:

• g(s_t​): Aggregated reward along the path so far.

• h(s_t): Heuristic estimate of the utility-to-go.

• λ: A scalar balancing the two terms.

To compute g(s_t), Q* uses an aggregation function of a reward function R_P​ that reflects task-specific preferences (e.g., correctness, coherence, confidence):

For the heuristic h(s_t), Q* uses the optimal Q-value:

Rather than exploring all possible continuations, Q* limits candidates to the top-K tokens proposed by the LLM.

The key challenge is estimating the optimal Q-value Q*(s_t,a_t) for a frozen LLM policy π_θ. To do this, Q* learns a proxy Q-value function Q̂ from offline trajectories:

where M is the number of actions in a trajectory and N is the number of trajectories. ŷ​(s_t​,a_t​) is a label approximating the true Q-value, computed using the following strategies:

1. Using Fitted Q-Iteration, labels are constructed recursively:

This process alternates between updating labels and training Q̂ for L iterations.

2. Learning from Rollouts: run random or MCTS rollouts from (s_t,a_t) and assign labels based on the best future trajectory:

3. Using a Stronger LLM: if a stronger LLM π_θ*​ is available (e.g., GPT-4), complete the reasoning with it to estimate the future reward:

where:

❌ While classic search-inspired frameworks like A* provide structured guidance for stepwise reasoning, they often rely on static heuristics and limited lookahead.

Beyond Search with Planning

To explore deeper decision spaces and balance exploration with exploitation, recent approaches turn to Monte Carlo Tree Search (MCTS)—a planning algorithm that enables LLMs to simulate multiple reasoning trajectories and refine their outputs through iterative rollouts. This shift empowers LLMs to deliberate more thoroughly, especially on complex tasks with long-horizon dependencies.

Although classical MCTS depends on access to an explicit environment model to simulate forward trajectories, this assumption is often unnecessary in language-based tasks. A key insight exploited by 👉Language Agent Tree Search (LATS) [13] is that for most language model (LM) environments—such as web navigation, multi-step reasoning, or tool use—we can reconstruct any prior state simply by restoring the full-text history. This removes one of the main barriers to applying MCTS in natural language settings and opens the door to powerful planning capabilities within LMs.

LATS adapts MCTS to language agents by interpreting each node in the tree as a structured tuple:

where x is the original task input, a_1:i is the action sequence taken so far (including reasoning traces), and o_1:i are the corresponding observations or feedback. The search tree is built incrementally using six operations:

1. Selection: Starting from the root, LATS uses the UCT criterion to traverse down the tree until it reaches a leaf node.

• V(s) is the estimated value of node s

• N(s) is the number of visits to node s

• N(p) is the number of visits to s's parent p

• w is the exploration weight

2. Expansion: At the selected node, the LM samples n next-step actions:

3. Evaluation: Each new node is scored with a novel value function:

• LM(s) is a scalar score predicted by the language model, prompted to end its reasoning trace with a self-assessed rating,

• SC(s) is a self-consistency score, aggregating agreement across multiple LM samples from the same node.

4. Simulation: LATS recursively samples and simulates actions from the current node until a terminal state is reached. The LM executes this by generating the remainder of the trajectory.

5. Backpropagation: Upon reaching a terminal state and obtaining a reward r, LATS updates all ancestor nodes along the path s₀​→s₁​→⋯→s_ℓ:

Here, the visitation count is updated as:

6. Reflection: When a trajectory fails, the LM is prompted to reflect on its reasoning and generate a natural language critique. This is stored and injected as additional in-context information in future rollouts, providing “semantic gradient” signals that improve decision-making without explicit gradient descent. The prompt for reflection is like this:

As the results speak, LATS improves on tasks that require reasoning and tool use (internet search):

As we’ve seen, many recent papers propose integrating search and planning with RL concepts to improve LLM reasoning, ranging from beam search guided by reward models to complex tree-based planning algorithms. However, this raises an important question: 🧠 Is more complexity always better?

A recent empirical study titled 👉Scaling LLM Test-Time Compute Optimally Can Be More Effective Than Scaling Model Parameters performed comprehensive evaluations that revealed many insights [14].

Core Idea: Test-Time Compute as a Resource:

• Instead of scaling parameters, we scale how we use compute at test time (e.g., number of samples, depth of revisions, search strategy).

• Given a prompt and a fixed compute budget N, different compute strategies yield different accuracies.

• Some questions benefit more from refinement (e.g., iterative revisions), while others require broader exploration (e.g., sampling diverse solutions in parallel).

Formally, we define Target(θ, N, q) as the output distribution, given a strategy θ, budget N, and prompt q.

• The goal is to find the optimal search strategy:

There are two main components in strategy:

1. Proposal Distribution: how candidate answers are generated (e.g., revision vs. independent sampling).

2. Verifier: how candidates are judged (e.g., majority vote, best-of-N, or scoring models like ORMs or PRMs).

The paper studies several options for strategies:

👀 The intuition is that, for easy problems, simple search may suffice, yet for hard problems, it is better to sample diverse reasoning paths up front.

This requires approximating question difficulty to guide how computing should be allocated. 🧠 How to estimate problem difficulty?

• Two ways:

  • Oracle difficulty: uses ground-truth labels (ideal but impractical).

  • Model-predicted difficulty: uses verifier scores to estimate difficulty without access to answers.

We often assume more sophisticated search equals better performance. But the empirical results say: not always. In particular,

• Beam search wins early: At low-generation budgets, beam search significantly outperforms best-of-N. But this advantage vanishes as the budget increases.

• Lookahead search underperforms: Despite being more powerful, it performs worse due to:

  • Costly rollouts.

  • Over-optimization of the verifier (PRM), leading to repetitive or trivial outputs.

• Diminishing returns from search are real: As shown in some failure cases, stronger search just amplifies bad verifier signals, generating:

  • Repetitive steps.

  • Overly short or spurious solutions.

Using difficulty-conditioned compute allocation—i.e., scaling search per question difficulty—improves efficiency:

• With oracle difficulty bins, the compute-optimal strategy nearly matches the best-of-64 using only 16 generations.

• Even with model-predicted bins, the gains largely hold.

Last but not least, the paper also explores improving the model’s proposal distribution by teaching it to revise its answers. This technique can be combined with parallel search approaches like best-of-N:

Unlike naive self-correction (which often fails), fine-tuning a model to iteratively refine previous answers yields clear improvements. So, they do supervised fine-tuning (SFT) on trajectories composed of a sequence of incorrect answers followed by a correct one. This trains the model to identify and correct mistakes in context rather than discard prior attempts and restart.

At inference time, the finetuned model generates sequential revisions, each conditioned on the latest few attempts (up to four). Empirically, performance improves with each revision step, showing the model effectively learns to self-correct over time.

Key insights: They find a tradeoff between sequential revisions and parallel sampling at test time. The optimal ratio depends on both the compute budget and the question’s difficulty:

• Easier questions perform best with purely sequential computation.

• Harder questions benefit from a balanced mix.

By adjusting this ratio per question, they can outperform standard best-of-N baselines using up to 4× less compute.

If you realize, all of these papers assume that the search process will eventually lead to a correct answer. In other words, they assume that LLMs already know the right answer, but just need the right trigger to reach it.

This is exactly the hypothesis behind 👉AlphaMath:

👀 Pre-trained LLMs already contain rich mathematical knowledge. The challenge is not learning from scratch, but activating this latent reasoning through better prompts and smarter search.

The paper follows prior works in modeling text generation as an MDP, where the LLM acts as a policy: At step t, state = partial solution s_t, action = next step a_t, and
the transition is deterministic:

Here, to find the optimal reasoning path, AlphaMath combines:

• A Monte Carlo Tree Search (MCTS) to balance exploration and exploitation, which is typical as in prior work:

• The novel bit is that a lightweight value model was added to the same LLM to judge intermediate reasoning quality

Not relying on the value estimated in MCTS, the paper proposes a value model V_ϕ(s) to estimate how likely a partial solution will lead to a correct final answer.

👀 A value model can generalize to new states, which will be useful to boostrap the value estimation of MCTS.

Training this model relies on Monte Carlo evaluation from multiple rollouts:

Then, we can optimize the value model via regression:

Given the value model, the MCTS phase uses both policy and value models to guide simulations:

Here, some steps introduce novelty. First, action selection uses a PUCT-style rule:

Each term here has a clear role:

• Q̂ : This is the MCTS’s average return (value) obtained when taking action a_t at state s_t​, based on past simulations. It encourages picking the best-known option so far (exploitation).

• π_θ​(a∣s_t​): the likelihood of taking action a according to the current policy. It helps bias the search toward actions the model thinks are promising — this is a key LLM-specific addition. It helps reduce the number of required simulations by starting from smart guesses (prior guidance).

• The counting term favors actions that haven’t been visited much.

  • If N(s_t,a)is low, the denominator is small, so the whole term becomes large → we explore new actions.

  • As N(s_t,a) increases (more visits), the term shrinks, reducing exploration pressure.

Second, the evaluation step is a mix of reward observation and value model estimation:

where i represents the i-th simulation. Given the current value estimation, we can do the backpropagation step to update the state-action value as:

After running N simulations with MCTS, we get a search tree with state-action values Q(s, a) stored at each node. Since the transition is deterministic, we assume for non-terminal nodes:

This implies that we can estimate the MCTS’s value as:

This “ground-truth“ value will be used to train the value model.

Finally, although MCTS works, it’s slow. During deployment inference, to speed up inference, AlphaMath introduces Step-level Beam Search (SBS):

• At each step, generate B₂​ next-step candidates with the LLM

• Evaluate each with V_ϕ, keep top B_1​

👀 No simulation, no tree—just fast reasoning with value guidance. MCTS is just the tool for training the value model. When B₁=1, SBS becomes a fast MCTS approximation for real-time use.

The results show that SBS is much faster while keeping competitive performance:

Conclusion & What’s Next

This blog covered test-time reinforcement learning methods that guide large language models without further gradient updates. By treating reasoning as a sequential decision process, we leveraged techniques like:

• Classic Tree Search

• Monte Carlo Tree Search

• Process Reward Modeling

But even with advanced test-time methods like MCTS and beam search, the model’s reasoning remains externally guided and self-supervised by final answers only. These methods help with step selection, but they don’t change the model’s internal reasoning ability. 🧠 So what if we want to go further? What if we aim to improve the LLM’s internal reasoning competence, not just steer it at test time?

To do that, we need more than search—we need to update the model’s weights. This is where reinforcement learning during post-training comes in: by assigning reward signals not only to the final answer, but also to the quality and structure of intermediate reasoning steps, we can directly fine-tune LLMs to reason better from the inside out.

🔜 Next up: We'll study how reinforcement learning frameworks can be adapted to fine-tune LLMs for better reasoning, using novel rewards and RL algorithms, and more. Stay tuned. Reasoning is just getting smarter. 🧩🚀

References

[1] Achiam, Josh, et al. "Gpt-4 technical report." arXiv preprint arXiv:2303.08774 (2023).

[2] Jaech, Aaron, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar et al. "Openai o1 system card." arXiv preprint arXiv:2412.16720 (2024).

[3] Ji, Yixin, Juntao Li, Hai Ye, Kaixin Wu, Jia Xu, Linjian Mo, and Min Zhang. "Test-time Computing: from System-1 Thinking to System-2 Thinking." arXiv preprint arXiv:2501.02497 (2025).

[4] Yao, Shunyu, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. "Tree of thoughts: Deliberate problem solving with large language models." Advances in Neural Information Processing Systems 36 (2024).

[5] Kambhampati, Subbarao, Kaya Stechly, and Karthik Valmeekam. "(How) Do reasoning models reason?" Annals of the New York Academy of Sciences (2025).

[6] Lightman, Hunter, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. "Let's verify step by step." In The Twelfth International Conference on Learning Representations. 2023.

[7] Wang, Peiyi, Lei Li, Zhihong Shao, Runxin Xu, Damai Dai, Yifei Li, Deli Chen, Yu Wu, and Zhifang Sui. "Math-Shepherd: Verify and Reinforce LLMs Step-by-step without Human Annotations." In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 9426-9439. 2024.

[8] Li, Wendi, and Yixuan Li. "Process reward model with q-value rankings." In The Thirteenth International Conference on Learning Representations, 2025.

[9] Setlur, Amrith, Chirag Nagpal, Adam Fisch, Xinyang Geng, Jacob Eisenstein, Rishabh Agarwal, Alekh Agarwal, Jonathan Berant, and Aviral Kumar. "Rewarding Progress: Scaling Automated Process Verifiers for LLM Reasoning." In The Thirteenth International Conference on Learning Representations. 2025.

[10] Kang, Jikun, Xin Zhe Li, Xi Chen, Amirreza Kazemi, Qianyi Sun, Boxing Chen, Dong Li et al. "Mindstar: Enhancing math reasoning in pre-trained llms at inference time." arXiv preprint arXiv:2405.16265 (2024).

[11] Meng, Silin, Yiwei Wang, Cheng-Fu Yang, Nanyun Peng, and Kai-Wei Chang. "LLM-A*: Large Language Model Enhanced Incremental Heuristic Search on Path Planning." In Findings of the Association for Computational Linguistics: EMNLP 2024, pp. 1087-1102. 2024.

[12] Wang, Chaojie, Yanchen Deng, Zhiyi Lyu, Liang Zeng, Jujie He, Shuicheng Yan, and Bo An. "Q*: Improving multi-step reasoning for llms with deliberative planning." arXiv preprint arXiv:2406.14283 (2024).

[13] Zhou, Andy, Kai Yan, Michal Shlapentokh-Rothman, Haohan Wang, and Yu-Xiong Wang. "Language agent tree search unifies reasoning, acting, and planning in language models." In Proceedings of the 41st International Conference on Machine Learning, pp. 62138-62160. 2024.

[14] Snell, Charlie, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. "Scaling llm test-time compute optimally can be more effective than scaling model parameters." arXiv preprint arXiv:2408.03314 (2024).

[15] Chen, Guoxin, Minpeng Liao, Chengxi Li, and Kai Fan. "AlphaMath Almost Zero: Process Supervision without Process." In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.

Now, a new wave of models is entering the scene: Large Language Models (LLMs). These models were originally built for language tasks—like chatbots, summarizing text, and answering questions. But recently, researchers have started using LLMs for time-series forecasting, too.

In this post, we’ll explore:

✔ How LLMs are being adapted to handle time-series data
✔ Some recent research and early results
✔ Key challenges and open questions

LLMs won’t replace every forecasting model, but they’re opening up new ideas about how we can approach time-series problems. Let’s take a look at where this is all heading.

Table of Contents

• Why Do We Need Large Language Models?

• Direct Application of Pre-trained LLMs

• Designing and Fine-tuning LLMs for Time Series Forecasting

• Building Foundational Models for Time Series: A New Era?

• Conclusion

Why Do We Need Large Language Models?

Time-series forecasting has been a long-standing challenge in data science, underpinning critical applications like financial market prediction, weather forecasting, and supply chain management. For the last 30 years, statistical models such as ARIMA and ETS have been extensively used for time-series tasks, followed by deep learning methods like LSTMs and Transformers. So, why are researchers now turning to Large Language Models (LLMs) for time-series forecasting?

Pattern Recognition Capability

LLMs excel at recognizing complex patterns across diverse data sources, including both structured numerical sequences and unstructured text. Their self-attention mechanism allows them to capture long-range dependencies and underlying trends in time-series numbers. For example, big LLMs like ChatGPT can find the number patterns easily:

Understanding Context Beyond Numbers

Additionally, LLMs can leverage their pre-trained knowledge to identify contextual signals—such as economic shifts, seasonal effects, or market sentiment—helping to enhance time-series forecasting beyond traditional statistical and deep learning approaches.

Long-term Forecasting

While traditional time-series models learn patterns from past data, they often struggle with long-term dependencies. LLMs, especially those equipped with long-term memory architectures, can store and recall historical insights more effectively, improving their ability to recognize seasonality, anomalies, and evolving trends over time.

👀 That said, since LLMs are primarily designed for text data, applying them to time-series data is not straightforward. In the following sections, we will explore recent works that address this challenge.

Direct Application of Pre-trained LLMs

A key approach is converting numerical time-series data into text, representing each value as a string, i.e., prompt-based forecasting. Because the sequence order is maintained in string form, LLMs can still capture temporal dependencies. Here, the choice of tokenization is crucial, as it determines how the model processes and learns patterns from the data.

Several notable examples illustrate the direct prompting approach. 👉PromptCast [1] represents a pioneering effort in this direction, converting numerical time series into natural language prompts using predefined templates. This framework frames the forecasting task as a sentence-to-sentence generation problem, where the LLM is prompted with historical data described in natural language and asked to generate a sentence representing the forecast.

The paper introduces simple templates to translate time series to prompts. For example, for the ECL dataset:

Despite being simple, the approach shows reasonable results compared to other deep learning methods, such as Autoformer and FEDFormer:

👀 The paper only uses small and old LLMs such as BigBird and Pegasus. The results can be even better with stronger and modern LLMs such as Chat-GPT.

Unlike PromptCast, which focuses on prompt engineering, the 👉LLMTime [2] paper demonstrates that LLMs can be used directly as forecasters without the need for additional text or prompt engineering, provided that the numerical values are carefully preprocessed.

For example, the choice of tokenization (with space or not) matters:

In particular, the paper proposes:

• Digits are spaced out for separate tokenization.

• Commas mark time steps.

• Decimal points are removed to save context (e.g., by multiplying the original value with 100)

• Example: 0.123, 1.23, 12.3, 123.0 → "1 2 , 1 2 3 , 1 2 3 0 , 1 2 3 0 0"

Other tricks include:

• Rescaling: The method rescales time series values so that the α-percentile of the rescaled data is 1, preventing extreme values from dominating while ensuring the LLM sees some examples where digit lengths change. For example, a time series has values [10, 50, 200, 1000], and we set α = 0.75. If the 75th percentile value is 200, we scale all values by 1/200, so the rescaled series becomes [0.05, 0.25, 1, 5],

• Sampling: To forecast, generate multiple samples (e.g., 20) from the LLM and use their statistics (e.g., median or quantiles) to create a point or probabilistic estimate.

Likelihood Estimation: To get good sampling, the paper refines the way to compute the probability of the LLM’s output. The approach allows LLMs to act like hierarchical softmax distributions over numerical values. A number u of n digits has probability:

To estimate the probability for a “continous“ value x by assigning a uniform probability within the bin that x falls into. For example, if x=0.5371 falls into bin k (0.537-0.538):

where B is the number base, e.g, B=10, each bin has the probability B⁻ⁿ.

The experiment results demonstrate that LLMTime is competitive against standard and deep forecasting methods:

LLMs are showing promise in time-series forecasting, but basic prompting limits them. A more recent work, 👉LSTPrompt, introduces a sophisticated Chain-oj-Thought prompt to break down forecasting, improving accuracy.

Some prompting techniques:

• Add task description

• Explain the long-term/short-term properties of the data

• Include “TimeBreath” trick prompt

Despite promising results, directly using LLMs “as-is“ for time-series forecasting has key limitations:

❌ No temporal inductive bias: LLMs lack built-in mechanisms for capturing time dependencies.

❌ Tokenization artifacts: Discretization may reduce numerical precision.

❌ High computational cost: More expensive than traditional models.

❌ Limited extrapolation: Struggles with long-term forecasts.

❌ Lack of domain constraints: Misses explicit trend/seasonality modeling.

In the next section, we will see more complicated solutions that address these concerns.

❌ High cost: Good results often require strong LLMs as GPT-4. However, this is expensive and not open-source.

Designing and Fine-tuning LLMs for Time Series Forecasting

To overcome these limitations, research has focused on repurposing and fine-tuning LLMs for time-series forecasting. This includes adapting their architecture and further training pre-trained models on time-series data.

One core technique is aligning the modalities of time series and language, making LLMs effective for time-series data. The obvious way is to choose some early layers of the LLMs to finetune on time-series forecasting tasks to align the encoding part of the LLMs to the time-series domain. For example, GPT4TS [4] only finetunes positional embeddings and layer normalization parts of the GPT model:

Similarly, 👉LLM4TS [7] develops an LLM framework for time-series forecasting. Compared to GPT4S, more engineering effort is spent on timestamp and time-series embedding to represent the tokens for the LLM engine. For example, the paper proposes to embed each time scale information separately with a final pooling to get the temporal representation e_temp:

Then, the final embedding is computed as:

where e_token is the time-series patch representation and e_pos is the positional encoding.

The paper also introduces a two-step training approach:

1. Alignment training, where the model learns to align representations through a next-token prediction task, similar to the pretraining phase in LLMs;

2. Forecasting fine-tuning, where the model is further trained on downstream time series forecasting tasks to specialize in prediction.

👀 Although these methods aim to align LLM text representations with time series data, they largely stop at simply retraining the encoding layers without deeper adaptation.

Digging deeper, 👉Time-LLM [5] exemplifies a reprogramming framework where the input time series is transformed into text prototype representations and natural language prompts are used to guide the LLM's reasoning process. This approach keeps the underlying LLM intact and trains a separate reprogramming layer to translate the observed time series into a language-based representation, as illustrated in Fig. c below:

👀Essentially, reprogramming framework finetunes the model's input and output layers while keeping its core frozen for efficiency.

Here, the paper proposes patch reprogramming to align time series with the LLM’s embedding space by transforming a subsequence of time series (patch) to a text-aligned token representation. As such, they start with pretrained embedding of text tokens:

where V is vocabulary size and D is embedding dimension. Direct mapping is inefficient, they propose to project the original embedding space to a smaller prototype space E’:

These prototypes encode phrases like "short up" or "steady down", keeping representations within the LLM’s text space, yet more specific. Then, they integrate prototypes to time-series representation via cross-attention. For an attention head k, the query Q is constructed from the time-series patch X_p while the key and value are generated from the prototype embedding:

The attention mechanism computes the alignment between a patch (i) and prototypes:

Aggregating across heads produces the reprogrammed representation, which is then linearly projected to match the hidden dimensions of the backbone model:

Given the reprogrammed patches, the paper proposes 2 ways to feed them to the LLMs:

• Patch-as-Prefix: Treat patch as input text followed by a simple prompt to trigger the LLM to predict time series values in natural language. However, this method faces challenges:

  1. LLMs struggle with high-precision numerals without external tools, making long-horizon forecasting difficult.

  2. Different pre-training corpora and tokenization strategies lead to inconsistencies in representing numeric values (e.g., ['0', '.', '6', '1'] vs. ['0', '.', '61']).

• Prompt-as-Prefix circumvents these limitations by structuring prompts with three key components:

  1. Use instruction prompt as input context, task instruction, and trends and lags statistics.

  2. The instruction precedes the patch, and we extract only the output segment corresponding to the patch for regression, similar to standard forecasting. A projection layer is needed to align the output.

With these architecture changes, the performance on long-term forecasting tasks is impressive:

Given these results, it seems that LLMs will have huge potential for boosting forecasting accuracy. However, a recent study [6] reveals that it is not that easy. The research question is: 👉Are Language Models Actually Useful for Time Series Forecasting?

The finding shows that:

• Despite the hype, large language models (LLMs) don’t help with time series forecasting—in fact, simpler models without LLMs often perform better.

• LLMs add extra cost without improving accuracy, struggling to understand time-based patterns, or helping in low-data situations.

To understand why, let’s look at the experimental setting:

Here, the paper considers the standard architecture of utilizing LLMs for forecasting tasks, which may involve a pre-trained LLM and other components:

• Input normalization

• Finetune encoding and projection layers

• Alignment training

To confirm the contribution of the LLM component, the paper removes or replaces the LLM module with other things. It turns out that without the LLM module, the performance tends to be better:

👀 The results suggest that improvements in previous papers is mainly attributed to normalization and encoding layer finetuning.

More results confirm that the pre-training knowledge stored in LLMs may not be useful for forecasting tasks. As shown below, even without using pre-training weights, the model can still perform best.

👀 Althouth the result is interesting, it is only tested on GPT-2, a weak LLM. It does not guarantee for other LLMs.

If LLMs do not help, 🧠where does the performance come from? The paper suggests that simple techniques used in current LLM forecasters are enough to build strong forecasters:

• Patching (channel independent)

• One-layer attention

• Linear projections

If this finding sounds grim for the future of LLMs in time series, there’s good news: a recent study shows that when used the right way, LLMs can significantly boost forecasting performance. It’s not that LLMs are useless—it’s that how we use them makes all the difference.

Concretely, most existing LLM-based time series models miss a key ingredient: autoregression—the core of how both LLMs and traditional forecasters make predictions. A new approach, 👉AutoTimes [8], brings this back by using LLMs in their natural autoregressive mode, enabling flexible, multi-step forecasts without needing to retrain for different input/output lengths.

AutoTimes leverages existing components and ideas:

• Patching to encode a segment of a time series as a token representation

• Timestamp representation modeling to improve token representation

• Only 0.1% of parameters are added by embedding time series as LLM tokens—keeping the LLM frozen and training efficient.

The main contribution here is that it trains the LLMs with autoregression style with the next token prediction task. The predicted tokens are then projected back to time-series space for forecasting:

Under this approach, the results look promising:

More importantly, the paper shows that the LLM module is really useful, as indicated in the ablations study:

We can see that by designing the right method for LLMs, their full potential in time series forecasting can finally be unlocked. With approaches like AutoTimes, LLMs go beyond being just expensive add-ons—they become efficient and generalizable forecasters.

Building Foundational Models for Time Series: A New Era?

Instead of just adapting existing LLMs, researchers are now exploring specialized foundational models trained directly on vast time series data.

• Why Specialized Models? Training LLMs from scratch on time series could help them understand temporal dependencies just as language models learn from text.

• Challenges & Breakthroughs: Time series data is highly variable and non-stationary, making dataset creation and training complex, but models like TimeGPT and Chronos show it's possible. Let’s discover how these specialized LLMs are reshaping time series forecasting.

👉TimeGPT [9], the first time-series foundation model, employs Transformer architecture to learn time-series representations. It processes a sliding window of historical values, enriched with local positional encodings, and passes this through a deep encoder-decoder stack with residual connections and layer normalization. At the end, a linear head maps the decoder’s output to a forecasting window—predicting what comes next. The training objective is:

The paper applies Transformer-based architecture on time-series datasets:

👀 TimeGPT is not based on an existing language model. While it shares the “foundation model” spirit—training on vast data at scale—its architecture is customized for numerical sequences, not text.

To train TimeGPT, researchers compiled what’s likely the largest public collection of time series ever used, with over 100 billion data points. This dataset spans a rich mix of domains: finance, economics, healthcare, demographics, weather, web traffic, transport, IoT sensors, and more. These series bring with them a wide spectrum of behaviors—multiple seasonalities, varying cycle lengths, non-linear trends, noise, and sudden anomalies.

👀 Instead of overly sanitizing the data, the creators kept most of it raw, simply standardizing formats and filling in missing values. This decision ensured the model would learn from real-world messiness—a critical factor in achieving robustness and generalization.

Training TimeGPT required multiple days on NVIDIA A10G GPU clusters. Extensive tuning was conducted to optimize batch sizes, learning rates, and other parameters. Consistent with earlier large-scale training studies, larger batch sizes, and smaller learning rates helped stabilize training.

One standout feature of TimeGPT is its support for probabilistic forecasting—not just predicting what’s most likely to happen, but quantifying uncertainty around those predictions. This is achieved using conformal prediction, a model-agnostic, non-parametric method that doesn’t assume a specific distribution. During inference, TimeGPT performs rolling forecasts on the latest data to estimate and calibrate the model’s prediction intervals, making it a more trustworthy tool for decision-making under uncertainty.

Sharing the same spirit with TimeGPT, 👉MOIRAI [10] prepares a large-scale time-series dataset and trains a Transformer-based model with the data. However, MOIRAI provides more sophisticated encoding and decoding pipelines. First, it proposes a Patch-based Masked Forecasting framework where it divides time series into non-overlapping patches and learns to forecast the masked future patches using a Transformer encoder.

• Each patch captures temporal chunks of data, and multiple patch sizes are used to adapt to both high- and low-frequency signals.

• During training, future patches are replaced with a special [mask] token — a learnable embedding — signaling the model to predict them.

Most models assume fixed-length multivariate inputs. MOIRAI aims to work with any number of variates, even if unseen during training. To this end, it flattens multivariate time series into a single long sequence and applies attention to it. To preserve variate identity and ensure robustness to variate permutation:

• It uses variates encodings (not positional encodings).

• Applies binary attention biases that help the Transformer:

  • Distinguish between intra- and inter-variate interactions.

  • Remain equivariant to variate order and invariant to variate IDs.

In particular, define the rotary positional encoding matrix as R and let u(1), and u(2) be learnable scalars, we have the attention score:

Here, the binary attention bias is controlled by u as the two scalars bias the attention scores based on whether the variate (m,n) indices match:

• u(1) enhances attention between elements of the same variate.

• u(2) modulates cross-variate attention.

They are learned to best fit the training data. The final attention weights are:

MOIRAI adopts several modern LLM techniques to stabilize and improve the Transformer encoder:

• RMSNorm replaces LayerNorm

• SwiGLU activation used in FFNs

• Query-key normalization for attention

• Pre-normalization layout for better gradient flow

• Biases are removed for simplicity and regularization

Finally, instead of predicting a single value, MOIRAI outputs parameters of a mixture of parametric distributions — allowing flexible and accurate uncertainty modeling. In particular, the output reads:

Here, Y_t:t+h is the time-series patch. Distribution choices:

• Student’s t – for general robustness

• Negative Binomial – for count-based data

• Log-normal – for right-skewed distributions (e.g., prices)

• Low-variance Normal – for high-confidence predictions

This setup supports both sampling and likelihood-based training with minimal overhead. The loss function represents the negative log-likelihood:

The whole pipeline is given below:

In terms of performance, MORAI can outperform some deep learning forecasters. This is remarkable because the foundation model does not train on each dataset as the other baselines.

Approaching differently, 👉Chronos [11] repurposes standard language models—like T5 or GPT-2—to handle time series data with minimal changes, by converting continuous time series values into discrete tokens. These tokenized sequences are then treated as a “language of time series” that LMs can ingest and model.

To make time series digestible for LLMs:

• Scaling: Each time series is normalized, particularly via mean scaling (preserving zero values, e.g., "zero sales" days).

• Quantization: The scaled values are mapped to discrete bins (e.g., 1024 total). They use uniform binning (equal bin widths) to generalize better across unseen data distributions.

👀 This converts a real-valued time series into a discrete sequence:
[x1, ..., xC] → [token1, ..., tokenC]. This is similar to fine-tuning LLM approaches mentioned above.

Then, the LLMs are employed to process the data

• Off-the-shelf LMs (like T5 or GPT-2) are used without architecture changes.

• Only modification: Adjusting vocabulary size to match the quantized token space.

• No time or frequency features (e.g., no day-of-week encoding), which surprisingly doesn’t degrade performance.

The training Objective is the standard cross-entropy loss over the token vocabulary:

• Predicts next token z_C+h+1 given previous tokens z_1:C+h

• Trained like standard language modeling

During inference, Chronos samples token autoregressively:

• These tokens are dequantized and unscaled back to real numbers.

• Multiple samples are drawn to get probabilistic forecasts, forming a distribution over possible futures.

Chronos tackles the data scarcity problem with two strategies:

1. TSMixup – Combines time series from different datasets via convex interpolation.

2. KernelSynth – Uses Gaussian processes to generate synthetic time series with random kernel compositions.

These data augmentation techniques improve generalization and robustness, especially for zero-shot settings.

Chronos was trained and evaluated across 55 datasets and:

• Outperformed traditional time series models and specialized deep nets.

• Achieved state-of-the-art zero-shot performance—i.e., good generalization without fine-tuning, i.e., better than MORAI.

• Performed competitively even with modest model sizes, making it computationally efficient.

The benchmark mainly involves 2 metrics:

• WQL evaluates the accuracy of probabilistic forecasts (i.e. when your model predicts a distribution or multiple quantiles, not just a single value).

• MASE is a scale-independent metric used to evaluate point forecasts. It compares your forecast to a naive baseline (like the previous value or seasonal naive).

While models like Chronos and Moirai have made notable progress towards universal forecasting, they face limitations:

❌ They rely on moderate and fixed context lengths or handcrafted heuristics that limit their flexibility and scalability.

❌ They follow dense training, i.e., being computationally expensive and requiring large GPUs to accommodate.

To address these issues, researchers recently introduced 👉TIME-MoE [12]—a scalable, sparse, and general-purpose time series foundation model built to mirror the success of LLMs and vision transformers in their respective domains. In particular, the model features:

• Decoder-only Transformer with a Mixture-of-Experts (MoE) backbone.

• Operates in an auto-regressive manner, enabling support for any forecasting horizon.

• Capable of handling long contexts (up to 4096 tokens), which is critical for long-term temporal dependencies.

• Sparse activation: Only a tiny subset of expert networks are activated per token, making the model highly efficient.

• Allows the model to scale up (up to 2.4 billion parameters) without linear increases in compute cost.

Below is the overall architecture and details of each processing step.

①② The first step is to embed time series data into tokens. Unlike the mainstream that uses patch-based tokenization, here they use point-wise tokenization, i.e., each time-series step is projected to high-dimensional embedding:

where:

• W and V are learnable projection weights

• ⊗ denotes element-wise product

③ TIME-MoE builds on the decoder-only transformer, with common tricks for time series:

• RMSNorm for stable training

• Rotary Positional Embedding (RoPE) to better generalize to long-range sequences

• Bias-free Layers (except in self-attention) for better extrapolation

The twist: each transformer block replaces the standard feedforward layer with a Mixture-of-Experts module. Here's what happens at a layer l step-by-step:

i. Self-Attention + Residual:

ii. Normalization:

iii. Mixture-of-Experts Activation:

In this step, the Mixture layer routes each token to a subset of K out of N possible experts, plus a shared expert to capture global knowledge:

Where the gating weights g_i,t are determined via softmax over router scores s_i,t​. In particular, the gating weights for each specialized expert are:

with the router’s score is computed as:

On the other hand, the gate for the global expert is:

④⑤ Another contribution of the approach is that instead of predicting at a single horizon, TIME-MoE predicts at multiple resolutions simultaneously. Each projection head targets a different forecast length, enabling the model to:

• Generalize across short and long horizons

• Learn a richer latent structure of the future

• Ensemble multi-scale predictions for better robustness

The total loss averages over all resolutions, using the Huber loss for each prediction:

Last but not least, MoE models suffer from a known issue: routing collapse, where only a few experts get chosen repeatedly. To fix this, TIME-MoE introduces an auxiliary expert balance loss:

where f_i​ is the fraction of tokens routed to expert i:

r_i​ is the average router probability for expert i:

👀 Minimizing the balance loss encourages uniform expert assignment. This is because the frequency f and the routing vector r are positively correlated (i.e., approximately proportional), allowing us to apply Chebyshev’s sum inequality.

The result is impressive. Time-MOE is now one of the SOTA in time-series forecasting when compared with other foundation models:

as well as other specialized methods:

Conclusion

LLMs aren’t the final answer to time-series forecasting, but they’ve changed the conversation. What started as a tool for language tasks is now being reimagined for sequences of all kinds. Researchers are finding creative ways to adapt pre-trained LLMs to time-series problems, fine-tune them for specific domains, or even build new foundational models trained purely on temporal data. Some of these early results are promising.

The takeaway? LLMs are giving us a new lens on forecasting problems. They open up ideas that weren’t obvious before—ideas about scale, transferability, and reasoning over time. That doesn’t mean we throw out everything we’ve learned from classical and deep learning models. It just means we now have another tool to work with—and a new frontier to explore.

Reference

[1] Xue, Hao, and Flora D. Salim. "Promptcast: A new prompt-based learning paradigm for time series forecasting." IEEE Transactions on Knowledge and Data Engineering 36, no. 11 (2023): 6851-6864.

[2] Gruver, Nate, Marc Finzi, Shikai Qiu, and Andrew G. Wilson. "Large language models are zero-shot time series forecasters." Advances in Neural Information Processing Systems 36 (2023): 19622-19635.

[3] Liu, Haoxin, Zhiyuan Zhao, Jindong Wang, Harshavardhan Kamarthi, and B. Aditya Prakash. "LSTPrompt: Large Language Models as Zero-Shot Time Series Forecasters by Long-Short-Term Prompting." In Findings of the Association for Computational Linguistics ACL 2024, pp. 7832-7840. 2024.

[4] Zhou, Tian, Peisong Niu, Liang Sun, and Rong Jin. "One fits all: Power general time series analysis by pretrained lm." Advances in neural information processing systems 36 (2023): 43322-43355.

[5] Jin, Ming, Shiyu Wang, Lintao Ma, Zhixuan Chu, James Y. Zhang, Xiaoming Shi, Pin-Yu Chen et al. "Time-LLM: Time Series Forecasting by Reprogramming Large Language Models." In The Twelfth International Conference on Learning Representations, 2019.

[6] Tan, Mingtian, Mike Merrill, Vinayak Gupta, Tim Althoff, and Tom Hartvigsen. "Are language models actually useful for time series forecasting?" Advances in Neural Information Processing Systems 37 (2024): 60162-60191.

[7] Chang, Ching, Wei-Yao Wang, Wen-Chih Peng, and Tien-Fu Chen. "Llm4ts: Aligning pre-trained llms as data-efficient time-series forecasters." arXiv preprint arXiv:2308.08469 (2023).

[8] Liu, Yong, Guo Qin, Xiangdong Huang, Jianmin Wang, and Mingsheng Long. "Autotimes: Autoregressive time series forecasters via large language models." Advances in Neural Information Processing Systems 37 (2024): 122154-122184.

[9] Garza, Azul, Cristian Challu, and Max Mergenthaler-Canseco. "TimeGPT-1." arXiv preprint arXiv:2310.03589 (2023).

[10] Woo, Gerald, Chenghao Liu, Akshat Kumar, Caiming Xiong, Silvio Savarese, and Doyen Sahoo. "Unified Training of Universal Time Series Forecasting Transformers." In International Conference on Machine Learning, pp. 53140-53164. PMLR, 2024.

[11] Ansari, Abdul Fatir, Lorenzo Stella, Ali Caner Turkmen, Xiyuan Zhang, Pedro Mercado, Huibin Shen, Oleksandr Shchur et al. "Chronos: Learning the Language of Time Series." Transactions on Machine Learning Research, 2024.

[12] Shi, X., Wang, S., Nie, Y., Li, D., Ye, Z., Wen, Q., & Jin, M. (2025). Time-MoE: Billion-Scale Time Series Foundation Models with Mixture of Experts. In Proceedings of the Thirteenth International Conference on Learning Representations, 2025.

• Finance: Accurate market predictions drive billion-dollar trades and risk management.

• Healthcare: Monitoring patient vitals helps detect early warning signs, enabling life-saving interventions.

• Energy: Power grid demand forecasting prevents blackouts and optimizes renewable energy integration.

• Climate Science: Weather and climate modeling help mitigate the impact of extreme events.

AI has already revolutionized many fields, outperforming humans in complex games, generating realistic images, and crafting coherent text. However, when it comes to time-series forecasting, AI still struggles to keep up.

In this blog, we’ll explore the evolution of time-series forecasting:
✔ From classic statistical methods like ARIMA and exponential smoothing,
✔ To deep learning breakthroughs with Transformers.
✔ To the non-attention methods that challenge the dominance of Transformers

We’ll dive into the latest research, uncover the biggest challenges, and explore future AI forecasting options. Because in a world that never stops moving, seeing what’s coming next is more valuable than ever.

Table of Contents

• Why Is Time-series Forecasting So Hard?

  • Time Series Are Unlike Others

  • The Last Stronghold of AI?

• Classical Time-Series Models

  • ARIMA (AutoRegressive Integrated Moving Average)

  • Exponential Smoothing (ETS)

  • Seasonal Decomposition and Trend Analysis

• Deep Learning for Time-Series Forecasting

  • Modern Formulation of Time-Series Forecasting

  • Transformer-Based Time-Series Models

  • When Attention is Not All You Need for Time Series

• Conclusion

Why Is Time-series Forecasting So Hard?

It's a question that echoes through research labs and boardrooms alike, a testament to the persistent challenge of predicting the future from the flow of time. The difficulty doesn't only stem from a lack of data, but also from the very nature of temporal sequences themselves.

• Dynamic and evolving patterns: Unlike static data, time-series sequences constantly shift over time.

• External disruptions: Economic shocks, pandemics, and supply chain issues create sudden, unpredictable changes.

• Long-range dependencies: Small fluctuations in early data points can have cascading effects later.

• Noisy and sparse data: Real-world time-series datasets are often incomplete or heavily affected by anomalies.

Time Series Are Unlike Others

In addition to these issues, time series data are unlike others due to their special properties. When analyzing time series data, it's crucial to understand the properties of trends and seasonality.

Trend:

• A trend represents the long-term movement or direction of a time series. It indicates whether the data generally increases, decreases, or remains constant over an extended period.

• Trends can be linear (straight line) or non-linear (curved).

Seasonality:

• Seasonality refers to recurring, predictable patterns that occur regularly within a time series. These patterns are typically influenced by seasonal factors like time of year, day of the week, or time of day.

• Seasonal patterns repeat with a fixed and known period.

Accurately modeling these properties remains a significant challenge and an active area of research.

The Last Stronghold of AI?

Amidst the rapid advancements in artificial intelligence, a crucial domain continues to resist its transformative touch. Despite the impressive capabilities demonstrated by large language models (LLMs) across various fields, a particular challenge persists, hinting at a deeper complexity within the nature of prediction.

• While LLMs have made huge strides in code generation, document understanding, and even creative writing, time-series forecasting remains one of the last frontiers AI has yet to conquer.

• The question is: 🧠 Can AI truly learn to predict the time-series future?

To examine this empirically, let's analyze recent state-of-the-art (SOTA) time-series forecasting results.

The results reveal a significant challenge: 👉while capturing broad trends is achievable, accurately predicting the nuanced fluctuations of time-series data remains a hurdle. This inability to capture detail can lead to significant financial loss in applications like stock market prediction.

Classical Time-Series Models

Before diving into deep learning and transformer-based forecasting models, it’s essential to understand the traditional statistical methods that have been widely used for decades. These methods remain competitive in many applications, especially when data is limited or interpretability is crucial.

Classical methods consider a time series as a sequence of observations indexed by time, assuming that past values contain sufficient information to model future values. Mathematically, a time series can be defined as:

Classical time-series approaches aim to find a precise mathematical model of y_t, representing it as a closed-form function of past data.

ARIMA: AutoRegressive Integrated Moving Average

ARIMA is a powerful model for univariate time-series forecasting that captures patterns in past values and forecast errors. Three parameters define it:

• p (Autoregression - AR): Number of past values used for forecasting.

• d (Differencing - I): Number of times the series is differenced to remove trends.

• q (Moving Average - MA): Number of past forecast errors included in the model.

The model is often written as 👉ARIMA(p, d, q) [2]. We will go through each step of the method.

Step 1: Differencing for Stationarity

Most real-world time series are non-stationary, meaning their statistical properties change over time. To make the series stationary, we apply d levels of differencing. For example, d=1:

For second-order differencing (d=2):

Step 2: Autoregressive (AR)

The autoregressive part models the relationship between a time step y_t and its previous values:

where ϕ_i are the AR coefficients and ϵ_t​ is a white noise error term. We aim to learn ϕ_i.

Step 3: Moving Average (MA)

Instead of relying only on past values, the MA component incorporates past forecast errors:

where θ_i are the MA coefficients and need to be learned. The errors ϵ should be determined after an initial prediction.

👀 Moving Average term here is a bit misleading, because this component works on error level, not the original signal.

Step 4: Combining Everything – The ARIMA Model

The final ARIMA equation defines the relationships between the time series and ARIMA parameters satisfying the 3 components above:

where:

• B is the backshift operator By_t=y_t−1

• Φ(B) is the AR polynomial

• Θ(B) is the MA polynomial

There are many ways to learn the parameters to satisfy the equation, and the procedure has been integrated into standard time-series libraries. For example, we can use Python:

import pandas as pd
from statsmodels.tsa.arima.model import ARIMA
import matplotlib.pyplot as plt

# Sample Time Series (replace with your data)
data = [50, 52, 55, 58, 61, 63, 65, 68, 70, 72]
time_series = pd.Series(data)

# Fit ARIMA(1, 1, 1) model (adjust p, d, q as needed)
model = ARIMA(time_series, order=(1, 1, 1))
model_fit = model.fit()

# Forecast next 3 values
forecast = model_fit.get_forecast(steps=3)
forecast_values = forecast.predicted_mean

# Plot results
plt.plot(time_series, label='Original')
plt.plot(pd.Series(range(len(time_series), len(time_series) + 3)), forecast_values, color='red', label='Forecast')
plt.legend()
plt.show()

print(forecast_values) # print the forecast values.

Exponential Smoothing (ETS Models)

Exponential Smoothing methods forecast future values by giving exponentially decreasing weights to past observations. This helps to capture trends and seasonality.

A basic model for smoothing a time series prediction can be simple:

where α is the smoothing parameter and ŷ_t is the forecasted value for step t.

👉Holt’s Linear Trend Model [3]

Holt's Linear Trend model is a powerful tool for forecasting time series data that exhibits both a level and a linear trend. It's an extension of simple exponential smoothing designed to capture the direction and magnitude of changes over time.

These core equations define the model:

Role of l_t: The level component, denoted as l_t, represents the smoothed average of the time series at time t. It's essentially our estimate of the "current" value, stripped of short-term fluctuations.

Role of b_t: The trend component, denoted as b_t, represents the estimated slope of the time series at time t, embedded in the term l_t-l_t-1. It captures the rate of change (increase or decrease) in the data.

The final forecast equation combines the level and trend to project future values of h steps ahead. This model separates the time series into its underlying level and trend components. This allows for more accurate forecasting of data with linear trends. By adjusting the smoothing parameters α and β, we can control the model's responsiveness to recent changes and fine-tune its performance.

As shown in the table below [10], it's surprising that simple methods like ARIMA remain competitive with complex deep learning models, explaining their continued widespread use in industrial applications.

Yet, they still have several inherent limitations:

❌ Stationarity Requirement: ARIMA models primarily work with stationary time series. If the data is not stationary (i.e., it has trends or seasonality), it needs to be transformed through differencing, which can sometimes lead to loss of information.

❌ Linearity Assumption: ARIMA and ETS models assume that the relationships within the time series are linear. They may not effectively capture non-linear patterns or complex dependencies.

❌ Parameter Selection: Determining the optimal values for these model's hyper-parameters (e.g., p, d, q, α, β) can be challenging. It often involves a trial-and-error process and requires expertise in time series analysis.

❌ Limited to Univariate Data: These models are designed for univariate time series, meaning they can only model a single variable. They cannot directly handle multiple related time series.

Seasonal Decomposition and Trend Analysis

Classical seasonal decomposition breaks a time series into three components:

1. Trend (T_t​): Long-term pattern.

2. Seasonality (S_t​): Repeating periodic variations.

3. Residual (R_t​): Unexplained random variations or fluctuations.

If we assume an additive model for the decomposition, we can model the time-series step as:

Additive models assume constant variance. For time series with increasing variance, multiplicative models are more appropriate.

The next task is to determine each component in the model. For example, we can estimate trends using a moving average:

where m is the seasonal period. Then, we can remove trends and estimate seasonality:

Or for multiplicative decomposition:

Finally, we can compute the residual:

Because seasonality is a fundamental characteristic of many time-series datasets, where patterns repeat at regular intervals. Classical time-series analysis leverages autocorrelation techniques to detect and quantify seasonal effects. Autocorrelation measures the linear relationship between a time series and its lagged versions, helping identify repeating cycles and dependencies over time.

The classic autocorrelation function (ACF) for a lag k in a time-series y_t of length N is defined as:

where ȳ is the mean of the time series. A significant peak in the ACF at lag k=s suggests a seasonal pattern of period s.

As we will see later, these basic analyses will be adopted in the modern deep learning approach and play a crucial role in ensuring good forecasting performance.

Deep Learning for Time-Series Forecasting

Traditional time-series models like ARIMA, ETS, and seasonal decomposition assume that the underlying process is stationary or can be transformed into a stationary form. These models rely on linear dependencies and handcrafted features, making them effective for simple and well-structured time series. However, as real-world applications grow in complexity, these assumptions often break down.

Modern time-series forecasting embraces a more data-driven and high-dimensional approach, leveraging machine learning and deep learning techniques. Instead of relying on predefined structures, modern models learn complex dependencies directly from data, making them more adaptable to dynamic, non-linear patterns. This shift also accommodates multivariate time series, where multiple interdependent variables evolve together over time.

Modern Formulation of Time-Series Forecasting

A time series is a sequence of data steps, considering that each step consists of D correlated variables (i.e., multivariate time series):

👀 There are other specific kinds of time series. For example, univariate time series: D=1, or multi-modal time series with M as the number of modalities:

Theoretically, we can consider all the past steps to forecast future values of the time series. However, it is hard and inconvenient to learn with long, undetermined sequences. Thus, we construct training samples by sliding a window of fixed size over the time series. For each valid index t, we define a data sample:

Time-series forecasting now involves learning a function that maps a historical window of L time steps to a future window of T time steps. Specifically, our goal is to predict the next T steps given the past L steps:

Objective: We want to learn a function f_θ​ that predicts future values given past input time series:

Transformer-Based Time-Series Models

It is natural for deep learning to reshape time-series forecasting, but the recurrent paradigm—RNNs, LSTMs—simply doesn't cut it for complex, long-horizon problems. Long-term dependencies? Computational overhead? These are well-documented issues. That's why the Transformer, a model born in the NLP domain, has become a serious contender. Its self-attention mechanism provides a powerful tool for capturing those elusive long-range dependencies.

Anyone who's attempted to apply vanilla Transformers to a time series will attest to the inherent scalability issues. Beyond that, a fundamental mismatch exists: Transformers were built for the discrete, data-rich world of text, while time-series data is continuous and often comparatively scarce. To overcome this, the community has rallied, producing a wave of specialized Transformer architectures tailored to the nuances of time-series data

👉PatchTST: Simple Extension of Transformer to Time Series [6]

PatchTST directly applies Transformer to time-series data with two modifications:

1. Segmentation into Patches: The method segments time-series data into subseries-level patches, which serve as input tokens to the Transformer model. This patching mechanism ensures that local semantic information is preserved within each segment, enabling the model to capture fine-grained details that may otherwise be missed in traditional models.

2. Channel-Independence: Each channel in PatchTST contains a single univariate time series. This setup allows all channels to share the same embedding and Transformer weights, promoting computational efficiency and reducing the complexity of the model. By minimizing the number of parameters that need to be learned, this approach ensures that the model remains scalable across different time-series datasets.

The ideas can be summarized in the diagram below:

👉Autoformer: Decomposition-Based Transformer for Long-Term Forecasting [4]

Inspired by classical methods, Autoformer introduces series decomposition into the Transformer architecture to separate trend and seasonality, improving efficiency and interpretability. Instead of relying on direct self-attention over raw time-series data, it models time series X_t as:

where:

This decomposition forms a basic computation block named SeriesDecomp. Another important block of Autoformer is the Auto-Correlation Block, which is designed to efficiently capture period-based dependencies in time-series data. Unlike traditional self-attention mechanisms that rely on pairwise dot-product comparisons, Autoformer leverages autocorrelation analysis to identify repeating patterns in the data and aggregate information across time-delayed sub-series. The Auto-Correlation function is defined simply as:

where R_XX​(τ) measures the similarity between the original series and its lagged counterpart. Peaks in R_XX​(τ) indicate potential periodic structures within the data. Autoformer selects the top-k most significant periods by computing:

where R_Q,K(τ)is the autocorrelation function computed between the query Q and key K corresponding to X. The number of selected periods, k, is defined as k = ⌊c × logL⌋ where c is a hyperparameter.

Once the top-k periodic dependencies are identified, Autoformer aligns and aggregates similar sub-series by shifting the time-series values according to the detected period lags:

• Roll(V,τ) shifts the value sequence by delay τ, ensuring elements that are shifted beyond the first position are re-introduced at the last position

• R̂ represents the softmax-normalized autocorrelation scores acting as attention weights.

👀 The process output is expected to amplify signals by effectively leveraging seasonal patterns.

Given the SeriesDecomp and Auto-Correlation Block, Autoformer stacks these blocks together, forming Encoder-Decoder architectures as follows:

👉FEDformer: Frequency Domain Transformer for Time-Series Forecasting [5]

Naive Transformers struggle with capturing global time-series temporal patterns and entail substantial computational costs. The FEDformer addresses these challenges by integrating frequency domain analysis into a computation framework similar to Autoformer’s.

The first improvement is the Decomposition block, where FEDformer introduces Mixture of Experts for Seasonal-Trend Decomposition. In this approach, the trend is computed as:

where F(·) applies a collection of average pooling filters, and L(x) calculates the weights used to combine the resulting trend features.

The second improvement is the Frequence-Enhanced blocks. Because time-series data frequently contain cyclic patterns that can be difficult to model directly, the paper proposes to use Discrete Fourier Transform (DFT) to decompose a signal into frequency components, making it easier to identify dominant trends.

Given a time-series sequence xnx_nxn​ with n=0,1,…,N−1, the DFT transforms it into the frequency domain:

where i is the imaginary unit, and X_l represents the frequency components of the sequence. The inverse DFT reconstructs the original signal:

👀 The strength of frequency analysis using DFT is that using Fast Fourier Transform (FFT), the computation can be reduced from O(N^2) to O(Nlog⁡N), and when selecting only a subset of frequencies, it can be further reduced to O(N).

Given these operators, the Frequency-Enhanced Block (FEB-f) applies FFT to transform inputs into the frequency domain, selects dominant modes, and applies a learned transformation before returning to the time domain to highlight the underlying dominant frequencies. The processing steps include:

1. Linear Projection: Transform time-series input x∈R^N×D using a learnable weight matrix w :

2. Fourier Transform & Mode Selection: Convert q to the frequency domain and retain only M dominant modes:

Here, Select() is the function that keeps M frequency values X_l.

3. Frequency Domain Processing: Apply a learnable transformation R (a parameterized kernel) to enhance the signal:

where ⊙ denotes element-wise multiplication.

Inverse Fourier Transform: Zero-pad and apply the inverse Fourier transform to return to the time domain:

The third novel module is Frequency-Enhanced Attention which aims to capture the relationship between frequency-enhanced signals (trend and seasonal). Given queries q, keys k, and values v, the computation steps are:

1. Fourier Transform & Mode Selection: Similar to Step 2 of FEB-f, just apply to all q, k, and v:

2. Frequency-Space Attention Calculation:

   where σ is an activation function (e.g., softmax or tanh).

3. Inverse Transform to Time Domain:

👀 The paper also proposes a similar procedure for FEB and FEA, just replacing Fourier Transform with Wavelet Transform, resulting FEB-w and FEA-f blocks, respectively.

The FedFormer framework looks similar to Autoformer’s with the new blocks:

When Attention is Not All You Need for Time Series

We've explored the allure of Transformers in time-series forecasting, marveling at their ability to capture intricate dependencies. However, a crucial question remains: 🧠 Does the complexity of Transformers always translate to superior performance? Recently, a growing body of research suggests that the answer might be a resounding 'no.'

While Transformers excel in capturing long-range dependencies through self-attention, they come with inherent drawbacks. These include:

• Computational Cost: The quadratic complexity of self-attention (O(n^2)) can be prohibitive for long time series, demanding significant computational resources and time.

• Data Hunger: Transformers, with their vast number of parameters, typically require massive datasets to train effectively. In many real-world time-series scenarios, such extensive data may not be readily available.

• Overfitting Risk: The flexibility of Transformers can lead to overfitting, particularly when dealing with noisy or short time series.

👉DLinear: A Surprisingly Simple Challenger to Transformer Models in Time-Series Forecasting [7]

DLinear takes the basic principles of forecasting and applies a clever yet straightforward approach to outshine some Transformer-based models in specific contexts. Concretely, each channel of the time-series input sequence is treated as a whole and provided as input to a simple model (e.g., linear transformation):

DLinear also proposes a decomposition strategy with linear layers, borrowed from models like Autoformer and FEDformer. Here’s how it works:

1. Decomposition: DLinear first decomposes the raw input time series into two components—trend and remainder (seasonal)—using a moving average kernel. This decomposition helps the model handle complex time-series data in a way that emphasizes the long-term trend and seasonal fluctuations separately.

2. Two Linear Layers: After decomposition, DLinear applies two separate one-layer linear models to each component (trend and seasonal). By doing so, it captures both the long-term trend and the short-term seasonality in a very efficient manner.

3. Combining the Components: Finally, the outputs of the trend and seasonal models are summed up to produce the final forecast.

👀 The model also has a variant, NLinear, designed to handle dataset distribution shifts. NLinear normalizes the input by subtracting the last value of the sequence, passing it through a linear layer, and then adding the subtracted value back.

The experiments reveal that DLinear and its variants show competitive performance against complicated Transformer-based methods:

❌ One obvious limitation of DLinear is that it treats each channel independently, which may ignore the correlations across variates that may be beneficial for forecasting.

👉iTransformer: A Step Beyond DLinear [8]

This work addresses the limitations of Transformers :

❌Falter with multivariate time series, compressing data at timestamps and losing crucial inter-variate correlations.

❌ The attention mechanisms also struggle with temporal dependencies.

iTransformer offers a solution by inverting the approach. It treats each variate as an independent token, allowing the attention mechanism to directly capture relationships between variates. This variate-centric view enables the model to effectively learn complex dependencies, leading to more accurate and robust time series forecasts. By shifting the focus from timestamps to variates, iTransformer unlocks a powerful new way to understand and predict time series data.

To this end, iTransformer uses DLinear to process temporal steps and Transformer’s attention to model channel relationships:

👉TimeMixer: A Multiscale Approach for Time Series Forecasting [9]

The TimeMixer method is designed to leverage the multiscale nature of time series data, where different scales exhibit unique properties. By focusing on modeling time-series specialty, the paper can achieve good results even without any attention mechanism.

To begin, TimeMixer first downsamples the input time series X into multiple scales using average pooling:

where x₀ is the finest scale, and x_M is the coarsest scale.

In TimeMixer, the past information is processed using stacked Past-Decomposable-Mixing (PDM) blocks. The core idea behind PDM is to separate seasonal and trend components at each scale and mix them separately across scales.

• Decomposition:

• Seasonal Mixing: In seasonal mixing, the paper uses a bottom-up approach to incorporate finer-scale data, enhancing the modeling of coarser scales and emphasizing the importance of detailed information for seasonal prediction.

  

• Trend Mixing: In contrast to seasonal components, detailed variations in trend data can introduce noise when capturing broader trends. Coarser scale time series offer clearer macro-level information than finer scales. Thus, they use a top-down mixing approach to leverage macro-level insights from coarser scales to guide trend modeling at finer scales.

• Seasonal and Trend Mix:

• The final prediction just combines prediction at all scales:

  

👉CycleNet: Leveraging Explicit Periodic Modeling for Long-Term Time Series Forecasting [11]

CycleNet introduces a novel approach to improve long-term time series forecasting by explicitly modeling periodic patterns present in the data. The core of CycleNet lies in its Residual Cycle Forecasting (RCF) technique, which decomposes the time series into periodic components and residuals. The residuals, representing the fluctuations that cannot be explained by periodic patterns, are then predicted to complete the forecast. This approach offers a clear distinction between cyclic and non-cyclic components.

Given a time series with D channels and a cycle length W, the goal is to model the inherent periodicity within each channel. This is achieved through learnable recurrent cycles, denoted by Q∈R^WxD, initialized to zero and trained to represent the cyclic components of the sequence. Some key points:

• These recurrent cycles are globally shared across channels

• The process of modeling the periodic patterns involves cyclic replications of the recurrent cycles to match the length of the input sequence.

For an input x_t-L+1:t the corresponding cyclic components c_t-L+1:t (past) and c_t+1:t+H (future) are generated by shifting and repeating the recurrent cycles Q as follows:

• Shift Q by t mod W positions to get Q(t), where t mod W represents the relative position of the current sequence within Q.

• Repeat Q(t) ⌊L/W⌋ times and concatenate the first L mod W elements of Q(t).

Given the cycles, we make predictions by:

• Subtract the past cycle from the original time series to get the residual

• Given past residual, predict the future residual

• Reconstruct the future time series by adding the future residual with the future cycles

Despite its simplicity, the method exhibits very good performance on standard time series benchmark:

That said, there are unavoidable limitations:

❌ This approach assumes that cyclic patterns with constant cycle length exist in the time series. The unit cycle length W must be pre-defined.

❌ Similar to DLinear, it does not model cross-channel relationships.

👉SOFTS: A Simple, More Efficient Time-Series Forecaster [12]

The new approach offers a fresh take on time-series forecasting that ditches the heavy attention mechanisms in favor of a lightweight MLP-based approach. Instead of treating each channel independently (losing useful correlations) or relying on expensive attention layers, SOFTS finds a sweet spot: it first extracts a global "core" representation of the time series and then fuses this core back into individual channels.

Concretely, there are 4 steps:

1. Embedding the Time Series: Each channel is projected into a hidden space of dimension d using a linear embedding function.

Here, C is the number of channels.

2. Channel Interaction via STAR: The STAR module refines the channel representations over multiple layers. Instead of direct pairwise interactions, STAR aggregates information into a core representation:

Concretely, at iteration/layer i, the method computes the function f as:

Here, Stochastic pooling (StochPool) aggregates representations from C series by randomly selecting values, effectively blending the characteristics of mean and max pooling.

3. Fusing Core and Channel Representation

The core representation is concatenated with each channel.

The fused representation is then projected back to the hidden space with a residual connection:

4. Prediction Layer:

MLP takes the fused representation to make predictions:

We can summarize the steps in the diagram below:

The results are impressive compared to prior works including Transformer-based methods:

Conclusion

Time-series forecasting remains one of the most challenging problems in AI, demanding models that can capture complex temporal dependencies, handle non-stationarity, and generalize across diverse patterns. Unlike standard machine learning tasks, time-series data is sequential, highly dynamic, and often influenced by external factors, making it resistant to one-size-fits-all solutions.

The Road Ahead: The increasing demand for long-range forecasting, adaptation to distribution shifts, and better uncertainty quantification calls for innovations beyond current architectures. One promising direction is the integration of Large Language Models (LLMs) into time-series forecasting, leveraging their pretrained knowledge and ability to process complex dependencies and adapt to diverse forecasting tasks.

🔥 Coming Next: In my next blog, I’ll explore how LLMs can be adapted as a new solution for time-series forecasting. Stay tuned! 🚀

References

[1] Shi, Xiaoming, Shiyu Wang, Yuqi Nie, Dianqi Li, Zhou Ye, Qingsong Wen, and Ming Jin. "Time-moe: Billion-scale time series foundation models with mixture of experts." ICLR, 2025.

[2] Shumway, Robert H., David S. Stoffer, Robert H. Shumway, and David S. Stoffer. "ARIMA models." Time series analysis and its applications: with R examples (2017): 75-163.

[3] Holt, C. C. (1957). Forecasting seasonals and trends by exponentially weighted averages (O.N.R. Memorandum No. 52). Carnegie Institute of Technology, Pittsburgh USA.

[4] Wu, Haixu, Jiehui Xu, Jianmin Wang, and Mingsheng Long. "Autoformer: Decomposition transformers with auto-correlation for long-term series forecasting." Advances in neural information processing systems 34 (2021): 22419-22430.

[5] Zhou, Tian, Ziqing Ma, Qingsong Wen, Xue Wang, Liang Sun, and Rong Jin. "Fedformer: Frequency enhanced decomposed transformer for long-term series forecasting." In International conference on machine learning, pp. 27268-27286. PMLR, 2022.

[6] Nie, Yuqi, Nam H. Nguyen, Phanwadee Sinthong, and Jayant Kalagnanam. "A time series is worth 64 words: Long-term forecasting with transformers." arXiv preprint arXiv:2211.14730 (2022).

[7] Zeng, Ailing, Muxi Chen, Lei Zhang, and Qiang Xu. "Are transformers effective for time series forecasting?." In Proceedings of the AAAI conference on artificial intelligence, vol. 37, no. 9, pp. 11121-11128. 2023.

[8] Liu, Yong, Tengge Hu, Haoran Zhang, Haixu Wu, Shiyu Wang, Lintao Ma, and Mingsheng Long. "iTransformer: Inverted Transformers Are Effective for Time Series Forecasting." In The Twelfth International Conference on Learning Representations.

[9] Wang, Shiyu, Haixu Wu, Xiaoming Shi, Tengge Hu, Huakun Luo, Lintao Ma, James Y. Zhang, and JUN ZHOU. "TimeMixer: Decomposable Multiscale Mixing for Time Series Forecasting." In The Twelfth International Conference on Learning Representations.

[10] Qiu, Xiangfei, Jilin Hu, Lekui Zhou, Xingjian Wu, Junyang Du, Buang Zhang, Chenjuan Guo et al. "Tfb: Towards comprehensive and fair benchmarking of time series forecasting methods." arXiv preprint arXiv:2403.20150 (2024).

[11] Lin, Shengsheng, Weiwei Lin, Xinyi Hu, Wentai Wu, Ruichao Mo, and Haocheng Zhong. "Cyclenet: enhancing time series forecasting through modeling periodic patterns." Advances in Neural Information Processing Systems 37 (2024): 106315-106345.

[12] Han, Lu, Xu-Yang Chen, Han-Jia Ye, and De-Chuan Zhan. "SOFTS: Efficient Multivariate Time Series Forecasting with Series-Core Fusion." In The Thirty-eighth Annual Conference on Neural Information Processing Systems.

• What are Memory-Augmented Neural Networks (MANNs)?

  • A Brief History of MANNs

  • What's Holding Back MANNs?

• The Rise of Memory in the LLM Era

  • Why Memory Craves LLMs?

  • Why LLMs Thrive with Memory?

  • Memory-Augmented Large Language Models (MA-LLM)

• Working Memory

  • String Memory: Enabling LLMs to Simulate Universal Turing Machines

  • Neural Memory: Long-term Storage and Generalization Power

• Episodic Memory

  • Rapid Knowledge Integration with Differentiable Memory

  • Prompt Optimization with Nearest Neighbor Memory

• The Future of Memory

What are Memory-Augmented Neural Networks?

Memory is the essence of intelligence. Thanks to memory, humans can recognize objects, recall events, plan, explain, and reason. It allows us to learn continuously, adapt to new environments, and apply past knowledge to unfamiliar situations. For AI, memory could be equally transformative. In neural networks, memory enables more than just storing patterns—it provides a way to connect past experiences to current tasks, to adapt across contexts, and to hold knowledge over time [1].

A Brief History of MANNs

Integrating memory into neural networks is not a new concept—it dates back to early models like recurrent neural networks (RNNs) and the Hopfield network, which introduced the idea of internal state retention. In these architectures, memory is embedded within the hidden states of the networks, allowing them to process sequences and retain a limited context. For example, classic RNN memory reads:

where h_t is the hidden state or memory at step t. Here, x_t is encoded in the current memory and h_t-1 represents the previous memory capturing previous inputs x₁, x₂, …, and x_t-1.

Long Short-Term Memory (LSTM) networks further developed this idea, adding mechanisms like gates to control memory flow and enable long-range dependencies. However, these internal states are still limited in how effectively they can store and retrieve large amounts of past information.

👀 The vector memory h_t fails to scale up when the number of dimension increases since W_h would require much more parameters, much bigger memory and is much slower to learn or compute. Therefore, vector-based memory is often low in capacity.

It was only with the development of memory-augmented neural networks (MANNs), such as Neural Turing Machines (NTM, [2]) and Differentiable Neural Computers (DNC, [3]), that memory became a specialized and external component of the model. Unlike RNNs, which are constrained by the fixed size of hidden states, NTMs, and DNCs use a matrix memory that allows the network to store vast quantities of data explicitly. This matrix memory functions like a writable memory bank, consisting of multiple slots, where information can be stored, retrieved, and updated independently, providing a dedicated structure for past knowledge. The role of memory in these models became clearer: it was not just a transient internal state but a long-term storage that could be accessed flexibly. Given previous input tokens, a Controller, which later becomes LLM, is trained to read, write the memory, and make predictions on the next tokens:

In MANNs, the matrix memory M can be updated recursively as the hidden states of RNNs:

where f is a general update function, which depends on the memory architectures and designs.

MANNs with matrix memory are now a promising direction for building neural networks with human-like memory capabilities. In these systems, memory is external to the network and accessed via learnable controllers, enabling the model to decide what to write and read from memory based on task needs. A complete memory model can capture both past data and relationships between data, resulting in a deliberate reasoning system that can simulate high-order reasoning [4].

Thus, the MANN approach allows for complex reasoning, long-term dependencies, and the handling of variable-length sequences, bringing neural networks closer to the adaptive and context-aware memory found in human cognition.

What's Holding Back MANNs?

While powerful, Memory-Augmented Neural Networks (MANNs) face key obstacles. First, their complex memory operations introduce high computational costs, making them slow and challenging to scale with large datasets. Let’s have a look at the write operator of the DNC memory to understand why:

where w, e, v are neural network functions, ⊙ and ⊗ are Hadamard (element-wise) and outer product, respectively.

MANNs also require a sophisticated controller to decide when and how to read from or write to memory. Training this controller is difficult, as no labeled steps guide memory access decisions. Without ground-truth reasoning labels, the controller must learn these strategies from scratch, a challenging process akin to meta-learning that can make training unstable and slow.

Finally, truly memory-demanding tasks are rare, and limited real-world data naturally requires such complex memory operations. This scarcity of relevant training data limits opportunities to unlock MANNs' full potential, holding them back from broader applications.

👀 The limitations outlined above highlight why, despite numerous Memory-Augmented Transformer designs, few are practical or capable of scaling to billion-parameter LLMs.

The Rise of Memory in the LLM Era

Things change when LLMs emerge. In the evolving landscape of AI, memory, and LLMs form a symbiotic relationship, each enhancing the capabilities of the other. This interdependence allows both memory and LLMs to reach their full potential, making them far more powerful together than when used independently, heading towards Artificial General Intelligence (AGI).

Why Memory Craves LLMs?

By leveraging vast amounts of pre-trained knowledge, LLMs can significantly reduce the reliance on task-specific data, making training memory models for memory-intensive tasks more feasible. Moreover, LLMs are sophisticated enough to act as controllers, managing memory access more intuitively and efficiently. This eliminates the need for complex, manually trained controllers, streamlining the training process.

LLMs with memory are also faster and more efficient. LLMs can access and store information quickly by integrating memory within the model architecture, reducing computational overhead. This makes them more scalable and suitable for handling large datasets and complex reasoning tasks.

Ultimately, LLMs provide a natural and efficient way to integrate memory into AI systems. Their inherent capacity for learning and adaptation allows them to seamlessly incorporate memory mechanisms, enhancing their ability to process long-term dependencies and generalize across tasks.

Why LLMs Thrive with Memory?

LLMs are powerful tools, but they often struggle with tasks that require long-term context or reasoning over multiple steps. This is where memory comes into play. By integrating memory mechanisms, LLMs can:

• Handle Long-Term Dependencies: LLMs can remember and utilize information from earlier parts of a text or conversation, improving their ability to generate coherent and contextually relevant responses.

• Facilitate Complex Reasoning: Memory allows LLMs to store intermediate results and refer back to them as needed, enabling more sophisticated reasoning processes.

• Enhance Creativity and Originality: LLMs can generate more creative and original content by accessing a vast knowledge base. Memory enables them to combine ideas from different sources and generate novel outputs.

By addressing these limitations and harnessing the power of memory, LLMs can become even more versatile and capable, opening up new possibilities for AI applications.

Theoretically, we cannot prove that a pre-trained LLM is computationally universal, meaning it can solve any computable problem or simulate a universal Turing Machine. This limitation arises because LLMs are trained on finite data, and their convergence after training cannot be guaranteed to meet the conditions of universality. However, by equipping an LLM with external memory, it can achieve computational universality under reasonable assumptions [5].

Now, you may ask 🧠 what is a universal Turing Machine? A Turing Machine (TM) is a theoretical computing model that can solve any computable problems given infinite memory. However, one TM only solves one specific problem. Universal Turing Machine (UTM) can simulate any TM given that the TM description is provided as the input. Therefore, one UTM, in theory, can solve all computable problems. UTM is realized as a general-purpose computer today, where the memory of the computers stores both programs and data.

👀 UTM can be simulated approximately using neural network architectures that are trained end-to-end [6].

The guarantee that an LLM with external memory can compute anything is significant, as it suggests that external memory can enhance LLM performance on complex tasks.

Memory-Augmented Large Language Models (MA-LLM)

Previous sections show that memory is vital in transforming LLMs from powerful tools into truly adaptive and intelligent systems. 🧠 But what kind of memory is suitable for LLMs? Just as the human brain relies on both working memory and episodic memory to navigate the world, LLMs benefit from these complementary forms of memory to handle diverse and complex challenges.

👉Working memory in LLMs is a short-term storage system, holding information relevant to the current task. Like a notepad for a writer, it is cleared and reset when the task concludes, making room for new data without interference from prior activities. This ensures that LLMs remain focused and efficient, adapting quickly to new tasks without being overwhelmed by the past.

In contrast, 👉episodic memory provides a longer-term perspective, capturing and storing knowledge that spans multiple tasks and events. It acts as a journal, retaining experiences, decisions, and outcomes, which can be revisited to improve understanding and performance over time. Episodic memory allows LLMs to learn from prior interactions and carry forward context, fostering continuity and depth in tasks requiring cumulative reasoning or personalized responses.

Working Memory

The bounded input length of LLMs, such as the typical 4096-token limit, restricts their capacity to handle complex, multi-step reasoning tasks. Augmenting LLMs with external read-write memory offers a promising solution by extending their computational abilities and enabling them to simulate algorithms beyond the scope of finite automata. This external memory acts as a working memory, dynamically storing intermediate computations, sub-problems, or task-specific data that can be retrieved and updated during a reasoning process.

Two key operators of working memory are:

• Read Operator: Retrieves relevant information based on the current input and stored entries.

• Write Operator: Updates the memory with new information generated during the task for future reference.

👀The working memory is often associative, in that it reads the data for the LLM based on the similarity of the current input and what is stored in the memory. Hence, it is also referred to as 👉associative memory.

String Memory: Enabling LLMs to Simulate Universal Turing Machines

A natural question arises: 🧠 What gets stored in the memory? One straightforward implementation is text—the memory can simply store raw text generated by the LLM or provided by users. This simplicity is sufficient to achieve universal computational capability, i.e., it can simulate a universal Turing Machine [5]. In the paper, the authors chose U_15,2-a well-known small Universal Turing Machine as the simulation target. The U_15,2 Turing Machine with a memory tape of an infinite number of slots can be described as follows:

1. States (Q): A finite set of possible configurations the TM can be in. In this case, U_15,2 has 15 states, denoted as: {A,B,C,…,I,J}.

2. Tape Alphabet (Σ): A finite set of symbols that can be written on the tape. U_15,2 has Σ={0,1}

3. Blank Symbol (b): A special symbol used to represent empty tape cells. Here, b=0

4. Start State (q₀): The initial state of the TM. q₀=A

5. Halting States (T): A set of (state, symbol) pairs that, when reached, halt the TM. T = {(J, 1)}.

6. Transition Function (f:Q×Σ→Σ×{−1,+1}×Q): A function that, given a current state and the symbol under the tape head, determines:

   • The symbol to write to the tape.

   • The direction to move the tape head (left -1 or right +1).

   • The next state to transition to.

The transition function or program of U_15,2 can be represented as a lookup table:

A Turing Machine starts in its initial state, reads the symbol under its tape head, and then, based on the current state and symbol, writes a new symbol, moves the head left or right, and transitions to a new state. This process repeats until the machine reaches a halting state.

Back to the contribution of the paper, the authors aim to simulate U_15,2 using LLM augmented with a string-based memory.

👀The LLM is like a CPU that can access the working memory (RAM) to fetch data and instructions.

The memory functions as a dictionary, mapping keys (variable names or addresses) to values: MEMORY[variable name] = "value". Unlike physical RAM, keys are strings for seamless interaction with the LLM, while values can be either strings or integers. Following the UTM principle, the string values can also be the instructions. For example, MEMORY[‘op‘] represents the current instruction, eg., ‘halt‘. Other important memory variables are:

• MEMORY[‘i‘] represents the current location of the Turing machine head. For example, MEMORY[‘i‘] =4 means the head is in the 4th slot.

• MEMORY[number] represents the value stored at the number-th slot where the number is a string representing a number. For example, MEMORY[4] = ‘0‘, indicating the symbol stored in the 4th slot is ‘0‘.

With that in mind, the LLM system executes in 3 steps:

1. Read: Retrieve and construct the next input prompt from the memory

2. Compute: Execute the prompt (instruction) using the LLM

3. Write: Parse the LLM's output to extract variable assignments, store them back in memory, and move to step 1.

👀Notably, this approach requires no additional training or weight modification of the LLM, relying purely on prompt engineering and memory management for universal computation. The authors can test this approach on Flan-U-PaLM 540B.

To guarantee computational universality stems from the system itself, interactions between the LLM and memory through Read and Write are limited to finite-state operations, like simple regular expression parsing. It is important to verify the power of LLM+String Memory to see if it is truly computational universality. Otherwise, during computation, the system can query another computer for answers. Now let’s look at each step in detail.

Read

The processing rule with the memory should be defined and executed by simple Python programs. For example, given a string, we can do Read by replacing the variable name that appears in the string (e.g., like @[variable_name]) with the value stored in the memory whose key equals the variable name:

Given this function, any string containing pattern char[] can be replaced by memory content easily.

Compute

In this step, the LLM tries to perform the U_15,2’s transition function above. It is provided with the string representing the current state and the program that reads like:

We are asking the LLM to do if-else generation conditioned on the head value using an instruction prompt. 🧠 How to design the prompt? The authors propose to use few-shot instruction to help the LLM familiarize with the syntax of generation output and if-else behavior. Some examples read:

result = " op="%[B]" %[i]="0" i+=1 "
if 0==1 then result = " op="%[A]" %[i]="1" i+=1 "
$result
" op="%[B]" %[i]="0" i+=1 "

result = " op="%[B]" %[i]="0" i+=1 "
if 1==1 then result = " op="%[A]" %[i]="1" i+=1 "
$result
" op="%[A]" %[i]="1" i+=1 "

All examples are stored in a variable boot stored in the memory. In other words, MEMORY[‘boot’] = “result = " op="%[B]" %[i]="0" i+=1 …".

Each transition rule corresponds to a state instruction prompt coupled with the few-shot examples. Each prompt is stored in the MEMORY as variables A, B, C …

A = """@[boot]result = " op="%[B]" %[i]="0" i+=1 "
if @[@[i]]==1 then result = " op="%[A]" %[i]="1" i+=1 "
$result
"""
B = """@[boot]result = " op="%[C]" %[i]="1" i+=1 "
if @[@[i]]==1 then result = " op="%[A]" %[i]="1" i+=1 "
$result
"""
...

For example, if the current state is A, MEMORY[‘op‘] =

"""@[boot]result = " op="%[B]" %[i]="0" i+=1 " if @[@[i]]==1 then result = " op="%[A]" %[i]="1" i+=1 " $result """

and we give that prompt to the LLM after loading the value head @[@[i]] from the memory using Read. For example, let's say the value is 1, then we expect that the LLM will generate: op="%[A]" %[i]="1" i+=1

Write

Given LLM’s raw output, we need to substitute the variable name (%[variable_name]) with its values stored in the memory, producing a post-processed output string. Then, the authors use a simple Python program to update the memory values with the assignment specified in the output string.

Putting together, the whole system simulating U_15,2 looks like this:

# Step 1: Initialize memory with predefined variables and values
Initialize MEMORY with:
    'boot', 'i', 'A', 'B', ... (values are defined earlier)

# Step 2: Set the initial operation
MEMORY['op'] = MEMORY['A']  # Start with operation defined in 'A'

# Step 3: Main execution loop
While True:
    op = MEMORY['op']  # Get the current operation
    
    # Check if the process should stop
    If op == 'halt':
        Exit loop  # End the program

    # Perform core actions
    Read  # Fetch required data
    Compute  # Perform calculations or operations
    Write  # Update memory 

Despite being theoretically powerful, the design of this kind of string memory faces difficulties:

❌ The system does not aim to solve the task directly. Rather, it simulates a UTM and then uses the program of the TM to solve the task. In theory, it can solve any task, but it still requires finding the program of TM that is suitable for the task. It is unclear how to find that program.

❌ Working on string or text level is slow because the LLM must generate texts and that involves a sequential sampling process, which is slow.

Tensor Memory: Long-term Storage and Generalization Power

In implementation, neural networks work with tensors. Therefore. it is convenient to store vector and tensor representations in the working memory, enabling the LLMs to communicate with LLMs in the representation-level, rather than text level. For example, Wang et al., (2024) proposed storing the LLM’s attention keys and values in external memory:

In particular, the input sequence is split into fixed-length (context_size) segments, each processed by a frozen LLM and their key-value pairs cached in memory. Current inputs use query vectors to retrieve memory content, fused with the local context for another trainable network (👉SideNet) to make predictions.

SideNet, implemented as a Transformer, uses the backbone LLM's embedding layer and frozen language modeling head:

The hidden state H will then be used to generate Q, K, and V to do the attention normally :

The catch is we will augment this attention with memory contents by a special memory retrieval and fusion.

Token-to-Chunk Memory Retrieval simplifies and accelerates memory operations by grouping tokens into fixed-size chunks. Instead of retrieving token-level key-value pairs, the system retrieves chunk-level pairs using mean-pooled vectors for efficient matching, then flattens the retrieved chunks back into token-level pairs for processing.

👀This approach reduces retrieval complexity, enhances accuracy, and allows adjustable granularity based on task needs, such as broader context for in-context learning.

Given the current token’s query, we can retrieve the top K/context_size chunks. After flattening the chunks, we obtain K key-value pairs as the memory contexts. Rewrite them in tensor form for all input token reads:

Memory Fusion combines local and retrieved memory contexts through a joint-attention mechanism in a specialized memory-augmented layer. Each token can attend to the tokens stored in the memory:

Then, we combine the memory attention with SideNet attention to produce the final output at layer l:

👀 During memory-augmented adaptation, only the SideNet's parameters are updated, while the backbone's pre-trained knowledge remains fixed. This streamlined method facilitates rapid convergence by effectively utilizing the existing expertise of the backbone model (405M GPT-2).

❌ The frozen LLM offers both advantages and drawbacks. While it keeps the system lightweight, it limits adaptability. For instance, if the pre-trained LLM has shortcomings, the quality of memory content may suffer, reducing the effectiveness of the augmentation.

❌ Chunk-level retrieval, while extending the memory span, cannot provide precise access to specific tokens, which limits its utility in applications requiring detailed reflection on past inputs, such as code generation or fine-grained question answering.

❌ The SideNet is trained to purely combine the memory content with its current input. This process relies on KNN retrieval without learning to read or write to the memory and offers no mechanism to ensure generalization.

To address these issues, we can consider a more flexible design where the memory can be differentiable:

The proposed mechanism, termed 👉Infini-attention, also stores Q, K, and V representations of the LLMs. However, unlike Wang et al., (2024), the author proposes to use a compressive memory that adopts a recurrent update using Linear Attention [9]:

or Linar Attention combined with Delta rule:

Then, the memory read performs a normalized matrix multiplication:

Like Wang et al., (2024), this is interpolated with LLM’s original attention A_dot:

👀 Each Infini-attention layer is trained via backpropagation through time using gradients of compressive memory states, similar to RNN training. This is expensive, yet allow better adaptiation to downstream tasks that requires finetuning LLM.

A different, more classical approach, for MA-LLM adopts the Controller Memory framework of NTM and DNC. In this MM-LLM architecture, the memory contains multiple slots. The LLM Controller reads and writes token representations to the slots. To save computation costs, the memory can be applied to the final layer of the LLM’s Transformer Encoder and Decoder as in Pointer-Augmented Neural Memory (👉PANM) [10]:

Unlike prior works that aim to use memory as long-term storage, PANM focuses on leveraging memory to empower LLM with length extrapolation capabilities. As such, PANM introduces a pointer-based mechanism with two principles:

1. Explicit Pointers as Physical Addresses: Incremental binary addresses replace softmax-based attention, ensuring scalable and predictable memory access for long sequences.

2. Decoupled Pointer Operations: Pointer manipulation is separated from input data, enabling abstract operations like copying or sorting independently of specific values.

Inspired by slot-based RAM, this design ensures reliable storage and retrieval across sequences of any length. The memory consists of:

1. An address bank (A)—a collection of binary memory addresses.

2. A Pointer Unit (PU)—a module that manipulates pointers to access memory efficiently.

👀 Apllying PANM to LLama-2 7B improves allows the LLM to generalize to a new sequence length 100 times longer than those seen in the training data. More details on PANM can be found in our previous blog post.

Episodic Memory

In deep learning and reinforcement learning, episodic memory refers to a memory mechanism that stores specific experiences or "episodes" for future reference. Unlike working memory, episodic memory can last longer, across tasks. These memories typically capture key events or knowledge in a task, in associative key-value pairs. For example:

• Eiffel Tower — Paris

• Friend’s DOB — 29/11/1995

• 9 AM yesterday — played soccer

A defining characteristic of episodic memory, distinguishing it from associative memories like semantic memory, is its ability to update or rewire associations rapidly. This is essential in the context of LLMs as although massive knowledge or facts have been embedded in the LLMs’ weights, they are more like semantic memory, and slow to update. This highlights the need for an additional storage mechanism that can rapidly adapt to new data or knowledge.

Rapid Knowledge Integration with Differentiable Memory

In recent work, Das et al., (2024) propose an MA-LLM named 👉Larimar that mirrors the hippocampus-neocortex interaction, where the memory rapidly captures factual updates as episodic memory, while the LLM encodes long-term patterns as semantic memory [11]. The episodic memory module serves as a global repository for storing the latest factual updates or edits, conditioning the LLM decoder to reflect this information.

Given a set of encoding Z_i computed by the Encoder, representing the i-th knowledge we want to add to the memory M, assuming that we can find the address (key) W that specifies where to write to the memory, the authors propose the following memory update rule:

This rule resembles the Linear Attention + Delta Rule mentioned above. Furthermore, the memory update rule ensures important theoretical properties. Assuming M_i-1 is the least-squares solution for Z_0:i-1:

• We want to remember Z_i, so we set α_i=1:

then M_i is the least-squares solution for Z_0:i

• Case 2: We want to forget Z_i, which is stored before at i_forget<i, so we set α_i=-1:

then M_i is the least-squares solution for Z_0:i-1 with Z_iforget removed from the data:

👀 Intuitively, the least-squares solution results in a good memory because the memory is computed such as we can minimize the reconstruction error, i.e., we can optimally reconstruct past data Z_j.

So far, we have assumed W_i is given. 🧠 How to compute W_i? The authors follow a prior work [12] in determining optimal values for W_i:

where X⁺ denotes the pseudo-inverse of X. Intuitively, the idea is still to minimize the reconstruction error. Given the memory, we can read or generate new content as follows:

Despite sound motivation for memory read/write design, Larimar has several drawbacks:

❌ The approach assumes no inherent order between episodes, neglecting temporal dependencies between knowledge across episodes. This limitation is evident in the memory update rule, where altering the order of Z does not affect the resulting memory values.

❌ Using pseudo-inverse operations to estimate memory and addresses can become computationally slow when the number of memory accesses is high.

❌ Training memory operations using backpropagation can be slow and require many training data samples.

To address the first limitation, Ko et al., (2024) propose to model the temporal dependencies between facts. This is critical for QA tasks that require advanced reasoning:

The proposed method, dubbed 👉MemReasoner leverages Larimar memory operations:

Furthermore. it introduces 2 new components:

• Temporal Encoding: Positional encodings are computed for each line of context using sine and cosine functions. Additionally, learnable encodings are explored using a bidirectional GRU, where input sequences generate ordered context embeddings through GRU outputs. These embeddings are then written to memory using Larimar’s write operation.

• Multi-step Reasoning with Query Rewriting: Multi-step reasoning tasks involve iterative "hops" between facts until the solution is reached. At each hop, z_q is processed by a simple linear transformation to align with memory content:

  Like Larimar, the memory readout is computed as:

The query is updated iteratively:

This new query will be used to get a new read value z’_r. This continues until convergence, i.e., ||z_r-z’_r||<τ or after a maximum number of iterations.

👀 Query rewritting is crucial when it requires multiple readings from the memory to find the relevant fact. For example, the orignial query is about A, and we want to know about E. If A-B, B-C, C-D, D-E facts are stored in the memory, we would need 4 reading “hops“.

Prompt Optimization with Nearest Neighbor Memory

An efficient way to enhance LLM performance is by refining the prompting process. Episodic memory can significantly improve prompt optimization, benefiting the LLM's outputs as effectively as augmenting the model itself.

By treating each training instance as an episode, the memory archives the experience of a seen prompt, which consists of input data, in-context learning (ICL) examples, and corresponding performance. During testing, we can refer to past experiences stored in the memory, to construct the prompt that potentially maximizes the testing performance.

Based on this principle, Do et al., (2024) propose an episodic memory for in-context example ordering optimization, called 👉POEM [14]. In their paper, the memory M is a set of tuples:

where s, a, and r are state, action, and reward, inspired by reinforcement learning notations and L is the size of the memory. In the context of LLM, they mean:

• State s: the input data, i.e., the question we want the LLM to answer

• Action a: the ordering of the in-context examples given we know the in-context examples

• Reward r: the accuracy of the LLM’s output when prompted with (s, a)

The paper models these components as follows:

• State representation: SentenceTransformer is used to encode the input to a state vector

• Example selection: The set of examples is simply selected using a nearest neighbor search on a given database. We do not need to optimize this process.

• Action encoding: The authors propose a clever way to allow generalization by representing the arrangement of in-context examples as a sequence of similarity ranks rather than their actual content. This rank-based representation captures relationships between examples, reducing overfitting and allowing the system to adapt better to new queries.

👀 This action encoding scheme allows discrete action space, which is convenient for later memory reading.

• Reward design: The reward changes depending on the tasks:

  • Exact Match: A reward is given if the LLM’s output perfectly matches the ground-truth answer:

  

  • Classification: The reward is the difference between the log probability of the correct class and the largest log probability of the other classes:

    

  • Generation: The reward is the difference between the log probability of the ground-truth sequence and the largest log probability of other sequences:

    

While building the memory, we sample the data—the state, from the training data, select the in-context examples, and explore possible actions—the ordering of in-context examples to construct the complete prompt. The prompt will be used for LLM to generate outputs and collect the rewards. Given the tuple (state, action, reward), POEM defines memory operations:

Memory writing: there are 2 scenarios:

• If the state-action pair is new to the memory, we just insert the tuple into the memory. If memory is overflow, the oldest tuple will be removed

• If the state-action pair already exists in memory, we update its stored reward to the current reward if the latter is higher. This maintains an optimistic estimation of the reward for the state-action pair.

Memory reading: We aim to estimate the value for taking action a for a new testing input s_t using nearest neighbor estimation with the state as the query:

where s_i, i = 1, ..., k are the k states with the highest similarity to the testing state s_t. CS is the Cosine Similarity measurement. The action that has the highest value estimation will be deemed optimal for the input s_t:

👀 The proposed approach is fast to train and do inference. The authors show that it works for an array of LLMs from e RoBERTa-large to Llama2-7B, 13B and 70B.

The Future of Memory

As LLMs continue to scale, integrating efficient memory mechanisms becomes critical for handling dynamic knowledge updates, long-term dependencies, and computational efficiency. Traditional transformer-based architectures face challenges in managing memory due to their quadratic complexity and reliance on attention mechanisms, which become computationally prohibitive for longer sequences. The future lies in fast, scalable memory mechanisms that exhibit linear complexity, parallelizability, and the capacity to manage long-term knowledge without degrading performance. Below, we highlight promising approaches shaping this evolution.

1. State Space Models (SSMs)

State Space Models represent a paradigm shift in sequence modeling. By combining linear state-space equations with deep learning, SSMs can efficiently process sequences with linear complexity. They inherently model long-range dependencies by operating in continuous time, making them well-suited for extending memory capabilities in LLMs. Their attention-free nature allows scalable memory management while retaining interpretability and robustness. Big SSMs such as Mamba have shown potential as a Transformer-based LLM alternative (see more in this blog post).

2. Linear Attention

Linear attention mechanisms reformulate traditional attention calculations to scale linearly with sequence length [9]. By approximating or reweighting the attention matrix using kernel methods or other simplifications, linear attention reduces the computational overhead without sacrificing the model’s ability to capture contextual dependencies. This method is inherently parallelizable and faster than standard self-attention, making it a practical choice for tasks requiring both speed and scalability. Linear attention also retains compatibility with existing architectures, providing an efficient path to extend memory capabilities.

3. xLSTM

xLSTM builds on traditional recurrent architectures, optimizing them for long-term memory retention and scalability [15]. Unlike vanilla LSTMs, xLSTM leverages architectural modifications to avoid vanishing gradients, allowing it to maintain information over extended sequences. It achieves near-linear complexity by avoiding redundant computations and parallelizing operations, making it a viable alternative for tasks requiring deep contextual understanding.

4. Stable Hadamard Memory Framework (SHM)

The Stable Hadamard Memory (SHM) framework employs the Hadamard product for memory updates and calibration, offering a robust and scalable memory solution [16]. Its core advantage is minimizing dependencies between time steps, stabilizing gradient flows, and preventing learning issues such as vanishing or exploding gradients. SHM is particularly adept at long-term reasoning, offering a linear complexity approach that is attention-free and inherently parallelizable.

Interestingly, these recent advancements converge under a unified Hadamard Memory Framework [16], where the matrix memory update process can be expressed in a linear formulation as:

For example,

• SSM: M, C, and U are vectors

• Linear Attention: C_t=1

• xLSTM: C_t is a scalar

Viewing these advancements through a unified lens opens opportunities for comprehensive theoretical analysis and technical enhancements, offering valuable insights into the behavior and potential of modern attention-free LLMs. We will dive deep into this model in a separate blog post. Stay tuned!

Reference

[1] Le, Hung. "Memory and attention in deep learning." arXiv preprint arXiv:2107.01390 (2021).

[2] Graves, Alex. "Neural Turing Machines." arXiv preprint arXiv:1410.5401 (2014).

[3] Graves, A., Wayne, G., Reynolds, M. et al. Hybrid computing using a neural network with dynamic external memory. Nature 538, 471–476 (2016).

[4] Le, Hung, Truyen Tran, and Svetha Venkatesh. "Self-attentive associative memory." In International conference on machine learning, pp. 5682-5691. PMLR, 2020.

[5] Schuurmans, Dale. "Memory augmented large language models are computationally universal." arXiv preprint arXiv:2301.04589 (2023).

[6] Le, Hung, Truyen Tran, and Svetha Venkatesh. "Neural Stored-program Memory." In International Conference on Learning Representations (2020).

[7] Wang, Weizhi, Li Dong, Hao Cheng, Xiaodong Liu, Xifeng Yan, Jianfeng Gao, and Furu Wei. "Augmenting language models with long-term memory." Advances in Neural Information Processing Systems 36 (2024).

[8] Munkhdalai, Tsendsuren, Manaal Faruqui, and Siddharth Gopal. "Leave no context behind: Efficient infinite context transformers with infini-attention." arXiv preprint arXiv:2404.07143 (2024).

[9] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and Franc¸ois Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pp. 5156–5165. PMLR, 2020.

[10] Le, Hung, Dung Nguyen, Kien Do, Svetha Venkatesh, and Truyen Tran. "Plug, Play, and Generalize: Length Extrapolation with Pointer-Augmented Neural Memory." Transactions on Machine Learning Research, 2024.

[11] Das, Payel, Subhajit Chaudhury, Elliot Nelson, Igor Melnyk, Sarath Swaminathan, Sihui Dai, Aurélie Lozano et al. "Larimar: Large Language Models with Episodic Memory Control." ICML, 2024.

[12] Kha Pham, Hung Le, Man Ngo, Truyen Tran, Bao Ho, and Svetha Venkatesh. Generative pseudo-inverse memory. In International Conference on Learning Representations, 2021.

[13] Ko, Ching-Yun, Sihui Dai, Payel Das, Georgios Kollias, Subhajit Chaudhury, and Aurelie Lozano. "MemReasoner: A Memory-augmented LLM Architecture for Multi-hop Reasoning." In The First Workshop on System-2 Reasoning at Scale, NeurIPS'24. 2024.

[14] Do, Dai, Quan Tran, Svetha Venkatesh, and Hung Le. "Large Language Model Prompting with Episodic Memory." ECAI, 2024.

[15] Beck, Maximilian, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. "xLSTM: Extended Long Short-Term Memory." NeurIPS, 2024.

[16] Le, Hung, Kien Do, Dung Nguyen, Sunil Gupta, and Svetha Venkatesh. "Stable Hadamard Memory: Revitalizing Memory-Augmented Agents for Reinforcement Learning." arXiv preprint arXiv:2410.10132 (2024).

• LLM Is Just Making Stuff Up

• Detecting Deception: Tools and Methods for Identifying LLM Falsehoods

• Score-based Approaches for Uncertainty Estimation in LLMs

  • Heuristic Uncertainty as a Clue

  • Quantifying Uncertainty with Information Theory

• Model-based Hallucination Detection

  • LLM as Evaluators

  • Simple Conformal Predictors

• Final Thoughts: The Future of LLM Hallucination Detection

LLM Is Just Making Stuff Up

Ever have a conversation with a large language model that sounds super confident, spitting out facts that seem...well, a little fishy? 🐟 You're not alone. One of the biggest challenges in working with Large Language Models (LLMs) is verifying the correctness of their output. Despite their advanced capabilities, LLMs can sometimes generate information that appears accurate but is fabricated. This phenomenon, known as 👉 hallucination, can lead to misinformation and erode trust in AI systems.

Hallucination in AI is not a new phenomenon. Deep learning models, in general, are notorious for their over-confidence in predictions. For instance, in classification tasks, these models can assign a very high probability to a label prediction, even when the prediction is incorrect [1]. Deep learning models can be misleading in how powerful they truly are.

In the context of AI-generated text, large language models (LLMs) can produce content that appears real and coherent, yet is irrelevant and unacceptable. Recent papers [2] categorize LLM hallucinations into 2 main types:

• Factuality Hallucination: This is like the LLM making stuff up entirely. It might sound convincing, but the information is just plain wrong. Think of it like telling you the capital of France is New York City (👉 factual inconsistency, i.e., simply wrong) or the Roman Empire was the first civilization to discover Antarctica (👉 factual fabrication, i.e., no evidence).

• Faithfulness Hallucination: This happens when the LLM strays from your topic or instructions. It might weave a good story but doesn't answer your question or follow the original idea. Imagine asking for a recipe and getting a poem about kitchens instead (👉 instruction inconsistency). Another example is that LLMs can summarize an input document, perpetuating any incorrect or unsupported information it contains (👉 context inconsistency) or perform wrong mathematics derivation (👉 logical inconsistency).

👀 The dangerous thing is that the generated text looks really smooth and confident, which make it hard to know if the content is hallucinated or not.

Imagine the catastrophic consequences of deploying LLMs in real-world applications without addressing their hallucination issues. For example, an LLM might mistakenly diagnose a benign condition when the symptoms indicate a serious illness, putting a patient's life at risk. The real-world impact of these errors is evident: Google's early LLM, BARD, cost the company $100 billion due to a critical hallucination error.

A reliable AI system should be upfront about its limitations. Imagine asking a friend a question, and they spout out an answer with zero hesitation, even if they're not entirely sure. Not ideal, right? The same goes for AI. The best AI systems are those that can signal their uncertainty when they're unsure. We don't want an AI that's either arrogantly confident about everything or so timid it never takes a guess. Hallucination or not, the key issue is identifying when an LLM’s output is unreliable. Now, the real questions are:

🧠 Can we detect when LLMs are generating misleading content? Or even better, can we mitigate the hallucination or dehallucinate LLMs’ output?

Today's focus is on the first question. We'll explore the second question in the next post.

Detecting Deception: Tools and Methods for Identifying LLM Falsehoods

We all know that feeling – you ask a language model a question, and it answers with booming confidence. So, 🧠 how can we tell if an LLM is just making stuff up? There are two key approaches to detect LLM "deception":

• Score-based Methods: One way to sniff out an LLM's fib is to look at how uncertain it is about its answer. Imagine a friend who gives you an answer with a shrug and a mumbled "maybe." That might raise a red flag, right? Similarly, LLMs that express high uncertainty about their output are more likely to be unreliable. By measuring this uncertainty as a score, we can get a sense of how trustworthy the information might be. In this vein, several approaches emerge, some inspired by heuristic methods in uncertainty estimation for deep learning, while others rely more on theoretical principles.

• Calling in the Backup: Another approach involves using external models called conformal predictors. Think of them as using AI to control AI. These models analyze the LLM's information and predict whether the output is real or fabricated (hallucinated). Two approaches to Conformal Prediction:

  1. LLM Evaluator: This method utilizes another LLM (or potentially the same one) to evaluate the generated text itself. This approach bypasses the need for handcrafted features but potentially introduces additional complexity.

  2. Simple Conformal Predictor: This approach leverages well-established methods like linear or logistic regression. However, it relies heavily on extracting informative features from the LLM's output.

The effectiveness of these methods depends on what information we have access to about the LLM itself. For example, the prediction will be more accurate if we can peek "inside" the LLM and see its internal workings (like a white box). However, if the LLM is a black box (we can't see its inner workings), we might need to ask it multiple times to get a clearer picture. Crucially, all methods assume that LLMs have some awareness of their uncertainty or confidence levels, meaning they have a rough idea of how accurate their outputs are. Without this self-awareness, estimating uncertainty or predicting correctness solely by observing the LLMs is impossible.

👀 Fortunately, recent evidence has pointed out that this assumption is practical and LLMs are aware of what they know or don’t know. All we need to do is find good ways to extract or trigger this information.

To sum up, in any method, we just need to craft a score to measure the uncertainty/confidence of LLMs. For methods that do not use external conformal predictors, the score can be computed as a scalar using different approaches. The score then can be calibrated with a training dataset to find a proper threshold for detection decision-making. For conformal prediction approaches, the score can be extracted from the predictor’s prediction logits. We also need a training dataset to train the predictor. Sometimes, if the logits or the score is good, we can just use a default threshold of 0.5 without calibration. The general framework for detection is depicted below:

Score-based Approaches for Uncertainty Estimation in LLMs

Heuristic Uncertainty as a Clue

Just like other powerful AI models, LLMs have a built-in tool to estimate how likely their answers are to be correct. This tool is embedded in the final layer, known as the softmax layer, which calculates the probability of each token in the vocabulary appearing at the current timestep. Typically, the token with the highest probability is the one you see in the LLMs' output.

Unfortunately, as discussed earlier, the built-in probability is not reliable. It does not reflect the reasonable confidence LLMs should have. For example, an LLM might assign a high probability to a factually incorrect answer, simply because the answer aligns with the patterns it has observed in its training data. Furthermore, the built-in probability only applies to individual tokens, not the overall coherence or accuracy of the entire response. This limitation hinders our ability to gauge the trustworthiness of a complete sentence. Ideally, we need methods to assess the probability of the entire content being factually sound and logically consistent.

Fortunately, the field of deep learning offers established methodologies for uncertainty estimation. As Huang et al. (2023) highlight in their comprehensive survey, three key approaches can be directly adapted from this literature to quantify the uncertainty associated with LLM responses [2].

👉 Probability Aggregation: This approach combines individual token-level probabilities (often in the form of log probabilities) to estimate a single probability score for the entire sentence and response.

👀 This approach is simple and economical, only requiring one forward pass of LLM to get the log probs. However, it needs access to softmax layer information, which may be unavailable for black box LLM services.

For example, max and average aggregation can be used to estimate the uncertainty of a sentence i:

where 𝑝_𝑖𝑗 is the probability of a token at position 𝑗 of a sentence i.

👀 Taking average of log probs is equivalent to measuring uncertainty as 👉Perplexity.

Since LLMs inherently predict probabilities for all possible tokens at each step, they produce a distribution p(x_j) for the j-th token. Hence we can leverage entropy, a well-established measure of uncertainty H(X_j) = -Σ p(x_j) * log(p(x_j)), to estimate sentence-level uncertainty as follows,

👉 Uncertainty through Voting: This trick involves generating multiple responses from the LLM for the same prompt. We then analyze the variance (how different or inconsistent the responses are from each other). The idea is that if the LLM keeps spitting out similar answers, it suggests a more consistent and potentially reliable thought process. The more the responses veer off course, the higher the uncertainty. Diving deeper, we can derive 2 metrics: (1) variation ratio (VR) and (2) variation ratio for original prediction (VRO). In particular, if we sample T sentence responses from the LLM and can measure the difference between 2 responses p_i and p_j via function dist(), we have:

The main difference between the VRO and VR formulas is that VRO only considers the variance between the original response and any additional generated responses (assigning a weight of 1 to the original response and 0 to the others). Here, the distance function can be the BLEU score which captures lexical matching:

We can also use BERT as the function to capture semantic similarity/difference between responses. SelfCheckGPT paper [3] proposes a similar uncertainty formula using BERT to measure the uncertainty of response r_i:

Here, for each of the N sampling iterations, we sample several sentences and select the one most similar to the original response in terms of BERT score then take the average.

👀 To sample different outputs from LLMs, we may need to access to the hyperparameter temperature t. t=0 means the generation will be deterministic and ends up with the same output. t>0 will enable more stochasticity in the generation process. Other than that, voting approach can work well with black box LLMs.

👉 Uncertainty through Perturbation: One fascinating aspect of LLMs is their inherent randomness during text generation. Like a chain reaction, a tiny change in one predicted word can ripple through the entire sequence, potentially leading to completely different meanings. This stochastic nature highlights the sensitivity of LLMs throughout the prediction process and we can measure it by:

1. Choose a token and replace (perturbed) it with other tokens (top-k highest probabilities). This leads to several responses.

2. Compute variance as the uncertainty score over the response as in Voting mechanisms above.

🧠 Which token should we mess with? The authors propose 3 ways:

• Most Uncertain Spot: This refers to the place in the generated token where the LLM itself seems unsure about what word to pick next (highest entropy, Max).

• Most Confident Spot: This is the opposite of point 1, where the LLM seems very certain about the word it chose (lowest entropy, Min).

• Biggest Shift: This focuses on the point where the LLM's confidence level changes the most compared to the previous word (maximum change in entropy. MaxDiff).

👀 This approach requires major interference into the LLM computation process, and thus more suitable for white box setting.

The research found that getting the LLM to vote on multiple responses is the best way to gauge uncertainty, followed by tweaking the text and looking at the changes, and lastly, simply looking at the probabilities the LLM assigns to each word.

🧪 Voting > Perturbation > Probability Aggregation

Quantifying Uncertainty with Information Theory

Continuing the line of reasoning that multiple samples help hallucination detection, we can investigate deeper into the hidden states of the LLMs instead of just probing outside. Concretely, Chen et al., (2024) sample responses multiple times, generating multiple hidden states and feature vectors, providing richer information about the LLM's confidence in its responses. The authors propose to use a metric based on the eigenvector of these feature vectors as an uncertainty metric, 👉 EigenScore [7].

In particular, given K hidden state vectors as composing a matrix Z, they compute the covariance matrix:

where J_d = I_d − 1_d 1_d^⊤ is the centering matrix and 1_d ∈ R^d is the all-one column vector. Then, EigenScore can be calculated as the logarithm of the determinant (LogDet) of the covariance matrix:

👀 The EigenScore represents the differential entropy in the sentence embedding space following Gaussian distribution. Hence, it is reasonable to use it for measuring uncertainty.

The authors also suggest clipping the features during the computation of EigenScore to reduce overconfidence estimation:

where h_min and h_max are hyperparameters that can be tuned or calibrated.

More principled research [4] throws shade on LLMs' hallucinations by revisiting the basics of uncertainty in machine learning. Turns out, there are two big types:

• Epistemic Uncertainty: When the LLM just doesn't know enough (think facts or grammar rules). This can happen because it hasn't seen enough training data or just isn't powerful enough yet.

• Aleatoric Uncertainty: This is when the question itself is tricky. Imagine there are multiple right answers, making it a guessing game even for the smartest LLM. Note that this kind of uncertainty is common in LLM settings because there can be many ways to generate reasonable responses.

So, the lower the epistemic uncertainty, the more likely the LLM's answer is on point. Since aleatoric uncertainty is not the fault of the model and we cannot do anything about it, it is important to differentiate the two sources of uncertainty.

❌ The problem with heuristic approaches is that they only measure LLM uncertainty as a whole, not inherent ambiguity in the problem itself (aleatoric uncertainty). This can be misleading. For example, a perfect predictor might have high aleatoric uncertainty, while a bad one might only have high epistemic uncertainty. Both would appear equally uncertain under heuristic methods.

Therefore, the authors propose to focus on identifying instances where only the epistemic uncertainty is large, which would suggest that the response is likely hallucinated. To this end, they propose 👉 epistemic uncertainty via an iterative prompting procedure.

Here's the trick to do that: first, they ask the model to respond to a query. Then, they ask for another response to the query plus the first response. After that, they request a third response to the query and the first two responses, continuing this process. If the LLM keeps changing its response across trials, it suggests a lack of confidence in its knowledge. In contrast, if the LLM consistently provides answers insensitive to the concatenation of its previous response, it indicates a stronger grasp of the topic.

👀 In other words, the responses should be independent. This means that the joint distribution of these responses, for a fixed query, must be a product distribution.

To illustrate the point, the authors observe the probabilities of LLM on the correct answers when prompting with the iterative procedure:

👀 Why? Intuitively, if the question is seen during training, the attention key and query weights of the LLM are tuned to be able to project the question to higher attention scores than other sentences. Thus, the question will be attended the most regard less of the context length, and the LLM will always have a chance to look at the quetion to give the right answer. On the contrary, if the question is novel, the weights can not do any thing and as the context get longer, the attention can be anywhere, leading to a wrong attended input for the LLMs to answer.

In short, the iterative prompting procedure gives us a hint at the uncertain behavior of the LLM. Given the right motivation, now we can derive a robust uncertainty score. Formerly, given a query x ∈ X and possible responses Y₁, . . . , Y_t, a family of prompts F = {F_t : X → X | t ∈ N} is defined with the prompt function F_t(x, Y₁, . . . , Y_t) as:

Then, we can model the distribution of the sequence of responses given the query x:

👀 The chain rule is approximated because of the use of the prompt function F_t to combine the random variables. Hence, it is pseudo join distribution.

Given the formulation of the joint distribution, it is intuitive to say that the response of LLMs Y₁,…,Y_n|x is wrong if the LLM’s probability of Y₁,…,Y_n|x is unlike the ground truth probability of Y₁,…,Y_n|x. Thus, a metric that can measure the truthfulness of LLM’s output is the KL divergence between LLM’s joint distribution and ground truth joint distribution. Yet, we don’t know ground truth distribution. Fortunately, we can replace the KL with an estimable lower bound:

Computing the exact mutual information requires evaluating Q over its entire support, which can be infinite. Therefore, the authors propose to estimate the term by sampling-based approximation. In particular,

1. Sample X₁, . . . , X_k sequence of responses from the LLM

2. Construct a set of indices of unique elements S =  i ∈ [k] : X_i ≠ X_j ∀j < i

3. Construct empirical distributions: for all i ∈ S:

4. Finally, compute the estimated mutual information:

where 𝛾 and k are hyperparameters.

Model-based Hallucination Detection

LLMs as Evaluators

It seems like a chicken-and-egg problem to use LLMs to detect LLMs’ falsehood – 🧠 how can an LLM identify a lie if it chose to generate the lie in the first place? Interestingly, early research has shown that they can do it [5]. It is reasonable since humans also exhibit this behavior. We often make mistakes and only realize them upon reflection. Similarly, when interacting with language models, they can acknowledge and correct their errors when shown their mistakes.

Enhancing validation reliability is possible by employing a more robust LLM to verify the outputs of a less powerful one. This approach, commonly used by the open-source community for benchmarking LLM improvements, involves leveraging a larger or more advanced model to assess the accuracy of a smaller or less-developed model. The detection framework is very simple:

👀 Simple evaluation prompts may not work well all the time, especially when the LLM Evaluator is not stronger than the main LLM.

Improving the accuracy of the LLM evaluator requires special methods [5]. One property the research found out is that LLMs excel at calibrating multiple-choice and true/false questions. Put simply, the probabilities they assign to the options are somewhat reliable.

👀 This property is more evident if the multiple choice has the suitable format. For example, if there is “None of the above” choice, the quality of calibration may be reduced. We may also need to tune the temperature t to have good probabilities.

Thus, they propose a simple trick to use 👉 LLM Prompting without finetuning to make the evaluation more accurate:

1. Present the response to the LLM Evaluator and ask if the response is True or False.

2. Measure the the probability P(“True”) that the LLM Evaluator assigns to the token “True”.

An example to illustrate the evaluation prompt:

To enhance accuracy, the authors suggest presenting the model with additional examples for comparison. For example, we can generate a total of 5 responses, and then ask the model to assess the validity of one of them—the original response of LLM.

👀 The result can be further improved with few-shot prompting techniques. In short, we can say that Comparison (Few-shot) > Comparison > One Proposed Answer.

In addition to prompting, the authors also propose 👉 finetuning LLMs for the detection task. Concretely, they train LLMs to predict whether they know the answer to any given free-form question, i.e., estimating P(IK) (“I know”), using 2 approaches:

• Value Head Integration: This approach introduces an additional "head" to the LLM architecture. This head is specifically trained to predict P(IK) as a logit value. A key advantage of this method lies in its flexibility. We can probe the value head at any point during text generation, allowing for dynamic uncertainty assessment.

• Natural Language Prompt-based Training: This approach leverages natural language processing (NLP) techniques. They train the LLM to respond to the prompt: "With what confidence could you answer this question?" The model's target output is a human-readable percentage value (0% - 100%) reflecting its estimated confidence level in answering the question. This method offers a more intuitive interpretation of uncertainty for users.

❌ Unfortunately, the Natural Language Prompt-based Training approach fails, so the authors only follow the Value Head Integration approach.

When training a model, it's essential to prepare training data. Similar to training other conformal predictors, we need data in a binary classification format:

• X: the input and response from the LLM

• Y: whether the response is correct or not.

In practice, they generated 30 response samples per question input. If 20 samples were deemed correct, they would have 20 positive-label data points in the training set, indicating the model "knew" the answer. Conversely, 10 incorrect samples resulted in 10 negative-label data points. The LLM is finetuned to output the value head following the ground-truth labels.

The results indicate that finetuning generally helps the model distinguish between correct and incorrect responses. In the in-distribution setting across datasets, the LLM's predicted P(IK) aligns somewhat with the ground truth. However, when generalized to a different dataset (from TriviaQA to Mixed-Arithmetic), this differentiation becomes less clear.

Despite the initial promising results, there are significant limitations with the approach:

❌ High Detection Cost: The cost of detection is high because it relies on LLMs. These models require substantial computational resources and energy, leading to increased expenses in terms of both hardware and operational costs.

❌ Insufficient Reliance on Textual Responses: Simply relying on the textual response to determine truthfulness is inadequate. Textual responses alone cannot comprehensively reveal the correctness of the information because LLMs are very good at making things look real.

Simple Conformal Predictor

The key to catching hallucinations might lie within the LLMs themselves. By peering deeper into their internal workings, we could extract valuable clues about their current state and what they "believe" to be true. This richer information would significantly boost the accuracy of hallucination detection. Think of it this way: with a clearer picture of the LLM's thought process, we wouldn't need such a complex detector. Even a simpler classifier could do the job if we have the right features to analyze. The workflow becomes:

This simple idea is attractive because a simple conformal predictor such as a feed-forward neural network can be used to perform the detection. However, the nature of these features can pose several challenges:

🤔 Is it easy and cost-effective to extract these features?

🤔 Are these features informative and can they generalize well to new prompts and different large language models (LLMs)?

Now, the main question is: 🧠 Which features should we extract? One candidate is 👉 the internal states of the LLMs. Recent works have investigated and declared that the internal states of LLMs are reliable sources of information for truthfulness detection on the final response [6,7].

👀 It is important to note that this must be the internal states, not the response text or the response embedding vector.

As quoted in their paper:

Ok, let’s use the internal states, which are represented by the hidden layers of the LLMs.🧠 Which layers should we use? Intuitively, the last hidden layer seems like a good candidate – it should theoretically hold all the processed information. But there's a catch: this layer is primarily focused on predicting the next word in the sequence, not necessarily retaining long-term context. Conversely, layers closer to the input are better at extracting basic features from the data, but might not capture the bigger picture.

To find out, the authors try out several hidden layers of the LLMs as the features. For each chosen layer, the feature vector can be simply the average across token timesteps or the last token’s hidden states. The results reveal that the middle layers perform best:

Recently, a more detailed investigation into the hidden states of LLMs aims to find out if these internal states can signal the risk of hallucination based on the given queries [8]. The goal is to see if we can reliably estimate this risk even before the LLM generates a response, i.e., 👉 self-awareness.

Self-awareness is the ability in humans that causes us to hesitate before responding to queries or making decisions in situations where we recognize our lack of knowledge (we know what we don’t know). The authors want to verify that ability in LLMs by studying LLMs’ internal states.

Concretely, they use internal states corresponding to the last token of queries, denoted as x_q. The conformal predictor, or estimator, employed is a variant of the multilayer perceptron (MLP) adapted from Llama’s. The estimator is mathematically formulated to predict the hallucination risk H as:

They also prepare a dataset containing both known and unknown queries for the LLMs. The LLMs are expected to be uncertain about the unknown queries, which they have never encountered before. They train the estimator on the dataset and compare it with simple baselines such as Perplexity and Prompting to illustrate the point that internal states are really good indicators of uncertainty for unknown queries.

🧪 Internal-State Conformal Predictor > PPL > ICL Prompt > Zero-shot Prompt.

Hidden states, while powerful tools within Large Language Models (LLMs), come with inherent limitations:

❌ Architecture Dependence: Extracting hidden states is intrinsically tied to specific LLM architectures and models. This creates a roadblock when transferring the extraction process across different LLMs. Each LLM architecture might require unique approaches to access and interpret its hidden states.

❌ Sensitivity and Generalizability: Hidden states are demonstrably sensitive to the input data and the specific LLM they are extracted from. This sensitivity poses a significant challenge to generalizability. Conformal predictors, for instance, trained in a particular dataset's hidden states might not perform well when applied to hidden states derived from a different dataset or LLMs.

In a promising new direction, researchers have proposed an alternative approach that bypasses hidden states altogether. This method, 👉 Lookback Lens, focuses on extracting features directly from the attention scores generated by LLMs. We focus on attention because it reveals how much the LLM considers the given context when generating text. This is especially valuable compared to other internal model workings. Since attention provides a human-understandable measure, it becomes a powerful tool for catching and fixing made-up information (hallucinations) in the generated text.

When a Transformer-based LLM performs attention, it attends to the context tokens and its newly generated tokens. The authors aggregate attention scores at attention head h and layer l, corresponding to the two types of attention:

Here N is the number of tokens in the context and t is the timestep of newly generated tokens. Hence, the lookback ratio LR(l,h t) for head h in layer l at time step t is:

👀 Intuitively, if the LLM focuses more on the context (the ratio is higher), it tends to be more reliable, less hallucinated.

Of course, we can combine different layers, heads, and timesteps to form a combined feature vector for a span of generated text Y={y_t , y_t+1, ..., y_{t+T −1}}. Given the feature vector, we can employ a simpler classifier (logistic regression) to detect if it is factual or hallucinated.

The experimental results are promising, at least for summarization tasks where attending to the context is crucial to summarize.

🧪 Lookback Lens > Hidden States > Prompt

The results also reveal that Lookback Lens might not always learn the training data perfectly, but it consistently performs better on completely different tasks (out-of-domain tasks). Lookback Lens analyzes attention maps (lookback ratio features), which is more robust than the hidden states, and thus, is powerful and adaptable, making it useful for a wider range of problems.

Given the Lookback ratio as a score, we can measure the factuality or confidence of different generated candidates:

Final Thoughts: The Future of LLM Hallucination Detection

While the methods explored here offer promising avenues for detecting LLM hallucinations, there's still much room for exploration. Future research directions include:

• Improved uncertainty estimation for LLMs: Refining techniques for LLMs to better quantify their uncertainty about generated content.

• Novel methods for leveraging internal LLM states: Exploring techniques to analyze internal LLM representations to glean deeper insights into the generation process and identify potential hallucinations.

• Integration with factual knowledge bases: Developing frameworks that seamlessly integrate LLM outputs with external knowledge sources to verify factual consistency and enhance detection accuracy.

• Benchmarking and interpretability: Establishing standardized benchmarks for evaluating hallucination detection methods and fostering interpretable models that provide clear explanations for their decisions.

By addressing these challenges, we can move towards a future where LLMs are reliable partners in human endeavors, offering creative and informative outputs while minimizing the risk of misleading information. This will be crucial for fostering trust and wider adoption of LLM technology across various domains.

References

[1] Guo, Chuan, et al. “On calibration of modern neural networks.” International conference on machine learning. PMLR, 2017.

[2] Huang, Lei, Weijiang Yu, Weitao Ma, Weihong Zhong, Zhangyin Feng, Haotian Wang, Qianglong Chen et al. "A survey on hallucination in large language models: Principles, taxonomy, challenges, and open questions." arXiv preprint arXiv:2311.05232 (2023).

[3] Manakul, Potsawee, Adian Liusie, and Mark JF Gales. "Selfcheckgpt: Zero-resource black-box hallucination detection for generative large language models." arXiv preprint arXiv:2303.08896 (2023).

[4] Yadkori, Yasin Abbasi, Ilja Kuzborskij, András György, and Csaba Szepesvári. "To Believe or Not to Believe Your LLM." arXiv preprint arXiv:2406.02543 (2024).

[5] Kadavath, Saurav, Tom Conerly, Amanda Askell, Tom Henighan, Dawn Drain, Ethan Perez, Nicholas Schiefer et al. "Language models (mostly) know what they know." arXiv preprint arXiv:2207.05221 (2022).

[6] Azaria, Amos, and Tom Mitchell. "The internal state of an LLM knows when it's lying." arXiv preprint arXiv:2304.13734 (2023).

[7] Chen, Chao, Kai Liu, Ze Chen, Yi Gu, Yue Wu, Mingyuan Tao, Zhihang Fu, and Jieping Ye. "INSIDE: LLMs' Internal States Retain the Power of Hallucination Detection." In The Twelfth International Conference on Learning Representations.

[8] Ji, Ziwei, Delong Chen, Etsuko Ishii, Samuel Cahyawijaya, Yejin Bang, Bryan Wilie, and Pascale Fung. "LLM Internal States Reveal Hallucination Risk Faced With a Query." arXiv preprint arXiv:2407.03282 (2024).

[9] Chuang, Yung-Sung, Linlu Qiu, Cheng-Yu Hsieh, Ranjay Krishna, Yoon Kim, and James Glass. "Lookback Lens: Detecting and Mitigating Contextual Hallucinations in Large Language Models Using Only Attention Maps." arXiv preprint arXiv:2407.07071 (2024).

• Reflection on Intrinsic Motivation

• Direct Exploration

  • Replay Memory Focused on Exploration

  • Performance-based Replay Memory

• Causal Exploration

  • What is Causality?

  • Dependency Test

  • Potential Outcome

  • Structural Causal Model

Reflection on Intrinsic Motivation

In the previous article, we reviewed an essential exploration framework called intrinsic motivation, which is widely used in deep RL due to its scalability. Within the framework, surprise and novelty are the medium for exploration. Regarding surprise, memory is often hidden within dynamics models, memorizing observed data to enhance predictive capabilities. This type of memory tends to be long-term, semantic, and slow to update, akin to a careful archivist meticulously preserving information.

On the other hand, novelty takes a more straightforward approach to memory. Memory is delineated here, resembling a slot-based matrix, a nearest neighbor estimator, or a simple counter. This memory is typically short-term, instance-based, and highly adaptive to environmental changes, acting more like a dynamic and responsive agent ready to adjust to new inputs swiftly.

They all begin with the memory origin, employing surprise or novelty mechanisms to create intrinsic rewards that guide the exploration of the RL agent. While being so convenient and easy to use, two major issues have hindered the ability of the framework to explore effectively:

1. Detachment: lose track of interesting areas to explore.

2. Derailment: prevent it from utilizing previously visited states.

🧠 But what other alternatives could there be to overcome these issues?

Direct Exploration

In this section, we will explore a new form of exploration that is more direct and mechanical. While it requires additional complexity, it proves to be quite effective. 🧠 What is direct exploration? In short, it is a systematic approach where an agent revisits previously discovered states to explore new actions and areas, enhancing efficiency and effectiveness in sparse reward environments.

Under this new scheme, the role of memory is simplified to:

1. Storing past states.

2. Retrieving states to explore.

3. Replaying past experiences.

We will refer to the memory as "replay memory" to distinguish it from semantic and episodic memory concepts.

The effectiveness of the direct exploration framework heavily relies on the sampling strategy, which can generally be categorized into exploration-focused and performance-driven.

Replay Memory Focused on Exploration

One of the pioneering works that paved the way for the direct exploration framework is 👉Go-Explore [1], the first paper to achieve superhuman performance on Atari games.

According to the paper, memory, and simulators offer effective solutions to address the challenges of detachment and derailment mentioned earlier. Detachment, where an algorithm loses track of interesting areas to explore, can be mitigated by using memory to group similar states into cells. This approach is akin to a hash count mechanism. Each state is mapped to a cell, with each cell having a score that indicates its sampling probability.

On the other hand, derailment, where exploratory mechanisms prevent the algorithm from utilizing previously visited states, can be tackled using a simulator, although this is only suitable for Atari environments. The simulator samples a cell's state from memory and resets the agent's state to that of the cell. Throughout the exploration process, the memory is continually updated with new cells, ensuring that the agent remains on track and fully utilizes past experiences.

The basic workflow extends the direct exploration as follows

Here, the replay memory stores three crucial elements: the state, the cell (corresponding to the state), and the cell's score. A higher cell score indicates a greater likelihood that the cell will be sampled for exploration from its state.

The cell is a compressed state representation, making storing efficient and informative against noise. The paper uses simple engineering tricks for the state-to-cell mapping process using downscaled cells with adaptive downscaling parameters to ensure robustness. It starts by grouping recent frames into cells and selecting hyperparameter values that optimize a specialized objective. This objective encourages frames to be distributed as uniformly as possible across cells, preventing the issues caused by overly non-uniform distributions. By avoiding excessive aggregation and insufficient aggregation of frames, the approach ensures balanced exploration and maintains tractability.

The probability of choosing a cell at each step is proportional to its count-based selection weight to enhance the selection of novel states: