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

Shaurya Omar

cs.LGSep 16, 2026

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

Authors: Shaurya Omar

Organizations: Indian Institute of Technology, Roorkee

Abstract

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.

Figures & tables

Figure 1: Same-State Causal Restoration: measure and freeze. From one compressed snapshot, exact SSR substitutes the original FFN on the same current branch input and follows the resulting continuation; the local FFN difference is only a local-error signal. Task-label-free source

U_{\mathrm{dense}}

effects select one global window, frozen before held-out evaluation. No dense hidden state is copied.

Work	Object / signal / intervention	Boundary relative to SSR
Quant-dLLM; DLLMQuant; Lin et al.	DLM compression / sensitivity-aware PTQ	optimizes or characterizes compression; no exact removed-FFN continuation effect
SureLock	converged-token compute / local posterior stability	skips query + FFN; no compressed-vs-original FFN contrast or exact restoration
Layer Collapse; Szatkowski et al.	FFN/layer compressibility, redundancy, activation sparsity	structural diagnostics; not causal downstream-effect measurement
Park & Jo	DLM factual localization / clean-run state or MLP restoration	transplanted mediator; not compressor-removed operator on current branch
TimeROME-DLM	temporal DLM causal tracing / editing	causal DLM state intervention; not compression-specific operator substitution
SSR	exact original FFN on current compressed/restored branch	same-current-input signed effect + frozen task-label-free transfer

Table 1: Closest conceptual neighbors and the intervention boundary relative to SSR . Existing work covers DLM compression, adaptive computation, or causal intervention; SSR targets the compression-specific same-current-input continuation effect.

Family / regime	Compressed FFNs	SSR $t/H$	NMSE $t/H$	Comp. $\rightarrow$ SSR saving
LLaDA c24	24/32	27/4	28/4	42.3 $\rightarrow$ 36.8%
LLaDA c14	14/32	27/4	28/4	24.7 $\rightarrow$ 21.5%
Dream c21	21/28	17/2	30/2	42.1 $\rightarrow$ 39.3%

Table 2: Frozen evaluation regimes.

t/H

denotes start transition / horizon. The last column reports estimated whole-model MAC saving relative to dense, from compressed-only to SSR .

Figure 2: Local error magnitude omits state and correction-structure information. Left: at fixed phase, negative release margin ranks downstream effect more strongly than local NMSE . Right: exact restoration is compared with a norm-matched structure control. State evidence is predictive; the structure test is interventional. Details are in Appendix C .

Figure 3: Closed-loop outcomes separate task repair from full-dense fidelity. (a) Frozen task accuracy. (b) Repair among dense-correct/compressed-wrong opportunities. (c) Dream token-agreement gains with paired 95% CIs. Aggressive c24 recovers 89.9% of the dense–compressed accuracy gap while retaining 36.8% estimated whole-model MAC saving; c14 is a task-equivalent boundary, and Dream improves fidelity without a reliable task effect.

Regime ( $n$ )	Dense	Comp.	SSR	NMSE	SSR $-$ Comp. [95% CI]	SSR $-$ NMSE [95% CI]
LLaDA c24 (297)	58.6	21.9	54.9	49.8	+33.0 [27.6,38.4]	+5.05 [2.36,8.08]
LLaDA c14 (72)	68.1	38.9	70.8	70.8	+31.9 [22.2,43.1]	0.00
Dream c21 (297)	83.8	70.7	72.7	71.7	+2.02 [ $-.34$ ,4.38]	+1.01 [ $-2.02$ ,4.04]

Table 3: Closed-loop task accuracy under frozen equal-horizon schedules. Accuracies and effects are pooled over prompts; effects are percentage points. The prespecified equal-dataset macro c24 effect is in the text; MAC savings are in Table 2 .

Question	Controlled comparison	Result
Does phase/local error explain ranking?	fixed phase, masks, layers, horizon; state vs. local NMSE	$\rho$ : LLaDA.456 vs. $-.041$ ; Dream.522 vs. $-.103$
Would any same-norm correction work?	exact vs. norm-matched cyclic permutation; sign flip in App. C	LLaDA $U_{\mathrm{dense}}$ 26/6/0; Dream trajectory AUC 20/9/3; task 2/28/2
Is c24 just a late-window / more-compute effect?	SSR 27–30 vs. NMSE /final-four 28–31; equal $H=4$ , identical est. MAC	SSR +5.05 $\,\mathrm{pp}$ (18–3, $p=.00149$ ); all-40 NMSE rescore keeps 28/4
Is one transition enough?	ARC-only, fixed $t=27$ ; $H=1$ vs. 4	1/12 vs. 7/12 repairs; six $H=4$ -only discordances, $p=.03125$
Could parsing create the Dream result?	parser-independent sequence/trajectory endpoints	token agreement +1.68 $\,\mathrm{pp}$ [.74,2.75]; trajectory AUC +.0082
Could selection leak held-out labels?	task-label-free objective; disjoint splits; MMLU unopened	frozen before held-out task evaluation; no per-prompt search/refit

Table 4: Claim-critical controls and robustness. Rows are ordered by the simplest reviewer alternatives; Appendix D gives additional sensitivities.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Endpoint	Effect	95% paired CI
Task correctness	.0101	[ $-.0202$ ,.0404]
Validity	.0236	[.0000,.0505]
Final-sequence token agreement to full dense run	.0143	[.0046,.0249]
Trajectory agreement AUC to full dense run	.0080	[.0052,.0109]
Released-aware trajectory agreement AUC	.0095	[.0064,.0129]

