PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yecheng Wu, Song Han, Han Cai

arXiv:2609.14708v2cs.AI

TL;DR

Improving reasoning accuracy while reducing inference cost is challenging because the objectives can favor different behaviors. Lightning Weave composes post-training policy shifts from independent specialists into a student target; across mathematics and code, it improves accuracy–efficiency trade-offs, including on Qwen3.5-4B.

  • Problem

    Reasoning traces can be costly, but shortening them can compromise accuracy, and accuracy and efficiency objectives can favor different behaviors.

    Source: §1 · p. 2

  • Method

    Lightning Weave represents specialists’ post-training capabilities as policy shifts and composes aligned shifts at shared student token states into a joint target.

    Source: §1 · p. 2, §3.3 · p. 5

  • Results

    Across evaluated students, composition improves accuracy–efficiency trade-offs; on Qwen3.5-4B, HMMT 2025 accuracy rises from 59.2% to 64.0% with 10.7% fewer response tokens.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    Adjusting anchor-signal strengths yields a family of students with an empirical accuracy–efficiency Pareto frontier.

    Source: §1 · p. 3

    Experiments cover models below 10B parameters, and the method assumes matched post-trained and pre-anchor models; benefits at larger scales are unestablished.

    Source: E LIMITATIONS AND FUTURE WORK · p. 27

Abstract

from arXiv · show

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.

1 INTRODUCTION

Lightning Weave composes policy shifts learned by accuracy- and efficiency-oriented specialists to improve a student's accuracy–efficiency trade-off. It addresses the challenge that jointly optimizing these objectives can favor conflicting reasoning behaviors.

1 INTRODUCTION

  • Each specialist’s capability is represented by the token-probability log-ratio between its post-trained model and its pre-anchor, aligned in student token space and composed into one target.This represents behavioral change acquired during post-training without requiring matching model parameters.Source: §1 · p. 2
  • Cached DOPD can retain a non-zero gradient after the student reaches the intended shifted policy, pushing it away from that solution.The cached-action regularizer may no longer equal the intended KL penalty as the student moves away from the behavior policy that generated the cache.Source: §1 · p. 2
  • Tilted-Target DOPD converts each cached shift into an explicit target distribution for student training.Lightning Weave composes weighted, aligned shifts at shared student token states before constructing the joint target.Source: §1 · p. 2
  • 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.Across fully evaluated students, composition improves the overall trade-off over the base and either single-anchor alternative.Source: §1 · pp. 2–3

2 RELATED WORK

Prior work improves reasoning efficiency through inference controls, training, distillation, and model merging, while capability composition also spans distillation and parameter-space approaches. Lightning Weave instead composes post-training policy shifts at shared student token states.

Efficient Reasoning

  • Efficient-reasoning methods target generated tokens, latency, memory, or throughput, whereas Lightning Weave targets the student’s accuracy–efficiency frontier through post-training capability composition.Prior approaches include token-budget control, length-aware objectives, chain-of-thought compression, adaptive reasoning modes, and inference-time systems optimization.Source: §2 · p. 3

Capability Composition

  • Capability integration includes policy distillation, output-space fusion, parameter-space merging, multi-objective alignment, and multi-teacher on-policy distillation.Lightning Weave composes aligned post-training policy shifts at shared student token states rather than merging endpoint parameters or directly matching specialist policies.Source: §2 · p. 3

On-Policy Distillation

  • On-policy distillation supplies supervision on student-generated trajectories; Lightning OPD caches teacher feedback, while Lightning OPD 2.0 corrects cross-teacher style bias.Tilted-Target DOPD adds an explicit stable target on cached student states for capability composition without serving multiple live anchor models during training.Source: §2 · p. 3

3 METHOD

Lightning Weave represents post-training capabilities as anchor-pair policy shifts aligned at shared student-generated token states. It composes them into a joint tilted target and trains the student to match that target.

3.3 MULTI-ANCHOR DISTILLATION: LIGHTNING WEAVE

  • At each shared student token state, Lightning Weave takes a weighted sum of anchor log-ratio shifts and applies it once to the common behavior policy.Compatible changes reinforce one another, while competing changes are reconciled locally; the weights set the capability balance, with equal composition using 1/K.Source: Figure 2 · p. 4, §3.3 · p. 5

