AsymFlow:

用秩-非对称速度参数化, 把 latent 流匹配模型微调到像素空间

AsymFLUX.2 klein 生成示例
Fig. 1 — AsymFLUX.2 klein 从 FLUX.2 klein 9B 微调而来的像素空间生成。论文用这些样例论证 "pixel 空间的细节质感比 latent 强" — 砂砾、毛发、皮肤纹理这些 VAE decoder 通常磨平的高频细节直接在像素流上生成。

1. 出发点 (Motivation)

2025-2026 的 pixel-space diffusion 复兴 (JiT、PixelGen、PixNerd、DeCo、PixelDiT、DiP、L2P) 都在攻击同一个观察: VAE 是有损的, latent 扩散的细节天花板被 decoder 卡死。所以社区开始问: 既然 plain DiT 已经 scale 到 9B+ 参数, 能不能直接在像素上用?

但直接把 DiT 塞像素有一个具体的技术障碍, 论文 §1 讲得很清楚:

"Velocity target \(u = \boldsymbol\epsilon - \mathbf x_0\) requires predicting a high-dimensional noise component in addition to structured data. ... per-patch noise dimension can pollute the network's internal states, creating a bottleneck."

翻成大白话: 在像素空间, 一个 16×16×3 = 768 维的 patch token 里, 一半是要被 transformer 内部表示和传递的高斯白噪声。但 transformer 的 hidden state 是有限带宽 (e.g., 1280 维), 你逼它"既要学语义结构, 又要在 hidden 里携带 768 维白噪声", 就会爆。U-Net 用 skip connection 让噪声"从 input 短路到 output"绕过这个瓶颈; plain transformer 没有 skip, 撞墙。

历史上社区有两条退路, AsymFlow 都不满意:

  1. \(x_0\)-prediction (JiT、PixelGen、PixelREPA): 直接让网络预测干净图 \(\hat{\mathbf x}_0\), 然后用 \(\hat u = (\mathbf x_t - \hat{\mathbf x}_0) / \sigma_t\) 换算 velocity。\(\sigma_t \to 0\) 时分母趋零, 数值不稳, low-noise 阶段质量崩 (JiT 论文里也承认)。
  2. 架构改造 (U-ViT、PixNerd、PixelDiT、DDT 的 decoder head): 在 plain transformer 上加 skip/decoder head 让噪声有路径短路。能 work, 但破坏了 plain DiT 的简洁性, 也不太方便复用 latent DiT 的预训练权重。

AsymFlow 的赌注: 不改架构, 不改 training/sampling 协议, 只改 velocity 的参数化目标 — 把 noise 部分限制到 patch 内的低秩子空间。两个好处一次拿到:

2. 方法 (Method) — 高中生友好 + 数学严谨

核心思想 (类比)

想象你是个画家, 任务是从一张被白噪声覆盖的画布逐步"擦掉噪声"还原一幅清晰画作。标准 flow matching 要你预测一个"擦除方向"向量, 这个向量同时编码 (a) 画作内容 和 (b) 噪声形状。在低维 latent 空间这不难, 因为画作和噪声的总维度都不算大。但跳到全像素 (768 维 per patch), 你的大脑 (transformer hidden state) 同时要在 working memory 里塞下 整张画的几何 和 每个像素的噪声值, 装不下, 内部表示开始打架。

AsymFlow 的招: 把"噪声形状"那部分降维, 只让你预测噪声在某 8 个"重要方向"上的投影。剩余 760 个方向的噪声怎么办? — 不让你预测, 直接从 \(\mathbf x_t\) 里减出来 (因为 \(\mathbf x_t = \alpha_t \mathbf x_0 + \sigma_t \boldsymbol\epsilon\), 知道 \(\mathbf x_0\) 和 \(\mathbf x_t\) 就能反推 \(\boldsymbol\epsilon\))。

