Parallel Rollout Approximation for Pixel-Space Autoregressive Image Generation¶

会议: NeurIPS 2026
arXiv: 2606.27978
OpenReview: https://openreview.net/forum?id=D55NaBnwtO
代码: https://github.com/MangataX/PRA
领域: 图像生成
关键词: 像素空间自回归, 连续token生成, 教师强制偏差, 并行rollout近似, 流匹配

一句话总结¶

针对像素空间自回归图像生成中高维patch单步误差大以及教师强制训练与自回归推理不匹配导致的跨步误差累积难题，本文提出并行rollout近似（PRA），通过端到端低维上下文感知中间态生成与并行的退化-解码像素反馈构造，在保持“像素输入-像素输出”原生接口的同时实现极具竞争力的生成与理解性能。

研究背景与动机¶

大语言模型的巨大成功推动了自回归范式向视觉等连续领域的延伸。然而，主流图像自回归模型普遍建立在离散量化代码本（如 VQ-VAE/VQGAN）或预训练的高维连续潜空间（如 VAE/autoencoder）之上。这种两阶段分工虽然规避了直接建模连续信号的复杂性，却使最终图像质量严格受限于外部 tokenizer 的重构上限，并割裂了端到端像素建模与底层视觉表征的统一。直接在原生像素空间对连续图像块进行 patch-wise 自回归建模虽然理念纯粹，却长期受到严峻性能瓶颈的掣肘，大幅落后于同时期的像素级扩散模型。

深入剖析该退化现象可以发现，像素空间自回归面临着一对紧密耦合的系统性矛盾：在输出端，每个原始像素 patch 均为高维连续向量（如 $16 \times 16 \times 3 = 768$ 维），单步条件生成难度极高，极易产生显著的局部预测误差；在输入端，标准自回归训练采用教师强制（Teacher Forcing），模型在前缀中始终接收完全干净的真实图像块，而在实际推理时却只能依赖模型自身先前生成的带噪图像块。这构成了自增强的恶性循环——高维预测引入严重的局部误差，而因果自回归循环又将误差作为后续步骤的历史上下文，引发误差在跨步生成中的非线性累积与放大。

现有改进方案如 $x$-prediction 或简单的输入高斯噪声注入只能局部缓解问题，仍与像素扩散模型存在显著性能鸿沟；而真正消除曝光偏差的精确在线 rollout 训练，由于每生成一个连续 token 均需运行多步流扩散采样，具有不可承受的串行时间开销。本文的切入角度是：解耦高维空间单步预测难度，并在不执行时序自回归采样的前提下，通过相同解码路径并行模拟出具有真实生成伪影的输入前缀。核心 idea：提出并行 rollout 近似（PRA），以端到端自适应学习的紧凑中间态替代高维像素块作为自回归生成目标，并在训练阶段利用相同因果像素解码器对加噪中间态并行重构出逼近推理分布的像素输入，在完全保留像素进、像素出接口的同时化解误差累积。

方法详解¶

整体框架¶

PRA 框架由四个紧密协作且完全端到端联合训练的模块构成：因果自回归 Transformer 主干 $f_\theta$、上下文感知中间态编码器 $g_\psi$、因果像素解码器 $D_\phi$ 以及基于整流流（Rectified Flow）建模的 token 级扩散头 $v_\omega$。对于分辨率为 $256 \times 256$ 的 RGB 图像，系统将其划分为 $T = 256$ 个互不重叠的 $16 \times 16$ 原始像素 patch，光栅扫描展平为 $d = 768$ 维的连续 token 序列 $\mathbf{x} = (\mathbf{x}_1, \dots, \mathbf{x}_T)$。