3.1 PRELIMINARIES

  • Standard OPD minimizes KL to a teacher on student-generated states, while DOPD represents a specialist with a post-trained anchor and its pre-anchor.The sequence-level policy shift decomposes into dense token-level rewards; this ratio isolates post-training behavior change rather than the endpoint policy alone.Source: §3.1 · p. 4

3.2 SINGLE-ANCHOR DISTILLATION: TILTED-TARGET DOPD

  • Naive offline DOPD caches anchor shifts on behavior-policy trajectories, but its cached-action regularizer may fail to cancel the shift once the student departs that policy.The cached surrogate can keep pushing after the shift has been absorbed, although the cached states remain an offline approximation.Source: §3.2 · p. 4
  • Tilted-Target DOPD exponentially tilts the behavior policy using the cached shift, then minimizes student-to-target KL on cached states.Its loss and gradient vanish when the student matches the target, giving offline transfer a well-defined fixed point.Source: §3.2 · p. 5

3.3 MULTI-ANCHOR DISTILLATION: LIGHTNING WEAVE

  • Each anchor pair scores shared cached trajectories once; after shifts are aligned and cached, training serves only the student, regardless of anchor count.The method aligns the anchor scores over a common student action space.Source: §3.3 · p. 5

4 EXPERIMENTS

Across five student backbones, Lightning Weave improves the accuracy–efficiency balance, and varying anchor weights traces empirical Pareto frontiers. Ablations and additional anchor combinations support its offline composition approach and show how added anchors shift the trade-off.

4.1 EXPERIMENTAL SETUP

  • Experiments cover five student models, mathematics and code benchmarks, and report accuracy, response tokens, and five-benchmark AES.Students generate four responses per prompt from 3,200 prompts in each training setting; multi-anchor weights are normalized to sum to one.Source: §4.1 · pp. 6–7, Table 1 · p. 6

4.2 MAIN RESULTS

  • Across five backbones, Lightning Weave achieves the highest AES and consistently improves the accuracy–efficiency trade-off over base and single-anchor policies.On Qwen3.5-4B, average accuracy rises by 9.70 points while response tokens fall by 5.3%; HMMT 2025 accuracy rises from 59.2% to 64.0% with 10.7% fewer tokens.Source: Table 1 · p. 6, §4.2 · p. 7

4.3 ACCURACY–EFFICIENCY FRONTIERS

  • On Qwen3-4B-Thinking-2507, non-dominated Lightning Weave weight settings trace an empirical Pareto frontier against efficient-reasoning baselines.The frontier is evaluated across three mathematical reasoning benchmarks.Source: Table 2 · p. 7, §4.3 · p. 7

4.4 GENERALIZATION TO ADDITIONAL ANCHOR COMBINATIONS

  • On Qwen3-4B, adding a third anchor shifts the operating point toward that capability, and the three-anchor composition achieves the highest AES of 0.42.Klear+MiMo-RL achieves the highest overall accuracy but substantially increases inference tokens; adding DECS improves efficiency, while adding an accuracy anchor to a DECS pair improves accuracy.Source: §4.4 · p. 8, Table 3 · p. 8

4.5 ABLATION STUDY

  • Tilted-Target performs comparably to Online DOPD and improves the trade-off over Naive Offline; in the two-anchor setting it uses 17.6% fewer tokens with slightly higher accuracy.Average accuracy improves by 1.51 and 3.85 points over Naive Offline with Klear and DECS, respectively; training avoids concurrently serving the four anchor-pair models.Source: §4.5 · pp. 8–9, Table 4 · p. 9
  • Lightning Weave improves accuracy over Data Mixture on all five benchmarks while using 3.8–8.7% fewer tokens, and achieves the highest AES of 0.41.Sequential Composition is order-sensitive: the last applied anchor determines whether the result favors accuracy or efficiency.Source: §4.5 · p. 9, Table 4 · p. 9

5 CONCLUSION

The conclusion presents capability composition as a route to efficient reasoning, using independently learned policy shifts in a student. The accompanying theoretical analysis specifies the local target’s scope and assumptions.

