PaperlayerView exact revision on arXiv

Source-linked AI summary

What does FFN compression change downstream? Same-state causal restoration in diffusion language models

Shaurya Omar

arXiv:2609.31685v1cs.LGcs.AI

TL;DR

Local approximation error does not show which removed FFN computations change a DLM’s downstream denoising trajectory. SSR restores the original FFN on the compressed model’s current input and follows the resulting continuation. Across two DLM families, state and correction structure add information beyond local magnitude; under aggressive LLaDA, a frozen window recovers most of the accuracy gap while retaining estimated MAC savings.

  • Problem

    Local FFN approximation error does not directly measure the downstream trajectory effect of restoring removed computation on the states a compressed DLM actually visits.

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

  • Method

    SSR restores the original FFN on the current compressed-trajectory branch state and follows the resulting continuation, using task-label-free source effects to freeze a restoration window.

    Source: §3.1 · p. 4, §3.3 · p. 6

  • Results

    Across two DLM families, compressed-side state adds ranking information beyond local magnitude; under aggressive LLaDA, SSR recovers 89.9% of the dense–compressed accuracy gap with an estimated 36.8% whole-model MAC saving.

    Source: §1 · p. 1, §5 · p. 9

  • Takeaways & Limitations

    DLM compression should distinguish local mismatch, downstream consequence, and task repair rather than treating them as interchangeable objectives.

    Source: §5 · p. 9

    Evaluation covers two DLM families, deterministic 32-transition length-32 generation, and mainly multiple-choice endpoints; MAC savings are analytical estimates, not measured latency.

    Source: §5 · p. 9

Abstract

from arXiv · show

Diffusion language models (DLMs) enable flexible, parallel generation, but their iterative denoising remains computationally expensive, motivating increasingly aggressive compression. Existing compression objectives largely measure how well compressed computation approximates the original locally, but local error does not reveal which removed computations actually matter to the downstream denoising trajectory. We introduce Same-State Causal Restoration (SSR), which restores the original FFN on the exact current input reached by the compressed model and measures how the resulting trajectory changes. To our knowledge, this is the first direct measurement of the same-current-input closed-loop effect of removed FFN computation in DLM compression. Across LLaDA-8B-Instruct and Dream-v0-Instruct-7B, compressed-side state ranks this downstream effect substantially better than local NMSE at fixed denoising phase, while controlled interventions show that correction structure matters beyond magnitude. Using task-label-free calibration, SSR freezes a single restoration window for held-out inference. Under aggressive LLaDA compression, restoring only four transitions recovers 89.9% of the lost accuracy while retaining an estimated 36.8% whole-model MAC saving and outperforming an equal-budget local-error baseline. Dream further shows that restoring dense behavior and repairing the final task are distinct outcomes.

1 INTRODUCTION

The introduction identifies a gap in measuring how restoring compression-removed FFN computation changes the trajectory actually visited by a compressed DLM. SSR evaluates that restoration on the same current input; a frozen window recovers much of aggressive LLaDA’s accuracy loss while retaining estimated compute savings.

1 INTRODUCTION

  • Prior work had not directly measured the downstream effect of reinstating the exact FFN computation removed by compression on the compressed model’s visited trajectory.Iterative FFN compression can change both current computation and the states used in later denoising transitions.Source: §1 · p. 1
  • SSR compares original and compressed FFNs on the same current restoration-branch input, isolating operator substitution from trajectory-induced input shift.Source: §1 · p. 1
  • 89.9% of the dense–compressed accuracy gap is recovered with an estimated 36.8% whole-model MAC saving, exceeding the equal-horizon local-NMSE/final-four schedule by 5.05 pp.The window is globally frozen before held-out evaluation.Source: §1 · p. 1

2 Freeze

SSR distinguishes immediate local mismatch from the downstream change caused by restoring an FFN on the current branch, and from whether that change repairs the final task. It uses task-label-free source effects to freeze one global restoration window for held-out evaluation.

