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 fuses all 24 DINOv3-L layers as norm-capped residuals; ImageNet rFID drops from 0.299 to 0.209 vs RAEv2, and post-SFT GenEval from 84.86 to 87.70.
Representation Autoencoders (RAE) freeze a pretrained vision encoder, train a decoder on top, and run diffusion in that feature space. RAEv2 already does this with DINOv3-L, but it aggregates a hand-chosen subset of layers with fixed weights. Final-layer features are semantically tidy; stroke-level detail and local color live in the intermediate layers.
Using those layers has a cost. DecQ adds shallow detail queries: reconstruction improves, generation gets worse. DRoRAE learns full-depth fusion, but only after a fusion-only stage before the decoder is adapted. Reconstruction loss will pull the latent toward shallow detail. Pull too far and the diffusion model cannot fit the distribution.
HiRAE keeps the deepest feature H23 as a fixed anchor. Other layers never replace it; they add a residual that is capped by a norm budget.
HiRAE-24 freezes all 24 DINOv3-L layers. Each layer gets a token-wise MLP expert (1024→4096→1024). A router, conditioned only on the deepest feature, produces signed, ℓ2-normalized weights at every spatial token. Layers are split into three depth groups; weighted expert outputs form one residual per group:
Each residual is dropped out, then scaled down if its Frobenius norm exceeds the cap. The triangle inequality keeps the total correction within 0.25 times the anchor. A token-wise LayerNorm yields a 16×16×1024 latent, same token count and width as RAEv2.
The tight shallow budget is the point. Reconstruction will grab high-frequency cues from early layers if nothing stops it. DRoRAE stages fusion then decoding to contain that drift. HiRAE writes the cap into the forward pass and jointly trains fusion and decoder (pixel L1, perceptual loss, adversarial loss, decoder-input noise τ=0.8) with the encoder frozen. The tokenizer is then frozen and a DiT is trained for 80 epochs under RAEv2's internal-guidance recipe.
HiRAE-7 is the compact sibling: the same seven layers RAEv2 already selected, about 59M expert parameters versus about 202M for HiRAE-24.
On ImageNet-256, same 16×16×1024 latent:
| Method | rFID ↓ | 5K PSNR ↑ | 5K LPIPS ↓ | Guided 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 |
After 80 generator epochs with internal guidance, Inception Score moves from 255.300 to 257.823 and the six-space aggregate FD6r from 2.170 to 1.856. Unguided gFID is 2.129, better than RAEv2's K=23 variant at 3.010 and worse than RAEv2's main 1.650; unguided IS is 210.339, also below 228.000. Reconstruction and guided generation can improve together. Unguided generation does not.
Text-to-image uses the same 875M DiT, 100K pretraining steps, CFG=6 at inference with internal guidance off:
| Stage | Metric | RAEv2 | HiRAE-24 |
| Pretrain | GenEval | 56.42 | 60.94 |
| Pretrain | DPG-Bench | 81.22 | 82.73 |
| After SFT | GenEval | 84.86 | 87.70 |
| After SFT | DPG-Bench | 84.90 | 86.35 |
GenAI-Bench after SFT goes from 71.69 to 72.66. Under the same protocol, a FLUX autoencoder paired with this DiT scores 84.13 GenEval after SFT. HiRAE-7 beats RAEv2 on every reported T2I column; HiRAE-24 goes further.
The ablation is the mechanism. Same 24 layers, same 24 experts, swap the group budgets for DRoRAE-style global interpolation at β=0.2: rFID falls to 0.065, well below 0.209, while guided gFID rises to 1.551 and FD6r to 3.375. Replace layer-wise experts with depth-mode experts and drop residual regularization: rFID hits 0.023, and guided gFID at 20 epochs is 7.905; with regularization that 20-epoch gFID is 2.410. Reconstruction can be driven very low. Generation collapses with it.
After fusion, same-class 10-NN stays at 79.536% versus 79.634% on the deep anchor. Spatial effective rank on 100 images rises from 130.360 to 154.893, in every image. At 10% relative latent perturbation, decoded LPIPS change is about 21%–22% of RAEv2's across Gaussian, low-frequency, and high-frequency directions.
This is a concrete regularizer on the RAE stack: full-depth fusion is learnable if shallow layers are budgeted. Anyone already on RAEv2 with DINOv3-L keeps the same latent shape and generator interface. If the layer subset is already chosen and the goal is reconstruction, HiRAE-7 matches the 1.038 guided gFID with seven experts.
The guided gFID move from 1.060 to 1.038 is marginal. The gaps that matter are reconstruction (rFID down about 30%, +3.710 dB PSNR on the paired 5K set) and T2I alignment. Unguided generation is worse than RAEv2. This is not a blanket generation win.
There is no Limitations section. The numbers still constrain the claim.
Unguided gFID 2.129 versus 1.650 means internal guidance is doing a lot of the generation work; without it, the fused latent is not easier to model. DRoRAE-style full-depth fusion reconstructs more accurately. HiRAE spends reconstruction quality to keep generation intact. How the 0.025 / 0.075 / 0.150 caps and dropout rates were chosen is not reported as a search.
Only DINOv3-L is tested, at 256 resolution. Tokenizer inference runs 24 experts. T2I comparisons use one in-house DiT; the FLUX-VAE row is the FLUX autoencoder with that DiT, not the FLUX generator. Reconstruction figures were screened by pixel error; the 24 guided samples were picked by eye from the first 384.
Three depth groups remove DRoRAE's fusion-only stage. That result holds. Other backbones, higher resolution, and whether the budgets need retuning are open.