BiGain: Unified Token Compression for Joint Generation and Classification

AI summary

Overview

Research area: Computer vision — efficient diffusion models, specifically training-free token compression for diffusion backbones used simultaneously for image generation and image classification.

Technical level: Intermediate. The paper assumes familiarity with diffusion models, self-attention, and token-merging baselines such as ToMe and ToDo, but its core argument (preserve both frequency bands rather than only optimizing synthesis) is explained in accessible terms.

Scope: The paper introduces BiGain, a training-free, plug-and-play framework of two frequency-aware operators that compress tokens in accelerated diffusion models while improving diffusion-based classification and maintaining or improving generation quality.

What This Paper Is About

Diffusion-model acceleration methods such as token merging and token downsampling are almost always evaluated on synthesis quality alone — keeping FID stable while cutting compute — which ignores the fact that the same backbone is routinely reused for recognition, either through linear probes, feature distillation, or diffusion-classifier protocols. In practice, the authors observe that accelerations which barely hurt generation can sharply degrade classification, because naive compression removes exactly the high-contrast, edge- and texture-bearing structures that recognition depends on. The goal of BiGain is to reframe token compression as a joint, bi-objective problem and to solve it with a training-free framework that preserves generative fidelity while markedly improving discriminative performance.

Key Contributions

1. Reframing token compression as a bi-objective problem. BiGain treats compression for diffusion models as a joint optimization that must safeguard generative fidelity and discriminative utility, and it offers a practical, training-free solution for it.
2. The first framework to jointly study and advance both tasks under acceleration. To the authors' knowledge, BiGain is the first framework to jointly study and advance both generation and classification under acceleration of diffusion models.
3. Two frequency-aware operators. Laplacian-gated token merging (L-GTM) merges spectrally smooth tokens while discouraging merges of high-contrast tokens, and Interpolate-Extrapolate KV-Downsampling (IE-KVD) downsamples keys and values via a controllable interpolation/extrapolation between nearest and average pooling while keeping queries intact.
4. Design guidance for a frequency-aware regime. Beyond throughput and recognition gains, the study provides a practical design rule — balanced spectral retention — for future compression in deployable, dual-purpose generative models.

Main Findings

• Synthesis-tuned compression is much harder on classification than on generation. On Oxford-IIIT Pets at roughly 10% FLOPs reduction (Table 1), the no-acceleration baseline is 81.03% accuracy; ToMe drops 8.07 points (72.96%), SiTo drops 12.19 points (68.84%), DiP-GO drops 4.50 points (76.53%), and MosaicDiff drops 3.65 points (77.38%). BiGain's Laplacian-gated merging retains 78.38% (a 2.65-point drop), which the authors describe as cutting the loss by 27 to 78% relative to these methods at matched FLOPs.

• Under downsampling at 14.2% FLOPs, BiGain is the best overall on Pets. ToDo slightly decreases accuracy (−1.88 points, to 79.15%), while BiGain's Interpolate-Extrapolate KV-downsampling reaches 79.90% (only a 1.13-point drop).

• Gains grow with compression aggressiveness. On ImageNet-1K with a 70% token merging ratio on Stable Diffusion 2.0, BiGain increases classification accuracy by 7.15% while also improving generation FID by 0.34, which the paper reports as 1.85%.

• The token-merging advantage widens as the merge ratio rises. In the Stable Diffusion 2.0 token-merging comparison across 10% to 70% merge ratios, the reported accuracy deltas over ToMe climb from 1.06 to 8.87 points on Pets, 0.34 to 4.90 on ImageNet-100, and 1.75 to 7.15 on ImageNet-1K, with FID also lower at virtually every ratio.

• DiT-XL/2 behaves very differently from U-Net backbones. On DiT-XL/2 token downsampling (ImageNet-100), ToDo performs very poorly — dropping to 8.46, 4.74, and 3.32 accuracy at 3x, 4x, and 5x — while BiGain remains more stable on this backbone. With relatively small downsampling (2x), the method can match or slightly surpass the unaccelerated model in both classification and generation.

• Reduction should be applied to self-attention only. The location ablation on Stable Diffusion 2.0 (Pets) shows self-attention-only reduction consistently delivers the best quality-efficiency trade-off: it preserves prompt adherence that cross-attention compression degrades and avoids compounding bias through MLP compression. At merge ratio 0.7, BiGain's SA-only setting scores 74.63 versus 73.89 for SA+CA, 68.27 for SA+MLP, and 68.25 for SA+CA+MLP.

