NeurIPS 2026

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

Ruilin Tong Dong Gong (corresponding author)
Artificer AI Lab, University of New South Wales (UNSW Sydney)
arXiv PDF Code BibTeX

A one-minute teaser, followed by a technical walkthrough: routing by answer agreement, writing memory, the memory unit, coarse-to-fine reading, the learning objective, and the results. Turn the sound on for the score. Example units and values in the walkthrough are illustrative.

Video made with Opus 5.5 + a bit of design (DG: Yes. I love Interstellar).

TL;DRMILES lets a frozen LLM improve itself at test time from its own reasoning trajectories, without ground-truth labels. From problems it answers confidently, it extracts sub-goals and the sub-skills (sub-instructions) that achieve them into a modular memory, and it learns online when to read each memory unit, so that later, uncertain problems are solved by composing these sub-skills step by step.

Please check out our other works MoRAM LLaVA-DyMoE MemAE SEMA CSReL

Abstract

Self-improving at test time from its own inference trajectories

Large language models (LLMs) increasingly improve their reasoning at test time via additional computation, yet most existing works treat each problem in isolation. When problems arrive sequentially, accumulating reusable experience across them can further improve performance. Existing memory-based methods either store whole-solution templates that generalize poorly to novel problems or use heuristic step-level selection that is not optimized for final-answer correctness. Learning selection policies requires large-scale training data and fixed action spaces, making such approaches unsuitable for test-time settings where memory expands incrementally and only limited supervision is available.

We propose MILES (Modular Instruction Memory with LEarnable Selection for self-improving LLM reasoning), a framework that dynamically expands step-wise memory and applies correctness-optimized memory composition under realistic test-time constraints. MILES maintains modular memory units consisting of asymmetric pairs of sub-goal embeddings and sub-instructions, each associated with a learnable selection head. This memory structure enables a coarse-to-fine retrieval mechanism: The coarse level enables memory expansion and collects supervision for training selection heads from confident samples, while the fine stage applies learned selection heads to rerank coarse-level candidates and guide reasoning for uncertain samples.

MILES consistently matches or outperforms prior methods while achieving superior accuracy–efficiency tradeoffs. Extensive experiments demonstrate its effectiveness, robustness, and transferability.

At a glance

Why MILES?

MILES treats the model's own inference trajectories as experience. It distills confident trajectories into step-level memory and learns, online, when each memory unit helps. Six characteristics summarize the design.

Learns from its own trajectories

No ground-truth labels. When the K = 10 initial responses to a question agree, their majority answer serves as the pseudo-label for scoring rollouts.

Step-level, modular memory

Each memory unit pairs a sub-goal embedding, what a step should achieve, with a sub-instruction, how to achieve it. Unlike whole-solution templates, units can be composed across problems.

Selection learned for correctness

Every memory unit has a lightweight selection head, trained on whether applying the unit led rollouts to the pseudo-labeled answer.

Two-layer, coarse-to-fine reading

Layer 1 shortlists candidate units by sub-goal similarity. Layer 2 reranks them with the learned heads for the current reasoning state, and the chosen sub-instruction guides the next step.

Grows during testing, or offline

Each new memory unit comes with its own selection head, and heads are trained independently, so memory expands along the test stream without a fixed action space. The memory can also be built offline, ahead of time, from labeled data.

Works with frozen LLMs

MILES needs no access to the LLM's weights and never updates them, so it works with frozen, API-only models; only lightweight per-unit selection heads (small MLPs) are trained. The textual memory can also be built by another model.

Method

How MILES works

Each incoming question is routed by answer agreement. Confident questions write: they expand the memory and supply training data for the selection heads. Uncertain questions read: they retrieve memory units coarse-to-fine and use them to guide a tree search.

Overview of MILES: confident questions write memory and train selection heads; uncertain questions read memory through memory-guided tree search.
Animated walkthrough. A question whose K = 10 answers agree goes to WRITE: MCTS rollouts scored against the pseudo-label add a memory unit and train the selection heads. A question whose answers disagree goes to READ: Layer 1 retrieves units with similar sub-goals, the learned heads rerank them, and the chosen sub-instruction guides the next step. The numbers in the animation are illustrative. Use the controls to pause or to step through the phases.
Overview of MILES: confident questions write memory and train selection heads; uncertain questions read memory through memory-guided tree search.
Figure 1. Overview of MILES. MILES improves test-time reasoning by constructing memory and learning memory selection from confident samples, then applying the learned memory to guide reasoning on uncertain samples.
Technical diagram of MILES: routing by answer agreement, memory writing on confident questions, coarse-to-fine memory reading on uncertain questions, and the selection-head objective.
Technical overview. Answer agreement routes each question. On the WRITE path, confident questions grow the memory and train per-unit selection heads from pseudo-labeled MCTS rollouts. On the READ path, uncertain questions compose sub-instructions through coarse-to-fine retrieval inside a memory-guided tree search. The bottom strip shows the selection objective and the per-unit BCE surrogate used to train the heads. Notation and hyperparameters follow the paper.
Open the diagram at full size

