Structured State-Space Regularization for Generation-Friendly Image Tokenization • Jinsung Lee’s Personal Homepage

Find the whitepaper here: arXiv preprint

This paper introduces a sound way to extend 1D state-space models (SSMs) to 2D. Our claim is that the new update rule defined on the image space should be enforced in a form of regularization loss , and it happens to improve image generation tasks when applied to image tokenizers. — This paper introduces a sound way to extend 1D state-space models (SSMs) to 2D. Our claim is that the new update rule defined on the image space should be enforced in a form of **regularization loss**, and it happens to improve image generation tasks when applied to image tokenizers.

TL; DR

We introduce a novel way to extend State-Space Models (SSMs) to 2D. It naturally translates the SSM’s implementation from model architecture to regularization function, which shows effectiveness when integrated to image generation and restoration.

What we came up with & what you can do with it

The hidden state update rule we newly derived in the paper is expressed in a quite simple way:

\begin{equation*} h_t = \overline{\mathbf{A}}h_{t-1}, \end{equation*}

where $\overline{\mathbf{A}}$ is a discretized matrix formed by a structured matrix $\mathbf{A}$ .

This update rule is realized in a form of regularization function $\mathcal{L}$ :

\begin{equation*} \mathcal{L} = d\bigl(\overline{\mathbf{A}}f(\mathbf{I}_{t-1}), f(\mathbf{I}_{t})\bigr), \end{equation*}

where $f(\cdot)$ is a feature extractor that takes images with different levels of blur $\{\mathbf{I}_t\}$ , and $d(\cdot, \cdot)$ is a distance measure.

With this loss being applied properly, you can see the channels of image feature $f(\mathbf{I})$ behaving like image frequency components. For example, you can do something like below just by manipulating channels of an image feature.

This loss happens to improve generation quality of a tokenizer too.

If you are interested in how it works, we encourage you to…

check out the related blog posts:

read the whitepaper here: arXiv preprint.

About the project name

Although does not appear in the paper, we call this project WHippo, in a sense that it extends the idea of the pioneer SSM paper HiPPO (NeurIPS 2020) to 2D (Width and Height).

What makes even more sense is that the biological classification of hippo is called Whippomorpha, which is a class of mammals that includes whales and hippos.
So, Whippo can be viewed as a superclass (or generalization) of Hippo!

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.