PaperlayerView exact revision on arXiv

Source-linked AI summary

Preserving Mathematical Reasoning in Compressed Diffusion Language Models via Trajectory-Aware Low-Rank Approximation

Tian Liang, Zishan Shao, Yiran Chen

arXiv:2610.03326v1cs.AI

TL;DR

dLLM compression calibrates on clean activations even though inference uses partially masked intermediate states, leaving open how calibration states affect low-rank optimality and reasoning preservation. The paper defines a trajectory-aware objective and estimates it with Traj-MC. Under matched compression settings, trajectory-aware calibration improves generation-state reconstruction and preserves more mathematical reasoning than clean-state calibration.

  • Problem

    dLLM low-rank compression typically calibrates on clean, fully visible activations despite inference traversing partially masked states, raising questions about trajectory-state optimality and reasoning preservation.

    Source: §1 · p. 1

  • Method

    The paper formulates a trajectory-aware low-rank objective over corruption levels and masking realizations and estimates its second moment with Monte Carlo sampling in Traj-MC.

    Source: §1 · p. 1, §3.1 · p. 2, §5 · p. 4

  • Results

    Under matched compression settings, trajectory-aware calibration preserves more mathematical reasoning and improves generation-state reconstruction, while knowledge and commonsense changes are smaller or mixed.

    Source: §4.2 · pp. 3–4, §5 · p. 5

  • Takeaways & Limitations

    Low-rank compression quality in dLLMs depends on the state distribution defining reconstruction geometry, making calibration-state design a component of compression.

    Source: §5 · p. 5

    The experiments focus on activation-aware low-rank compression with SVD-LLM and layerwise reconstruction, leaving other solvers and model families for future work.

    Source: C Limitations and Discussion · p. 8

Abstract

from arXiv · show

Diffusion language model (dLLM) compression faces a known challenge because calibration is typically performed on clean, fully visible activations, whereas inference traverses partially masked intermediate states. For low-rank compression, this raises two questions. First, can low-rank optimality still be characterized when approximation quality is measured over trajectory-distributed states, and second, does the choice of calibration states affect mathematical reasoning preservation under compression? We address these questions by formulating a trajectory-aware low-rank objective over corruption levels and masking realizations. To estimate this objective efficiently, we propose Traj-MC, which estimates the trajectory second moment through Monte Carlo sampling and yields exact sampled-state optimality and population consistency. Under matched compression budgets, trajectory-aware calibration improves reconstruction over the generation trajectory and preserves substantially more mathematical reasoning than clean calibration on mathematical reasoning benchmarks. Our results connect trajectory-aware low-rank optimality to the reasoning capability retained after dLLM compression. Our code is available at: https://github.com/Zishan-Shao/traj-mc.git.

1 Introduction

dLLM low-rank compression calibrates on clean activations even though inference traverses partially masked states, raising questions about trajectory-state optimality and reasoning preservation. The paper formulates a trajectory-aware objective and reports improved generation-state reconstruction and mathematical reasoning preservation under matched compression.

1 Introduction

  • Clean, fully visible calibration differs from the partially observed states traversed during iterative dLLM denoising.This mismatch motivates asking whether low-rank optimality can be characterized over trajectory states and whether calibration states affect reasoning preservation.Source: §1 · p. 1
  • At 20% compression, trajectory-aware calibration preserves substantially higher mathematical reasoning performance, with similar understanding accuracy.Source: Figure 1 · p. 2, §1 · p. 2
  • The paper formulates a trajectory-aware low-rank objective, establishes its fixed-rank optimum, and develops Traj-MC with consistency guarantees.The objective is defined over trajectory states, and the introduction reports improved matching to the generation-state second moment.Source: §1 · pp. 1–2

2 Related Work

Prior activation-aware low-rank methods use calibration activations, while diffusion compression research recognizes mismatch between calibration and iterative-generation states. This paper asks which state distribution should define low-rank approximation for dLLMs.