1Modular memory units

A memory unit is an asymmetric pair. The sub-goal \(g_i\) is an embedding of what the step should achieve and serves as the retrieval key. The sub-instruction \(u_i\) is a short natural-language directive that conditions generation. Each unit has its own selection head, so new units can be added without retraining existing ones.

Units are summarized from (reasoning step, step-goal) pairs collected on confident trajectories. The pairs are clustered with k-means, the LLM writes one generalized sub-instruction per cluster, and the mean step-goal embedding becomes the sub-goal.

\[\mathcal{M}=\{(g_i,u_i)\}_{i=1}^{N},\qquad g_i\in\mathbb{R}^{d}\] \[d=1-\cos(\mathbf{s}_1,\mathbf{s}_2)+\alpha_{\text{sum}}\big[1-\cos(\mathbf{o}_1,\mathbf{o}_2)\big]\] \[R^{(i)}_{\text{effect}}=\alpha_{\text{eff}}\cdot\bar{d}_i\]

\(\mathbf{s},\mathbf{o}\): step and step-goal embeddings. \(N\) grows as experience accumulates. \(\bar{d}_i\), the mean distance within unit \(i\)'s cluster, sets its effective threshold, so units from looser clusters apply more broadly.

2Coarse-to-fine selection

Layer 1 (coarse). The LLM proposes step-goals for the next step. Each step-goal embedding \(\hat g\) retrieves the stored sub-goals with the highest cosine similarity, and units beyond their effective threshold are excluded. If no unit matches, the LLM writes the step directly and the pair is kept as a seed for new units.

Layer 2 (fine). Each candidate's learned head scores how applicable it is in the current state \(x\). The state is split into overlapping chunks, and the unit's context codebook turns it into a fixed-length feature vector by max-similarity pooling.

\[\mathbf{m}_a(x)=[s_{a1},\dots,s_{an}],\qquad s_{ai}=\max_{j\in[1,T]}\cos(\mathbf{q}_{ai},\mathbf{c}_j)\] \[\pi_{\theta_a}(a\mid x)=f_{\theta_a}\big(\mathbf{m}_a(x)\big)\in[0,1]\] \[\operatorname{score}(a,x)=(1-\beta_{\text{c2f}})\cos(\hat g,g_a)+\beta_{\text{c2f}}\,\pi_{\theta_a}(a\mid x)\]

\(\mathbf{q}_{ai}\): context prototypes for unit \(a\). \(\mathbf{c}_j\): chunk embeddings of \(x\). The top-scoring unit's sub-instruction guides the next step. \(\beta_{\text{c2f}}=0\) on math benchmarks (heads only) and 0.4–0.6 on academic ones.

3Learning the selection heads

Choosing unit \(a\) at state \(x\) should maximize the probability of a correct final answer. With a KL penalty toward the Layer-1 sampling policy \(\pi_{\text{sim}}\), the optimal policy has the same closed form as in advantage-weighted regression (Peng et al., 2019).

MILES approximates this optimum with one lightweight one-vs-rest classifier per unit, an ensemble of 10 three-layer MLPs, trained with binary cross-entropy on labels from MCTS rollouts on confident samples. Positives are states where applying \(a\) reached the pseudo-label answer. Negatives are states where \(a\) led to a wrong answer, or where a different unit succeeded.

\[\max_{\pi}\;p_m(y^{*}\mid x)=\sum_a\pi(a\mid x)\,V(x,a)\] \[\pi^{*}(a\mid x)=\frac{1}{Z(x)}\,\pi_{\text{sim}}(a\mid x)\,\exp\!\big(V(x,a)/\beta\big)\] \[\ell_a=-\big[r_a\log\pi_{\theta_a}(a\mid x)+(1-r_a)\log\big(1-\pi_{\theta_a}(a\mid x)\big)\big]\]

\(V(x,a)\): expected correctness of the final answer after applying \(a\) at \(x\). \(r_a\in\{0,1\}\): whether the rollout reached the pseudo-label. The one-vs-rest design lets each head be trained independently as memory grows.

Test-time pipeline

Confident question → write

The initial responses agree on an answer.

  1. The majority answer becomes the pseudo-label \(\hat y\). It is also the output for this question.
  2. MCTS explores reasoning states (20–40 rollouts). Expansions choose units by Layer-1 similarity only, so the collected data is not biased toward the current heads.
  3. Each expansion is labeled by whether its rollout reaches \(\hat y\), and the (state, label) pair is added to that unit's dataset \(\mathcal{D}_a\).
  4. Seeds are summarized into new memory units, and the affected heads are retrained on all data collected so far.