"8 个重要方向"是哪 8 个? — PCA 选出来的、自然图像 patch 上方差最大的 8 个方向 (从头训用); 或者把预训练 latent 空间通过 Procrustes 对齐"提升"到像素空间的 8/128/d 个方向 (微调用)。这两种构造让网络只需要在"语义相关的低秩子空间"里学 \(u\)-prediction, 在"高频纯噪声方向"上自动退化为 \(x_0\)-prediction (后者不要求网络精确预测噪声, 只要求预测干净数据)。

2.1 Flow Matching 基础 (Eq. 1–2)

记 \(\mathbf x_0 \in \mathbb R^D\) 是干净数据, \(\boldsymbol\epsilon \sim \mathcal N(0, \mathbf I)\)。线性 flow 调度: \(\alpha_t = 1-t,\ \sigma_t = t,\ t \in (0,1]\):

ODE 速度目标 (条件后验均值 → 用样本速度做 unbiased estimator):

—— 翻译: velocity 就是从干净指向噪声的向量。学一个网络 $\hat u = G_\theta(\mathbf x_t, t)$ 来逼近它。

Flow matching loss:

关键问题: \(u\) 里包含 \(\boldsymbol\epsilon\) — 一个 \(D\) 维白噪声。在像素空间 \(D\) 很大 (768/patch), 这部分占用了网络内部 features 大量带宽, 而它本应能从输入 \(\mathbf x_t\) 里"读出"。

2.2 非对称速度 (Eq. 3)

设 \(A \in \mathbb R^{D \times r}\) 是 rank-\(r\) 子空间的正交基 (\(A^\top A = I_r\)), \(P := AA^\top\) 是对应的正交投影矩阵。\(\text{Im}(P)\) 是 rank-\(r\) 子空间, \(\text{Im}(I-P)\) 是其正交补。

AsymFlow 把要预测的目标改成非对称速度:

—— 翻译: 噪声项 $\boldsymbol\epsilon$ 被替换成它在 $r$ 维子空间上的投影 $P\boldsymbol\epsilon$ (维度从 $D$ 降到 $r$ 的有效自由度), 数据项 $-\mathbf x_0$ 保持 full-rank。训练时网络输出 $\hat u_A = G_\theta(\mathbf x_t, t)$ 来拟合这个 $u_A$。

核心好处: 网络 hidden state 不再需要把 \(D\) 维噪声"端到端"传递, 只需要保留 \(r\) 维的有意义信息。

AsymFlow 参数化与还原图
Fig. 2 — (a) 标准 $u$-prediction vs AsymFlow 参数化的对比: 上排 $u$ 里的噪声是 full-rank, 杂乱; 下排 $u_A$ 里的噪声是 $P\boldsymbol\epsilon$, 限制在 $r$ 维子空间, 整个目标更"结构化"、更可预测。(b) 把网络预测的 $\hat u_A$ 还原成 full-rank $\hat u$ 的几何: $P$ 分量直接用 $P\hat u_A$, 正交分量用 $x_0$-to-$u$ 关系 $(\mathbf x_t - \hat{\mathbf x}_0)/\sigma_t$ 换算。

2.3 正交分解 + Full-rank 还原 (Eq. 5)

把 \(u_A\) 分解到 \(\text{Im}(P)\) 和 \(\text{Im}(I-P)\):

—— 翻译: 在子空间 $\text{Im}(P)$ 内, AsymFlow 就是标准 $u$-prediction; 在正交补 $\text{Im}(I-P)$ 内, AsymFlow 退化成 $x_0$-prediction (符号相反)。这是 Fig. 3 那张"族谱图"的数学根据 — 调 $r$ 就是在两端连续滑动。

AsymFlow 参数化族
Fig. 3 — 调 rank ratio $r/D$ 形成一族参数化。最左 $r/D{=}0$ 是 $x_0$-prediction, 最右 $r/D{=}1$ 是 full $u$-prediction; 中间是 AsymFlow。上排是非对称速度 $u_A$ 的可视化, 中排是 $\text{Im}(P)$ 分量 (是 $u$-style), 下排是 $\text{Im}(I-P)$ 分量 (是 $-x_0$-style)。 读法: $r$ 越大, 噪声"全维出现"; $r$ 越小, $u_A$ 越接近 $-\mathbf x_0$, 几乎是纯 $x_0$-prediction。

