💬 NLP

MILES: Modular Instruction Memory with Learnable Selection for Self-Improving LLM Reasoning

MILES is a novel framework for self-improving LLM reasoning that dynamically expands a modular memory of step-wise instruction pairs and employs a coarse-to-fine, learnable selection mechanism to optimize memory composition and reasoning accuracy under realistic test-time constraints.

Original authors: Ruilin Tong, Dong Gong

Published 2026-07-09

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Ruilin Tong, Dong Gong

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

The Big Picture: The "Smart Notebook" for AI

Imagine you have a brilliant but stubborn friend (the AI) who is great at solving puzzles but has a terrible memory. Every time you ask them a new math problem, they act like they've never seen one before. They solve it, but if you ask a similar question later, they have to start from scratch.

Usually, to make this friend smarter, you'd have to send them back to school for years to relearn everything (this is called "training" or updating the model's parameters). But what if your friend is a closed book that you can't send back to school? You can't change their brain.

MILES is a new way to help this friend get smarter while they are working, without changing their brain. It gives them a dynamic, learnable notebook that they can write in and consult in real-time.

The Problem with Old Notebooks

Previous attempts to give AI a "memory" had two main flaws:

The "Whole Solution" Notebook: Some methods stored entire answers to past problems. If you asked a question that was 90% similar to an old one, it worked great. But if the question was slightly different (a "novel" problem), the old answer was useless. It was like trying to use a recipe for "Apple Pie" to bake a "Blueberry Muffin."
The "Heuristic" Notebook: Other methods stored small steps (like "divide by 2" or "check the units"). But the AI had to guess which step to use next by just looking at how similar the words were. It was like a chef grabbing a random spice jar because the label looked similar to the one they needed, hoping for the best.

The MILES Solution: A "Modular" and "Learnable" System

MILES changes the game by treating reasoning like a Lego set rather than a single block of concrete.

1. The Building Blocks: "Sub-goals" and "Sub-instructions"

Instead of storing whole answers, MILES breaks problems down into tiny, reusable Lego bricks.

The Sub-goal (The Label): A short description of what needs to be done next (e.g., "Simplify the fraction").
The Sub-instruction (The Brick): A specific, natural-language tip on how to do it (e.g., "Divide the top and bottom numbers by their greatest common divisor").

Think of this as a toolbox where every tool has a clear label.

2. The "Smart Selector" (The Learnable Head)

This is the magic part. In the past, the AI just grabbed the tool that looked most similar to the current problem. MILES gives the AI a personal assistant (a "selection head") for every single tool in the toolbox.

How it learns: When the AI solves a problem confidently (it gets the answer right), MILES looks back at the steps it took. It asks: "Hey, when we used the 'Simplify Fraction' tool, did it help us get the right answer?"
The Training: If the answer is "Yes," the assistant learns to trust that tool for similar situations. If the answer is "No," the assistant learns to avoid it.
No School Needed: The AI's brain stays frozen. Only these tiny "assistants" (which are very small computer programs) get updated.

3. The Two-Step Search Process

When the AI faces a new problem, MILES uses a "Coarse-to-Fine" search, like looking for a book in a library:

Layer 1 (The Rough Search): The AI quickly scans the labels (Sub-goals) to find a few promising tools. It's like walking down the "Math" aisle and grabbing a few books that might be relevant.
Layer 2 (The Expert Review): The "assistants" (selection heads) take those few candidates and give them a score based on their past experience. "Wait, this tool worked great for algebra but failed for geometry. Let's pick the other one."

How It Works in Real Time (The "Confident vs. Uncertain" Flow)

The system operates in two modes depending on how sure the AI is:

The "Confident" Mode (Learning Phase):
If the AI solves a problem easily and gets the right answer, it treats this as a "training session." It breaks down its successful steps, saves them as new Lego bricks in the notebook, and teaches its assistants which tools worked best. The notebook grows richer with every success.
The "Uncertain" Mode (Helping Phase):
If the AI is stuck or unsure, it stops guessing. Instead, it opens the notebook, asks its assistants for advice, and uses the best tools to guide its thinking. It's like a student raising their hand to ask the teacher for a hint when they are stuck.