Uncertain question → read

The initial responses disagree, so there is no reliable pseudo-label and no MCTS.

  1. A multi-stage tree search keeps a pool of responses. Each stage samples existing reasoning chains and truncates them, using longer prefixes at later stages.
  2. The LLM proposes step-goals. Layer 1 retrieves 16 candidate units, and the Layer-2 heads rerank them.
  3. The top unit's sub-instruction guides the next step, the LLM completes the reasoning, and the response joins the pool.
  4. The final answer is the majority vote over all responses in the pool.

Animated overview

A short animation of how MILES incrementally constructs memory and applies learned memory selection during reasoning.

Results

Main results

Across the 22 benchmark–backbone settings, MILES has the best accuracy in 17 and ties for the best in 4. The one exception is AIME 2025 with GPT-4.1, where Dynamic Cheatsheet (DC) scores 50.00 against 40.00 for MILES. DC's strong AIME results come mainly from its use of code execution.

Dataset Model ZS-CoT SC DC BoT MILES
MATH-500GPT-4.1-mini88.0089.4087.0086.8092.60
GPT-4.187.2088.0083.2084.8091.60
AIME 2024GPT-4.1-mini46.6756.6763.3350.0066.67
GPT-4.140.0050.0053.3350.0053.33
GPT-OSS-20B83.3393.3343.3350.0093.33
Qwen3-30B-Instruct70.0083.3323.3363.3386.67
AIME 2025GPT-4.1-mini46.6753.3360.0040.0060.00
GPT-4.133.3336.6750.0033.3340.00
GPT-OSS-20B86.6790.0036.6733.3393.33
Qwen3-30B-Instruct60.0070.0030.0053.3373.33
GPQA-DiamondGPT-4.1-mini70.2071.7263.6466.6773.23
GPT-4.168.9468.6964.6561.1170.71
GPT-OSS-20B68.1871.7261.6237.3774.24
Qwen3-30B-Instruct70.2072.2342.4267.6873.74
MMLU-Pro PhysicsGPT-4.1-mini82.0083.5083.5083.0085.00
GPT-4.181.0083.0078.0084.0084.50
GPT-OSS-20B85.5086.5079.0042.0087.00
Qwen3-30B-Instruct87.0088.5063.5086.0088.50
MMLU-Pro EngineeringGPT-4.1-mini65.5073.5067.5070.5075.00
GPT-4.175.5078.5072.0067.5080.00
GPT-OSS-20B68.5070.0066.5041.5071.00
Qwen3-30B-Instruct74.5078.0052.0075.5081.00

Table 1. Final-answer accuracy (%). Bold marks the best result in each row, including ties. The same backbone is used for reasoning and for memory construction within each row. MMLU-Pro subsets use 200 randomly sampled questions each; the other benchmarks are evaluated in full.

Accuracy–efficiency tradeoff

Scatter plots of accuracy versus response tokens (millions) on AIME 2024 and AIME 2025 with GPT-4.1-mini. MILES points lie above CoT-SC, ToT, rStar, and DORA at every budget and keep rising with more tokens.
Figure 2. Accuracy vs. response-token cost on AIME 2024 and AIME 2025 with GPT-4.1-mini.

On both datasets, every MILES run is at least as accurate as every baseline run. MILES also keeps improving as the token budget grows, while Self-Consistency, Tree-of-Thoughts, rStar, and DORA level off.

More compute helps MILES because extra MCTS rollouts on confident questions build additional memory units (Table 5). Reusing cross-problem memory captures knowledge that per-problem scaling alone cannot.

Transferability

Memory transfers across models and from offline data

The memory is textual, so it does not have to be built by the reasoning model itself. It can also be prepared ahead of time from a separate labeled set.

Cross-model transfer

GPT-4.1-mini is the reasoning backbone and produces the first 10 self-consistency rollouts used for routing. An auxiliary model produces the next 40 rollouts, which are used to build the memory.

DatasetAuxiliary modelAccuracy
AIME 2024no auxiliary model60.00
Qwen3-30B-Instruct66.67
Qwen3-4B-Instruct63.33
GPT-4.1-mini66.67
AIME 2025no auxiliary model50.00
Qwen3-30B-Instruct60.00
Qwen3-4B-Instruct56.67
GPT-4.1-mini60.00

Table 2. Accuracy (%). Memory built by each auxiliary model, including the 4B-parameter Qwen3-4B-Instruct, improves on the no-auxiliary baseline.

Pre-built memory

An auxiliary model builds the memory ahead of time on a separate labeled set: AIME 2020–2024 for AIME 2025, and the non-Diamond GPQA split for GPQA-Diamond. The memory is then applied to the 13 and 31 most uncertain test questions.

