HiRAE: Hierarchical Representation Autoencoding with Residual Budgets
Xuanyu Zhu, Yan Bai, Yang Shi, Yihang Lou, Yuanxing Zhang, Tengfei Liu, Jing Jin, Yuan Zhou
cs.CV, cs.AI
2026-09-29
HiRAE把冻结的DINOv3-L全部24层分成浅中深三组,用相对最深层2.5%/7.5%/15%的范数预算做残差融合,重建rFID从0.299降到0.209,文生图SFT后GenEval从84.86升到87.70。
Representation Autoencoder(RAE)把冻结的视觉编码器接到可训练解码器上,扩散模型直接在这个表征空间里生成。RAEv2 已经能用 DINOv3-L 做这件事,但它对中间层是固定加权,层选哪些要靠反复训练、反复评估。最后一层语义整齐,细笔画和局部颜色却更多留在中间层。
把中间层拿来用有代价。DecQ 加浅层 query,重建变好、生成变差。DRoRAE 学全部层的融合,但要先单独训融合、再适配解码器。核心张力就一条:重建损失会把 latent 往浅层细节拽,拽过头,扩散模型就学不动。
HiRAE 的锚点永远是最深层特征 H23。其他层不准替换它,只允许加一段被范数封顶的残差。
主配置 HiRAE-24 冻结 DINOv3-L 全部 24 层,每层一个 token-wise MLP expert(1024→4096→1024),把该层特征变成 Uℓ。路由器只看最深层,在每个空间位置给出带符号、ℓ2 归一化的层权重。24 层按深度切成三组,组内加权求和得到残差 Rg:
残差先按组 dropout,再按整图 Frobenius 范数封顶;超了就缩小。三组加起来,总修正被三角不等式卡在锚点的 0.25 倍以内。最后做 token-wise LayerNorm,输出仍是 16×16×1024,token 数和通道数都不增加。
浅层预算最紧,是故意的。重建目标会本能地去抢浅层高频,不设上限,分布会漂成「好解码、难生成」。DRoRAE 靠分阶段适配压这个漂移;HiRAE 把上限写进前向计算,融合模块和解码器一次联合训练(像素 L1、感知损失、对抗损失,解码器输入加噪 τ=0.8),编码器始终冻住。随后冻 tokenizer,按 RAEv2 的内部引导框架训 DiT,生成器 80 epoch。
另有 HiRAE-7,只用 RAEv2 已选的 7 层,专家参数大约 59M,对比 HiRAE-24 的约 202M。
ImageNet-256,同一套 16×16×1024 latent:
| 方法 | rFID ↓ | 5K PSNR ↑ | 5K LPIPS ↓ | 引导 gFID ↓ |
| RAEv2 | 0.299 | 22.667 | 0.074 | 1.060 |
| HiRAE-7 | 0.217 | 25.068 | 0.049 | 1.038 |
| HiRAE-24 | 0.209 | 26.377 | 0.043 | 1.038 |
引导设置下训 80 epoch,HiRAE-24 的 Inception Score 从 255.300 到 257.823,六表征聚合 FD6r 从 2.170 到 1.856。无引导 gFID 是 2.129,好于 RAEv2 的 K=23 变体 3.010,但差于 RAEv2 主结果 1.650;无引导 IS 210.339,也低于 228.000。重建和引导生成可以同时变好,无引导生成没有一起跟上。
文生图在同一套约 875M 的 DiT 上比,预训练 100K step,推理 CFG=6、关掉内部引导:
| 阶段 | 指标 | RAEv2 | HiRAE-24 |
| 预训练 | GenEval | 56.42 | 60.94 |
| 预训练 | DPG-Bench | 81.22 | 82.73 |
| SFT 后 | GenEval | 84.86 | 87.70 |
| SFT 后 | DPG-Bench | 84.90 | 86.35 |
SFT 后 GenAI-Bench 从 71.69 到 72.66。同一协议下 FLUX autoencoder 配这套 DiT,SFT 后 GenEval 是 84.13。HiRAE-7 在每个能比的栏里也高于 RAEv2,HiRAE-24 再往上走一截。
消融把机制钉死了。同样 24 层、同样 24 个 expert,把分组预算换成 DRoRAE 式全局插值(β=0.2),rFID 能到 0.065,比 0.209 好看很多,引导 gFID 却是 1.551,FD6r 是 3.375。把逐层 expert 换成深度 mode expert 并去掉残差正则,rFID 能到 0.023,20 epoch 引导 gFID 升到 7.905;加上正则后 20 epoch gFID 回到 2.410。重建数字可以很低,生成跟着崩。
融合后同类 10 近邻从锚点的 79.634% 变成 79.536%,类邻域几乎不动;100 张图的空间有效秩从 130.360 升到 154.893,张张都升。10% 相对范数的 Gaussian / 低频 / 高频扰动下,解码 LPIPS 变化大约是 RAEv2 的 21%–22%。
这是 RAE 路线上一档很具体的正则:全层融合可以学,浅层必须加预算。已经在 RAEv2 和 DINOv3-L 上做 tokenizer 的人,HiRAE-24 不改 latent 形状,生成器接口不用动。只想提高重建、层已经选好,HiRAE-7 用 7 个 expert 就能把引导 gFID 打到同样的 1.038。
引导 gFID 从 1.060 到 1.038,是边际改进。拉开差距的是重建(rFID 降约 30%,配对 5K 上 PSNR +3.710 dB)和文生图对齐。无引导生成还不如 RAEv2。别把它读成生成全面更好。
论文没有单独的 Limitations 节。数字自己会说话。
无引导 gFID 2.129 对 1.650,内部引导抹平了大部分生成差距;关掉引导,这个 latent 并不更好建模。DRoRAE 风格全层融合重建明显更准,HiRAE 是主动拿重建换生成。0.025 / 0.075 / 0.150 这组预算和 dropout 怎么定下来的,正文没有搜索记录。
主干只测了 DINOv3-L,分辨率停在 256。tokenizer 推理要跑 24 个 expert。文生图对照的是同一套自训 DiT;表里的 FLUX-VAE 是 FLUX 的 autoencoder 配他们的 DiT,不是 FLUX 原系统。重建示例经过像素误差筛选,引导生成图从 384 张里目视挑了 24 张。
三组划分能去掉 DRoRAE 的融合-only 阶段,这些结果成立。换编码器、上更高分辨率、预算要不要重搜,没有答案。