Structured State-Space Regularization for Generation-Friendly Image Tokenization

Jinsung Lee; Jaemin Oh; Namhun Kim; Dongwon Kim; Byung-Jun Yoon; Suha Kwak

cs.CVApr 13, 2026

Structured State-Space Regularization for Generation-Friendly Image Tokenization

Authors: Jinsung Lee, Jaemin Oh, Namhun Kim, Dongwon Kim, Byung-Jun Yoon, Suha Kwak

Organizations: POSTECH · Brown University · KAIST · Texas A&M University ·

Abstract

Image tokenizers play a central role in modern generative models, where the structure of the latent space critically determines the downstream generation performance. A key but underexplored property of effective latent representations is spectral organization, the ability to encode information across frequency components. In this work, we introduce structured state-space regularization, a principled approach to inducing spectral structure in latent spaces. We derive a regularization objective by revisiting state-space models (SSMs) as systems mimicking a basis function's behavior. This perspective reveals that hidden states of SSMs are induced to capture the frequency components, resulting in a novel regularizer that enforces the latent space to capture spectral structure of images. Experiments demonstrate that our regularizer improves the generative performance of image tokenizers while incurring only minimal loss in their reconstruction fidelity.

Figures & tables

Panel of Figure 1 — Figure 1: Different choice of $\bigl(\mathbf{c}(\cdot),\theta(\cdot)\bigr)$ and the resulting update rules. Existing SSMs update their hidden state based on Eq. ( 4 ), which is derived from the coefficient dynamics based on the choice (a) . One can choose a different combination of $\mathbf{c}(\cdot)$ and $\theta(\cdot)$ such as (b) , to derive new dynamics and formulate a new SSM framework.

Figure 2: State-space regularization applied to an image tokenizer. With probability $\alpha$ , the update Eq. ( 21 ) is applied to the latent representation of the input images to produce $\hat{\mathbf{z}}_{2}$ and $\hat{\mathbf{I}}_{\tau_{2}}$ . The network is trained to match $(\hat{\mathbf{z}}_{2},\hat{\mathbf{I}}_{\tau_{2}})$ to $(\mathbf{z}_{2},\mathbf{I}_{\tau_{2}})$ .

Table 3

Figure 3: Generation results comparison using the Flux tokenizer. With only marginal loss in reconstruction quality, our method improves the generative performance of the image tokenizer.

Panel of Figure 4 — Figure 4: Effect of spatial mean-centering function $\mathcal{C}.$ The one trained without $\mathcal{C}$ shows latents with higher latent norm, and exhibits tremendous degree of KL loss.

Figure 5: Latent channels encoding different frequency bands. The regularized tokenizer is able to reconstruct recognizable images solely from using three low-frequency channels.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 6 — Figure 6: Visualization of derived ${\mathbf{A}}$ matrices. We set $H=W=8$ . Every matrix we derived is very sparse, which enables efficient matrix-vector multiplication. To maintain training stability, we normalize the matrix to have maximum absolute value of 1.

Tokenizer	Num	batch	KL	Warmup init.	Base	Min	$\mathbf{A},\Delta$	Gradient
Tokenizer	params	size	weight	learning rate	learning rate	learning rate	learning rate	clipping norm
Flux	83.8M	32	$10^{-6}$	$1.0\times 10^{-7}$	$1.0\times 10^{-5}$	$1.0\times 10^{-6}$	$1.0\times 10^{-6}$	50.0
Cosmos	81.5M	32	-	$1.0\times 10^{-7}$	$1.0\times 10^{-4}$	$1.0\times 10^{-5}$	$1.0\times 10^{-5}$	50.0

Appendix

Table 3: Implementation details of the reconstruction experiments from Sec. 4 .

Tokenizer	$\tau_{\textrm{max}}$	$\Delta_{\mathbf{I}}$	$\Delta$ init.	Max $\mathbf{A}$ init.	$d_{\mathbf{z}}(\cdot,\cdot)$	$d_{\mathbf{I}}(\cdot,\cdot)$	$\lambda_{\mathbf{z}}$	$\lambda_{{\mathbf{I}}}$	$\lambda_{\textrm{MAE}}$	$\lambda_{\textrm{LPIPS}}$
Flux	8.0	4.0	0.1	16.0	MSE	MAE & LPIPS	0.0	1.0	1.0	0.1
Cosmos	8.0	4.0	0.1	1.0	MSE	MAE & LPIPS	5.0	1.0	1.0	0.1

Appendix

Table 4: Hyperparameters of the regularization framework

Figure 7: Image displayed with respective $\tau$ values.