2 Related Work

  • Activation-aware low-rank methods use calibration activations, and their approximation quality depends on the activations defining reconstruction.ASVD uses activation statistics, SVD-LLM uses whitening to align truncation with activation reconstruction, and MoDeGPT extends reconstruction across coupled Transformer modules (2; 3; 1).Source: §2 · p. 2
  • Diffusion quantization addresses timestep-aware calibration and dLLM quantization models masked denoising states, while this work studies their implications for low-rank optimality (7; 9; 8; 6).The paper derives a trajectory-aware objective and fixed-rank solution, approximates it with Traj-MC, and evaluates mathematical reasoning preservation.Source: §2 · p. 2

3 Methodology

The methodology defines low-rank reconstruction over corrupted trajectory states, derives an exact sampled-state solution, and estimates the population objective with Traj-MC.

3.1 Trajectory-aware low-rank optimality

  • The trajectory-aware objective integrates reconstruction across examples, corruption levels, and masking realizations using a single rank-k operator.Corrupted calibration windows provide a tractable surrogate for partially observed denoising states.Source: §3.1 · pp. 2–3
  • For sampled states with positive-definite second moment, Cholesky whitening reduces reconstruction to rank-k SVD, yielding an exact sampled-state optimum without independence assumptions.The identity expresses sampled reconstruction error as the Frobenius norm of the whitened residual.Source: §3.1 · p. 3

3.2 Traj-MC: Monte Carlo approximation

  • Traj-MC samples corruption levels uniformly and independently masks tokens conditional on each level, then accumulates the second moment and applies the existing compression procedure.It avoids dense-model rollouts and leaves compressed layers and rank allocation unchanged; the implementation uses the whitening-based SVD solver of SVD-LLM (3).Source: §3 · p. 2, §3.2 · p. 3
  • Under the stated sampling assumptions and singular-value gap, the estimated rank-k operator converges to a global minimizer of trajectory reconstruction error.The estimator is unbiased under a finite fourth-moment condition on activations.Source: §3.2 · p. 3

4 Experiments

Under matched compression settings, trajectory-aware calibration improves mathematical reasoning preservation and generation-state reconstruction. The experiments also test uniform corruption mixing and the Monte Carlo approximation’s behavior.

4.1 Experimental Setup

  • The evaluation compares LLaDA-8B-Instruct and Dream-7B-Instruct at 20% and 40% parameter reduction, using matched settings and the same C4 calibration windows.SVD-LLM uses clean-state activations, Traj-MC uses trajectory-masked states, and Vanilla Weight SVD is a no-activation reference.Source: §4.1 · p. 3

4.2 Main Results

  • Across GSM8K, MATH-500, SVAMP, and ARC-C, Traj-MC consistently outperforms clean-state calibration on both dLLMs, with advantages persisting under more aggressive compression.Knowledge and commonsense benchmarks show relatively minor or mixed changes.Source: §4.2 · p. 3, Table 1 · p. 4
  • Traj-MC consistently lowers generation-state error compared with clean-state calibration using only corrupted calibration windows.The reported gains in generation-state reconstruction coincide with downstream gains.Source: §4.2 · p. 4
  • Uniform sampling across corruption levels achieves lower generation-state error than every fixed-t alternative, without selecting or tuning a specific corruption level.Increasing a fixed mask probability also progressively improves generation-state reconstruction.Source: §4.2 · p. 4, Figure 2 · p. 4
  • The estimated second moment approaches its reference value at approximately the expected N^-1/2 rate, while improved covariance estimation accompanies smaller low-rank solution deviations.Figure 2 reports the estimation rate and the relationship between covariance error and subspace deviation.Source: §4.2 · p. 4, Figure 2 · p. 4

5 Conclusion

The conclusion presents Traj-MC as a practical approximation to an exact trajectory-distributed low-rank objective. Under matched settings, trajectory-aware calibration improves generation-state reconstruction and preserves more mathematical reasoning than clean-state calibration.

5 Conclusion

  • The method formulates low-rank dLLM compression over trajectory-distributed states and derives the corresponding exact fixed-rank optimum.Source: §5 · p. 4
  • Traj-MC approximates the population objective with Monte Carlo sampling without requiring dense generation rollouts.Source: §5 · p. 4
  • Under matched settings, trajectory-aware calibration improves held-out generation-state reconstruction and preserves substantially more mathematical reasoning than clean-state calibration.Changes on knowledge and commonsense tasks are comparatively smaller or mixed.Source: §5 · p. 5

