Structured State-Space Regularization for Generation-Friendly Image Tokenization¶

会议: NeurIPS 2026
arXiv: 2604.11089
OpenReview: https://openreview.net/forum?id=ZD298GwlMx
领域: 图像生成
关键词: 图像分词器 (Image Tokenizer)、状态空间模型 (SSM)、频域规整 (Spectral Organization)、潜在扩散模型 (LDM)、生成友好性

一句话总结¶

本文将状态空间模型（SSM）重新解释为“基底投影加输入变换”的动态系统，推导了高斯模糊下 2D 正交基系数的演化规律并作为轻量正则项约束编码器，使图像分词器隐空间具备天然的频域组织结构，在几无重建损失的前提下大幅提升下游扩散模型的生成质量。

研究背景与动机¶

现代视觉生成模型（如基于潜在扩散模型 LDM、流匹配与自回归架构的系统）已全面转向潜在空间建模。与高维像素空间相比，紧凑低维的潜在表示计算开销更低，且能更好地捕获高层语义布局。在此范式下，图像分词器（Image Tokenizer，如 VAE 或自编码器）所构建隐空间的几何与概率分布特性，直接决定了下游生成模型采样的难易程度与生成上限。

然而，如何构建一个“对生成模型友好（Generation-Friendly）”的隐空间仍是一个悬而未决的核心难题。现有探索主要围绕几何等变性（如像素与隐空间间的尺度/旋转等变性）或特征分布对齐（如减少高斯混合成分、引入视觉大模型判别特征），却普遍忽视了自然图像固有的基本先验——频域组织特性（Spectral Organization）。自然图像的频谱能量服从经典的幂律衰减分布，但将频域显式引入隐空间的传统尝试往往因频域与空间域表征的显著差异而破坏训练动态，导致生成与重建性能受损。

本文的切入视角是：如何在不改变分词器现有卷积/Transformer架构、不直接拼接频域特征的前提下，为隐空间注入频域感知归纳偏置？研究者从序列状态空间模型（SSMs）的基本数学机理中获得启发——SSM 的隐藏状态本质上是在通过正交多项式基底投影连续压缩输入，其内部演化追踪的是基底系数的变化动态。核心 idea：将状态空间模型抽象为「正交基投影 + 输入动态变换」系统，基于热方程推导二维正交基底在连续高斯模糊演化下的系数动态矩阵 $\bar{\mathbf{A}}$，构造结构化状态空间正则项（$S^3\text{R}$）强制分词器隐特征模拟基底系数衰减行为，赋予潜在通道明确的由低至高频域解耦结构。

方法详解¶

整体框架¶

分词器采用标准的自编码器结构 $f = \mathcal{D} \circ \mathcal{E}$。整个训练过程在标准自编码重建目标与结构化状态空间正则化目标之间以概率 $\alpha$ 进行交替切换。当触发正则化时，输入清晰图像经过不同程度的高斯模糊变换生成成对视图，正则项强制两级模糊图像编码特征之间的跃迁严格遵循理论推导出的状态转移矩阵 $\bar{\mathbf{A}}$，并辅以空间均值中心化与解码器防塌陷重构约束。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入清晰图像 I"] --> Blur["高斯模糊序列采样<br/>得到两级模糊图像 I(τ1) 与 I(τ2)"]
    Blur --> Enc1["编码器 E 提取浅度模糊特征<br/>z1 = E(I(τ1))"]
    Blur --> Enc2["EMA 编码器提取深度模糊特征<br/>z2 = E(I(τ2)) (截断梯度)"]
    Enc1 --> S3R["1. 结构化状态空间动力学推导<br/>应用离散转移矩阵: z_hat2 = A_bar * z1"]
    S3R --> Centering["2. 空间均值中心化校准<br/>校准特征范数爆炸: C(z) = z - mean(z)"]
    Centering --> DualLoss["3. 潜在-像素双重正则化与防塌陷优化<br/>联合隐空间 MSE 与像素空间解码重建损失"]
    Enc2 --> DualLoss
    DualLoss --> Opt["更新分词器编码器 E 与解码器 D 权重"]

关键设计¶

1. 结构化状态空间动力学推导：将模糊退化映射为频域状态转移

