BudgetPix：面向像素空间图像扩散的计算自适应分词框架 BudgetPix: Compute-Adaptive Tokenization for Pixel-Space Image Diffusion

完整流程分五步。第一步，布局构建：图像转 luma 并做 512-bin 直方图；patch 尺度为 $s = 2^j p$（$j = 0,1,\dots,N-1$）。从最粗的 $np \times np$ 网格开始，对每个 cell $R$ 计算熵 $H(R) = -\sum_{k=1}^{K} q_k(R)\log_2 q_k(R)$；若 $H(R) \le \tau_s + \delta$ 则保留为单 token，否则四等分递归细分，直到最细的 $p \times p$。阈值取 $\tau_{2p}=6.0$、$\tau_{4p}=4.0$ bits（沿用 APT 不调参），$\delta$ 逐图逐步随机。训练时熵来自真值图 $x$；推理时前 $n_w$ 个 warm-up 步（JiT 用 5 或 10 步）跑均匀网格，之后从当前预测 $\hat{x}$ 估计布局，每 $m=5$ 步重建；'预算模式'下对 $\delta$ 做二分搜索，硬性保证每图 token 数不超过预算。第二步，编码：尺度 $s=np$ 的 patch $P_i$ 在 $n=1$ 时直接 $x_i = E(P_i) + e_{\text{scale}}(s)$；$n>1$ 时 $x_i = E(\text{resize}_p(P_i)) + e_{\text{scale}}(s) + P_{As}(\text{split}(P_i))$，其中 Patch Aggregator 把大 patch 切成 $n^2$ 个子块分别过 embedder $E$、加 patch 内相对位置编码 $e_{pos}$，再经卷积+MLP 的 scale mixer 融合成高频修正向量。第三步，中央 Transformer：用预训练权重初始化，位置编码（RoPE）放在 token 中心的全局坐标系，照常处理变长序列。第四步，解码：$n=1$ 时用预训练头直接输出 $p\times p$ patch；$n>1$ 时先解码再双线性上采样到 $np\times np$，Patch Refiner 把预测块切成 4×4 子块、经两层自注意力（AdaLN 以该 token 的隐状态 $h_i$ 为条件）产生高频残差，即 $\hat{P}_i^{(s)} = \text{resize}_{np}\big(\hat{P}_i + P_{Rs}(\text{split}(\hat{P}_i), h_i)\big)$，最终所有 patch 拼成完整的 $\hat{x}$。第五步，训练与采样：损失 $\mathcal{L}(\theta) = \mathbb{E}_{x,y,L,t,\epsilon}\, w(t)\lVert x-\hat{x}\rVert_2^2$，$w(t)$ 沿用各基模型自身的调度权重；全模型微调，新增模块学习率为骨干的 8 倍，EMA 0.9999；采样时warm-up 后每 5 步重估布局，布局在期间冻结。

与推理时 token 合并（ToMe/FeatSim）的本质区别在于：合并方法在冻结权重上按特征相似度把 token 粗暴平均，模型从未被训练过去'读'粗 token，预算一降就系统性崩溃（JiT-L/32 在 50% 预算下 FID 从 2.66 恶化到 49–110），而且合并操作本身有额外开销，实测 ToMe 的墙钟加速比反而低于直接截断；BudgetPix 通过微调让模型真正掌握多尺度 token 这门语言。与 FlexiDiT 等均匀变 patch 尺寸的方法相比：后者只能全图统一换粒度，无法在同一张图里同时保住平坦区与细节区。与并发工作 RTI 相比：RTI 用 LoRA 挂在冻结模型上消费 region token，BudgetPix 则完全微调、从第一个 block 起改变输入分布，且布局来自模型自身对干净图像的预测而非内部特征。与重训 latent tokenizer（如 CAT）相比：那些方法要整体替换 tokenizer，BudgetPix 保持像素空间架构的纯粹性，新增参数仅 2.26M–4.77M（约占骨干 1%）。工程上的新颖点还包括：零初始化保证训练起点严格等于发布模型；预算模式通过二分 $\delta$ 给出每图 token 上限的硬保证；Patch Refiner 只在单个粗 token 内部做短序列注意力，成本不随图像分辨率增长。