5 CONCLUSION

  • Lightning Weave composes independently learned policy shifts into a student and improves the accuracy–efficiency trade-off across diverse models and mathematics and code benchmarks.The conclusion also reports generalization to additional anchor combinations and a controllable empirical Pareto frontier as composition weights vary.Source: §5 · p. 9

A THEORETICAL ANALYSIS AND PROOFS

  • The theoretical results concern the zero-discount, token-level DOPD surrogate at fixed states, not conditionals of the idealized sequence-level optimum.The appendix derives a local target and a joint target for weighted compositions of multiple anchor-pair shifts.Source: A THEORETICAL ANALYSIS AND PROOFS · p. 17

A.1 SETUP AND ASSUMPTIONS

  • The setup assumes a finite student-token action space, positive behavior-policy probabilities, finite aligned anchor scores, and cached quantities independent of optimized parameters.The behavior-policy state distribution is represented in implementation by the empirical distribution of cached states.Source: §A.1 · p. 17
  • The variational results require finite anchor scores, while the exact product-of-policy-ratios interpretation additionally requires the stated shared-support condition.The paper treats cross-tokenizer-aligned scores separately as finite scores.Source: §A.1 · p. 17

A.2 TILTED-TARGET DOPD: OBJECTIVE EQUIVALENCE AND FIXED POINT

Tilted-Target DOPD makes the cached loss equivalent to a regularized objective with an explicit target, while the appendix distinguishes this guarantee from the behavior of the naive cached sampled-KL surrogate. The multi-anchor objective extends the equivalence to a joint target.

A.2 TILTED-TARGET DOPD: OBJECTIVE EQUIVALENCE AND FIXED POINT

  • For a fixed state, the tilted distribution is the unique maximizer of the single-anchor local objective.For the multi-anchor objective, the supplied proposition identifies a unique solution, and the single-anchor and zero-weight reductions are specified separately.Source: §A.2 · p. 18, §A.3 · pp. 19–20
  • Minimizing Ld and maximizing Jd are equivalent over a common policy class; if the class can realize Tr on cached states, Tr is a global minimizer with zero loss.Otherwise, training minimizes expected reverse KL to the explicit target.Source: §A.2 · p. 18
  • At πθ0 = πb, the Tilted-Target loss and zero-discount DOPD surrogate have the same expected per-state update direction, but this does not extend to later online state distributions after policy departure.A finite cache still has ordinary sampling error.Source: §A.2 · p. 19
  • The expected cached sampled-KL objective is generally not DKL(p ∥ bs), and a two-action example shows the intended tilted target need not be stationary.This establishes no general fixed-point guarantee, but does not imply every finite run of the naive surrogate fails.Source: §A.2 · p. 19

A.3 LIGHTNING WEAVE: JOINT TARGETS FOR CAPABILITY COMPOSITION

  • For Lightning Weave, the loss equals the negative joint regularized objective on fixed cached states up to a student-policy-independent constant.The equivalence holds over any common policy class; one anchor recovers Tilted-Target DOPD, while zero anchor weights recover the behavior policy.Source: §A.3 · p. 20

A.4 STRUCTURAL PROPERTIES OF THE JOINT TARGET

The joint target combines anchor shifts through weighted odds, with explicit invariances that do not imply order-invariant sequential training. The appendix also bounds the formal guarantees and describes the implemented models, data, and cached trajectories.

A.4 STRUCTURAL PROPERTIES OF THE JOINT TARGET

  • When positively weighted anchors raise one action’s score relative to another, the joint target raises their odds relative to the behavior policy; disagreement is resolved by the signed weighted sum.The target is invariant to simultaneous permutation of anchors and weights, state-only score baselines, and proportional rescaling of weights and the KL coefficient.Source: §A.4 · p. 20
  • These invariances apply to the constructed target, not sequential training, where stages can change the model, visited states, and approximation error.The finite-candidate squared-log-residual surrogate retains zero loss and gradient at its finite target but not the full-distribution KL objective equivalence.Source: §A.4 · p. 21
  • The guarantees assume finite aligned scores, full action distributions, and fixed cached states; they do not ensure exact target representation, optimizer success, or coverage after large policy changes.Finite-support scoring and cross-tokenizer alignment are practical approximations to the analyzed full-distribution setup.Source: §A.4 · p. 21

