PaperlayerView exact revision on arXiv

Source-linked AI summary

HiRAE: Hierarchical Representation Autoencoding with Residual Budgets

Xuanyu Zhu, Yan Bai, Yang Shi, Yihang Lou, Yuanxing Zhang, Tengfei Liu, Jing Jin, Yuan Zhou

arXiv:2609.37775v1cs.CVcs.AI

TL;DR

Pretrained visual representations can omit fine detail needed for faithful reconstruction, while learned fusion must also remain compatible with generation. HiRAE uses depth-dependent residual controls to fuse the full encoder hierarchy. It improves reconstruction and text-to-image alignment over RAEv2 while maintaining competitive guided generation.

  • Problem

    Final-layer visual representations can omit reconstruction detail, while using intermediate layers for learned fusion raises concerns about compatibility with generation.

    Source: §1 · p. 2

  • Method

    HiRAE fuses all encoder layers through depth-group residual corrections to the deepest representation, with tighter norm budgets for shallower groups.

    Source: §1 · p. 2, §3 · p. 3, §3.3 · p. 5

  • Results

    HiRAE-24 reduces reconstruction FID by 30% relative to RAEv2, improves text-to-image alignment across three benchmarks before and after fine-tuning, and maintains competitive guided generation.

    Source: §1 · p. 2, §5.1 · p. 8

  • Takeaways & Limitations

    HiRAE-24 improves reconstruction over RAEv2 while achieving competitive guided generation, supporting its use for both reconstruction and generation.

    Source: §5.1 · p. 8, §6 · p. 10

Abstract

from arXiv · show

Pretrained visual representations support image generation, but may not fully preserve the fine-grained details needed for faithful reconstruction. Meanwhile, intermediate encoder layers contain complementary visual details, but learning to fuse them for reconstruction can produce a latent distribution that is difficult to model. Existing fusion methods require empirical tuning of layer selection or staged optimization of fusion and decoding, increasing configuration effort or training complexity. We introduce HiRAE (Hierarchical Representation Autoencoder), which learns a hierarchical fusion framework over the full encoder hierarchy to improve reconstruction fidelity while maintaining compatibility with generative modeling. HiRAE groups encoder layers by depth and learns residual corrections to the deepest representation. Group-wise norm caps bound these corrections relative to the deep anchor, with tighter budgets for shallower groups. Our HiRAE-24 preserves the latent token count and channel dimension. On ImageNet-256, HiRAE-24 reduces reconstruction FID from 0.299 to 0.209 relative to RAEv2 while maintaining competitive guided generation quality. For text-to-image generation, HiRAE-24 improves alignment over RAEv2 on GenEval, DPG-Bench, and GenAI-Bench both before and after supervised fine-tuning. Under the same generator-training and evaluation protocol, post-fine-tuning GenEval increases from 84.86 to 87.70.

1 Introduction

Pretrained encoders offer semantic representations but can omit reconstruction detail, while using intermediate features raises concerns about compatibility with generation. HiRAE learns depth-aware fusion with residual controls; it improves reconstruction and text-to-image alignment while retaining competitive guided generation.

1 Introduction

  • Intermediate layers retain reconstruction detail missing from final-layer features, but learned fusion can favor shallow detail without accounting for its effect on generation quality.Prior approaches instead use predefined aggregation, additional detail pathways, or staged fusion and decoder adaptation; selecting a suitable layer subset can require repeated training and evaluation.Source: §1 · p. 2
  • HiRAE-24 learns spatially varying contributions from all 24 layers of a frozen DINOv3-L encoder, using depth-group residual corrections around the deepest representation.The corrections have distinct norm budgets that increase with depth, and the fusion module and decoder are trained jointly without a separate fusion-only adaptation phase.Source: §1 · p. 2
  • HiRAE-24 reduces reconstruction FID by 30% relative to RAEv2 while maintaining competitive guided ImageNet generation.Under the shared text-to-image protocol, it also improves GenEval, DPG-Bench, and GenAI-Bench before and after supervised fine-tuning, with a 2.84-point post-fine-tuning GenEval gain.Source: §1 · p. 2
  • HiRAE-24 improves GenEval, DPG-Bench, and GenAI-Bench scores before and after supervised fine-tuning, with a 2.84-point post-fine-tuning GenEval gain.The comparison uses the paper's shared text-to-image protocol.Source: §1 · p. 2
  • Hierarchical fusion adds spatial detail while largely preserving class neighborhoods, and the learned tokenizer has lower decoding sensitivity to the tested latent perturbations.The fused representation preserves the original latent token count and channel dimension.Source: §1 · p. 2