full-rank velocity 的还原公式 (Eq. 5):

—— 翻译: 子空间内直接用 $Pu_A$; 正交补内用 $\hat{\mathbf x}_0 = -(I-P)u_A$ 代回 $x_0$-to-$u$ 公式。结果是 full-rank velocity, 可以直接喂进 standard flow matching loss 和 standard ODE/SDE sampler。所以 training/sampling 协议完全不动, 只换了一个 head 的语义。

代码里这个还原逻辑就是 asymflow_velocity 函数:

repo/lakonlab/models/architectures/asymflow/common.py:L42-L72 — 正交分解 + Eq. 5 的实现

@staticmethod
def orthogonal_decomposition(full_rank_state, proj_buffer):
    subspace = full_rank_state @ proj_buffer @ proj_buffer.T
    complement = full_rank_state - subspace
    return subspace, complement

def asymflow_velocity(
        self,
        u_a_packed,
        x_t_packed,
        calibration: AsymFlowCalibration):
    with torch.autocast(device_type='cuda', dtype=torch.float32, enabled=False):
        sigma_min = self.train_sigma_min if self.training else self.sigma_min
        u_a_packed = u_a_packed.float()
        x_t_packed = x_t_packed.float()
        proj_buffer = self.proj_buffer.float()
        # orthogonal decomposition
        u_a_subspace, u_a_complement = self.orthogonal_decomposition(u_a_packed, proj_buffer)
        x_t_subspace, x_t_complement = self.orthogonal_decomposition(x_t_packed, proj_buffer)
        # read calibration output
        sk = calibration.s * calibration.k
        sigma_clamped = calibration.sigma.clamp(min=sigma_min)
        # low-rank subspace (Pu_A)
        u_subspace = (
            sk * u_a_subspace
            + (1 - sk) / sigma_clamped * x_t_subspace
        )
        # orthogonal complement: (x_t + s * u_A) / sigma_clamped
        u_complement = (x_t_complement + calibration.s * u_a_complement) / sigma_clamped
        # full velocity
        return u_subspace + u_complement

对照公式: u_subspace 就是 Eq. 5 的左半 \(Pu_A\) (加上 scale calibration 的 \(s, k\) 校正项, 见 §2.5 / Appendix A.2), u_complement 就是右半 \((I-P)(\mathbf x_t + u_A)/\sigma_t\)。注意 sigma_clamped = max(sigma, sigma_min) 防止低噪声时分母崩 — 但 因为这只发生在正交补里, AsymFlow 比纯 \(x_0\)-prediction 对 \(\sigma_{\min}\) 的依赖小得多 (Table 1: JiT 关掉 clamp FID 从 1.90 掉到 3.27, AsymFlow 只从 1.76 掉到 2.28)。

2.4 子空间怎么选: PCA (from-scratch) vs Procrustes (latent→pixel)

整个机制的"形状"由投影矩阵 \(A\) 决定。论文给两种构造:

(a) PCA 子空间 (训练 from scratch)

对训练集 patches 做 Gram 矩阵特征分解, 取 top-\(r\) 主方向:

直观: PCA 保留了自然图像 patch 上"方差最大"的 \(r\) 个方向 — 这些是低频结构 (光照、色块、整体颜色), 也是 velocity prediction 最需要"精细"控制的部分。剩下 \(D-r\) 个方向 (高频纹理、噪声样态) 推到 \(x_0\)-prediction 通道里, 不要求精确预测。Fig. 5 消融: 同一 rank \(r{=}8\), PCA 比 random subspace 好将近 0.4 FID — 子空间的方向比维度更重要。

repo/tools/asymflow_subspace_pca_dit.py:L94-L123 — PCA 子空间构造主函数

