Frequency-Aware Token Compression: Reconciling Generation and Discrimination in Diffusion Models

Diffusion models have evolved from pure generative tools into dual-purpose computational substrates. In medical imaging workflows, the same backbone that synthesizes counterfactual lesions for data augmentation also provides diagnostic predictions through feature probing. In autonomous perception systems, diffusion-based classifiers leverage the exact same denoising network used for generative augmentation to provide robust safety-critical classifications. This functional convergence creates a tension that current acceleration methodologies have largely ignored. Most training-free compression techniques, including token merging and spatial downsampling, optimize exclusively for synthesis fidelity measured by FID or perceptual metrics. Yet as Liu et al. demonstrate in BiGain: Unified Token Compression for Joint Generation and Classification, such myopic optimization systematically destroys the high-frequency signal components essential for discriminative performance. Their work introduces a frequency-aware compression framework that challenges the assumption that generation quality and classification accuracy must trade off under computational constraints.

The Spectral Bias of Dual-Use Models

The central observation driving BiGain is that generative and discriminative objectives exhibit complementary, rather than overlapping, spectral dependencies. Image generation fundamentally requires coherent low and mid-frequency structures; global layout, color consistency, and semantic scene composition depend on smooth spatial variations. In contrast, visual recognition relies disproportionately on high-frequency content: edge discontinuities, textural micro-patterns, and high-contrast boundaries that distinguish object categories. When token compression methods apply uniform similarity metrics, such as cosine similarity on raw feature vectors, they conflate these spectral regimes. The result is a compression schedule that preserves sufficient low-frequency energy to maintain plausible generation while decimating the high-frequency cues that anchor classification.

This spectral asymmetry explains the empirical findings in the paper's motivating experiments. On COCO-2017 and ImageNet-100, baseline token merging methods maintain generation quality (FID) even at aggressive compression ratios while classification accuracy collapses precipitously. At extreme sparsity, accuracy can degrade to near chance levels while generated images remain visually coherent. This divergence reveals that standard compression techniques treat high-frequency information as redundant noise rather than task-critical signal. BiGain addresses this by formalizing compression as a frequency separation problem, mapping hidden states into a representation that disentagles fine detail from global semantics before any reduction operation occurs.

Technical Mechanisms: From Principle to Practice

BiGain implements frequency-aware compression through two complementary operators that operate without retraining or architectural modification. The first, Laplacian-gated Token Merging (L-TM), replaces similarity-based merging with a spectral selection criterion. The method applies a Laplacian filter to hidden-state tokens to compute local frequency scores, quantifying the magnitude of spatial variation within each token's receptive field. Tokens exhibiting low Laplacian magnitude represent spectrally smooth regions suitable for merging, while high-magnitude tokens indicate edge or texture content that must be preserved.

The merging process proceeds through a constrained bipartite matching. Within each spatial stride, the lowest-scoring (smoothest) token is designated as a destination, while remaining tokens form the source set. A global bipartite matching then selects top source-destination pairs based on feature similarity, but with the critical constraint that high-frequency sources are strongly discouraged from merging. This preserves the fine-grained spatial structure that classification relies upon while still achieving the computational savings of reduced token counts. Unlike aggressive uniform merging, which homogenizes local detail into smoothed averages, L-TM maintains a sparse set of high-frequency anchors that retain discriminative capacity.

The second operator, Interpolate-Extrapolate KV Downsampling (IE-KV), addresses the attention bottleneck in diffusion transformers. Standard downsampling of keys and values in the attention mechanism typically employs either nearest-neighbor selection (preserving detail but potentially noisy) or average pooling (smooth but detail-obliterating). IE-KV introduces a controllable interpolation between these extremes, allowing the model to extrapolate beyond simple averaging when high-frequency preservation is critical. Importantly, queries remain at full resolution, maintaining attention precision at the output layer while compressing the KV cache. This asymmetry proves crucial; by keeping queries uncompressed, the model retains the capacity to attend to the fine-grained details preserved in the partially compressed key-value representations.

Empirical Validation and the Spectral Budget

The authors evaluate BiGain across both Diffusion Transformer (DiT) and U-Net architectures on ImageNet-1K, ImageNet-100, Oxford-IIIT Pets, and COCO-2017. The results substantiate the frequency separation hypothesis. On ImageNet-1K using Stable Diffusion 2.0 with a 70% token merging ratio, BiGain increases classification accuracy by 7.15% over baseline token merging while simultaneously improving FID by 0.34 points (a 1.85% relative improvement). This dual improvement is remarkable; it demonstrates that respecting spectral structure does not merely mitigate the accuracy degradation of compression but can actually enhance generative fidelity by preventing the smooth, over-merged artifacts that plague aggressive compression.

The consistency across datasets suggests a general design rule: balanced spectral retention. Compression methods should explicitly budget computational resources to preserve high-frequency components alongside low-frequency semantic content. The failure modes of baseline methods on COCO-2017, where classification accuracy collapses entirely under high compression while generation remains viable, illustrate the cost of ignoring this budget. BiGain's ability to maintain both metrics indicates that much of what standard compression discards as "redundant" is actually task-critical information, merely filtered through the wrong spectral lens.

Limitations and Broader Implications

Despite its efficacy, BiGain operates within the constraints of training-free acceleration, which limits its capacity to reshape the underlying feature space. The Laplacian filter, while mathematically principled, represents a hand-crafted prior on what constitutes "important" frequency content. Certain semantic categories might rely on subtle high-frequency patterns that differ from the edge-detection assumptions inherent in Laplacian filtering, suggesting that learned spectral decomposition could yield further gains at the cost of requiring fine-tuning. Additionally, computing Laplacian scores across hidden states introduces marginal overhead, though the authors demonstrate this is amortized by the subsequent reduction in attention computation.

The framework's implications extend beyond diffusion models. Any vision system employing attention over spatial tokens, from video transformers to multimodal large language models, faces analogous compression pressures. The insight that different downstream tasks occupy distinct spectral niches suggests that future compression schemes should be task-adaptive or multi-objective by default. The current dichotomy between "generation-optimized" and "classification-optimized" compression reflects an artificial separation that unified frameworks like BiGain can bridge.

From a signal processing perspective, BiGain reconnects deep learning compression with classical image processing foundations. The explicit use of frequency-domain analysis within the latent space of a neural network represents a promising hybrid approach, combining the representational power of learned features with the interpretability and structural guarantees of traditional filtering. This suggests opportunities for incorporating other harmonic analysis tools, such as wavelet decompositions or steerable filters, into token reduction pipelines.

Concluding Thoughts

BiGain establishes that token compression in generative models need not be a zero-sum game between synthesis quality and discriminative utility. By recognizing that generation and classification occupy different, but compatible, spectral regimes, the framework achieves simultaneous acceleration for both tasks. This opens several avenues for future inquiry. Can we learn the frequency separation end-to-end, training models with explicitly factorized spectral pathways that are inherently compressible? How does the spectral budget extend to video diffusion, where temporal coherence introduces additional frequency dimensions? And can similar principles apply to the compression of non-visual modalities where "frequency" must be defined in abstract feature spaces?

As diffusion models increasingly serve as general-purpose visual processors rather than single-purpose generators, compression methods must evolve correspondingly. BiGain provides both the technical machinery and the conceptual framework for this evolution, demonstrating that respecting the spectral structure of visual information is not merely an engineering optimization but a fundamental requirement for building efficient, dual-capable intelligent systems.