现有 SSM（如 HiPPO、S4、Mamba）通常针对 1D 序列因果输入，将新到达的 token 拼接推导为正交多项式基系数的连续常微分方程（ODE）。本文突破性地将 SSM 推广为广义框架：由任意基底投影 $\mathbf{c}(\cdot)$ 与输入连续变换 $\theta(\cdot)$ 共同决定。针对 2D 图像域，本文选择 2D 正交基底（如二维傅里叶基、Chebyshev、Legendre 或 Hermite 多项式）与连续各向同性高斯模糊过程 $\mathbf{I}_\tau = G_\tau * \mathbf{I}$。由于高斯模糊等价于热传导方程的解析解 $\frac{\partial \mathbf{I}_\tau}{\partial \tau} = \frac{1}{2}\nabla^2 \mathbf{I}_\tau$，基底系数向量 $\mathbf{c}(\mathbf{I}_\tau)$ 对连续扩散时间 $\tau$ 的导数可解析推导为： $$\frac{\mathrm{d}}{\mathrm{d}\tau}\mathbf{c}_k(\mathbf{I}_\tau) = \frac{1}{2}\sum_{n=1}^N \mathbf{c}_n(\mathbf{I}_\tau) \langle \nabla^2 \phi_n, \phi_k \rangle \implies \frac{\mathrm{d}}{\mathrm{d}\tau}\mathbf{c}(\mathbf{I}_\tau) = \mathbf{A}\mathbf{c}(\mathbf{I}_\tau)$$ 其中矩阵 $\mathbf{A} \in \mathbb{R}^{N \times N}$ 的元素完全由基底函数间的拉普拉斯内积 $\frac{1}{2}\langle \nabla^2 \phi_n, \phi_k \rangle$ 严格决定（对于傅里叶基，因正交性 $\mathbf{A}$ 精确为对角阵）。通过欧拉离散化或零阶保持（ZOH），得到跨时间步长 $\Delta$ 的转移矩阵 $\bar{\mathbf{A}} = \delta(\Delta, \mathbf{A})$。这为潜在特征沿通道维度的演化赋予了严格的物理与数学依据。

2. 空间均值中心化校准：消除由于基底系数范数衰减引发的训练不稳定性

在连续高斯模糊算子作用下，高频细节快速衰减导致系数模长满足 $\|\bar{\mathbf{A}}\| < 1$。如果直接强制编码器隐特征满足 $\mathcal{E}(\mathbf{I}_{\tau_2}) = \bar{\mathbf{A}}\mathcal{E}(\mathbf{I}_{\tau_1})$，反向推导清晰图像特征 $\mathcal{E}(\mathbf{I}_0) = \bar{\mathbf{A}}^{-N}\mathcal{E}(\mathbf{I}_N)$ 将不可避免地导致清晰图像潜在表示的范数出现严重膨胀爆炸。这种极不均衡的特征方差会极大地恶化变分自编码器的 KL 散度与训练稳定性，甚至破坏下游扩散模型的先验假设。为此，本文设计了空间均值中心化算子 $\mathcal{C}$： $$\mathcal{C}(\mathbf{z})_{c,h,w} = \mathbf{z}_{c,h,w} - \frac{1}{HW}\sum_{h'=1}^H \sum_{w'=1}^W \mathbf{z}_{c,h',w'}$$ 通过对中心化残差特征施加正则化，促使图像在退化模糊时潜在表示自然收敛于各通道的空间均值偏置，而非绝对零向量，从而在保持不同频带相对衰减率的同时，彻底稳固了隐空间的数值分布尺度。

3. 潜在-像素双重正则化与防塌陷优化：避免模式崩溃并确保解码保真度

若仅在隐空间最小化 $d_\mathbf{z}(\mathbf{z}_2, \hat{\mathbf{z}}_2)$，自编码器极易退化至常数映射或零映射 $\mathcal{E}(\mathbf{I}) \mapsto \mathbf{0}$（平凡解崩溃）。为解决此问题并兼顾特征的解码重构能力，本文设计了双重目标空间损失函数：一方面计算预测隐特征 $\hat{\mathbf{z}}_2 = \bar{\mathbf{A}}\mathbf{z}_1$ 与目标隐特征 $\mathbf{z}_2$ 的均方误差（MSE）；另一方面将预测的隐特征输入解码器 $\mathcal{D}(\hat{\mathbf{z}}_2)$，计算其与真实模糊图像 $\mathbf{I}_{\tau_2}$ 在像素层面的 MAE 与感知相似度 LPIPS 损失。同时，在提取参考目标特征 $\mathbf{z}_2$ 时使用指数移动平均（EMA）编码器权重并截断其反向传播梯度，完全切断两个隐特征相互迎合的退化途径，确保状态跃迁正则化稳定施加于主干网络。

