Lightning Weave: Improving the Accuracy-Efficiency Frontier of Reasoning Models through Capability Composition

AI summary

Overview

Research area: Efficient reasoning for large reasoning models (post-training, on-policy distillation, and capability composition).

Technical level: Intermediate. The paper assumes familiarity with language-model post-training and KL-divergence-based distillation, but its central idea — combining what two separately trained specialists have learned — is explained through a clear analogy of "policy shifts."

Scope: A post-training framework that extracts the behavioral changes acquired by independently trained specialist models and composes them into one student model, improving accuracy and reducing response length simultaneously across five student models and five math/code benchmarks.

What This Paper Is About

Large reasoning models solve hard problems well but often produce very long reasoning traces, and simply forcing them to be shorter tends to hurt accuracy. The authors observe that some independently post-trained models are already better at accuracy while others are already better at brevity, and ask whether those two capabilities can be extracted and woven together into a single student model. Lightning Weave does this by treating each specialist as a "policy shift" relative to its own pre-training checkpoint, then combining those shifts into one learning target the student can be trained against offline.

Key Contributions

1. Capability composition as a route to efficient reasoning. Rather than jointly optimizing competing accuracy and efficiency objectives from scratch, the paper formalizes each specialist as an "anchor pair" — a post-trained model and its pre-anchor checkpoint — and represents what it learned as the log-ratio of their token probabilities.

2. Tilted-Target DOPD. The authors identify that naively reusing DOPD's sampled-token objective on cached trajectories does not preserve the intended stationary target: the cached regularizer can keep producing non-zero gradients even after the student reaches the desired shifted policy. They fix this by tilting the frozen behavior policy into an explicit target distribution q and minimizing KL divergence to it, so loss and gradient vanish when the student matches the target on cached states.

3. Statewise multi-anchor composition (Lightning Weave). Multiple aligned shifts are combined by a weighted sum of log-ratios (weights non-negative and summing to one), and that composed shift constructs one joint tilted target. Composition happens before target construction, so all capability sources supervise each student decision, and each anchor pair scores the cached trajectories only once — removing the need to serve up to 2K anchor models concurrently during student training.

4. Empirical validation of a controllable accuracy–efficiency frontier. Across five student models, composition beats the base and both single-anchor alternatives, generalizes to additional anchor combinations (including a three-anchor composition), and produces a family of non-dominated operating points as the anchor weights are varied.

Main Findings

• Composition beats both specialists separately. On Qwen3-4B, Lightning Weave improves average accuracy over the base model by 4.02 points while reducing response tokens by 19.7%. On Qwen3-4B-Thinking-2507 it improves average accuracy by 4.5 points while reducing tokens by 9.0%. On Qwen3.5-4B it gains 9.70 accuracy points while reducing tokens by 5.3%. On OLMo-3-7B-Think it gains 3.74 points while reducing tokens by 12.9%.

• Headline benchmark numbers. On Qwen3.5-4B, HMMT 2025 accuracy rises from 59.2% to 64.0% with 10.7% fewer response tokens, and LiveCodeBench v5 accuracy rises from 41.7% to 54.2% with 9.6% fewer tokens. On Qwen3-4B, AIME 2024 accuracy rises from 72.9% to 76.0% while response tokens drop by 21.3%.

• Highest AES across all five backbones. Lightning Weave achieves the highest Accuracy–Efficiency Score (AES) among the compared policies for every student model tested (for example, AES 0.61 on Qwen3.5-4B versus 0.25 for Klear-only and 0.06 for DECS-only).

• A controllable Pareto frontier. Sweeping the relative Klear–DECS weight on Qwen3-4B-Thinking-2507 produces five interior compositions (ρ_DECS = 0.250, 0.375, 0.500, 0.625, 0.750). Responses generally shorten as the DECS fraction increases while intermediate compositions preserve or improve accuracy. At ρ_DECS = 0.625, AIME 2024 reaches 84.9% accuracy with 15,048 tokens, exceeding every listed external baseline and both single-anchor policies in accuracy while using fewer tokens. The non-dominated points trace an empirical Pareto frontier.