2 Related work

Related work connects pretrained visual representations to generation through latent diffusion, alignment, and learned decoders. HiRAE extends this landscape with full-hierarchy learned fusion and depth-dependent residual budgets.

Visual representations for image generation

  • RAE pairs a frozen vision encoder with a learned decoder and trains diffusion in the encoder representation space; HiRAE adds learnable full-hierarchy fusion jointly trained with the decoder.Other alignment approaches supervise diffusion features, regularize tokenizer latents, or align jointly optimized autoencoders and diffusion models.Source: §2 · p. 3

Hierarchical fusion and detail enrichment

  • RAEv2 uses fixed aggregation of selected layers, while DRoRAE learns full-depth routed fusion before decoder adaptation; HiRAE adds depth-dependent residual budgets for joint fusion and decoder training.IDEAL, DecQ, and LV-RAE provide other detail-enrichment designs.Source: §2 · p. 3

Latent structure and generative modeling

  • Reconstruction detail changes the representation a generator must model, so improved reconstruction alone does not establish better generation.Prior work studies latent properties across tokenizer families; HiRAE analyzes spatial detail, class neighborhoods, and decoding response to latent perturbations.Source: §2 · p. 3

3 HiRAE: controlled hierarchical composition

HiRAE fuses the full frozen encoder hierarchy into a latent with the original token count and channel dimension. It combines layer-wise experts and spatial routing with depth-group residual controls, then jointly trains fusion and decoding before training a generator on frozen tokenizer latents.

3.1 Layer-wise experts

  • Separate token-wise MLP experts process encoder-layer features, while a router assigns expert contributions at each spatial location.HiRAE-24 uses all 24 layers of a frozen DINOv3-L encoder.Source: §3.1 · p. 3, §3 · p. 3

3.2 Routing

  • A shared projection of the deepest feature produces spatial routing scores, which are L2-normalized while retaining their signs.The router's weights assign layer contributions at each spatial location.Source: §3 · p. 3, §3.2 · p. 4

3.3 Residual regularization

  • Residual controls convert each depth group's routed feature sum into a correction to the deepest-layer anchor, then layer normalization produces the fused latent.Group-wise norm caps bound each correction relative to the anchor; shallower groups receive tighter budgets and stronger dropout, while middle and deep groups allow progressively larger corrections.Source: §3.3 · pp. 4–5

3.4 Training the tokenizer and generator

  • Tokenizer training updates fusion and decoder together while keeping the pretrained backbone frozen; generator training then freezes the tokenizer and trains a DiT on its latents.This removes DRoRAE's separate fusion-only adaptation phase.Source: §3.4 · p. 5

4 Reconstruction and generation

The section specifies datasets, metrics, and evaluation protocols for ImageNet reconstruction and generation and text-to-image alignment. It evaluates HiRAE-24 using EMA generators for ImageNet and follows RAEv2 for text-to-image training.

4 Reconstruction and generation

  • ImageNet tokenizer training uses the training split at 256 × 256 resolution, with reconstruction evaluated on the validation split and 50,000 generated images per configuration for FID.Text-to-image pretraining uses JourneyDB and BLIP3o caption subsets, followed by supervised fine-tuning on BLIP3o-60k.Source: §4 · p. 5
  • ImageNet evaluation reports reconstruction FID and generation FID and Inception Score, with PSNR and LPIPS also measured on a matched reconstruction subset.Text-to-image alignment is assessed with GenEval, DPG-Bench, and GenAI-Bench.Source: §4 · p. 5
  • Main comparisons use HiRAE-24's 16×16×1024 latent and EMA generators with and without internal guidance, while retaining class conditioning in both settings.Text-to-image pretraining uses 100K optimizer updates, and supervised fine-tuning continues from the corresponding pretrained weights.Source: Implementation details · p. 5

Implementation details.

HiRAE-24 improves ImageNet-256 reconstruction while retaining competitive guided generation, and it improves text-to-image alignment over RAEv2 before and after supervised fine-tuning.

4.1 Reconstruction quality

  • With the same encoder and latent shape, HiRAE-24 reduces rFID from 0.299 to 0.209, while PSNR rises from 22.667 to 26.377 dB and LPIPS falls from 0.074 to 0.043.On the matched subset, the visual comparisons also show finer text strokes and local colors than RAEv2.Source: §4.1 · pp. 5–6, Table 1 · p. 6, Figure 4 · p. 6