Why This Matters (The Results)

The paper tested MILES on hard math and science exams (like AIME and MATH-500).

Better Accuracy: It consistently solved more problems correctly than other methods that use memory or search.
Efficiency: It didn't waste as much computing power trying random guesses; it used the "learned" memory to pick the right path faster.
Transferability: The notebook built by one AI model could actually help a different AI model solve problems. It's like if a master chef wrote a cookbook that a junior chef could pick up and use immediately, even if they had never met.

Summary Analogy

Imagine you are playing a video game.

Old Way: You try to beat every level by guessing. If you die, you restart.
Previous Memory Way: You keep a list of "Winning Strategies" for specific levels. If you see Level 5, you use the Level 5 strategy. If you see Level 5.1, you are stuck.
MILES Way: You keep a dynamic strategy guide that updates itself. When you beat a level, the guide automatically writes down exactly which move worked best for that specific situation. Next time you face a tricky jump, the guide doesn't just show you a list; it has a "smart filter" that says, "Based on your past wins, the 'Double Jump' move is 90% likely to work here."

MILES allows AI to accumulate experience and get smarter over time, not by retraining its brain, but by learning how to better use its own growing library of tips and tricks.

Technical Summary: MILES (Modular Instruction Memory with Learnable Selection)

Problem Formulation

Large Language Models (LLMs) have demonstrated significant reasoning capabilities through test-time computation techniques such as Chain-of-Thought (CoT), Self-Consistency, and Tree Search. However, existing approaches largely treat each problem in isolation. In practical scenarios where problems arrive sequentially, models fail to accumulate and reuse experience across related queries. Updating model parameters for every new problem is often infeasible, particularly for closed or frozen models.

While external memory methods exist, they suffer from critical limitations:

Poor Generalization: Methods storing whole-solution templates or strategies often fail when new problems possess novel structures.
Heuristic Selection: Step-level memory approaches typically rely on heuristic selection (e.g., similarity retrieval or prompting) rather than optimizing for final-answer correctness.
Training Constraints: Learning selection policies usually requires large-scale labeled datasets, fixed action spaces, or parameter updates, which are incompatible with test-time settings where memory expands incrementally and supervision is limited.

The core challenge is to organize, maintain, and reuse an external memory that improves reasoning accuracy as experience accumulates, without updating the underlying LLM parameters.

Methodology: MILES Framework

MILES (Modular Instruction Memory with LEarnable Selection) addresses these challenges by enabling test-time self-improvement through the accumulation and composition of reusable step-level reasoning modules. The framework operates on a frozen LLM and consists of four core components:

1. Modular Asymmetric Memory Structure

MILES stores memory as a collection of modular units, defined as asymmetric pairs $(g_i, u_i)$ :

Sub-goal Embedding ( $g_i$ ): A dense vector summarizing the objective of a reasoning step. It serves as a retrieval key.
Sub-instruction ( $u_i$ ): A natural-language directive used to condition the LLM's generation when the unit is applied.
Selection Head: Each memory unit is associated with a lightweight, learnable selection head ( $\pi_{\theta_a}$ ) that predicts the applicability of the unit given the current reasoning state.

This structure allows for incremental expansion; new units and their corresponding heads can be added without retraining existing components.

2. Coarse-to-Fine Retrieval Mechanism

To balance retrieval efficiency with selection accuracy, MILES employs a two-layer selection process:

Layer 1 (Coarse Retrieval): Based on sub-goal similarity. The LLM generates candidate step-goals, and their embeddings are compared against stored sub-goals using cosine similarity. A per-unit effective threshold filters out irrelevant units. This layer facilitates memory expansion and collects training data.
Layer 2 (Learned Applicability Scoring): A learned selection head scores the applicability of retrieved candidates based on the full reasoning state. To handle variable-length states, MILES uses a context codebook (clustered embeddings of past step-goals and conditions) to create fixed-length feature vectors via max-similarity pooling.
Combined Score: The final selection combines the semantic similarity from Layer 1 and the learned applicability score from Layer 2.

