跳转至

Spanning Tree Autoregressive Visual Generation

会议: ECCV 2026
论文: ECCV 原文
代码: https://github.com/oddqueue/star
领域: 图像生成
关键词: 自回归视觉生成、均匀生成树、广度优先遍历、图像修复、排列自回归

一句话总结

STAR 将图像 patch 网格建模为网格图,利用基于随机边角根节点的均匀生成树 BFS 遍历序列作为自回归生成的结构化随机生成顺序,既保留了天然的局部连通性与中心先验以保证图像生成质量,又通过拒绝采样支持任意连通观测作为前缀的原生图像修复。

研究背景与动机

基于纯解码器(Decoder-only)Transformer 架构的自回归(AR)模型通过下一 token 预测在语言建模领域取得了巨大成功,这一范式随后被广泛引入基于图像 token 的视觉生成任务。传统视觉 AR 模型大多采用固定的光栅扫描顺序(raster-scan order),将二维图像按先行后列的光栅顺序扁平化为一维序列。这种固定的单向因果依赖虽然符合二维网格的局部邻近性和从角落到中心的先验分布,训练收敛快且生成质量优异,却带来了严重的架构局限:在推理阶段无法根据任意空间掩码进行灵活的条件前缀补全,使得传统的自回归模型极难原生支持局部区域擦除、外扩或重绘等图像编辑任务。

为了打破固定光栅扫描的单向约束,近期的随机排列自回归(Permutation AR,如 RandAR 等)方法尝试在训练阶段随机打乱 token 顺序,并通过向模型注入待预测 patch 的位置信息来强迫模型适应任意生成顺序。然而,盲目的全排列随机化引入了高达 \(N!\) 的巨大搜索空间。二维图像与自然语言不同,具有强烈的局部连续性(locality)和中心偏置(center bias)。全随机打乱完全破坏了相邻 token 之间的拓扑连通,使得模型在没有周围上下文连通支持的情况下极难预测远距离的孤立 patch,导致条件模型熵剧烈波动、训练收敛极为缓慢,最终在相同算力预算下的图像生成 FID 显著劣于传统光栅扫描模型。

因此,视觉自回归模型的核心矛盾在于:如何在保持解码器极简架构的前提下,既获得支持任意局部编辑的灵活生成顺序,又不破坏图像固有的局部连续性与中心偏置先验。本文提出不再面向无约束的任意排列做随机化,而是将图像 patch 视作网格图(taxicab lattice),采样以边角为根节点的均匀生成树(Uniform Spanning Tree)。核心 idea:利用网格图上以随机边角为根节点的均匀生成树的广度优先遍历(BFS)序列作为自回归生成顺序,在几何结构上保证局部连通性与由外向内的先验,并通过 BFS 深度单调性配合拒绝采样高效实现图像前缀连接与后缀补全。

方法详解

整体框架

STAR 的核心理念是将原本全排列随机化的生成顺序约束在由图像网格拓扑所定义的均匀生成树集合中。首先将分辨率为 \(h \times w\) 的图像划分为 \(N = h \times w\) 个 patch token,构建正则出租车网格图 \(G = (V, E)\)。每个 token 对应图中的一个顶点 \((i, j)\),相邻 token 之间连有无向边。STAR 在训练和生成时,随机选取图像的四个边角之一作为树根 \(r\),通过 Wilson 算法采样一棵均匀生成树 \(T\),并提取该生成树的 BFS 遍历顺序 \(\tau_s\)。自回归 Transformer 按照此遍历顺序逐个预测下一个 token,并通过额外的可学习位置编码注入下一步预测的目标位置。在下游图像修复任务中,通过在未掩码区域做拒绝采样构造最大深度顶点位于掩码边界的生成树,无缝衔接掩码区域树,将未掩码观测天然排在前缀(prefix),将待修复区域排在后缀(postfix),实现原生条件生成。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["图像 Patch 离散化<br/>构建网格图 G=(V,E)"] --> B["均匀生成树采样<br/>选定边角根节点 r 运行 Wilson 算法"]
    B --> C["BFS 遍历序列提取<br/>按深度由浅入深保证局部连通"]
    C --> D["自回归 Transformer 预测<br/>因果掩码 + 目标位置嵌入条件"]
    D --> E["双模式生成输出"]
    E -->|无条件/类别条件生成| F["沿树遍历逐 Token 采样"]
    E -->|局部图像编辑/修复| G["前缀约束拒绝采样<br/>观测区置于前缀,修复区置于后缀"]