2 Freeze

  • Local error, downstream effect, and task repair are distinct; at fixed phase, compressed-side state ranks measured effect better than local NMSE, and norm-matched interventions test correction structure beyond magnitude.Source: §2 · p. 2
  • SSR measures the continuation after restoring the original FFN on the compressed trajectory’s current endogenous branch state.The figure describes substituting the original FFN on the same current input without copying a dense hidden state.Source: §2 · p. 2, Figure 1 · p. 2
  • Across the two DLM families, task-label-free source effects freeze one global restoration window; aggressive LLaDA shows strong task repair, while Dream improves dense fidelity without a conclusive task effect.Source: §2 · p. 2

2 RELATED WORK AND NOVELTY BOUNDARY

Related work studies DLM efficiency and causal interventions, but the paper locates its novelty in measuring the effect of reinstating the exact compression-removed FFN on the current compressed-trajectory input. It also uses task-label-free source effects to freeze a selective window for held-out execution.

2 RELATED WORK AND NOVELTY BOUNDARY

  • DLM efficiency work covers timestep-dependent masking and quantization, adaptive computation, and non-uniform layer or activation structure, but does not measure the downstream effect of reinstating the exact removed computation.The passage cites Quant-dLLM, DLLMQuant, and Quantization Meets dLLMs (Zhang et al., 2026a; Lin et al., 2026); adaptive methods (Zhu et al., 2026; Oba et al., 2026); and structure studies (Conzelmann et al., 2026; Szatkowski et al., 2026).Source: §2 · p. 3
  • DLM causal tracing restores clean-run hidden states or MLP contributions, or edits temporally localized states, rather than evaluating the removed FFN on the current compressed/restored branch.The passage names Park & Jo (2026) and Yao et al. (2026).Source: §2 · p. 3
  • The stated novelty is the same-current-input closed-loop effect of the exact compressor-removed FFN, paired with task-label-free source effects that freeze a restoration window for held-out execution.The table caption frames SSR against neighboring work on DLM compression, adaptive computation, and causal intervention.Source: §2 · p. 3, Table 1 · p. 3

3 SAME-STATE CAUSAL FFN RESTORATION

SSR restores original FFN computation on the current compressed/restoration-branch input, then measures how the continuation changes. The section also tests information beyond local error magnitude and describes how a restoration window is selected and transferred to held-out prompts.

3.1 MEASURE: DOWNSTREAM EFFECT ON THE COMPRESSED TRAJECTORY

  • SSR evaluates original and compressed FFNs on the same current input, isolating operator substitution while allowing later states to diverge.The same-state control applies at each treated event; a restoration window is a joint intervention because earlier restorations can change later inputs.Source: §3.1 · p. 4
  • The restoration and compressed continuations start from one saved compressed snapshot, with no dense hidden state copied and non-FFN computation unchanged.Original and compressed FFNs share the residual-stream interface; representation-changing compressors requiring an alignment map are outside this estimand.Source: §3.1 · pp. 4–5

3.2 WHAT LOCAL ERROR MAGNITUDE MISSES: STATE AND CORRECTION STRUCTURE

  • After denoising phase and other listed factors are fixed, compressed-side state observables provide ranking information about downstream effect beyond local NMSE.The observables are computed from the compressed branch; this is controlled predictive evidence, not a causal manipulation of the observable.Source: §3.2 · p. 5
  • Norm-matched feature-permuted corrections produce different downstream outcomes, showing that correction magnitude alone is insufficient.The frozen cyclic permutation preserves the correction's L2 norm while changing its feature-space direction or alignment; this does not identify a complete structural causal model.Source: §3.2 · p. 5