def main():
    args = parse_args()
    cfg = Config.fromfile(args.config)
    # ... build dataset, dataloader ...
    denoising_cfg = cfg.model.diffusion.denoising
    patch_size = denoising_cfg.patch_size
    rank = denoising_cfg.base_rank  # e.g., 8 for AsymFlow-H/16

    torch.set_grad_enabled(False)
    pixel_gram = compute_pixel_gram(cfg, args, patch_size, device)  # accumulate X^T X
    dim = pixel_gram.size(0)                                          # patch dim (e.g., 768 for 16x16x3)
    if rank > dim:
        raise ValueError(f'Rank {rank} exceeds patch dimension {dim}.')

    _, evecs = torch.linalg.eigh(pixel_gram)                          # eigendecomp, ascending
    proj_mat = evecs[:, -rank:].contiguous().float().cpu()            # top-r eigenvectors = PCA basis

    write_checkpoint_to_file({f'proj_mat_p{patch_size}': proj_mat}, out)

(b) Procrustes 子空间 (latent→pixel 微调)

用于把预训练 latent flow 模型 lift 成像素模型。设 \(Z \in \mathbb R^{d \times N}\) 是 latent token, \(X \in \mathbb R^{D \times N}\) 是配对的 pixel patch (来自同一张图, 经 VAE encoder 得 \(Z\), 直接 patchify 得 \(X\))。求正交 lift:

—— 翻译: 找一个 正交 线性映射 $A$, 让 lift 后的 latent $Az$ 尽量像配对的 pixel patch $x$。这是经典 Procrustes 问题, 闭式解就是 cross-Gram 矩阵 $XZ^\top$ 的 SVD。

repo/tools/asymflow_subspace_procrustes.py:L141-L171 — Procrustes 主循环 + SVD 闭式解

    for data in dataloader:
        # ... resize images, encode latents ...
        latent_tokens, pixel_tokens = build_token_pairs(
            encoder, latents, images, latent_patch_size, patch_size)

        if cross_gram is None:
            d_latent = latent_tokens.size(1)
            d_pixel = pixel_tokens.size(1)
            cross_gram = torch.zeros((d_latent, d_pixel), dtype=torch.float64, device=device)
            pixel_gram = torch.zeros((d_pixel, d_pixel), dtype=torch.float64, device=device)

        cross_gram += latent_tokens.T @ pixel_tokens     # Z X^T 累加
        pixel_gram += pixel_tokens.T @ pixel_tokens
        latent_norm_sq += (latent_tokens * latent_tokens).sum()
        num_seen += latents.size(0)

    # Orthogonal Procrustes closed-form: A = V U^T from SVD of cross-Gram
    u, _, vh = torch.linalg.svd(cross_gram, full_matrices=False)
    rank = cross_gram.size(0)
    proj_mat = (vh[:rank].T @ u[:, :rank].T).contiguous()
    scale = fit_scale(pixel_gram, proj_mat, latent_norm_sq)   # 尺度校正 (Appendix A.2)

    write_checkpoint_to_file({
        f'proj_mat_p{patch_size}': proj_mat.float().cpu(),
        f'scale_p{patch_size}': scale.float().cpu(),
    }, out)

SVD 给闭式解 \(A = V U^\top\) (经典 Schönemann 1966), 配上一个 scale factor \(s\) (Procrustes 只校方向不校尺度)。结果是一个 把 \(d\) 维 latent token 1-to-1 映射到 \(D\) 维 pixel patch 的正交矩阵。

2.5 Latent→Pixel 初始化 (§5 — 论文最有意思的部分)

有了 \(A\) 之后, 论文证明: 预训练的 latent \(u\)-prediction 模型可以被精确地"提升"为一个 rank-\(d\) 的像素 AsymFlow 模型, 零额外训练就能生成有意义的 lift 像素。