B.1 MODELS, ANCHOR PAIRS, AND TRAINING DATA

  • Experiments use five student models and three anchor pairs; Klear and MiMo-RL supply accuracy-oriented shifts, while DECS supplies an efficiency-oriented shift.The described supplementary setup distinguishes the full-distribution method from its finite-candidate implementation.Source: §B.1 · p. 21, B SUPPLEMENTARY EXPERIMENTAL DETAILS · p. 21
  • Mathematics and code use separate 3,200-prompt datasets, and four initial-student responses per prompt are cached for reuse throughout training.The code prompts are screened against LiveCodeBench v6 for overlap; all anchors score the same cached trajectories.Source: §B.1 · p. 21

B.2 TRAINING HYPERPARAMETERS

Training replays cached trajectories for two epochs, using shared run settings for single- and multi-anchor configurations and separate checkpoints for mathematics and code.

B.2 TRAINING HYPERPARAMETERS

  • Each run replays 12,800 cached trajectories for two epochs, totaling 400 optimizer updates, with mathematics and code checkpoints trained separately.The runs use 100 rounds of 256 trajectories and four updates per round at a global batch size of 64.Source: §B.2 · p. 21
  • Single- and multi-anchor runs share the hyperparameters summarized in Table 5.The caption identifies the table’s scope as shared training hyperparameters for both configurations.Source: §B.2 · p. 21, Table 5 · p. 22
  • All runs use α = 2.0, and multi-anchor weights sum to one; equal composition assigns wi = 1/K, so a single weight of one recovers the corresponding single-anchor configuration.That endpoint uses the same target and training configuration as the single-anchor run.Source: §B.2 · p. 22

B.3 TRAINING COST

Lightning Weave reduces measured GPU-hour costs and avoids concurrent anchor serving during student training. Its cross-tokenizer alignment and finite-candidate target construction define how cached shifts are scored and fitted.

B.3 TRAINING COST

  • The cost comparison uses four mathematics settings, 100 rounds of 256 trajectories, and eight H100 GPUs for all GPU stages.Offline training replays 12,800 unique trajectories for two epochs.Source: §B.3 · p. 22
  • 1.70–2.18× lower GPU-hour cost per composition across four Qwen3 settings, with sequential anchor scoring followed by student-only training.Both online and offline pipelines score each anchor pair; the reduction need not grow with anchor count.Source: §B.3 · p. 22, Table 6 · p. 22
  • 4.58× projected amortized cost reduction for five Qwen3-4B compositions, assuming unchanged training cost per composition.Reusing cached trajectories, student-reference scores, and anchor shifts makes additional configurations student-only runs.Source: §B.3 · p. 22

B.4 CROSS-TOKENIZER POLICY-SHIFT ALIGNMENT

  • Student prefixes are losslessly projected into each anchor tokenizer; added tokens are accepted only when decoding and re-encoding preserves their text.This provides a tokenizer-independent view of the text prefix, including when a student token decomposes into multiple anchor tokens.Source: §B.4 · p. 22
  • A student candidate is scored only when it maps to one anchor token for every cached candidate; invalid boundaries are masked without fallback scores.The multi-anchor target intersects validity masks, and 93.8–100% of generated positions survive in the reported Klear–DECS experiments.Source: §B.4 · p. 23
  • The finite-candidate target preserves behavior mass for 16 cached candidates and aggregates remaining vocabulary mass as “other,” then fits the target with a squared log-residual surrogate.The surrogate is a finite-support approximation to the full-vocabulary KL, with value and gradient vanishing when the student matches the target.Source: §B.4 · p. 23

B.5 BENCHMARKS AND TEST SETS

The evaluation covers competition mathematics and recent competitive-programming problems, using task-specific sampling and answer checks. Reported metrics aggregate accuracy and response length across benchmarks, with implementation and training settings specified.

B.5 BENCHMARKS AND TEST SETS

  • Mathematics evaluation uses AIME 2024, AIME 2025, and HMMT February 2025, each with 30 competition-level problems.The problems cover algebra, geometry, combinatorics, and number theory.Source: §B.5 · p. 23
  • LiveCodeBench v5 and v6 provide 279 and 131 recently released problems, respectively, and all methods use the same problem sets.The recent-problem filters reduce the risk of test-set contamination.Source: §B.5 · p. 23