在整体流程上，PRA 在单步生成时并不强求模型直接预测 768 维的像素值，而是让自回归模型仅负责预测紧凑的低维中间态 $\mathbf{z}_i \in \mathbb{R}^{d_z}$（默认 $d_z = 16$），再由因果像素解码器将其瞬时映射为像素 patch 输出 $\hat{\mathbf{x}}_i$，从而维持纯正的“像素输入、像素输出”自回归闭环。在训练阶段，模型通过两趟并行的自回归前向传递完成计算：第一趟在真实像素序列上提取因果隐状态并生成监督目标 $\mathbf{z}_i$，对其施加中间扰动后由像素解码器并行还原为具有推理阶段伪影的近似像素输入 $\bar{\mathbf{x}}_i$；第二趟以前缀 $\bar{\mathbf{x}}_{<i}$（截断梯度）驱动自回归主干，监督流匹配扩散头预测目标 $\mathbf{z}_i$。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入真实像素序列 x_1:T + 类别前缀"] --> B["前向提取因果隐状态 h_i-1"]
    B --> C["端到端上下文感知中间态构建"]
    C --> D["并行扰动与因果像素解码重构"]
    D --> E["并行自回归前向与流匹配训练"]
    E --> F["推理期像素级逐步因果 Rollout 输出"]

关键设计¶

1. 端到端上下文感知中间态构建：压缩预测维度并与因果前缀对齐