• Best performance under tight budgets. When AIME 2024 accuracy is evaluated under different token budgets, Lightning Weave outperforms both single-anchor policies at every budget, with the largest gains under tight budgets.

• Generalization to other anchor combinations. Adding MiMo-RL as an accuracy-oriented anchor on Qwen3-4B, Klear+MiMo-RL reaches the highest overall accuracy (64.38 average) but substantially increases inference tokens, while DECS+MiMo-RL improves accuracy over the base while reducing tokens. A three-anchor composition (Klear+DECS+MiMo-RL) achieves the highest AES of 0.42.

• Tilted-Target matches online supervision cheaply. Tilted-Target performs comparably to Online DOPD across the Klear, DECS, and Klear+DECS settings while avoiding concurrently serving the four models that constitute the anchor pairs in the two-anchor setting. It also beats Naive Offline: average accuracy improves by 1.51 and 3.85 points with Klear and DECS respectively, and the two-anchor setting uses 17.6% fewer tokens with slightly higher accuracy.

• Better than data mixture and sequential composition. Compared with Data Mixture (interleaving equal data scored separately by Klear and DECS), Lightning Weave improves accuracy on all five benchmarks while using 3.8–8.7% fewer tokens. Sequential Composition is strongly order-sensitive, with the last applied anchor determining whether the result favors accuracy or efficiency.

Methodology in Plain English

Treat each specialist as a change, not as a destination. The authors do not copy the specialist's outputs directly. Instead, for each specialist they take two checkpoints — the model before the post-training stage and the model after it — and compute, token by token, the log-ratio of their probabilities. That ratio captures what changed during post-training rather than what the endpoint model happens to say.

Score once, train offline. Rather than running the specialists live on every new student rollout, the student (the "behavior policy") generates trajectories once, and every anchor pair scores those same cached trajectories once. After that, training only needs the student plus the cached scores.

Fix the cached-objective mismatch. Applying DOPD's sampled-token objective directly to cached actions is unreliable because once the student drifts away from the behavior policy, the cached regularizer no longer matches the intended KL penalty, so training can keep pushing even after the shift has been absorbed. The fix is to turn each cached shift into an explicit target distribution — the behavior policy reweighted exponentially by the shift, controlled by a temperature-like parameter α — and then train the student to minimize KL divergence to that target. The loss and gradient vanish exactly when the student matches the target on cached states.

Combine shifts before building the target. With K anchor pairs, the shifts are added as a weighted sum (weights non-negative and summing to one; equal composition uses 1/K each), and the summed shift tilts the behavior policy into one joint target. Because the composition happens in log-ratio space, anchors that support the same action reinforce each other and conflicting changes are reconciled locally. Setting one weight to one and the rest to zero exactly recovers the corresponding single-anchor objective.

Experimental setup. Five students were tested (Qwen3-1.7B, Qwen3-4B, Qwen3-4B-Thinking-2507, Qwen3.5-4B, OLMo-3-7B-Think). The accuracy-oriented anchor pair is Qwen3-8B-Base → Klear-Reasoner-8B, and the efficiency-oriented pair is DeepSeek-R1-Distill-Qwen-1.5B → DECS-1.5B. Training used 3,200 prompts from the math split of Skywork-OR1-RL-Data and 3,200 prompts from KlearReasoner-CodeSub-15K, with four responses per prompt and a 2,048-token maximum response length; optimization used Adam, a global batch size of 64, a constant learning rate of 1.0×10⁻⁶, and α = 2.0. Evaluation covered AIME 2024, AIME 2025, HMMT February 2025, and LiveCodeBench v5 and v6, with 64 samples per math problem and four per code problem at temperature 0.6 and top-p 0.95. Where anchor and student use different tokenizers, anchor scores are projected into the student token space via cross-tokenizer alignment.

Why This Matters