B.6 EVALUATION PROMPTS AND ANSWER EXTRACTION

  • Mathematics evaluation samples 64 responses per problem and scores the last balanced boxed answer by exact comparison, without computer-algebra equivalence.Length-truncated responses are scored when they contain a valid extracted answer.Source: §B.6 · p. 24
  • Code evaluation uses LiveCodeBench’s generation format, samples four responses per problem, and counts a sample correct only when its extracted program passes every test.Runtime errors, timeouts, extraction failures, and incomplete programs score as incorrect.Source: §B.6 · p. 24

B.7 METRICS AND AGGREGATION

  • Code accuracy is Avg@4, not pass@4; mean response length counts generated tokens over the same samples, and reported averages are unweighted across five benchmarks.AES compares each trained policy with its corresponding base student and rewards accuracy gains and token savings while penalizing accuracy loss more strongly than an equal relative gain.Source: §B.7 · p. 24

B.8 IMPLEMENTATION DETAILS

  • Training uses one eight-GPU node with bfloat16; each configuration is trained once, so repeated evaluation generations measure decoding variability rather than training-seed variance.The reported checkpoint is the final checkpoint of each run.Source: §B.8 · p. 25

C ADDITIONAL EXPERIMENTAL RESULTS

Additional experiments examine controllable Klear–DECS trade-offs, transfer of MiMo-RL compositions across students, and response-level examples. Together, they show how composition behaves across operating points and individual cases.

C ADDITIONAL EXPERIMENTAL RESULTS

C.1 KLEAR–DECS OPERATING POINTS

  • Increasing ρDECS generally shortens responses, while lower and intermediate fractions preserve higher accuracy across the Qwen3-4B mathematics and code sweep.The five standardized compositions make the fraction a controllable accuracy–efficiency axis.Source: §C.1 · p. 25, Table 7 · p. 25

C.2 GENERALIZATION WITH MIMO-RL

  • On Qwen3.5-4B, Klear+MiMo-RL improves average accuracy by 7.56 points with 20.0% fewer tokens, while DECS+MiMo-RL improves it by 7.17 points with 19.4% fewer tokens.On Qwen3-1.7B, DECS+MiMo-RL improves average accuracy over the base by 2.08 points with 7.3% fewer tokens; Klear+MiMo-RL has the highest average accuracy and AES.Source: §C.2 · pp. 25–26, Table 8 · p. 25

D CASE STUDIES

  • The case studies select mathematics examples from shared Qwen3-4B AIME 2024 evaluation and a code example from LCB v5, using fixed criteria before inspecting responses.The examples illustrate aggregate trends rather than establish causal effects.Source: D CASE STUDIES · p. 26

D.1 COMBINING COMPLEMENTARY ACCURACY SIGNALS

  • In two complementary accuracy cases, MiMo-RL alone solves Case A and Klear-only solves Case B, while Klear+MiMo-RL solves both.The composition preserves behavior contributed by either accuracy-oriented anchor.Source: §D.1 · p. 26, Figure 4 · p. 26

D.2 COMBINING ACCURACY AND EFFICIENCY

  • In mathematics and code examples, Klear-only is correct but long, DECS-only is shorter but incorrect, and Klear+DECS is correct and shortest.The figure covers one case in each domain.Source: §D.2 · p. 27, Figure 5 · p. 27

E LIMITATIONS AND FUTURE WORK

The authors identify limits in model scale, anchor availability, and efficiency gains, and plan broader model and task evaluations.

E LIMITATIONS AND FUTURE WORK

  • Experiments cover models below 10B parameters, so whether the benefits persist at larger scales is not established.The authors plan evaluation on larger models and more architectures.Source: E LIMITATIONS AND FUTURE WORK · p. 27
  • Capability shifts require both a post-trained model with the desired capability and its corresponding pre-anchor, which may be unavailable for new tasks.The planned extension to agentic reasoning and tool-use tasks is future work.Source: E LIMITATIONS AND FUTURE WORK · p. 27
  • Lightning Weave improves the accuracy–efficiency balance, but its efficiency gains remain smaller than those of methods optimized exclusively for efficient reasoning.Source: E LIMITATIONS AND FUTURE WORK · p. 27
Loading 2609.14708v2…