具体地, 给一个预训练 latent flow \(\hat u_z = G_\phi(\mathbf z_t, t)\), 定义"lift 后的低秩像素" \(\mathbf x_0^L := A\mathbf z_0\) 和它的扩散 \(\mathbf x_t^L := \alpha_t \mathbf x_0^L + \sigma_t \boldsymbol\epsilon\)。论文 Thm. 1 给出 latent ODE 和 lifted pixel ODE 的 轨迹耦合:

—— 翻译: 一个 $d$-维 latent $u$-prediction 网络 $G_\phi$ 等价于一个 rank-$d$ 的像素 AsymFlow 网络 $AG_\phi(A^\top \cdot, t)$。在实现里, $A^\top$ 和 $A$ 直接融进 latent DiT 的 input/output 线性层权重, 模型 forward 就变成像素流。

这一招在论文里的视觉证据是 Fig. 4 — 初始化好的像素 AsymFlow 跑 sampling 出来的图, 跟原 latent 模型 decode 出来的图构图、语义、风格几乎一致, 只是低频细节有 lift 误差。微调因此变成"补一个 \(\mathbf x_0 - \mathbf x_0^L\) 的低频 gap", 而不是"从 random init 重新学像素生成"。这是 AsymFlow 文中最 elegant 的一步。

2.6 方差缩减微调 loss (Eq. 7)

初始化后, 微调目标本可以就用标准 FM loss 回归 \(\mathbf x_0\)。但因为 lift 出来的 \(\mathbf x_0^L\) 跟真实 \(\mathbf x_0\) 已经很接近, 论文用 control variates 思想引入方差缩减项:

—— 翻译: 引入"冻结的初始化模型" $\hat{\mathbf x}_0^L$ 作为对照变量, 减去 $\lambda(\mathbf x_0^L - \hat{\mathbf x}_0^L)$ 这一项不影响期望但能降方差。$\lambda$ 是 patch-wise 自适应权重 (选成"使 loss 梯度范数最小"的解析最优, 见下面代码的 vr_coef)。

具体实现:

repo/lakonlab/models/diffusions/asymflow.py:L83-L113 — AsymFlowVR 的 forward_train 核心

    # reference low-rank forward (frozen teacher initialized via Procrustes lift)
    with torch.no_grad(), module_eval(teacher):
        ref_u_low_rank = teacher(return_u=True, x_t=x_t_low_rank, t=t, **teacher_kwargs)
        ref_x_0_low_rank = self.u_to_x_0(ref_u_low_rank, x_t_low_rank, sigma=sigma)

    # full-rank forward (student)
    pred_u = self.pred(x_t, t, **kwargs)
    pred_x_0 = self.u_to_x_0(pred_u, x_t, sigma=sigma)

    # adaptive variance reduction coefficient lambda (Eq. 7)
    low_rank_diff = self.denoising.patchify(
        x_0_low_rank - ref_x_0_low_rank, self.denoising.patch_size, pack_channels=False
    )  # (bs, c, patch_numel, *)
    full_rank_diff = self.denoising.patchify(
        x_0 - pred_x_0.detach(), self.denoising.patch_size, pack_channels=False
    )

    num = (full_rank_diff * low_rank_diff).mean(dim=2, keepdim=True)
    den = low_rank_diff.square().mean(dim=2, keepdim=True).clamp(min=eps)
    vr_coef = (num / den).clamp_(0.0, 1.0)               # lambda, per-patch in [0,1]

    if self.loss_shift is None:
        shifted_signal_ratio = 0.0
    else:
        shifted_signal_ratio = calc_shifted_signal_ratio(sigma, self.loss_shift)

    tgt_x_0 = x_0 - (1 - shifted_signal_ratio) * self.denoising.unpatchify(
        vr_coef * low_rank_diff, self.denoising.patch_size, packed_channels=False)
    mse_loss = F.mse_loss(pred_x_0, tgt_x_0, reduction='none')
    mse_loss = 0.5 * self.mse_loss_weight * (mse_loss / sigma_clamped.square()).mean()