3. Learnable Memory Composition

MILES treats memory selection as a sequential decision-making problem. The goal is to maximize the probability of a correct final answer.

Training Data Collection: For "confident" samples (where multiple LLM rollouts agree on an answer), the system runs Monte Carlo Tree Search (MCTS) using Layer 1 sampling. This generates trajectories where sub-instructions are applied, and the resulting correctness serves as a pseudo-label.
Optimization: The selection heads are trained using binary cross-entropy on these collected trajectories. The objective approximates an optimal policy derived under a KL-divergence constraint, encouraging the heads to select units that lead to correct answers while avoiding overfitting through one-vs-rest classification and ensemble methods.

4. Inference Pipeline

The system routes test samples based on confidence:

Confident Samples: Used to expand the memory and train selection heads. The LLM generates responses; if consistent, the trajectory is used for MCTS data collection and memory summarization.
Uncertain Samples: Routed to a memory-guided tree search. The system retrieves candidates via Layer 1, reranks them via Layer 2, and applies the best sub-instruction to guide the next reasoning step, repeating until a final answer is reached.

Key Contributions

MILES Framework: A novel framework for test-time self-improvement that incrementally builds modular instruction memory from confident reasoning traces and applies learned memory to uncertain cases without parameter updates, external supervision, or fixed action spaces.
Modular Asymmetric Memory: A structure of (sub-goal, sub-instruction) pairs with per-item selection heads, enabling a coarse-to-fine mechanism that supports both incremental memory expansion and learnable composition.
Learnable Memory Composition: A method to train lightweight selection heads from rollout outcomes on confident samples, enabling self-supervised, correctness-optimized memory composition under limited test-time data.
Empirical Validation: Extensive experiments demonstrating that MILES matches or outperforms prior methods across six reasoning benchmarks and four backbone models, achieving superior accuracy-efficiency tradeoffs.

Experimental Results

The authors evaluated MILES on six benchmarks (MATH-500, AIME 2024/2025, GPQA-Diamond, MMLU-Pro Physics/Engineering) using four backbones (GPT-4.1, GPT-4.1-mini, GPT-OSS-20B, Qwen3-30B-Instruct).

Accuracy: MILES consistently matches or exceeds baseline methods (including Zero-shot CoT, Self-Consistency, Buffer-of-Thoughts, and Dynamic CheatSheet). Notable improvements were observed on smaller open-weight backbones where other memory-based methods failed to outperform zero-shot CoT.
Efficiency: MILES dominates the accuracy-vs-token frontier. Unlike baselines that plateau with increased compute, MILES continues to improve as the token budget grows by constructing additional memory units.
Transferability: The framework supports cross-model transfer. Memory constructed from rollouts of auxiliary models (e.g., Qwen3-30B) effectively improves the performance of a different, stronger frozen backbone (e.g., GPT-4.1-mini).
Ablation Studies: Removing either the coarse retrieval layer or the learned selection head significantly degrades performance, validating the necessity of the two-layer mechanism. The learned heads outperform alternative rerankers based on state similarity or LLM prompting.

Significance and Claims

The paper claims that MILES demonstrates that modular memory composition is a viable and effective mechanism for the test-time self-improvement of frozen LLMs.

Key significance points include:

Complementarity: The step-wise memory composition approach is complementary to existing prompting and test-time search methods, providing an additional avenue for improvement through accumulated experience.
Robustness: The method is robust across different model families and dataset distributions, and it does not rely on external Process Reward Models (PRMs) or curated training datasets.
Scalability: By decoupling memory selection from parameter updates, MILES offers a path for continual improvement in settings where model weights cannot be modified.

The authors conclude that learning reusable memory-selection knowledge is a promising direction for building continually improving reasoning systems, though they acknowledge limitations regarding the computational overhead of rollout-based data collection and the reliance on the underlying LLM's instruction-following capabilities.