损失函数 / 训练策略¶

分词器训练的总目标为标准重构/KL 损失与状态空间正则项的组合。在正则化步骤中，从预设区间 $[0, \tau_\text{max}]$（$\tau_\text{max}=10.0$）中依线性衰减概率采样初态时间 $\tau_1$，并取固定间隔步长 $\Delta_\mathbf{I}=4.0$ 得到 $\tau_2 = \tau_1 + \Delta_\mathbf{I}$。正则化损失定义为： $$\mathcal{L} = \lambda_\mathbf{z} d_\mathbf{z}(\mathcal{C}(\mathbf{z}_2), \bar{\mathbf{A}}\mathcal{C}(\mathbf{z}_1)) + \lambda_\mathbf{I} \big( \lambda_\text{MAE} \|\mathbf{I}_{\tau_2} - \mathcal{D}(\hat{\mathbf{z}}_2)\|_1 + \lambda_\text{LPIPS} \mathcal{L}_\text{LPIPS}(\mathbf{I}_{\tau_2}, \mathcal{D}(\hat{\mathbf{z}}_2)) \big)$$ 在网络更新方面，冻结编码器前两层和解码器后两层以保留预训练权重的底层表达并免除额外判别器对抗训练；将转移矩阵 $\mathbf{A}$ 的非零元素与离散化步长 $\Delta$ 设为可微参数，采用为主干网络 0.1 倍的学习率微调更新，确保主干网络遵循严格的动力学约束收敛。

实验关键数据¶

主实验¶

在 ImageNet-1K（$256 \times 256$）上评估分词器重构性能（PSNR、rFID），并使用主流潜在扩散模型 LightningDiT（包含 B/2 与 L/2 两种规格，训练 80K steps，CFG 采样）评估下游生成质量（gFID、IS）。所有对比基线均基于统一训练设置在 ImageNet-1K 上对 Flux 和 Cosmos 进行微调。

分词器类型	优化方法	PSNR ↑	rFID ↓	gFID ↓ (-B/2 / -L/2)	IS ↑ (-B/2 / -L/2)
Flux (83.8M)	Baseline (Finetuned)	31.70	0.87	16.08 / 5.56	89.32 / 184.64
Flux	+ EQVAE (CVPR 2025)	31.53	0.94	13.95 / 5.37	97.81 / 195.64
Flux	+ FT-SE (ICML 2025)	31.62	0.90	13.87 / 5.08	99.79 / 201.12
Flux	+ Ours ($S^3\text{R}$)	31.54	0.90	13.17 / 5.25	105.99 / 207.14
Cosmos (81.5M)	Baseline (Finetuned)	25.16	2.88	13.52 / 5.49	112.48 / 230.59
Cosmos	+ EQVAE (CVPR 2025)	25.20	2.85	11.66 / 5.45	123.54 / 231.70
Cosmos	+ FT-SE (ICML 2025)	25.37	2.87	12.01 / 5.36	117.97 / 234.18
Cosmos	+ Ours ($S^3\text{R}$)	25.18	2.82	11.20 / 5.23	128.91 / 237.69

消融实验¶

在基于 Flux 分词器的设定下，探究不同正交基底诱导的转移矩阵类型（Table 5）以及空间均值中心化算子 $\mathcal{C}$ 对隐空间与生成指标的影响（Table 4(b)）：