• Local frequency cues beat global and spectral alternatives for scoring tokens. Laplacian Filter with an ℓ1 reduction is best at all merge ratios, outperforming global statistics (norms, channel variance), DFT spectral measures, and cosine similarity by 0.3 to 1.9%. At merge ratio 0.7 on Pets it reaches 74.63, versus 72.96 for global mean deviation, 73.04 for channel variance, 73.75 for DFT spectral centroid, and 74.00 for cosine to neighbors.

• Faster variants trade a little accuracy for more FLOP savings. Measured at merge ratio r = 0.7 on Stable Diffusion 2.0, Laplacian Gated Merge costs 704.99 GFLOPs, Cached Assignment Merge 698.88, and Adaptive Block Merge 695.08, while all three closely track each other across 10 to 70% merge ratios on Pets, ImageNet-100, ImageNet-1K, and COCO-2017.

• Frequency awareness is necessary, not incidental. The authors report that removing Laplacian gating disproportionately hurts classification, and that downsampling KV in the frequency-balanced way is helpful for generation. Visualization on Stable Diffusion 2.0 at a 90% merge ratio at denoising step t = 200 shows Laplacian-gated merging preserves more class-discriminative structure (for example, a cat's edges) than standard ToMe.

Methodology in Plain English

The central insight is frequency separation: when intermediate features are mapped into a frequency-aware representation, fine high-frequency detail (edges, textures) separates from low- and mid-frequency content (shapes, layouts, semantics). The design rule that follows is balanced spectral retention — keep the high-frequency components that anchor recognition while maintaining the low and mid bands that support coherent generation.

BiGain implements this with two training-free operators:

1. Laplacian-gated token merging (L-GTM). Hidden-state tokens are reshaped into a spatial grid and scored by the magnitude of a 2-D Laplacian filter (a 3x3 kernel with 1s on the four neighbors and −4 in the center), aggregated across channels. In each spatial cell, the lowest-scoring (spectrally smoothest) tokens become destination anchors and the rest become sources; a bipartite matching then merges the top r% most similar source-destination pairs by equal-weight averaging. Merges are encouraged among flat regions and discouraged among high-contrast tokens, so edges and textures survive. When merging applies to queries, keys, and values, attention cost drops from O(N²d) to O(N′²d). Class tokens in DiT and time/text conditioning tokens in U-Net cross-attention are never touched. A tiled variant, Adaptive Block Merging (ABM), pools an s×s block only when its maximum Laplacian score falls below a quantile threshold τ, and serves as a drop-in replacement at high-resolution stages.

2. Interpolate-Extrapolate KV-Downsampling (IE-KVD). Given a stride s and a reduced grid, each site's output is a convex-style combination of nearest-neighbor pooling and average pooling controlled by a factor α. Keys and values are downsampled this way while queries remain at full resolution, so attention costs O(N Ñ d) but the output length stays N. Keeping queries intact preserves fine-grained receptive fields for every output token, which stabilizes synthesis and retains the edge/texture cues that per-token attention precision matters for in diffusion classification.

Why this works with diffusion classifiers. The diffusion classifier's decision rule draws a shared Monte Carlo set of (t_s, ε_s) pairs for all classes and picks the class whose average squared noise-prediction error is lowest. BiGain's operators are timestep-local, deterministic given the input, and rely on no cross-timestep caching, so every class receives identical (t_s, ε_s) samples and identical compression schedules. The paired-difference estimator therefore stays valid, while per-class FLOPs fall and accuracy rises. The paper also uses staged pruning of classes: computation is allocated over N_stages rounds with cumulative budgets (T₁, …, T_Nstages) and keep-counts (K₁, …, K_Nstages), retaining only the lowest-scoring classes at each stage, which reduces wall-clock compute without changing the final argmin decision.

Experimental setup. Two backbones are tested with official pretrained weights: Stable Diffusion v2.0 (U-Net latent diffusion with text conditioning) and DiT-XL/2 (Transformer backbone). Classification is evaluated on ImageNet-1K, ImageNet-100, Oxford-IIIT Pets, and COCO2017, using a 2,000-image validation subset for ImageNet-1K and full validation splits elsewhere. Top-1 accuracy is reported for single-label datasets; Top-1 precision and macro mAP for multi-label COCO; FID for generation on COCO2017 captions, ImageNet-100, and ImageNet-1K. DiT-XL/2 is evaluated only on ImageNet datasets (class-index conditioning, no free-form prompts). The same token-reduction policy is applied at every denoising step, with no cached merge pairings or pooling indices across timesteps. For Stable Diffusion 2.0 classification, the interextrapolation factor is fixed at 0.9 across all timesteps for stability, while for generation it is varied linearly from 0.8 in early steps to 1.2 in later steps; for DiT-XL/2 it is fixed at 0.1 across all timesteps for both tasks.

Why This Matters

Impact on research. The paper challenges the default evaluation protocol for diffusion acceleration, which treats FID (or perceptual quality) as the only target. It shows that a single diffusion backbone is a dual-purpose system, and that acceleration policies optimized for synthesis alone systematically discard the high-frequency structure that recognition needs. The frequency-separation lens — and the empirical finding that Laplacian scoring of local contrast is the strongest token-salience heuristic tested — offers a reusable design principle for future compression, pruning, and distillation work on diffusion backbones.

Real-world applications (these are the domains the paper itself cites as already using a single diffusion backbone for both tasks):

• Medical imaging: class-conditional diffusion models used for diagnostic prediction alongside counterfactual or uncertainty-aware reconstruction within one generative prior.
• Safety-critical perception: robust diffusion classifiers that use the same denoiser for likelihood-based classification and image generation.
• Industrial visual inspection: diffusion backbones supporting both defect identification and defect reconstruction.
• Remote sensing: diffusion models performing cloud-removal or super-resolution synthesis while also improving land-cover and object classification.

Industry relevance. Because BiGain is training-free and plug-and-play, it can be inserted into existing inference pipelines for pretrained checkpoints without retraining or fine-tuning. The claimed benefit — lower per-class FLOPs and higher accuracy relative to synthesis-only baselines, across both DiT and U-Net architectures — directly targets the cost of deploying dual-purpose generative systems, where a single model must serve both an image-generation endpoint and a recognition endpoint.

Future Directions

• How far can balanced spectral retention scale? The authors establish the principle as a reliable design rule but leave open how finely the high/low/mid frequency budget should be tuned per task, per timestep, and per backbone.
• Closing the remaining gap at aggressive compression. At the highest merge ratios the method still degrades noticeably (for example on DiT-XL/2 at 70% merging), raising the question of whether stronger frequency-aware operators or learned schedules can extend the usable compression range.
• Beyond the two tested backbones and four datasets. The paper covers DiT-XL/2 and Stable Diffusion v2.0 across ImageNet-1K, ImageNet-100, Oxford-IIIT Pets, and COCO2017; whether the frequency-separation principle transfers to video diffusion, other transformer diffusion families, and larger conditioning vocabularies is untested here.
• Unifying with orthogonal accelerations. The paper positions its intra-step token compression as orthogonal to step-reduction methods such as DDIM, DPM-Solver, and progressive distillation, but does not report experiments combining the two families.

Target Audience

Researchers and engineers working on efficient diffusion inference, token reduction, and model acceleration who care about more than FID. It is also relevant to practitioners building dual-purpose generative systems that must classify as well as synthesize, to those studying diffusion classifiers and diffusion-based feature extraction, and to readers interested in frequency-domain interpretations of transformer and U-Net representations. Readers without background in diffusion sampling or attention mechanics will need to consult the cited baselines (ToMe, ToDo) first, since the paper's comparisons are framed against them.

Authors’ abstract

Acceleration methods for diffusion models (e.g., token merging or downsampling) typically optimize synthesis quality under reduced compute, yet often ignore discriminative capacity. We revisit token compression with a joint objective and present BiGain, a training-free, plug-and-play framework that preserves generation quality while improving classification in accelerated diffusion models. Our key insight is frequency separation: mapping feature-space signals into a frequency-aware representation disentangles fine detail from global semantics, enabling compression that respects both generative fidelity and discriminative utility. BiGain reflects this principle with two frequency-aware operators: (1) Laplacian-gated token merging, which encourages merges among spectrally smooth tokens while discouraging merges of high-contrast tokens, thereby retaining edges and textures; and (2) Interpolate-Extrapolate KV Downsampling, which downsamples keys/values via a controllable interextrapolation between nearest and average pooling while keeping queries intact, thereby conserving attention precision. Across DiT- and U-Net-based backbones and ImageNet-1K, ImageNet-100, Oxford-IIIT Pets, and COCO-2017, our operators consistently improve the speed-accuracy trade-off for diffusion-based classification, while maintaining or enhancing generation quality under comparable acceleration. For instance, on ImageNet-1K, with 70% token merging on Stable Diffusion 2.0, BiGain increases classification accuracy by 7.15% while improving FID by 0.34 (1.85%). Our analyses indicate that balanced spectral retention, preserving high-frequency detail and low/mid-frequency semantics, is a reliable design rule for token compression in diffusion models. To our knowledge, BiGain is the first framework to jointly study and advance both generation and classification under accelerated diffusion, supporting lower-cost deployment.