DatasetAuxiliary modelAccuracy
AIME 2025no memory23.08
Qwen3-30B-Instruct30.77
GPT-4.1-mini38.46
GPQA-Diamondno memory48.39
Qwen3-30B-Instruct58.06
GPT-4.1-mini54.84

Table 3. Accuracy (%) on the most uncertain test questions, with GPT-4.1-mini as the reasoning backbone.

Ablation

Removing either selection layer lowers accuracy

DatasetExp.Layer 1Layer 2Accuracy
AIME 20241✕✕60.00
2✓✕63.33
3✕✓63.33
4✓✓66.67
5✓State sim.63.33
6✓LLM rerank60.00
AIME 20251✕✕53.33
2✓✕53.33
3✕✓56.67
4✓✓60.00
5✓State sim.56.67
6✓LLM rerank53.33

Table 4. Accuracy (%) with GPT-4.1-mini. Exp. 5 and 6 replace the learned heads with state-similarity reranking or LLM-prompt reranking.

Combining both layers (Exp. 4) gives the best accuracy on both datasets, and disabling either layer (Exp. 2–3) lowers it.

State-similarity reranking (Exp. 5) recovers part of the learned heads' gain. LLM-prompt reranking (Exp. 6) does not improve over using Layer 1 alone. The learned heads therefore capture signal that neither prompting nor surface similarity of states provides.

Bar charts for AIME 2024 and AIME 2025 showing correct-answer proportion and sub-instruction correct rate for Exp. 1 to 6. Exp. 4 is highest on both metrics: 39.64 and 57.14 on AIME 2024, 35.54 and 47.27 on AIME 2025.
Figure 3. Correct-answer proportion and sub-instruction accuracy on uncertain samples for each experiment in Table 4. The full method (Exp. 4) is highest on both metrics.

Analysis

Scaling, robustness, and what a sub-instruction does

Rollout budget and memory size

More MCTS rollouts per confident question give a larger memory with better coverage of the step-goal space, and higher accuracy.

DatasetRolloutsUnitsAcc.
AIME 202420960.00
301566.67
402066.67
AIME 2025201053.33
301956.67
402560.00

Table 5. GPT-4.1-mini. Units: memory size.

Effect of sample order

With 20 rollouts, easy-to-hard order gives the highest accuracy on both datasets, which the paper attributes to useful memory accumulating before hard questions arrive. With 40 rollouts, the three orders are within 3.33 points of each other.

DatasetOrder2040
AIME 2024default60.0066.67
easy→hard66.6766.67
hard→easy63.3366.67
AIME 2025default53.3360.00
easy→hard60.0060.00
hard→easy56.6763.33

Table 6. GPT-4.1-mini, accuracy (%) with 20 or 40 rollouts per confident question.

Effect of a selected sub-instruction

Under matched reasoning states, trajectories that apply a selected sub-instruction are more accurate than trajectories without one, in all 12 dataset–backbone cells, by 0.63 to 10.80 points.

Dataset4.1-mini4.1
MATH-500+6.48+2.23
AIME 2024+10.80+3.77
AIME 2025+9.12+2.31
GPQA-Diamond+4.36+3.21
MMLU-Pro Phys.+4.12+8.73
MMLU-Pro Eng.+0.63+3.87

Table 14 (appendix). Accuracy gain in percentage points, GPT-4.1-mini and GPT-4.1.

Case study: one sub-instruction changes the solution path

MATH-500 (algebra). Let \(x, y, z\) be positive real numbers such that \(xyz = 2\). Find the minimum value of \(x^4 + 4y^2 + 4z^4\).

Zero-shot CoT

The model applies AM–GM to \(x^4, 2y^2, 2y^2, 2z^4, 2z^4\), then works through a long equality-case computation that ends at an expression, not a number.

Final answer: \(2^{8/5} + 4\cdot 2^{3/5} + 4\cdot 2^{6/5}\) ✕ incorrect

With a selected sub-instruction

Sub-instruction from memoryRewrite or simplify algebraic expressions or equations by substituting variables, factoring, or using properties of roots and exponents to transform them into more manageable forms.

The next step substitutes \(z = \tfrac{2}{xy}\), which gives \(x^4 + 4y^2 + \tfrac{64}{x^4y^4}\). The model then sets \(a = x^4\), \(b = y^2\) and minimizes \(a + 4b + \tfrac{64}{ab^2}\).

Final answer: \(16\) at \(x = y = \sqrt{2},\ z = 1\) ✓ correct

From the paper's case studies (MATH-500, memory-guided tree search). Both trajectories share the first reasoning step.

Citation

BibTeX

@inproceedings{tong2026miles,
  title     = {MILES: Modular Instruction Memory with Learnable Selection for Self-Improving LLM Reasoning},
  author    = {Tong, Ruilin and Gong, Dong},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year      = {2026}
}