4.2 Image generation

  • After 80 generator-training epochs, HiRAE-24 achieves guided gFID 1.038 versus 1.060 for RAEv2, with IS 257.823 versus 255.300.The reported guided-generation comparison uses the same training duration.Source: §4.2 · p. 6, Table 3 · p. 7

4.3 Text-to-image generation

  • After SFT, HiRAE-24 gains 2.84, 1.45, and 0.97 points over RAEv2 on GenEval, DPG-Bench, and GenAI-Bench, respectively.The alignment advantage is also reported after pretraining; the cited table covers both stages.Source: §4.3 · pp. 7–8, Table 4 · p. 8

5 Ablation Studies and Analysis

The ablations favor layer-wise experts and residual regularization, while depth-group and latent-space analyses examine complementary detail and the fused representation’s structure and decoding behavior.

5.1 Ablation studies

  • Reconstruction fidelity alone does not determine guided-generation quality; HiRAE-24 improves reconstruction over RAEv2 while maintaining competitive guided generation.The ablations favor retaining layer-wise experts and residual regularization to balance the two outcomes.Source: §5.1 · pp. 8–9
  • Replacing 24 raw-layer experts with seven depth-mode experts raises rFID from 0.209 to 0.230 and guided gFID from 1.038 to 1.067 at epoch 80.Without residual regularization, the depth-mode configuration reaches rFID 0.023, while guided gFID and FD6-r rise to 7.905 and 14.722 at epoch 20, versus 2.410 and 2.650 for the regularized configuration.Source: §5.1 · pp. 8–9

5.2 Latent-space analysis

  • Fusion raises spatial effective rank from 130.360 for the deep anchor to 154.893 while mean spatial CKA remains 0.985 and same-class neighborhoods are largely preserved.Same-class 10-neighbor fraction is 79.634% for the deep anchor and 79.536% after fusion.Source: Figure 6 · p. 9, §5.2 · p. 9
  • At 10% relative latent perturbation, HiRAE-24 produces 21%–22% of RAEv2’s output LPIPS change across Gaussian, low-frequency, and high-frequency directions.The generator predicts HiRAE-24 latents with lower clean-latent mean squared error at five interior noise levels, but slightly higher error at both endpoints.Source: Figure 6 · p. 9, §5.2 · pp. 9–10

5.1 Ablation studies

  • Three depth groups yield the lowest guided gFID, with rFID close to the two-group result; removing any group raises reconstruction LPIPS.Equal-norm removals damage high-frequency detail more in the shallow group and low-frequency detail more in the deep group; the middle-group difference is small, with a paired 95% interval including zero.Source: §5.1 · p. 9, Table 7 · p. 9, Table 8 · p. 9

6 Conclusion

HiRAE uses depth-dependent residual bounds to fuse encoder features with the deepest representation, aiming to improve reconstruction while retaining the original latent dimensions. The conclusion reports higher reconstruction fidelity than RAEv2, competitive guided generation, and matched-image analyses linking added detail to depth groups while largely preserving class neighborhoods.

6 Conclusion

  • HiRAE-24 learns all 24 encoder layers without manual layer-subset selection, using depth-dependent residual bounds to control fusion and jointly training fusion with the decoder.It preserves the original latent dimensions and requires no separate fusion-only adaptation phase.Source: §6 · p. 10

A.1 HiRAE-24 architecture

  • Shallow, middle, and deep groups add residual corrections to the deepest anchor, with tighter norm caps for shallower groups.The three groups contain layers 0–7, 8–15, and 16–23; their caps sum to 0.250.Source: §A.1 · p. 14, Table 9 · p. 14

A.2 ImageNet training

  • Stage 1 jointly trains the decoder and fusion module using pixel, perceptual, and adversarial objectives.The diffusion generator is trained in a separate Stage 2.Source: §A.2 · p. 15

A.3 ImageNet sampling

  • ImageNet generation evaluation uses 50,000 class-conditioned 256 × 256 samples, with guided and unguided settings specified separately.Guided evaluation uses IG=1.78 and CFG=1; unguided evaluation sets both scales to 1 while retaining the class input.Source: §A.3 · p. 15

A.4 Text-to-image training and evaluation

The text-to-image protocol compares tokenizers with a shared generator-training and evaluation setup, then assesses alignment on three benchmarks. The appendix also reports Fréchet distances across six representation spaces and describes an aggregate over normalized distances.