Appendix

Table 5: Dream SSR versus NMSE at essentially matched whole-model MAC cost on 297 prompts.

Endpoint	SSR $-$ NMSE	Paired evidence
Task accuracy	+5.05 $\,\mathrm{pp}$	18 vs. 3 discordances; exact $p=.00149$ ; bootstrap CI [2.36,8.08] $\,\mathrm{pp}$
Whole-model MAC saving	0.00 $\,\mathrm{pp}$	both 36.76% under the common ledger

Appendix

Table 6: LLaDA c24 SSR versus equal-horizon NMSE on the frozen 297-prompt population.

Family / endpoint	Counts	Mean exact $-$ control
LLaDA same-snapshot downstream effect	26/6/0	.0713
LLaDA task correctness	12/20/0	.3750 indicator
Dream trajectory agreement AUC	20/9/3	.0529
Dream final-sequence token agreement	5/23/4	.0410
Dream task correctness	2/28/2	.0000

Appendix

Table 7: Correction-structure specificity. Counts are original-FFN restoration better / tie / structure-control better.

Comparison	Result	Interpretation
LLaDA $H=4$ vs. $H=1$ , ARC-sensitive failures ( $n=12$ )	7 vs. 1 rescues; 6 $H=4$ -only discordances; exact $p=.03125$	one restored transition is insufficient on this ARC cohort
Exact vs. permuted correction	positive exact-control downstream effect	extra correction norm alone does not explain the effect
Broad late $-$ early LLaDA macro	positive point estimate; CI crosses zero	do not universalize a late-transition preference
Dream task specificity	2/28/2	correction structure need not change final task correctness

Appendix

Table 8: Horizon and placement controls.

Transition	LLaDA c24	Dream c21
Invalid $\rightarrow$ correct	98	7
Valid-wrong $\rightarrow$ correct	1	2
Correct $\rightarrow$ wrong	0	2
Correct $\rightarrow$ invalid	1	1
Invalid $\rightarrow$ valid-wrong	26	1
Correct $\rightarrow$ correct	64	207

Appendix

Table 9: Compressed

\rightarrow

SSR outcome transitions on matched populations.

Family / condition	Mean saving
LLaDA c24 compressed	42.26%
LLaDA c24 SSR	36.76%
LLaDA c24 NMSE	36.76%
LLaDA c14 SSR	21.45%
Dream c21 compressed	42.08%
Dream c21 SSR	39.29%

Appendix

Table 10: Estimated whole-model MAC saving under one common accounting convention.

LLM Compression

Large language model compression reduces the memory, compute, and context costs of LLMs while aiming to preserve their capabilities.

See all 273 papers

Prompt Compression in Diffusion Large Language Models: Evaluating LLMLingua-2 on LLaDA
May 18, 2026Sterling Huang, Abigayle Brown, Jiyoo Noh +4LLM Evaluation LLM Compression
University of Toronto Toronto, Canada · King Mongkut’s University of Technology Thonburi Bangkok, Thailand
Preserving Mathematical Reasoning in Compressed Diffusion Language Models via Trajectory-Aware Low-Rank Approximation
Oct 2, 2026Tian Liang, Zishan Shao, Yiran ChenDiffusion Model Compression Low-Rank Compression
Department of Electrical and Computer Engineering, Duke University
Layer Collapse in Diffusion Language Models
May 7, 2026Alexander Conzelmann, Albert Catalan-Tatjer, Shiwei LiuLLM Compression Diffusion Language Models
Tübingen AI Center · Max Planck Institute for Intelligent Systems · ELLIS Institute Tübingen

May 18, 2026J/K move · Enter open · S save

Prompt Compression in Diffusion Large Language Models: Evaluating LLMLingua-2 on LLaDA

Sterling Huang, Abigayle Brown, Jiyoo Noh, Jiakang Xu+3

University of Toronto Toronto, Canada · King Mongkut’s University of Technology Thonburi Bangkok, Thailand

LLM Evaluation LLM Compression Diffusion Language Models Prompt Compression

Open paper AskPDF arXiv

Prompt compression reduces inference cost and context length in large language models, but prior evaluations focus mainly on autoregressive architectures. This study examines whether LLMLingua-2 transfers effectively to diffusion large language models (DLLMs), specifically LLaDA-8B-Instruct. We evaluate GSM8K, DUC2004, and ShareGPT using 250 prompts per dataset at an approximate 50% compression ratio, covering mathematical reasoning, prompt reconstruction, and summarization. Outputs from original, compressed, and reconstructed prompts are compared using exact-match accuracy, BLEU, ROUGE, and BERTScore. Results show that high semantic preservation does not necessarily ensure stable downstream behavior in diffusion models. Summarization remains relatively robust, while mathematical reasoning degrades substantially despite high semantic similarity. Reconstruction further shows that semantically similar prompts may omit reasoning-critical information needed for stable denoising. Overall, compression failures are mainly driven by information omission rather than semantic drift, suggesting that autoregressive prompt compression methods may not transfer uniformly to DLLMs. These findings motivate diffusion-aware compression strategies.

Abstract

LLM Compression

Explore similar work

Prompt Compression in Diffusion Large Language Models: Evaluating LLMLingua-2 on LLaDA

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

Layer Collapse in Diffusion Language Models