对齐潜空间几何的球面流匹配图像生成 Aligning Latent Geometry for Spherical Flow Matching in Image Generation
潜码投到定半径超球面并改用 slerp 路径,ImageNet FID 大幅提升且无需改动扩散架构
前置知识
潜空间流匹配 (Latent Flow Matching)
在 VAE 潜码上学习速度场，沿线性插值 $\mathbf{z}_t=(1-t)\mathbf{z}_0+t\mathbf{z}_1$ 把高斯噪声输送到数据端点，监督速度 $\mathbf{z}_1-\mathbf{z}_0$。
全文围绕"线性路径"假设展开。理解 $\mathbf{z}_t$ 的直线插值与速度目标，才知道为何能用 slerp 替换。
VAE 潜码与 token-wise 表示
VAE 把图像编码为空间位置上的 d 维 token 集合，每个 token 含范数 $\|\mathbf{z}\|$ 与方向 $\hat{\mathbf{z}}=\mathbf{z}/\|\mathbf{z}\|$。
径向/角度分解建立在 token-wise 表示上，"潜码是位置上的小向量"是后续投影与 slerp 的前提。
球面线性插值 (Slerp) 与测地线
球面测地线公式 $\mathbf{z}_t=\sqrt{d}(\sin((1-t)\omega)\hat{\mathbf{z}}_0+\sin(t\omega)\hat{\mathbf{z}}_1)/\sin\omega$，$\omega$ 为端点夹角。
方法是 slerp 几何核心。看懂公式与切投影 $\Pi_{\mathbf{z}_t}$，才懂 slerp 消径向监督且点始终在球上。
高维测度集中 (Concentration of Measure)
标准高斯概率质量集中在半径 $\approx\sqrt{d}$、宽度 $O(1)$ 的薄球壳上，形式化为 $P(\|\mathbf{z}\|-\sqrt{d}>t)\le 2\exp(-ct^2)$。
论证"VAE 潜码与高斯噪声同处薄球壳"的理论基础，是本文几何修正的前提。
SiT 与分类器无关引导 (CFG)
SiT 是流匹配专用 Transformer 骨干；CFG 用条件差 $\tilde{v}_\theta=(1-w)v_\theta(\cdot|y)+w\,v_\theta(\cdot|\varnothing)$ 加权样本。
实验均以 SiT 在 ImageNet-256 做 class-conditional 评测，读懂 CFG 数字必须先懂 CFG 放大条件信号的原理。
FID-50K 与 rFID
FID-50K 用 5 万生成样本相对训练参考的 Fréchet 距离；rFID 用验证集衡量 VAE 重建。FLUX.2 原 0.186、FT-vanilla 0.486、spherical 1.072。
两者分衡量"生成"与"重建"，才能解读"重建变差但生成 FID 反而变好"的 trade-off。
研究动机
线性潜空间流匹配默认用欧氏直线连接噪声 $\mathbf{z}_0\sim\mathcal{N}(\mathbf{0},\mathbf{I}_d)$ 与数据 $\mathbf{z}_1=E(\mathbf{x})$，但高维测度集中告诉我们两端点都集中在半径 $\approx\sqrt{d}$、宽度 $O(1)$ 的薄球壳上。FLUX.2、VA-VAE、REPA-E FLUX.1 处理后 $\bar{r}/\sqrt{d}$ 分别为 0.94、0.95、0.35，CV 都 $\le 0.23$。把两端连成直线时，中点范数向内塌陷到 $\approx R/\sqrt{2}$，路径相对最近球壳偏离达到 1.4σ、1.8σ、2.5σ，把监督信号投到训练分布几乎不生产的区域。component-swap 探针还证明 decoder 对方向远比半径敏感：与同类邻居交换方向几乎等同完全替换潜码，而交换半径几乎不改变图像，因此约 50–90% 的速度目标耗在 decoder 不在乎的径向运动上。
本文的目标是论文的目标是把"潜码长在球壳上"这一结构性质显式地编码进流匹配路径，使得两端点和插值轨迹始终共处同一个半径为 $\sqrt{d}$ 的超球面 $\mathbb{S}^{d-1}(\sqrt{d})$，从而速度目标天然是纯切向的、且训练分布与生成分布严格匹配。具体目标包括：在 FLUX.2 + SiT-B/2 上把 ImageNet-256 FID-50K 从 26.35（vanilla-linear）降到 20.55（spherical-slerp），同时让模型用约 2.2 倍更少的训练步达到 FID=30；并把这一收益横跨 FLUX.2、VA-VAE、REPA-E FLUX.1 三类 tokenizer 和 SiT-B 至 SiT-XL 两种规模、CFG=1.0/1.5/2.0 三档引导下都稳定复现。所有这一切都必须在不改动 SiT 扩散架构、不引入训练或推理阶段辅助编码器的前提下完成。
与已有工作不同的是，现有的几何感知替代方案都不同时满足"用现有预训练 VAE 的潜码"与"不改扩散架构"这两条。Riemannian flow matching 已经在 DINOv2 LayerNorm 诱导的球面上成功，却要求训练与推理每步都跑辅助编码器；Davidson et al.、Xu & Durrett、Ke & Xue、Yue et al. 等超球面 VAE 工作要么从零训练带 von Mises-Fisher 先验的 encoder，要么在自回归解码时做投影，无法与流匹配共用；REPA 一类表示对齐方法仍需在 VAE 或扩散训练中加入 DINOv2 等特征对齐损失与辅助编码器，同样带来部署成本。本文恰好落在这个交叉点上：通过对既有 VAE 施加一次性 token-wise $\sqrt{d}\,\mathbf{z}/\|\mathbf{z}\|$ 投影并冻结 encoder，只对 decoder 做少量 finetune，就把"球面潜码"和"流匹配"两条独立研究的优势缝合在一起，且全程不引入辅助编码器。
核心方法
核心直觉：VAE 潜码与高斯噪声都集中在细球壳上，decoder 几乎只读取方向，线性路径却把 50%+ 监督花在径向位移上。修正是把潜码硬投影到定半径超球面 $\mathbb{S}^{d-1}(\sqrt{d})$。路线四步：(1) encoder 与 decoder 间插 token-wise L2 投影 $\pi(\mathbf{z})=\sqrt{d}\,\mathbf{z}/\|\mathbf{z}\|$，强制 $\|\mathbf{z}_{i,j}\|=\sqrt{d}$；(2) 冻结 encoder，对 decoder+discriminator 复用原 VAE 重建目标 finetune 5 epoch；(3) 噪声端用 $\mathbf{z}_0=\sqrt{d}\,\boldsymbol{\epsilon}/\|\boldsymbol{\epsilon}\|$ 得球面均匀先验；(4) 插值换为 slerp，目标与输出都投到切空间 $\Pi_{\mathbf{z}_t}$，推理用指数映射保点始终在球。SiT 零改动，可对接 FLUX.2/VA-VAE/REPA-E。
区别于一切已有工作的本质在于把"方向比半径重要"这一非对称性质转化为几何约束，而非靠损失或调度去平衡两条监督。具体而言，本文做了三件新事：(1) 提出"token-wise 定半径硬投影"——既不是 Riemannian flow matching 要求的"目标潜码天然在流形上"，也不是表示对齐方法在 VAE 目标里加 auxiliary loss，而是直接在 encoder 输出做 $\sqrt{d}\,\mathbf{z}/\|\mathbf{z}\|$；(2) 在上述投影下用 slerp 作为训练时路径——slerp 不再只是 Riemannian FM 里的几何积分器，而是替代了线性 $\mathbf{z}_t$ 的监督目标，速度天然切向、半径恒等于 $\sqrt{d}$，CV=0；(3) 通过 decoder-only finetune 5 epoch 弥补被投影截断后的重建损失，而不修改 encoder 或下游 SiT 任何参数，让"球面潜码"对接标准扩散训练流程成为可能。
方法步骤详情
**球面投影**：每位置执行 $\mathbf{z}_{i,j}\leftarrow\sqrt{d}\,\mathbf{z}_{i,j}/\|\mathbf{z}_{i,j}\|$。**Decoder finetune**：冻结 $E$，L1+LPIPS+GAN 训 $D$ 5 epoch。spherical rFID: 1.072/1.064/0.824，对 FT-vanilla 0.486/1.116/0.759 有得有失。**Slerp 训练**：沿球面测地线监督 $v_\theta$，目标与输出都投到 $\Pi_{\mathbf{z}_t}$。噪声 $\mathbf{z}_0=\sqrt{d}\,\boldsymbol{\epsilon}/\|\boldsymbol{\epsilon}\|$ 得 $\mathrm{Uniform}(\mathbb{S}^{d-1}(\sqrt{d}))$。**球面采样**：推理亦投 $\Pi_{\mathbf{z}_t}$，指数映射 250 步匹配训练。SiT 超参：80 epoch、batch 256、bfloat16、AdamW。
技术新颖性
**几何命题精确化**：Tab. 1 的 $\bar{r}$/CV 与 Fig. 2b off-shell 距离（最大 2.5σ）量化"线性路径走偏"，Fig. 3 component-swap 把"decoder 对方向敏感"做成可复现现象。**极简干预**：单一 $\sqrt{d}\,\mathbf{z}/\|\mathbf{z}\|$ 投影 + decoder-only 5 epoch finetune，把任意 VAE 转成"球面 VAE"，无需 DINOv2、新 loss、新架构。**训练时路径 slerp 化**：Davis et al. 仅把 slerp 用于采样期，本文用 slerp 完整替换 flow matching 监督路径 $\mathbf{z}_t$，$\|\mathbf{u}_t\|$ 沿弧长恒定。**decoder-swap 控制**：Tab. 9 中 spherical 与 vanilla decoder 互换都使 FID 退化。Tab. 2 显示 spherical-linear 22.85、spherical-slerp 20.55 验证双层增益解耦。
实验结果
**(1) 主结果**：SiT-B/2 + FLUX.2 + CFG=1.0，spherical-slerp FID=20.55 vs vanilla-linear 26.35（−22.7%），FID=30 在约 2.2× 更少训练步达到（Fig. 6）。**(2) 跨 tokenizer**：CFG=1.0 下 FLUX.2 B/2 26.14→20.21、XL/2 13.00→10.71、VA-VAE B/1 26.99→21.96、XL/1 12.16→10.36、REPA-E B/2 38.00→26.07；CFG=1.5 下 REPA-E 13.83→6.88。**(3) 预算扩展**：200 epoch 下 CFG=1.0 8.35 vs 9.15、CFG=1.5 2.91 vs 3.22，排除"短预算优势"。**(4) 路径对比**：vanilla-linear 26.35→spherical-linear 22.85→spherical-slerp 20.55；vanilla+shell 反 27.26；slerp radial share 全程 0（Fig. 4）。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ImageNet-256 class-conditional FID-50K @ 80 epoch | FID-50K (lower is better) | Spherical-slerp SiT-B/2 + FLUX.2 VAE, CFG=1.0: 20.21 | Vanilla-linear SiT-B/2 + FLUX.2 VAE, CFG=1.0: 26.14 | −22.7% (绝对下降 5.93) |
| ImageNet-256 class-conditional FID-50K @ 80 epoch, REPA-E FLUX.1 | FID-50K (lower is better) | Spherical-slerp SiT-B/2 + REPA-E FLUX.1, CFG=1.5: 6.88 | Vanilla-linear SiT-B/2 + REPA-E FLUX.1, CFG=1.5: 13.83 | −50.2% (绝对下降 6.95) |
| Scaling to SiT-XL | FID-50K (lower is better) | Spherical-slerp SiT-XL/2 + FLUX.2, CFG=1.0: 10.71 | Vanilla-linear SiT-XL/2 + FLUX.2, CFG=1.0: 13.00 | −17.6% (绝对下降 2.29) |
| Extended training (200 epoch) | FID-50K (lower is better) | Spherical-slerp SiT-B/2 + FLUX.2, CFG=1.5: 2.91 | Vanilla-linear SiT-B/2 + FLUX.2, CFG=1.5: 3.22 | −9.6% (绝对下降 0.31)，且 gap 在更长预算下不闭合 |
| Convergence speed to FID=30 | 训练 epoch 数 | Spherical-slerp 约 35 epoch 达到 FID=30 | Vanilla-linear 约 80 epoch 达到 FID=30 | ≈2.2× fewer steps to FID=30 (Fig. 6) |
| Decoder-swap 控制实验 (Tab. 9) | FID-50K (lower is better) | Spherical VAE → spherical decoder: 20.55 | Vanilla VAE → vanilla decoder: 26.35 (交叉互换都退化) | decoder-swap 双方向均出现明显 FID 退化，证明增益不来自 decoder 容量 |
局限与改进
局限主要集中在三个层面。**重建代价**：spherical FLUX.2 rFID 从原 0.186 升到 1.072、相比 FT-vanilla 也从 0.486 升到 1.072（Tab. 3），重建质量近乎减半，必须用 FT-vanilla 控制确认增益不来自重建退化。**验证场景单一**：实验只覆盖 ImageNet-256 class-conditional 设置，未在 text-to-image、高分辨率或视频/3D 验证；球面几何在 slerp 退化（$\omega\to\pi$）时插值不稳定，作者在 Eq. 8 后给出 $\omega<\pi$ 的边界假设但无系统 ablation。**辅助编码器仍是 token-level**：球面 VAE decoder finetune 仍需 5 epoch，主实验只能在自建小规模模型上展开。Tab. 5 限制在 CFG=1.0/1.5/2.0 三档，CFG 越大提升越被压缩（FLUX.2 XL/2 CFG=1.5 仅 −8.6%），提示引导信号本身正在补偿几何失配。
独立分析的弱点
独立审视识别四点。**球面半径锁定 $\sqrt{d}$**：作者只验"匹配高斯集中半径"一个选择，未 sweep 不同 $R$。若取 $\sqrt{d-1/2}$ 或更大，重建/生成 trade-off 可能更友好，但缺数据。**Slerp 极端 $\omega$ 退化**：$\langle\hat{\mathbf{z}}_0,\hat{\mathbf{z}}_1\rangle\to-1$ 时 $\sin\omega\to0$，数值不稳；虽给 $\omega<\pi$ 边界但无 fallback，可对 $\omega>\pi/2$ 子区切回"线性+再投影"或球面随机游走。**vanilla+shell 退化诊断不足**：Tab. 2 中 vanilla+shell 27.26 比 vanilla 还差，未测 shell+slerp 等组合，也没加 radial schedule 消融。**球面 VAE 训练成本被低估**：decoder 5 epoch finetune 非无成本，跨 3 tokenizer × 多 patch size 都要重训，论文缺 batch/step 细节。
未来方向
作者 Sec. 5 已点名"更强压缩比检验球面约束"与"重建匹敌无约束 VAE 的球面 tokenizer"。可延伸：(1) spherical shell family——多离散半径分配 token 或球面调和基替换笛卡尔潜码，恢复被投影破坏的重建。(2) few-step distillation——slerp 解析 $\mathbf{u}_t$ 与 $\omega$ 闭式可能为 consistency model 提供结构，使 2.2× 训练加速叠乘为采样加速。(3) 推广到视频/3D/自回归图像——VDVAE、MagViT 同样有集中现象，球面 slerp 与 frame-time slerp 可组合。(4) 与 representation alignment 联用——VA-VAE 已对齐 DINOv2，可在 frozen sphere 上做 DINO 特征层 contrastive 微调。(5) 在 CFG scaling laws 下重新评估，研究引导强度与几何 gain 的耦合。
复现评估
**代码**：项目页 https://aligning-latent-geometry.github.io 但无仓库链接；需自实现 $\sqrt{d}\,\mathbf{z}/\|\mathbf{z}\|$ 投影、slerp loss 与 exponential map，约 200 行 PyTorch。**数据**：仅 ImageNet-256 train；latent 在 5 epoch decoder finetune 后一次预计算。**算力**：SiT-B/2 单卡 8×A100 80G 约一周跑完；XL 规模加 200 epoch 扩展与消融需数十到上百 A100-day；论文未公开精确 GPU 数与 wall-clock。**难度**：Eq. 8/9/10 给显式公式，App A 补推导与旋转不变性证明；按公式实现无隐藏 trick。decoder-swap (Tab. 9) 与多规模多 tokenizer 多 CFG 一致复现，证明结果鲁棒。
论文图表
三个子图：(a) 在二维示意把 VAE 数据壳与高斯噪声壳按欧氏直线相连，直线穿过两壳之间极不常到达的区域；(b) 把数据与噪声都投影到固定半径球面上，再用 spherical arc（slerp）相连，使整条路径始终在球壳上；(c) 训练曲线 SiT-B/2 + FLUX.2 VAE，spherical-slerp 约在 35 epoch 触到 FID=30，约为 vanilla-linear 80 epoch 的 2.2× 加速。
这张图是全文最核心的视觉锚点，把"线性 vs 球面"几何差异和"2.2× 收敛加速"这条最显眼的实验数字放在同一画面里，是把握论文论证结构的关键。
(a) 三种 tokenizer 下，沿 2048 token 配对的 linear/shell/slerp 路径，画出每步 token 平均范数的时间曲线与 $\pm\sigma$ 带，水平参考线是定球面半径；线性路径在中段明显向内塌陷至 $R/\sqrt{2}$ 量级。(b) 同一组路径上的 off-shell 距离，以"距离最近端点球壳的标准差"衡量；FLUX.2 1.4σ、VA-VAE 1.8σ、REPA-E 2.5σ，而 slerp 全程为 0。
这是把"几何失配"从理论命题转化为可量化实验的核心证据，配合 Tab. 1 一起支撑 motivation 中的"线性路径走偏"论点。
对每个 anchor token 取同类的邻居 token，把方向与半径分别互换再 decode；在 LPIPS 与 DINOv2 cosine 两个指标下，"keep direction"几乎不变（保留邻域半径即可），而"keep radius / full swap"则把图像拉向邻域，证明 decoder 主要消费方向信息。
这是后续"投影到定半径球面相当于丢弃 decoder 不敏感的成分"这一论证的实验根据，决定了方法的合理性。
在三种 tokenizer 的 flow-training 坐标系里，把每个 token 的速度目标 $\mathbf{u}_t$ 分解为径向和切向分量，画出径向份额随 $t$ 的变化。FLUX.2、VA-VAE 在两端 $\approx 50\%$，REPA-E FLUX.1 在噪声端 $\approx 90\%$；spherical-slerp 在所有 $t$ 上恒为 0。
定量说明 vanilla 线性流匹配有多少监督被浪费在 decoder 不在乎的径向运动上，并直观比较 slerp 在此指标上的完美解。