3.3 CLOSE THE LOOP: FREEZE AND TRANSFER A RESTORATION WINDOW

  • Task-label-free source measurements select one global restoration window, which is frozen before held-out evaluation without per-prompt search or refitting.The protocol uses separate source sets for screening, confirmation, and horizon selection; task correctness is not used before the freeze.Source: §3.3 · pp. 5–6
  • Frozen execution uses original FFNs only inside the selected window, retaining original weights; it measures selective operator execution, not parameter compression or latency speedup.For LLaDA, the equal-horizon local-NMSE comparator is also the final-four-transition schedule.Source: §3.3 · p. 6
  • Full-dense fidelity and task repair are separate endpoints, because a positive downstream effect need not cross the final task-correctness boundary.The evaluation separately reports parser validity and conditional repair of prompts that were dense-correct but compressed-incorrect.Source: §3.3 · p. 6

4 EXPERIMENTS

The experiments test whether compressed-side state and correction structure reveal downstream effects beyond local error, then evaluate frozen restoration windows on task repair and dense fidelity. Under aggressive LLaDA compression, a four-transition SSR window recovers most of the accuracy gap while retaining estimated compute savings.

4.1 SETUP AND FROZEN PROTOCOL

  • Frozen deterministic 32-transition, length-32, temperature-0 tests span LLaDA and Dream, with disjoint source cohorts and 297 held-out prompts.Compression construction is held fixed to isolate the causal question; savings are analytical whole-model MAC estimates, not measured latency, memory, or parameter storage.Source: §4.1 · pp. 6–7

4.2 EVIDENCE BEYOND LOCAL ERROR MAGNITUDE

  • At fixed phase, negative release margin ranks downstream effect better than local NMSE: LLaDA ρ=.456 vs −.041 and Dream .522 vs −.103.A fresh LLaDA ARC-Challenge confirmation gave ρ=.500 (n=32); the authors characterize this as predictive evidence, not a causal intervention on release margin.Source: Figure 2 · p. 7, §4.2 · p. 7
  • Norm-matched controls favor original-FFN restoration in 26/32 LLaDA prompts with zero reversals, and 20/32 Dream prompts with 3 reversals.Dream task specificity is null (2/28/2); the controls show magnitude alone is insufficient but do not identify a complete structural causal model. Predictor generalization remains weak, so the authors make no router claim.Source: Figure 2 · p. 7, §4.2 · pp. 7–8

4.3 FROZEN CLOSED-LOOP OUTCOMES: FULL-DENSE FIDELITY AND TASK REPAIR SEPARATE

  • 89.9% of the dense–compressed accuracy gap is recovered by c24 SSR, with 36.8% estimated whole-model MAC saving.SSR reaches 54.9% accuracy versus 21.9% for compressed and 49.8% for the equal-horizon local-NMSE/final-four schedule, a +5.05 pp advantage.Source: Figure 3 · p. 8, §4.3 · p. 8
  • At c14, SSR and matched local-NMSE both reach 70.8% prompt-level correctness across 72 prompts, marking a task-equivalent moderate-compression boundary.Parsed labels differ once when both methods are wrong; partial restoration above dense accuracy is not evidence of superiority over the original model.Source: §4.3 · p. 8
  • Dream c21 task effect is inconclusive (+2.02 pp, 95% CI [−.34,4.38]), while final-sequence token agreement improves +1.68 pp (95% CI [.74,2.75]).Trajectory agreement AUC improves +.0082; at near-matched MAC, SSR exceeds NMSE by +1.43 pp token agreement and +.0080 AUC.Source: §4.3 · p. 8

4.4 CLAIM-CRITICAL CONTROLS AND ROBUSTNESS

  • Claim-critical controls challenge fixed-phase/local-NMSE and magnitude-only accounts, while c24 beats literal final-four at identical estimated MAC.Selection is frozen before held-out evaluation; the all-40 rescore does not equalize calibration budgets.Source: §4.4 · p. 9

5 DISCUSSION, LIMITATIONS, AND CONCLUSION