配置 / 基底类型	PSNR ↑	rFID ↓	KL 散度 ↓	gFID ↓ (-B/2)	IS ↑ (-B/2)	状态说明
Fourier 基底 (Ours 默认)	31.54	0.90	203.2K	13.17	105.99	对角转移阵，频域解耦最佳，生成指标最优
Legendre 多项式	31.57	0.92	-	13.22	104.70	稀疏下三角矩阵，综合生成性能居次
Chebyshev 多项式	31.56	0.90	-	13.28	102.11	逼近最优，生成指标显著优于 Baseline
Hermite 多项式	31.60	0.92	-	13.34	101.49	高阶高斯衰减特性，显著优于未正则化基线
随机随机矩阵 (Random A)	N/A	-	-	-	-	无法收敛 (训练崩塌，证明结构化基底必要性)
去掉中心化 (w/o $\mathcal{C}$)	31.53	0.85	10.45M	17.23	88.14	特征范数暴涨导致 KL 激增 50 倍，生成质量严重恶化

关键发现¶

生成友好性显著提升：在两种顶尖分词器（Flux 与 Cosmos）上，本文正则化均在生成指标上大幅超越了专为重建旋转/尺度等变性设计的对比方法（EQVAE 与 FT-SE），Flux-B/2 的 gFID 优化幅度达 2.91，IS 提升超 16 分，印证频域组织特性是扩散生成建模更直接的结构诱导偏置。
重建与训练开销近乎无损：在 RTX 6000 Ada GPU 上实测显示，迭代耗时仅从 165.1 ms 微增至 165.3 ms，吞吐量维持 48.1 img/s（Cosmos 维持 34.7 img/s 不变）；重建 PSNR 与 rFID 下降极其微弱（Flux PSNR 仅微跌 0.16 dB），有效化解了分词器“保真度与生成友好性不可兼得”的矛盾。
隐特征通道呈现由低到高的频域涌现排列：通过依次解掩码通道可视化发现，前 3 个通道即包含了图像的低频轮廓与整体结构，后续通道逐级叠加高频纹理边缘细节；而未经正则化的基线分词器通道间杂乱交织，缺乏明确的频带对应规律。

亮点与洞察¶

将连续动力系统作为非侵入式隐空间正则化器：巧妙避开了显式在隐空间执行离散余弦变换（DCT）或傅里叶变换所引起的训练动态冲突，通过自编码网络对连续模糊过程的响应拟合，将频域紧致性内化为神经网络的隐式先验。
基于热方程实现 2D SSM 的确定性推导：将 1D SSM 的时序拼接更新推广至基于热传导算子的 2D 平滑演化，利用偏微分方程的解析解推导拉普拉斯内积矩阵，为多模态/多维 SSM 在连续尺度空间的发展提供了范式参考。
可迁移性强：该正则化方法作为通用的损失项即插即用，无须改动任何模型主干（无论是 CNN 还是 ViT 架构的分词器），未来可直接拓展至视频分词器（3D 热传导方程）、多视角 3D 资产隐空间建模等跨模态生成场景。

局限与展望¶

作者承认的局限：正则化设计以生成友好性为核心导向，在重建保真度上相较于纯重构微调存在微小的指标折损（如 Flux 的 rFID 上升 0.03）；此外，目前实验主要针对图像扩散生成任务，未在自回归视觉建模（VAR）、视频生成或表征预训练等更宽广的下游任务上做全方位展开。
评测与应用层面的延伸思考：转移矩阵中假设模糊核为标准各向同性高斯模糊，在真实复杂场景下（如运动模糊、非均匀退化）其先验假定可能存在失配；此外，将通道硬性绑定到固定低频到高频序列可能一定程度限制了隐空间的自由表征维度。
改进方向：探索引入非各向同性或自适应扩散系数的热传导方程；将该频域分解思想与自回归分词器（如 Next-Scale Prediction）相结合，实现更高效的分层多尺度由粗到细生成。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从广义 SSM 视角将图像高斯模糊与热传导基底系数演化建立桥梁，构思极为巧妙。
实验充分度: ⭐⭐⭐⭐☆ 覆盖 Flux 与 Cosmos 双分词器，消融实验涵盖 4 种正交多项式及稳定性关键算子，但生成端缺少非扩散模型（如自回归）对比。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨自洽，附录中各类正交基底内积解析解详实，逻辑层层递进。
价值: ⭐⭐⭐⭐⭐ 为解决图像分词器“重构与生成权衡困境”提供了轻量优雅的新范式，具高度工程与学术复用价值。