弱点一：布局信号低级且与语义脱钩。像素强度熵无法区分'高熵但可牺牲的纹理'（草地、胶片噪点）与'低熵但承载主体的区域'（纯色墙上的一只猫），文生图场景下布局完全不感知文本提示——模型想画什么不影响算力往哪花。改进方向：引入可学习的布局头或语义显著性（DINO 特征、文本交叉注意力图）与熵联合打分。弱点二：warm-up 阶段的布局不稳与错误锁定。早期预测模糊时熵估计噪声大，Figure 41 显示无 warm-up 时布局与最终布局的差异比例可达 20–30%；构图复杂的场景可能在低细节阶段被错误布局锁定。改进：不确定性感知的布局更新（对预测做方差估计，高方差区域保守用细 token），或前若干步强制细节区域细粒度。弱点三：全分辨率阶段瓶颈。PixelDiT 的像素阶段只做条件修正而非结构性加速，改进方向是对像素级解码也做空间自适应采样，只对合并区域做轻量解码。弱点四：布局在 $m=5$ 步内冻结，而生成后期细节逐渐清晰，可能出现'该细分却没细分'的时间错配。改进：引入只分裂不合并的单向细化调度，或按噪声水平调整粒度。弱点五：单步模型退化为类别级布局先验，改进方向是在蒸馏目标中显式加入布局条件，或用本文的 random layout 增强训练单步模型使其对任意布局鲁棒。弱点六：批处理内每图同预算导致 GPU 利用率次优，论文 D.9 的 pooled budget 实验已显示批内共享预算能提升保真度（PSNR 26.35 vs 24.57 @70%），值得发展为默认调度器。

作者明确提出方向：学习'在哪里放 token'——用可学习模块取代熵启发式，从数据中学出与语义和生成过程对齐的布局，这是最自然的下一步。基于本文成果可延伸的方向还包括：其一，视频与时空扩散——视频帧间背景高度冗余，熵四叉树可自然扩展为时间-空间联合自适应 tokenization，收益可能比图像更大；其二，与正交加速路线叠加——论文已证明 BudgetPix 与一步蒸馏（MeanFlow）兼容，还可继续与少步求解器、特征缓存、量化组合；其三，语义/条件感知布局——用文本嵌入或交叉注意力调制各区域的熵阈值，让'模型想画什么'决定'哪里多花算力'；其四，批级预算调度——把 pooled budget 从消融实验变成产品特性，按批次动态分配 token 以提升硬件利用率；其五，迁移到其他像素空间任务——超分辨率、图像编辑、inpainting 的空间复杂度差异更大，自适应 tokenization 的收益可能更显著；其六，协同设计布局估计与少步采样器，降低 warm-up 开销；其七，探索布局随生成阶段演化的动态调度，让 token 在去噪后期按需分裂。

复现条件相当好。论文有公开项目页（karaozgur.com/BudgetPix），所有基座模型（JiT、MiniT2I、PixelDiT、pMF）均为已发布权重；附录 C 给出罕见完整的训练配方：熵阈值（$\tau_{2p}=6.0$、$\tau_{4p}=4.0$ bits，512-bin luma 直方图）、三种布局比例（$a=0.15$、$b=0.1$）、逐模型的优化器设置（AdamW、骨干 lr 3.5e-5、新增模块 8 倍 lr、EMA 0.9999 fp32、bf16 autocast）、训练量（JiT-L 50 epoch/6410 万图像、MiniT2I-L 4 万步/1020 万图像、PixelDiT 3 万步）、数据混合（MiniT2I-L 用 BLIP3o-60k + DALL-E 3 + ShareGPT-4o-Image 共 12 万图；JiT 用 ImageNet-1k；PixelDiT 用 PD12M 20 万图）以及 Table 16 逐家列明的模块架构。新增参数仅 2.26M–4.77M。算力方面：微调多数模型用 4 GPU（pMF 用 7 GPU），评估计时在单张 H100 上完成，实验室级集群即可承受。主要不确定性是官方代码是否随项目页开放（论文正文未明说）；若未开源，熵四叉树 + 双路径编解码器 + 变长序列注意力 + 预算模式二分约需 1–2 周工程实现，难度中等偏上；若开源则可直接复用其逐模型配置。评估协议（FID-50k、GenEval、DPG-Bench、PartiPrompts、用户研究与 VLM 判官提示词全文）在附录 C.5–C.7 中均可照做，复现门槛主要在工程而非科学。