A.4 Text-to-image training and evaluation

  • RAEv2, HiRAE-7, and HiRAE-24 use DINOv3-L/16 and 16 × 16 grids of 1024-dimensional latent tokens, with tokenizer, decoder, and text encoder frozen during generator pretraining and fine-tuning.Pretraining uses JourneyDB and BLIP3o caption subsets; supervised fine-tuning uses BLIP3o-60k.Source: §A.4 · pp. 15–16
  • The shared text-to-image generator is evaluated after 100K pretraining updates, then fine-tuned for 2,850 updates from the corresponding 100K checkpoint.Fine-tuning resets the optimizer and scheduler.Source: §A.4 · p. 16
  • Text-to-image alignment is evaluated on GenEval, DPG-Bench, and GenAI-Bench using EMA weights, BF16, 50 Euler steps, and CFG=6.The evaluation generates one image per prompt; GenAI-Bench uses CLIP-FlanT5-XL VQAScore as a continuous alignment measure.Source: §A.4 · p. 16

B.1 Metrics and aggregation

  • The evaluation reports Fréchet distances in Inception, ConvNeXt, DINOv2, MAE, SigLIP, and CLIP representation spaces and describes an aggregate over normalized distances.Text-to-image benchmark scores are reported after multiplying the benchmark mean by 100.Source: §B.1 · p. 17

B.2 Baseline configurations and result sources

The appendix documents the sources and settings behind external baseline comparisons, alongside additional generation comparisons and a matched-image reconstruction protocol. It notes a missing RAEv2 checkpoint result and reports HiRAE-24 generation scores under guided and unguided evaluation.

B.2 Baseline configurations and result sources

  • The baseline tables combine external measurements from cited papers with HiRAE evaluations, and distinguish complete-system comparisons from comparisons within the RAEv2 family.Some entries have unavailable metrics, which are left unfilled.Source: §B.2 · p. 17, Table 12 · p. 17, Table 3 · p. 18

B.3 Official-checkpoint unguided evaluation

  • The checked public RAEv2 release had a K=23 decoder and statistics but no Stage-2 checkpoint, so the missing FDr is not replaced with a different generator.The passage lists six normalized representation distances from the official-checkpoint evaluation.Source: §B.3 · p. 18

B.4 Additional generation comparisons

  • HiRAE-24 reaches guided gFID 1.038 and unguided gFID 2.129, while RAEv2 reaches guided gFID 1.060 and unguided gFID 1.650.These results use the stated 80-epoch comparison context; external systems retain their reported sampling settings.Source: §B.4 · p. 18, Table 13 · p. 19

C.1 Matched-image evaluation

  • Matched reconstruction metrics use the same 5,000 validation images from 100 fixed ImageNet classes, with fixed image IDs and preprocessing across tokenizers.These subset measurements complement the separate 50K rFID benchmark.Source: §C.1 · p. 19

C.2 Paired improvements and spatial-detail errors

Matched-image analyses report lower spatial-derivative errors for HiRAE-24 and reconstruction gains across texture quartiles, with a measurable contribution from its shallow group. A selected-layer extension, HiRAE-7, has similar rFID and guided gFID to HiRAE-24, while HiRAE-24 has stronger PSNR and LPIPS results.

C.2 Paired improvements and spatial-detail errors

  • HiRAE-24 has lower Sobel and Laplacian reconstruction errors than the comparison, consistent with improved local detail.The derivative-error measurements use the matched 5K subset.Source: §C.2 · p. 20, Table 16 · p. 20

C.3 Reconstruction gains across texture strata

  • Across texture quartiles Q1–Q4, HiRAE-24’s absolute LPIPS gains rise from 0.022 to 0.039, while relative reductions fall from 47.419% to 37.312%.Removing the shallow group increases LPIPS by 0.004 to 0.007 across the same quartiles; the authors report an association between stronger texture, larger absolute gains, and larger shallow-group contribution.Source: §C.3 · p. 20, Figure 7 · p. 20

D.1 HiRAE-7: selected-layer extension

  • HiRAE-7 achieves reconstruction rFID 0.217 and, on 50K guided samples, FD6r 1.913 and gFID 1.038; HiRAE-24 has similar rFID and guided gFID, higher PSNR, lower LPIPS, and guided FD6r 1.856.PSNR and LPIPS use the common preprocessing pipeline on the matched 5K subset.Source: §D.1 · p. 21