The conclusion distinguishes local FFN mismatch, downstream consequences, and task repair, while presenting SSR as a diagnostic rather than a dense-free deployment selector. The section also bounds the evidence to two DLM families and specified evaluation settings.

5 DISCUSSION, LIMITATIONS, AND CONCLUSION

  • The evaluation covers two DLM families, deterministic 32-transition generation at length 32, and mainly multiple-choice endpoints; long-form generation and dense-free selection remain open.Cross-compressor generalization also remains open, and original FFN weights stay resident; compute savings are analytical MAC estimates rather than latency, memory, or parameter-compression results.Source: §5 · p. 9
  • SSR measures downstream consequences that local FFN approximation error does not capture, and the paper argues that compression objectives should distinguish mismatch, downstream effect, and task repair.The diagnostic still uses original FFNs and same-snapshot dense continuations, so it is not a dense-free deployment selector.Source: §5 · p. 9

REPRODUCIBILITY STATEMENT

  • The appendix records model revisions, compression schedules, generation and population construction, and window candidates, with released manifests, raw predictions, and recomputation notebooks.Source cohorts use ARC-Challenge, ARC-Easy, SciQ, PIQA, and HellaSwag, while MMLU remains unopened during selection.Source: REPRODUCIBILITY STATEMENT · p. 10, §A.1 · p. 11

AI USE

  • The authors report using generative AI across literature search, experimental work, interpretation, and manuscript preparation, and say the authors checked AI-assisted outputs against project artifacts and source literature.The authors state that experiments were executed by them and that they take responsibility for the final content.Source: AI USE · p. 10

A.1 MODELS, COMPRESSION, AND FROZEN SCHEDULES

  • Calibration uses disjoint cohorts and leaves MMLU unopened until windows and source-side predictors are frozen; LLaDA and Dream use fixed compression schedules and transition-zero anchors.The LLaDA and Dream schedules use different frozen SSR windows, as detailed in the appendix.Source: §A.1 · p. 11

A.2 STATISTICS AND ENDPOINTS

  • Prompt-level paired inference uses bootstrap intervals and exact tests, while full-dense token and trajectory agreement are fidelity endpoints distinct from final multiple-choice correctness.The released-aware trajectory variant evaluates agreement on positions already released by the sampler at each transition.Source: §A.2 · p. 11

B HELD-OUT CLOSED-LOOP RESULTS

  • Table 5 covers Dream SSR versus NMSE at essentially matched whole-model MAC cost on 297 prompts, while Table 6 covers LLaDA c24 SSR versus equal-horizon NMSE on the frozen 297-prompt population.The supplied captions identify table scope but do not provide outcome values.Source: Table 5 · p. 12, Table 6 · p. 12

C.1 STATE RANKING

At fixed phase, compressed-side state observables add ranking information beyond local NMSE, while held-out prediction remains insufficiently calibrated for a learned router. Norm-matched interventions test whether correction structure matters beyond magnitude.

C.1 STATE RANKING

  • Under one common equal-dataset macro rule, state-versus-NMSE correlations are .456 vs. −.041 for LLaDA and .506 vs. −.185 for Dream, preserving the state>NMSE ordering across families.The authors claim directional consistency, not equality of absolute correlations.Source: §C.1 · p. 12
  • On 32 fresh LLaDA ARC-Challenge prompts, state ranking at the frozen setting yielded Spearman ρ = .500, one-sided permutation p = .0022, and 95% CI [.190,.743].The transition, horizon, remaining masks, compressed layer set, and negative release margin were fixed before outcomes.Source: §C.1 · p. 12
  • Negative or weak leave-dataset-out R2 supports ranking but not calibrated deployment, so the paper makes no learned-router claim.Compressed-side observables exclude dataset identity and gold answers.Source: §C.1 · p. 12