The work reframes efficient reasoning as a composition problem rather than a trade-off problem. Instead of fighting the tension between accuracy and brevity inside a single objective, it reuses specialists that were already trained separately and transfers their learned behaviors into one deployable student — with each anchor scored only once, so the approach scales to more anchors without proportionally increasing serving cost.

Real-world applications:

• Cost reduction in inference serving. Shorter response traces at equal or higher accuracy directly reduce token cost and latency for deployed reasoning models.
• Mathematics tutoring and problem-solving assistants. The reported gains on AIME 2024, AIME 2025, and HMMT February 2025 target exactly the kind of competition-style reasoning used in educational tools.
• Code generation assistants. Gains on LiveCodeBench v5 and v6 apply to coding copilots, where both correctness and response latency matter.
• Model development pipelines. Teams that already own several fine-tuned checkpoints can compose their capabilities without retraining a merged model from scratch or maintaining parameter compatibility.

Industry relevance: The offline caching design is the practically important part for production. Because each anchor pair scores the cached trajectories once, student training no longer requires serving up to 2K live anchor models concurrently, which the authors describe as making the training pipeline substantially more accessible. The weight sweep also gives engineers a single knob for choosing an operating point on the accuracy–efficiency curve, and the three-anchor results show the recipe extends beyond a single pair of specialists.

Future Directions

• Extending composition beyond two or three anchors. The paper demonstrates a three-anchor composition and shows it achieves the highest AES of 0.42, but the behavior of larger anchor sets is left open.
• Closing the gap between the token-level theory and sequence-level optimum. The theoretical analysis is stated for the zero-discount, token-level DOPD surrogate used in the method, and the authors explicitly note that the resulting local target is not identified with the conditionals of the idealized sequence-level optimum, whose conditionals can also depend on future shifts.
• Broadening the capability types composed. The study covers accuracy-oriented and efficiency-oriented shifts in mathematics and code; whether the same framework transfers other capability types is not reported.
• Improving cross-tokenizer alignment. Composition across models with different tokenizers currently relies on projecting anchor scores into the student token space, with details deferred to the appendix — a practical bottleneck as more heterogeneous specialists become available.

Target Audience

This paper is most useful to machine-learning researchers and engineers working on post-training and distillation for large reasoning models, particularly those interested in knowledge transfer between separately trained checkpoints. It also suits practitioners building production reasoning or coding systems who care about inference token cost, and graduate students studying efficient reasoning, multi-teacher distillation, or model composition. Readers wanting to reproduce the method should be comfortable with KL-divergence-based policy optimization, but the high-level idea of composing "what changed" rather than "what was learned" is accessible without that background.

Authors’ abstract

A core goal of efficient reasoning is to improve the accuracy-efficiency frontier. However, jointly improving reasoning accuracy and inference efficiency can be challenging, as the two objectives can favor different reasoning behaviors. Independently post-trained models already offer distinct strengths in accuracy and efficiency. We introduce Lightning Weave, a post-training framework that extracts and composes these independently learned capabilities in a single student through on-policy distillation. Each acquired capability is represented by the policy shift from the model before post-training to the resulting specialist. Lightning Weave combines aligned log-ratio shifts at shared student token states and uses Tilted-Target DOPD to convert the cached signals into a stable learning target. Each anchor pair scores the cached trajectories once, enabling subsequent student training without serving multiple live anchor models concurrently. Across diverse student models and benchmarks in mathematics and code, Lightning Weave substantially improves upon the base students and achieves a state-of-the-art accuracy-efficiency frontier. On Qwen3.5-4B, it raises HMMT 2025 accuracy from 59.2% to 64.0% with 10.7% fewer response tokens, and LiveCodeBench v5 accuracy from 41.7% to 54.2% with 9.6% fewer response tokens. Adjusting the relative strengths of the anchor signals yields a strong empirical accuracy-efficiency Pareto frontier. These results establish Lightning Weave as a new practical route to efficient reasoning through capability composition. Code is released at https://github.com/jet-ai-projects/Lightning-Weave.