Figure 8: (Left) Illustration of the channel-to-coefficient index mapping. (Right) Channels illustrated with colors, where darker color indicates that a channel encodes lower frequency.

$\mathbf{A}$ type	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	rFID $\downarrow$	gFID $\downarrow$	sFID $\downarrow$	IS $\uparrow$	pre. $\uparrow$	re. $\uparrow$
Fourier	31.54	0.9101	0.0370	0.90	5.25	6.99	207.14	0.85	0.45
Legendre	31.57	0.9091	0.0366	0.92	5.22	7.65	206.70	0.84	0.45
Chebyshev	31.56	0.9087	0.0368	0.90	5.28	7.27	199.11	0.83	0.46
Hermite	31.60	0.9092	0.0363	0.92	5.34	7.03	196.49	0.84	0.46
random	N/A (does not converge)				-	-	-	-	-

Appendix

Table 5: Ablation for the type of

\mathbf{A}

$\alpha$	PSNR $\uparrow$	LPIPS $\downarrow$	SSIM $\uparrow$	rFID $\downarrow$	gFID $\downarrow$	sFID $\downarrow$	IS $\uparrow$	pre. $\uparrow$	re. $\uparrow$
0	31.70	0.0354	0.9105	0.87	16.12	7.22	89.27	0.68	0.49
0.25	31.53	0.0373	0.9082	0.94	13.17	6.71	105.99	0.71	0.48
0.5	30.95	0.0459	0.9018	1.01	13.73	8.19	111.65	0.69	0.43
0.75	29.83	0.0649	0.8884	1.28	15.48	8.72	112.02	0.68	0.39

Appendix

Table 6: Ablation for the regularization strength

\alpha

Figure 9: Additional results from the channel revealing experiment

Panel of Figure 10 — Figure 10: Generation results from the Cosmos tokenizer

Panel of Figure 11 — Figure 11: Additional generation results from the Flux tokenizer

Image Tokenization

Image tokenization converts images into compact discrete or continuous visual units for efficient generation, compression, and multimodal reasoning.

See all 56 papers

Patch-Derived or Register Tokens? Comparing Variable-Length Image Tokenizers
Oct 24, 2026Philipp Schenkel, Tobias Schürmann, Omer Sahin TasImage Tokenization Visual Token Compression
FZI Research Center for Information Technology · Karlsruhe Institute of Technology
Semantic One-Dimensional Tokenizer for Image Reconstruction and Generation
Mar 17, 2026Yunpeng Qu, Kaidong Zhang, Yukang Ding +2Image Tokenization Image Reconstruction
Tsinghua University, Beijing, China · Alibaba Group, Beijing, China
Test-Time Registers as Global Priors for Tokenized Image Generation
Jul 18, 2026Cheng-Yao Hong, Yifan Wang, Yuewei Lin +1Image Tokenization Vision Transformer
Stony Brook University · Brookhaven National Laboratory

Oct 24, 2026J/K move · Enter open · S save

Patch-Derived or Register Tokens? Comparing Variable-Length Image Tokenizers

Philipp Schenkel, Tobias Schürmann, Omer Sahin Tas

FZI Research Center for Information Technology · Karlsruhe Institute of Technology

Image Tokenization Visual Token Compression

Open paper AskPDF arXiv

Image tokenizers convert pixels into compact latent sequences for efficient generative modeling. Most approaches arrange latents on a downsampled two-dimensional grid, tying sequence length to image resolution. We compare two variable-length strategies: computationally efficient patch-derived tokens, which inherit explicit spatial organization, and more demanding register tokens, which have no predefined spatial layout. Our reconstruction results reveal that patch-derived tokens perform substantially better under extreme compression, whereas register tokens achieve slightly better reconstruction quality at intermediate budgets. At high token budgets, both approaches achieve similar reconstruction quality. For downstream generative tasks, register tokens outperform patch-derived tokens. Adding representation alignment with external pretrained visual encoders substantially improves the generative performance of patch-derived tokens, allowing them to match or outperform register tokens depending on the token budget. We further find that both approaches learn spatially localized representations. Patch-derived tokens largely retain their explicitly imposed spatial organization, while register tokens develop an implicit two-dimensional spatial structure despite having no predefined layout. Representation alignment reduces this spatial organizational constraint, enabling individual tokens to exert a stronger non-local influence.

Abstract

Image Tokenization

Explore similar work

Patch-Derived or Register Tokens? Comparing Variable-Length Image Tokenizers

Semantic One-Dimensional Tokenizer for Image Reconstruction and Generation

Test-Time Registers as Global Priors for Tokenized Image Generation