C.2 NORM-MATCHED CORRECTION STRUCTURE

  • The correction-structure control holds correction norm and treated locations fixed while changing feature-space direction or alignment through orthogonal permutations or sign flips.The table caption identifies counts as original-FFN restoration better, tie, or structure-control better, but supplies no counts here.Source: §C.2 · p. 12, Table 7 · p. 12

D CLAIM-CRITICAL ABLATIONS AND SENSITIVITIES

Sensitivity checks report that the LLaDA NMSE selector’s chosen start is unchanged with a larger broad-screen subset, while the authors caution that the checks do not equalize calibration budgets.

D CLAIM-CRITICAL ABLATIONS AND SENSITIVITIES

  • Table 8 covers horizon and placement controls, but its caption alone supplies no control outcomes.The sensitivity checks do not establish a fully equal calibration budget between SSR and NMSE.Source: Table 8 · p. 13, D CLAIM-CRITICAL ABLATIONS AND… · p. 13
  • Recomputing the identical energy-weighted selector on all 40 broad-screen prompts leaves t = 28, H = 4 as the winner, with macro NMSE .12107 versus .11693 at t = 27 and .11360 at t = 26.The original executed NMSE baseline used a frozen six-prompt-per-dataset subset, with the SSR horizon fixed at four transitions.Source: D CLAIM-CRITICAL ABLATIONS AND… · p. 13
  • A post-hoc check on the disjoint 40-prompt confirmation cohort ranks t = 28 above t = 27, but cannot reproduce the exact energy-weighted aggregation because per-prompt SSE and energy components were not archived.The reported per-prompt NMSE values are .11457 at t = 28 and .10894 at t = 27.Source: D CLAIM-CRITICAL ABLATIONS AND… · p. 13

E FAILURE TYPES AND COMPUTE ACCOUNTING

The compute accounting is an analytical whole-model MAC estimate rather than a measurement of runtime or storage, and the supplied tables identify accounting and outcome-transition scopes without reporting their cells.

E FAILURE TYPES AND COMPUTE ACCOUNTING

  • Table 9 reports compressed-to-SSR outcome transitions on matched populations, but its caption alone does not establish transition counts or outcomes.Source: Table 9 · p. 13
  • Table 10 is scoped to estimated whole-model MAC saving under one common accounting convention, but the supplied caption provides no values.The estimates are not measurements of latency or memory.Source: Table 10 · p. 13, E FAILURE TYPES AND COMPUTE ACCOUNTING · p. 13
  • The ledger estimates arithmetic for an optimized direct-dispatch implementation, not wall-clock latency, memory, parameter storage, or hardware FLOPs; original FFN weights remain resident.Offline calibration is excluded from per-inference savings, while anchor construction and frozen restoration are included in the accounting convention.Source: E FAILURE TYPES AND COMPUTE ACCOUNTING · p. 13

F ARTIFACT AND REPRODUCIBILITY BOUNDARY

The section documents terminology mappings, the frozen LLaDA evaluation population, and checks supporting the source-to-freeze chronology.

F ARTIFACT AND REPRODUCIBILITY BOUNDARY

  • Historical CFR, global_cfr, and CFRController names refer to SSR, while historical step refers to a denoising transition.The archived final_token_agreement quantity is mean token equality across all generated output positions, not agreement of only the last token.Source: F ARTIFACT AND REPRODUCIBILITY BOUNDARY · p. 14
  • The 12-hour execution budget reduced the LLaDA final manifest from 594 prompts to 297 before inference, without inspecting outcomes or changing the evaluation policies.All headline LLaDA statistics use the frozen 297-prompt population.Source: F ARTIFACT AND REPRODUCIBILITY BOUNDARY · p. 14
  • Matched Stage59 and Stage59B artifacts, disabled selection during cross-family evaluation, and identity-audited populations support the source→confirm→horizon→freeze chronology.Stage60 disables temporal, layer, channel, width, predictor, and test-set selection.Source: F ARTIFACT AND REPRODUCIBILITY BOUNDARY · p. 14
Loading 2609.31685v1…