A.1 Exact optimality for sampled trajectory states

The proof applies a rank-preserving change of variables and the Eckart–Young theorem to characterize the minimum reconstruction error. Its argument is deterministic conditional on the collected states.

A.1 Exact optimality for sampled trajectory states

  • The proof introduces normalized, whitened activations Z_N as part of its optimality argument.Source: §A.1 · p. 6
  • Because S_N is invertible, the change of variables A = f W S_N preserves rank, allowing the Eckart–Young theorem to establish Eq. (3) and the minimum reconstruction error.Source: §A.1 · p. 6
  • The proof’s argument is deterministic conditional on the collected states and does not require independence across windows or tokens.Source: §A.1 · p. 6

A.2 Monte Carlo variance decay

The appendix gives a variance-decay result for moment estimation, a consistency argument for the trajectory-optimal solution, and the rationale for uniform corruption mixing across mask counts.

A.2 Monte Carlo variance decay

  • For i.i.d. window tuples, the mean-squared Frobenius estimation error decays as N^-1 and its root mean square as N^-1/2, without requiring independence among tokens within a window.Source: §A.2 · p. 6

A.3 Consistency of the trajectory-optimal solution

  • A strict spectral gap makes the best rank-k truncation continuous at the population solution, supporting consistency of the trajectory-optimal solution.The argument also invokes eventual positive definiteness and continuity of Cholesky factorization and inversion.Source: §A.3 · pp. 6–7

A.4 Why Uniform Corruption Mixing?

  • Uniform corruption mixing gives every mask count equal marginal probability and makes all subsets of a given size equally likely.It therefore assigns equal weight to mask-count-conditioned reconstruction risks rather than concentrating the objective around one masking severity.Source: §A.4 · p. 7

B Experimental Details

Experiments hold calibration windows, compressed layers, and rank allocation fixed while varying activation states, then assess generation-state reconstruction and moment and subspace diagnostics.

B Experimental Details

  • Generation-state reconstruction uses 256 held-out GSM8K prompts, one intermediate state per prompt, sampler steps 0–255 tiled across prompts, and generation-region error.When confidence intervals are reported, they use 2,000 paired bootstrap resamples over prompts.Source: B Experimental Details · p. 7
  • Fig. 1(a) projects normalized second-moment mismatches onto a Gram–Schmidt basis, placing generation at the origin while preserving all three pairwise distances.Its annotations report normalized Frobenius mismatch; smaller values indicate closer second-moment agreement.Source: B Experimental Details · p. 7
  • For Fig. 2(b), moment error uses an independently sampled reference with seed 31337; panel (c) instead uses a FullGrid16 reference over 16 midpoint masking levels and two realizations.The FullGrid16 reference averages over 256 fixed C4 windows.Source: B Experimental Details · p. 8
  • The subspace comparison uses the retained row-space projector with k = 1638, averages distances over two modules, and reports means and one standard deviation across five mask seeds.The protocol keeps calibration windows, compressed layers, and rank allocation fixed; clean calibration uses visible windows, while Traj-MC samples t uniformly and masks tokens independently with probability t.Source: B Experimental Details · pp. 7–8

C Limitations and Discussion

The paper’s experiments focus on activation-aware low-rank compression with SVD-LLM, while its discussion presents trajectory-aware calibration as a broader design direction for iterative generative models.

C Limitations and Discussion

  • Experiments focus on activation-aware low-rank compression instantiated with SVD-LLM; extending the principle to other low-rank solvers and model families remains future work.The calibration uses uniform corruption levels and independent token masking; more adaptive distributions may improve sample efficiency or alignment for particular dLLMs.Source: C Limitations and Discussion · p. 8
  • The discussion frames trajectory-aware calibration as a design paradigm for post-training optimization in iterative generative models, aligning optimization with intermediate-state dynamics rather than clean inputs.It suggests efficient deployment should jointly consider the generation process, compression objective, and execution substrate.Source: C Limitations and Discussion · p. 8

Related papers

Loading 2610.03326v1…