直接在 768 维高维空间进行单步扩散预测会导致极大的回归方差与局部误差。若直接引入外部 VAE 或冻结的离散 tokenizer，则会退化为传统的两阶段潜空间自回归，丧失原生像素接口的优势。PRA 创新性地引入可微分的轻量级目标编码器 $g_\psi$，将当前真实像素 patch $\mathbf{x}_i$ 与自回归主干先前生成的因果隐表征 $\mathbf{h}_{i-1} = f_\theta(\mathbf{x}_{<i})$ 联合投影为低维紧凑状态： $$ \mathbf{z}i = g\psi(\mathbf{x}i, \mathbf{h}, \quad d_z \ll d $$ 将因果前缀 }), \quad \mathbf{z}_i \in \mathbb{R}^{d_z$\mathbf{h}_{i-1}$ 引入目标构建的核心机理在于：单纯基于当前局部 patch 提取的潜码往往割裂了图像块间的全局上下文连续性；而借助已有的因果隐状态进行调制，能够使目标 $\mathbf{z}_i$ 与自回归上下文的语义预期高度对齐。为防止编码器完全依赖无损的当前真实 patch 而退化为无上下文的局部自编码，训练中对 $\mathbf{x}_i$ 引入了两级掩码机制（以概率 $p_{\text{sample}} = 0.9$ 选取样本，并以 $p_{\text{token}} = 0.5$ 将 patch 替换为可学习的 mask 嵌入），强迫编码器从因果上下文中汲取信息，从而构建出对自回归模型而言更易预测的中间状态。

2. 并行扰动与因果像素解码重构：零采样开销下的曝光偏差逼近

即使将生成空间降低至 $d_z = 16$，如果在训练时主干网络仍然接收完美的真实图像 patch $\mathbf{x}_{<i}$，推理阶段面对自身生成（包含解码伪影与近似偏差）的图像块时仍会瞬间崩溃。精确的多步 rollout 需要对 256 个 token 依次调用几十步数值求解器，计算上完全不可行。PRA 提出在中间状态空间施加插值扰动并重构：对每个目标中间态 $\mathbf{z}_i$，采样高斯噪声 $\mathbf{z}_i^0 \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ 及时间戳 $t \sim \mathcal{U}(t_{\min}, 1)$，构造加噪状态 $\tilde{\mathbf{z}}_i^t = (1-t)\mathbf{z}_i^0 + t\mathbf{z}_i$，随后将其送入具有因果结构的 Transformer 像素解码器 $D_\phi$ 中： $$ \bar{\mathbf{x}}i = D\phi(\tilde{\mathbf{z}}i^t, \tilde{\mathbf{z}}^t) $$ 由于所有位置的加噪与解码可以跨空间位置完全独立且并行地执行，PRA 彻底免除了时序采样的开销。至关重要的是，生成的 $\bar{\mathbf{x}}_i$ 是通过与推理期完全相同的“中间态 $\to$ 像素”解码路径生成的，天然蕴含了解码器重构损失、轻微几何形变与色彩量化等具有归纳偏置的结构化误差。将其截断梯度作为后续 AR 训练的输入前缀，使模型在训练期便习惯了推理分布，从根本上阻断了跨步误差的雪崩式累积。

3. 双向梯度解耦与混合监督机制：平衡生成质量与表征学习

由于自回归主干、目标编码器、像素解码器与扩散头全部联合端到端优化，梯度流向的设计对系统稳定性至关重要。若允许自回归扩散损失反向传播至目标编码器，将诱发生成目标向极小方差塌缩的模式崩溃。为此，PRA 严格冻结了定义 AR 训练目标时的目标编码器分支（通过无掩码的 clean pass 提取目标 $\mathbf{z}_i$ 并显式 detach），使目标编码器与像素解码器完全由图像重构及感知损失驱动。此外，编码器重构损失流入共享自回归主干的梯度被乘以恒定缩放因子 0.3，以防止底层像素重建破坏自回归顶层的全局语义表征。同时，在训练前缀中以 0.04 的微小概率随机用带噪真实 patch 替换重构 patch，极大提升了模型在初始收敛阶段的鲁棒性。

损失函数 / 训练策略¶

PRA 的整体端到端多任务训练损失由四部分线性组合而成，无需针对各项繁琐调权： $$ \mathcal{L} = \mathcal{L}{\text{AR}} + \lambda}} \mathcal{L{\text{rec}} + \lambda}} \mathcal{L{\text{lpips}} + \lambda}} \mathcal{L{\text{repr}} $$ 其中权重默认均设为 1.0。各项损失的具体定义与优化机制如下： 1. 自回归流匹配损失 $\mathcal{L}_{\text{AR}}$：对于每个位置，流匹配时间变量 $s_i \sim \sigma(\mathcal{N}(0, 1))$，在加噪目标 $\mathbf{z}_i^{s_i} = (1-s_i)\boldsymbol{\epsilon}_i + s_i \mathbf{z}_i$ 上监督扩散速度场： $$ \mathcal{L}}} = \frac{1}{T} \sum_{i=1}^T \left| v_\omega(\mathbf{zi^{s_i}, s_i, \bar{\mathbf{h}}_i) \right|_2^2 $$ 其中条件隐状态 }) - (\mathbf{z}_i - \boldsymbol{\epsilon$\bar{\mathbf{h}}_{i-1} = f_\theta(c, \operatorname{sg}(\bar{\mathbf{x}}_{<i}))$ 基于重构的退化像素前缀计算，并采用 stop-gradient 算子阻断对解码器的梯度。 2. 像素级重构损失 $\mathcal{L}_{\text{rec}}$ 与感知损失 $\mathcal{L}_{\text{lpips}}$：约束因果解码器在输入扰动中间态时，将拼接还原后的图像 $\mathrm{Unpatchify}(\bar{\mathbf{x}}_{1:T})$ 与原图 $\mathtt{IMG}$ 进行对齐，结合了 $\ell_1$ 空间像素误差与 VGG 特征空间的 LPIPS 感知距离。 3. 辅助表征损失 $\mathcal{L}_{\text{repr}}$：利用第一趟重构前向中的自回归隐藏状态，通过轻量级回归头预测每个 patch 的局部归一化像素值，增强 Transformer 底层对高频纹理的感知能力。

在优化配置上，全量模型在 ImageNet-1K 上采用 AdamW 优化器（$\beta_1 = 0.9, \beta_2 = 0.95$）训练 400 个 epoch，全局 batch size 为 512，学习率采用带有 20K 步预热的余弦衰减调度（峰值 $3 \times 10^{-4}$），并对模型权重施加衰减率为 0.9999 的指数移动平均（EMA）。噪声下界超参数 $t_{\min}$ 在前 350 个 epoch 设为 0.5，后 50 个 epoch 提高至 0.7 以让模型在微调期适应更高保真度的数据。

实验关键数据¶

主实验¶

在 Class-Conditional ImageNet-1K $256 \times 256$ 基准测试中，使用包含 50K 生成样本的 ADM 评测套件，通过对无分类器引导（CFG）强度进行步长为 0.1 的细粒度网格搜索评估最佳 FID 与 Inception Score（IS）。实验全面对比了两阶段扩散、两阶段自回归、像素级扩散以及现有像素级自回归四大主流范式。

范式 / 模型	参数量	FID $\downarrow$	IS $\uparrow$	核心特性说明
两阶段自回归基准
VAR-d20	600M+40M	2.57	302.6	多尺度离散潜代码预测
LlamaGen-XL	1.4B+70M	2.34	253.9	1D 光栅扫描离散 VQ-VAE
MAR-L	479+49M	1.78	296.0	连续潜空间掩码自回归
像素空间扩散基准
ADM-G	554M	4.59	186.7	U-Net 引导像素扩散经典基线
JiT-L/16	459M	2.36	298.5	像素空间纯注意力快速扩散
PixelFlow-XL/4	677M	1.98	282.1	像素级整流流模型
像素空间自回归基准
JetFormer	2.8B	6.64	–	超大参数量直接建模像素 patch
FractalMAR-H	848M	6.15	348.9	分形多尺度连续 patch 预测
FARMER-1.9B/8	1.9B	3.60	269.2	此前像素空间 AR 的最佳 SOTA
PRA (本文)
PRA-S	135M	2.58	273.9	仅 1.3 亿参数即超越 19 亿参数 FARMER
PRA-B	250M	2.21	276.9	超越像素级扩散旗舰 JiT-L (2.36)
PRA-L	511M	1.94	287.3	创像素空间自回归新记录，首次破 2.0

此外，针对生成式表征迁移能力的 ImageNet 线性探测（Linear Probing）主实验显示：冻结主干特征训练线性分类头，PRA-L 取得了 68.80% 的 Top-1 准确率，大幅领先于潜空间自回归模型 SphereAR-L（52.19%）、潜空间扩散模型 DiT-XL/2（43.28%）以及像素空间扩散模型 JiT-L（42.76%）。

消融实验¶

基于统一的 PRA-B 模型架构在 ImageNet-1K（100 epochs）上进行系统消融，分别验证输出端中间态设计与输入端训练前缀构造的有效性。

模块类别	消融配置变体	FID $\downarrow$	IS $\uparrow$	设计机理解析
输入端前缀构造	真实像素输入 (GT pixel)	42.36	72.50	严重曝光偏差引发推理跨步雪崩崩溃
	真实像素 + 高斯噪声 (GT pixel + noise)	32.60	110.60	独立高斯扰动无法模拟解码器的结构化伪影
	真实中间态 (GT Inter. states)	3.21	261.86	失去原生像素接口，且存在潜在分布漂移
	真实中间态 + 噪声扰动	3.05	268.34	依赖潜码输入，背离端到端像素原生目标
	解码像素输入 (Decoded Pixel, PRA)	2.88	279.05	通过真实解码路径并行构造，消除曝光偏差
输出端目标选择	冻结 LDM 编码器空间 (LDM enc.)	3.37	266.73	通用图像潜空间不适配因果时序自回归预测
	纯局部投影编码 (Local only, 无前缀)	3.08	271.96	缺乏因果上下文，目标与时序演进脱节
	前缀感知编码 (Prefix aware, PRA)	2.88	279.05	利用因果隐状态调制，目标更易被 AR 主干拟合
中间态维度 $d_z$	$d_z = 8$	3.36	285.85	维度过小导致像素解码器重构瓶颈明显
	$d_z = 16$ (默认)	2.88	279.05	在可预测性与像素重构保真度间取得最佳平衡
	$d_z = 32$	3.50	297.17	维度增大使得单步扩散建模方差增大
	$d_z = 64$	7.03	220.11	逼近高维像素预测瓶颈，性能发生断崖式下跌

关键发现¶

双重瓶颈协同效应明显：仅做输出端压缩而不处理输入端偏差时，模型 FID 高达 42.36；仅做输入端加噪而不做维度压缩时，FID 停滞在 7.68。两项机制缺一不可，只有同时解决高维单步预测与跨步误差累积，像素自回归才能发挥全部潜力。
结构化伪影优于独立噪声：单纯向输入像素注入高斯白噪声仅能将 FID 压至 32.60，而通过像素解码器生成的 $\bar{\mathbf{x}}_i$ 能将 FID 骤降至 2.88，证明暴露给模型与推理路径同源的退化模式是克服曝光偏差的决定性因素。
紧凑维度存在清晰极值点：中间态维度并非越宽重构越好。当 $d_z$ 从 16 扩展到 64 时，FID 从 2.88 急剧劣化至 7.03，印证了连续 token 自回归对目标维度的极端敏感性。
表征学习潜力巨大：像素级自回归主干在无监督生成训练中展现出了优越的线性分类迁移性能（68.80%），印证了直接在像素层级学习因果预测能够促使注意力层捕获远比潜空间扩散模型更为深厚底层的几何与语义表征。

亮点与洞察¶

并行展开的巧思化解时序死锁：巧妙利用因果像素解码器，在训练阶段独立且并行地对加噪中间态进行重构，既获得了如同真正自回归展开（Rollout）一般的推理分布输入，又完全保留了现代 Transformer 批处理并行的极致训练效率。
兼具两阶段优势与原生像素纯粹性：通过内部动态生成的端到端低维流匹配目标与即时因果解码，既拥有两阶段潜空间自回归的平滑收敛特性，又对外保持了“像素进、像素出”的原生接口，彻底摆脱了对外部复杂且固化的 VAE 代码本的强依赖。
开辟统一生成与理解的新通路：在线性探测中高达 68.80% 的准确率揭示了一个重要趋势——原生像素级自回归不仅是一种有竞争力的图像合成工具，其内在因果注意力表征天然蕴含着强大的视觉感知与判别能力，具备成为通用多模态统一基座的潜力。

局限与展望¶

训练显存与计算开销增加：PRA 需要在训练期执行两趟自回归前向传递（第一趟生成前缀与目标，第二趟进行截断梯度的流匹配训练），相比标准单趟自回归增加了约 30%~50% 的单步反向耗时与显存占用。
高分辨率直接生成尚未验证：当前验证主要聚焦在标准的 $256 \times 256$ ImageNet-1K 分辨率上，针对 $512 \times 512$ 及以上更高分辨率的大规模自然图像与图文跨模态生成效果仍需进一步验证。
推理期因果解码开销待优化：推理时每生成一个中间 token 均需调用因果像素解码器进行瞬时展开，虽然结合了 KV Caching，但在极长上下文（如大图或视频 token 序列）下的累计延迟仍需进一步通过推测解码或非自回归并行解码进行压缩。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙提出并行近似 rollout 机制，攻克了像素自回归领域长期存在的曝光偏差与训练效率不可调和的矛盾。
实验充分度: ⭐⭐⭐⭐⭐ 从单步高维维度瓶颈诊断、输入端机制消融、多参数量扩展对比到下游线性探测，论证逻辑严密完整。
写作质量: ⭐⭐⭐⭐⭐ 动机阐述深刻，问题归因精准，算法伪代码与附录实现参数披露详实清晰。
价值: ⭐⭐⭐⭐⭐ 为原生像素空间自回归模型确立了全新的 SOTA 范式，为下一代端到端视觉生成与理解大一统模型提供了极具启发性的技术路线。