Structured State-Space Regularization for Compact and Generation-Friendly Image Tokenization

1. Problem Statement

Modern generative vision models (e.g., Diffusion models, World models) increasingly rely on image tokenizers to map high-dimensional pixel data into a lower-dimensional latent space. An ideal latent space must satisfy two often conflicting objectives:

1. Compactness: The representation must be information-preserving yet efficient, compressing essential content while discarding redundant details.
2. Generation-Friendliness: The latent space must be easy to model by generative frameworks. Recent studies suggest that latent spaces supporting "diffusability" (smooth progression from noise to data) perform better.

The Core Challenge: Existing tokenizers often struggle to balance these goals. Strong reconstruction fidelity (capturing every pixel detail) can lead to latent spaces that are difficult for diffusion models to synthesize, while optimizing purely for generation can degrade reconstruction quality. The paper posits that the lack of frequency awareness in standard tokenizers is a key bottleneck.

2. Methodology: Structured State-Space Regularization

The authors propose a novel regularization technique that aligns the dynamics of image tokenizers with State-Space Models (SSMs), specifically leveraging the "frequency awareness" inherent in SSMs.

A. Theoretical Foundation: Generalized SSMs

The paper reframes SSMs not just as sequence models, but as a combination of two components:

1. Basis Projection ($c(\cdot)$): Compressing an input signal into coefficients using orthogonal basis functions (e.g., Fourier, Polynomials).
2. Input Transformation ($\theta(\cdot)$): A transformation applied to the input over time (e.g., token concatenation in 1D SSMs).

The authors argue that SSMs achieve frequency awareness because their hidden states mimic the dynamics of basis coefficients under specific transformations.

B. The Proposed Framework

To transfer this property to 2D image tokenizers, the authors define a new pair of $(c, \theta)$:

• Basis Projection ($c$): Projection onto 2D orthogonal basis functions (e.g., 2D Fourier bases). This decomposes the image into frequency components.
• Input Transformation ($\theta$): A progressive Gaussian blurring process. Blurring naturally removes high-frequency components, creating a smooth trajectory from sharp to blurry images.

Derivation of Dynamics:
By analyzing the heat equation (which governs Gaussian blur), the authors derive a differential equation describing how the basis coefficients evolve as an image blurs. This yields a state transition matrix $\mathbf{A}$ such that:
$$\frac{d}{d\tau} c(I_\tau) = \mathbf{A} c(I_\tau)$$
where $I_\tau$ is the image at blur level $\tau$.

C. Regularization Loss

The tokenizer encoder $\mathcal{E}$ is regularized to mimic these coefficient dynamics.

1. Process: Given an input image $I$, two versions are sampled at different blur levels ($I_{\tau_1}$ and $I_{\tau_2}$).
2. Latent Update: The latent representation of the sharper image, $z_1 = \mathcal{E}(I_{\tau_1})$, is updated via the learned transition matrix: $\hat{z}_2 = \bar{\mathbf{A}} z_1$.
3. Objective: The model minimizes the distance between the predicted latent $\hat{z}_2$ and the actual encoded latent $z_2 = \mathcal{E}(I_{\tau_2})$.
4. Stabilization: To prevent latent norms from collapsing to zero (as blur increases), the authors apply spatial mean-centering to the latent features before regularization.

The total loss includes the standard reconstruction loss plus this Structured State-Space Regularization (SSR) loss:
$$\mathcal{L} = \lambda_z d_z(z_2, \hat{z}_2) + \lambda_I d_I(I_{\tau_2}, \mathcal{D}(\hat{z}_2)) + \mathcal{L}_{recon}$$

3. Key Contributions

1. Novel Regularizer: Introduction of a theoretically grounded regularizer that forces tokenizers to learn latent spaces with explicit frequency-aware dynamics.
2. SSM-Image Tokenization Bridge: The first framework to extend SSM principles (specifically basis projection and input transformation) from 1D sequences to 2D image tokenization.
3. Efficient Implementation: The method uses pre-computed or learnable transition matrices (e.g., Fourier-based $\mathbf{A}$) that are sparse and efficient, requiring minimal computational overhead during training.
4. Empirical Validation: Demonstrated effectiveness on widely used tokenizers (Flux and Cosmos) without sacrificing reconstruction quality.

4. Experimental Results

The method was evaluated on ImageNet-1K using Flux (8x downsampling) and Cosmos (16x downsampling) tokenizers, followed by training a Latent Diffusion Transformer (DiT-B/2).

• Generation Quality: The proposed method significantly outperformed baselines (EQVAE, FT-SE) in generation metrics.
  • Flux + Ours: Improved gFID from 16.08 (baseline) to 13.33 and IS from 89.32 to 106.61.
  • Cosmos + Ours: Improved gFID from 13.52 to 11.20 and IS from 112.48 to 128.91.
• Reconstruction Fidelity: The method maintained competitive reconstruction quality (PSNR, LPIPS, SSIM), showing only marginal degradation compared to the baseline, and in some cases (Cosmos), even improved reconstruction metrics.
• Ablation Studies:
  • Basis Choice: Fourier bases yielded the best results compared to Chebyshev, Legendre, and Hermite polynomials.
  • Regularization Strength: An optimal strength ($\alpha=0.25$) was found; too high a strength degraded reconstruction, while too low failed to improve generation.
  • Mean-Centering: Removing the spatial mean-centering function caused latent norm inflation and a collapse in generation performance (gFID jumped from 13.33 to 17.23).

Qualitative Findings:

• Latent Dynamics: Visualizations showed that regularized latent channels evolve smoothly and exponentially as the image blurs, mirroring the theoretical behavior of Fourier coefficients.
• Frequency Ordering: Channels naturally organized by frequency. Unmasking channels from low-to-high frequency allowed the decoder to reconstruct images progressively from coarse shapes to fine textures, a property absent in non-regularized tokenizers.

5. Significance

This work provides a new paradigm for designing image tokenizers. By explicitly enforcing frequency-aware inductive biases through the lens of State-Space Models, the authors demonstrate that it is possible to create latent spaces that are simultaneously compact (efficiently encoding structure) and generation-friendly (aligning with the spectral synthesis nature of diffusion models).

The approach bridges the gap between control theory/SSMs and computer vision, suggesting that the mathematical properties of SSMs (specifically their ability to handle long-range dependencies and frequency decomposition) are highly beneficial for visual representation learning. This offers a promising direction for future generative AI systems, potentially reducing the need for complex architectural changes by simply improving the latent space geometry.