注意 vr_coef = (num/den).clamp(0,1) 是 patch-wise 解析最优 \(\lambda\) (论文 Appendix A.3 的正交投影解释), 把 control variate 的贡献按 patch 自适应缩放。当 student 的 \(\hat{\mathbf x}_0\) 和 teacher 的 \(\hat{\mathbf x}_0^L\) 高度相关时 \(\lambda \to 1\), 方差缩减最强; 不相关时 \(\lambda \to 0\), 退化为标准 FM loss。

感知校正 (LPIPS gating): 方差缩减带来一个副作用 — 在 low-noise 阶段 \(\mathbf x_t - \mathbf x_t^L\) 不严格落在 \(\text{Im}(I-P)\) 里, 引入有界近似误差, 表现为"图过分锐利, 出现噪点"。论文加了一个 LPIPS perceptual loss 来抑制 (gated by same \(\lambda\), schedule 见 Appendix A.4)。Table 3 的消融显示 +perceptual 把 pFID 从 27.8 拉到 22.5, HPSv3 从 12.99 拉到 13.06。

3. 结论 (Key Findings)

3.1 ImageNet 256×256 主结果 (训练 from scratch)

JiT-H/16 主干 + AsymFlow (\(r{=}8\), PCA 子空间), 同 setup 600 epochs:

ImageNet 256×256 像素扩散对比表
Tab. 2 — ImageNet 256×256 pixel diffusion 全谱对比。AsymFlow-H/16 (最后一行) 用 $P\boldsymbol\epsilon - \mathbf x_0$ 参数化, 953M 参数, 363 GFLOPs, 拿到 1.57 FID。注意 plain transformer 这一栏里 (DiT-like, 无 skip/decoder head), AsymFlow 是唯一进 1.6 以下的。还比 hierarchical / DDT-decoder 系列的最强 PixelDiT-XL/16 (1.61) 和 DeCo-XL/16 (1.62) 都低 — 不改架构地超过了"加架构改造"的同侪。

3.2 T2I: AsymFLUX.2 klein (从 FLUX.2 klein 9B 微调)

这是论文的"重磅展示"。从 FLUX.2 klein Base 9B (latent flow, \(d{=}128\), 16-ch latent + 2×2 patch) 用 Procrustes lift 初始化, 然后在 3M LAION-Aesthetics 上微调 (input/output 投影 + rank-256 LoRA, 中间冻结)。结果 (Table 4):

三个 benchmark 全面超过 latent base, 这是 pixel-space T2I 之前从未做到过的。论文还做了 controlled baseline: 一个用同样数据但只在 latent 上微调 FLUX.2 的 baseline (HPSv3 10.70) 几乎打平 AsymFLUX.2 (10.66 — 实际略输), 但是 AsymFLUX.2 在 HPSv3 提升的来源 是 "visual realism" 而非 "数据 bias", 说明像素化生成本身有质感优势 (Fig. 7 + Fig. 8 qualitative 很明显)。

T2I 定性对比
Fig. 7 — T2I 定性对比。从上到下: Qwen-Image (latent), FLUX.2 klein Base (latent), PixelDiT-T2I (pixel from scratch), AsymFLUX.2 klein (pixel from FLUX.2 lift)。最后一行: AsymFlow 保留了 latent base 的构图能力 (来自初始化的 lift), 但细节质感 (皮肤、毛发、纹理) 明显比 latent 更真; 比 PixelDiT-T2I 这种从头训练的像素模型, 细节更锐利、整体不"糊"。中间两行的 latent 模型出来的脸有典型的"AI 磨皮"感, 像素模型没有。

3.3 关键消融

4. 实现细节 (Implementation Notes)

代码里关键、论文未完全展开的细节:

5. 批判性总结 (Critical Assessment)

优点

不足 / 疑点

适用 vs 不适用

进一步阅读

讨论 / Comments

评论托管在本仓库的 GitHub Discussions, 需 GitHub 账号。