关键设计

1. 均匀生成树序列重排:平衡随机多样性与图像几何先验 全排列自回归将序列空间放宽到所有 \(N!\) 种可能,其中充斥着大量缺乏几何邻近性的离散跳跃,导致模型在预测孤立 token 时的条件熵大幅上升。相比之下,在正则网格图上,生成树的总数渐进满足 \(\exp(N \cdot z_G)\)(在二维网格图上常数 \(z_G \approx 1.166\)),样本空间既保证了组合爆炸式的充沛随机性以支持灵活的生成路径,又天然将每个生成步骤限制在网格的局部连通分支上。STAR 采用 Wilson 循环擦除随机游走算法(Loop-Erased Random Walk),能在 \(O(N \log N)\) 的极低复杂度内高效抽取无偏的均匀生成树 \(T \sim \mathcal{T}(G, r)\)。以随机选择的边角点作为根节点 \(r\),生成树的分支必然从非显著的图像边界向中心区域辐射展开,自然契合了自然图像中主体大多聚集在中央的中心偏置先验,同时保证了生成前缀集合在网格上的局部连通。

2. 基于 BFS 遍历的深度单调性与拒绝采样:高效实现 Inpainting 后缀补全 为了支持图像修复,序列顺序必须满足:未掩码区域 \(V \setminus V_M\) 中的所有 patch 必须先于被掩码区域 \(V_M\) 生成,即未掩码部分作为前缀 \(\tau_\text{pre}\),掩码区域作为后缀 \(\tau_\text{post}\)。若采用深度优先搜索(DFS),要求最后一个被访问的顶点必须恰好落在掩码边界上,其满足条件的概率极低。而广度优先搜索(BFS)天然具备由浅入深的层级性质。STAR 证明,只要未掩码子图 \(G' = V \setminus V_M\) 上的生成树 \(T_{G'}\) 的最大深度顶点集合 \(V_{G'}^\text{max}\) 中至少有一个顶点位于掩码边界 \(B_i\) 上,即可保证被掩码区域内各节点的 BFS 深度严格大于未掩码区域:

\[\text{depth}_T(v_M) \ge \text{depth}_T(v), \quad \forall v \in V \setminus V_M, \; v_M \in V_M\]

在推理时,算法先在未掩码区域中选择距离掩码边界最远的边角作为根节点 \(r_{G'}\),然后通过极少轮次的拒绝采样得到满足边界最大深度的子树 \(T_{G'}\);接着在掩码区域独立采样局部生成树,并通过单边连入最大深度边界点。这一设计使得图像修复无需微调或修改模型权重,仅凭序列重构即可达成端到端的高质量补全。

3. 极简因果 Transformer 适配:零侵入保留标量生成优势 STAR 严格保留标准纯解码器 Transformer 的因果自注意力架构与自回归训练范式。为了告知模型当前自回归步正在预测网格上的哪一个几何位置,STAR 引入与 RandAR 和 \(\sigma\)-GPT 一致的轻量化机制:除了当前 token 的内容嵌入和顺序位置编码外,额外引入待预测目标 token 的可学习二维空间坐标嵌入(Learnable Next-Token Positional Embedding)。模型在每一步预测输入序列中下一位置的概率分布:

\[\mathcal{L}_\text{STAR}(\theta) = - \mathbb{E}_{\tau \sim \tau(\mathcal{T}(G, r))} \left[ \sum_{i=1}^N \log p_\theta(x_{\tau_i} \mid x_{\tau_{<i}}) \right]\]

无需复杂的注意力掩码重塑,也无需分块结构或多流注意力,最大程度保留了解码器 Transformer 易于规模化扩展(Scaling)的工程红利。

损失函数 / 训练策略

模型采用标准自回归交叉熵损失进行端到端优化。训练阶段对每张图像动态采样根节点和生成树,生成对应的 BFS 遍历序列 \(\tau\)。分词器采用 MaskGIT 预训练的 VQGAN tokenizer(下采样率为 16,码本大小为 1024,在 \(256 \times 256\) 图像上生成 \(16 \times 16 = 256\) 个 token)。训练采用 AdamW 优化器,学习率配合线性预热与余弦衰减调度;采用 10-crop 数据增强,总批大小设为 2048,训练 250k 步(约 400 epochs)。在推理时结合 Classifier-Free Guidance(CFG),采用 Power-Cosine CFG 调度策略,无需截断 top-\(k\) 或 top-\(p\) 采样即可取得高质量图像生成。

实验关键数据

主实验

在 ImageNet-1k \(256 \times 256\) 类别条件图像生成任务上,STAR 与扩散模型(DiT)、固定光栅扫描自回归模型(LlamaGen、RAR)、多尺度自回归模型(VAR)以及随机排列模型(RandAR)进行了全面对比。结果如 Table 1 所示:

模型架构 模型名称 参数量 FID (↓) IS (↑) Precision (↑) Recall (↑)
Diffusion DiT-XL/2 675M 2.27 278.2 0.83 0.57
Raster-scan AR LlamaGen-XXL 1.4B 3.09 253.6 0.83 0.53
Raster-scan AR RAR-B 261M 1.95 290.5 0.82 0.58
Raster-scan AR RAR-XL 955M 1.50 306.9 0.80 0.62
Raster-scan AR RAR-XXL 1.5B 1.48 326.0 0.80 0.63
Block-wise AR VAR-d30 2.0B 1.92 323.1 0.82 0.59
Randomized AR RandAR-XL 775M 2.25 314.2 0.80 0.60
Randomized AR RandAR-XXL 1.4B 2.15 322.0 0.79 0.62
STAR (Ours) STAR-B 261M 2.24 295.3 0.82 0.57
STAR (Ours) STAR-L 461M 1.98 322.0 0.82 0.58
STAR (Ours) STAR-XL 955M 1.65 333.2 0.80 0.62
STAR (Ours) STAR-XXL 1.5B 1.55 338.8 0.81 0.62

消融实验

在相同算力预算(以 Base 规模配置为准)下对生成顺序策略、退火策略及推理遍历方式进行详细消融分析(原论文 Table 3),并测试各方法在跨 0.1~0.9 掩码比率下的平均图像修复性能(原论文 Table 2):

表 1:不同序列重排与退火策略消融(原论文 Table 3)

训练起始顺序 训练终止顺序 推理顺序 生成 FID (↓) 生成 IS (↑) 修复 FID (↓) 修复 IS (↑) 说明
光栅扫描 光栅扫描 光栅扫描 2.04 266.5 3.91 82.7 传统固定光栅扫描基线
全排列随机 全排列随机 全排列随机 3.57 291.4 2.57 108.7 RandAR 纯随机排列
全排列随机 全排列随机 生成树 BFS 3.58 234.9 2.43 105.8 随机训练后测试树遍历
全排列随机 生成树 BFS 生成树 BFS 2.31 285.2 2.35 107.9 排列退火至生成树
生成树 BFS 生成树 BFS 生成树 BFS 2.24 295.3 2.39 108.8 STAR 默认策略

表 2:图像生成与修复综合能力对比(掩码比率 0.1~0.9 平均,原论文 Table 2)

模型类型 模型 参数量 生成 FID (↓) 生成 IS (↑) 修复 FID (↓) 修复 IS (↑)
Diffusion DiT-XL/2 675M 2.27 278.2 4.58 50.4
Raster-scan AR RAR-XL 955M 1.50 306.9 3.40 88.2
Randomized AR RandAR-XL 775M 2.25 314.2 2.58 60.3
Randomized AR STAR-XL 955M 1.65 333.2 2.07 111.3

关键发现

  • 生成与修复的双赢:全排列模型(RandAR)虽具备修复能力但生成 FID 较差(XL 为 2.25);光栅扫描模型(RAR-XL)生成 FID 虽优(1.50)但在修复任务上严重失真(修复 FID 恶化至 3.40);STAR-XL 实现了生成 FID 1.65、修复 FID 2.07 的卓越平衡,Inception Score 高达 111.3。
  • 拒采样效率与遍历算法选择:原论文 Table 4 显示,采用 BFS 配合最远边角根节点策略时,在 0.1~0.9 掩码比率下所需的拒绝采样试验次数极低(0.1 掩码下仅需 4.41 次,0.9 掩码下仅 1.15 次),失败率恒为 0.0%;而 DFS 在 0.1 掩码下试验次数高达 89.37 次,失败率高达 82.3%。
  • 生成树额外计算开销可忽略:使用 Wilson 算法抽取生成树及计算 BFS 序列平均单图耗时仅 0.56 ms,在模型总推理延迟中占比低于 0.0004%,几乎没有运行时惩罚。

亮点与洞察

  • 将图论生成树精巧引入自回归视觉表征:避开了无约束全排列空间(\(N!\))与刚性固定单向扫描(1 种)两个极端,以指数级子集 \(\exp(N \cdot z_G)\) 恰到好处地覆盖了多样性与几何先验。
  • 巧妙利用 BFS 深度单调性解决 Inpainting 前缀约束:无需对掩码边界进行繁琐的启发式打补丁,只需保证未掩码子树的最深节点落在边界上即可确保拓扑层级正确,使通用 AR 模型兼具前缀补全功能。
  • 对多模态早融合架构具备极高迁移价值:由于未引入复杂的双向注意力修改,该序列重排思想可天然推广到文本-图像统一建模(如 Chameleon 等 early-fusion MLLM),赋予大语言模型原生、非侵入式的视觉局部重绘与交互式编辑能力。

局限与展望

  • 依赖未掩码区域与边角的连通性假设:当前算法假设未掩码区域在网格图上保持单连通,且未覆盖全部四个边角;当用户输入极度碎片化、断开的不连通掩码或完全覆盖四角的掩码时,必须退化或借助更复杂的虚拟连通图处理。
  • 离散 Tokenizer 重建瓶颈:当前基于 VQGAN 离散码本进行建模,相比连续潜空间的扩散损失模型(如 MAR、FLUX),在极致高频细节上可能受到离散词表保真度的限制。
  • 高分辨率下的扩展挑战:当特征图尺度进一步扩大(如 \(64 \times 64\) 或更高)时,尽管 Wilson 算法为 \(O(N \log N)\),但极大图上的生成树深度分布更加离散,拒绝采样的接受率在极端不规则掩码下可能会有所波动。

相关工作与启发

  • vs RAR (Randomized Autoregressive):RAR 在训练初期使用随机排列作为预训练任务,随后必须通过退火策略强行收敛到固定光栅扫描以恢复生成质量,这牺牲了推理阶段的自由编辑能力;STAR 始终在生成树子集上训练与推理,免除超参退火,原生保留修复能力。
  • vs RandAR:RandAR 直接采样无约束的随机排列,破坏局部连续性与中心偏置,模型预测熵居高不下;STAR 将排列空间收敛至网格图生成树,生成 FID 明显优于 RandAR(STAR-XL 1.65 vs RandAR-XL 2.25)。
  • vs VAR (Visual Autoregressive):VAR 采用由粗到细的尺度树结构,需要定制多尺度分词器与分块掩码注意力;STAR 仅在 patch 拓扑序列上做图遍历重排,完全沿用标配 Transformer。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从图拓扑生成树的角度重新审视自回归视觉生成的序列顺序,理论优雅且契合图像固有先验。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖生成、修复、消融分析、条件熵变化及表征学习线性探测,证据链极其扎实。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰,图表逻辑连贯,算法步骤与复杂度分析详实透彻。
  • 价值: ⭐⭐⭐⭐⭐ 为纯解码器视觉自回归模型的序列设计和原生区域编辑提供了极具实用性与通用性的新基石。