D.2 DRoRAE-style full-depth fusion

This block compares HiRAE-24 with adapted DRoRAE-style fusion and examines depth-group design, frequency contributions, representation structure, and decoding response. It reports generation and reconstruction comparisons alongside latent-space diagnostics.

D.2 DRoRAE-style full-depth fusion

  • At 80 generator epochs, HiRAE-24 reaches gFID 1.038 versus 1.551 for adapted DRoRAE-style fusion; gFID falls 53.7% versus 30.9% over 20–80 epochs.The comparison uses full-depth fusion configurations; Table 18 summarizes results across training budgets.Source: Table 18 · p. 22, §D.2 · p. 22

D.3 Expert inputs and residual regularization

  • Mode experts process 24 inputs through seven DCT-mode experts, with approximately 59M expert parameters compared with 202M for HiRAE-24.Source: §D.3 · p. 22

D.4 Number of depth groups

  • Three groups yield the best gFID and IS under both sampling settings, though two groups slightly improve reconstruction; a shared prior preserves total effective cap across partitions.The prior sums per-layer budgets for each new group and averages their dropout probabilities.Source: Table 19 · p. 22, §D.4 · pp. 22–23

D.5 Layer-group and frequency interventions

  • Removing any residual group worsens reconstruction LPIPS; equal-norm frequency removal hurts the shallow group more at high frequencies and the deep group more at low frequencies.The middle-group interval includes zero.Source: §D.5 · pp. 23–24, Table 21 · p. 24, Table 22 · p. 24

E.1 Analysis protocol

  • The analysis freezes tokenizers and generators, uses HiRAE-24’s epoch-80 EMA checkpoint, and requires no additional training.Source: §E.1 · p. 24

E.3 Cross-system spatial statistics

  • Adjacent-patch cosine similarity is 0.346/0.432/0.479 for RAEv2/HiRAE-7/HiRAE-24, while high-frequency energy fractions are 0.258/0.215/0.194.The figure reports class-mean adjacent-patch cosine similarity and high-frequency DCT energy fraction without DC.Source: Figure 8 · p. 25, §E.3 · p. 25

E.2 Class organization and spatial enrichment

  • Fusion raises mean spatial effective rank by 24.533 (95% CI [23.832, 25.264]) with mean spatial CKA 0.985, while spatial high-frequency content increases and class neighborhoods change less.The table compares HiRAE-24’s deep anchor and fused representation under shared anchor normalization and without additional standardization.Source: Table 23 · p. 25, §E.2 · p. 25

E.4 Decoding response to latent perturbations

  • At 10% relative latent perturbation, both HiRAE decodings change less than RAEv2, with HiRAE-7 showing the smallest response across tested direction types.The table reports output LPIPS at the 10% perturbation endpoint.Source: Table 24 · p. 26, §E.4 · p. 26

E.5 Generator prediction error

The generator prediction-error analysis finds lower HiRAE-24 error at interior noise levels and in frequency bands at log-SNR 0. The remaining sections provide reconstruction, generation, and text-to-image examples and visualization details.

E.5 Generator prediction error

  • HiRAE-24’s prediction-error advantage persists at five interior noise levels after target-energy normalization; at log-SNR 0, it also has lower error in each frequency band.The table gives per-coefficient clean-latent MSE, including DC, across seven coefficient log-SNR levels.Source: §E.5 · pp. 26–27, Table 25 · p. 26

F.1 Reconstruction comparisons

  • Figure 10 pairs inputs with official RAEv2 and HiRAE-24 reconstructions, including selected cases and cases at predetermined validation indices.Source: Figure 10 · p. 27, §F.1 · p. 27

F.2 Generation samples

  • Figure 11 shows the first eight stored unguided HiRAE-24 samples without quality filtering, using the epoch-80 EMA model with CFG=IG=1.Source: Figure 11 · p. 28, §F.2 · p. 28

F.3 Sample selection and visualization details

  • Pixel-error screening and visual inspection select reconstruction examples to illustrate local detail, while matched 5K metrics quantify average performance.The 24 guided examples were selected visually from the first 384 archive entries and are displayed uncropped.Source: §F.3 · p. 28

F.4 Text-to-image comparisons

  • Figure 12 compares outputs from FLUX-VAE, RAEv2, HiRAE-7, and HiRAE-24 after supervised fine-tuning on five selected GenEval prompts.Source: Figure 12 · p. 29
Loading 2609.37775v1…