PaperlayerView exact revision on arXiv

Source-linked AI summary

Hyperspherical Autoencoder for High-Fidelity Image Reconstruction and Generation

Hun Chang, Byunghee Cha, Jong Chul Ye

arXiv:2601.22904v2cs.CVcs.AIcs.LG

TL;DR

Existing VFM-based autoencoders provide semantic representations but can lose high-frequency detail and reconstruction fidelity. HAE combines detail-preserving patch embedding, directional feature alignment, and Riemannian Flow Matching on spherical latents, achieving gFID 1.96, rFID 0.78, and 25.2 dB PSNR.

  • Problem

    VFM-based autoencoders improve semantic expressiveness but often suffer degraded pixel-level reconstruction fidelity from lost high-frequency details.

  • Method

    HAE uses hierarchical convolutional patch embedding, cosine-based directional feature alignment, and Riemannian Flow Matching on a hyperspherical latent manifold.

  • Results

    HAE achieves 1.96 gFID, 0.78 rFID, and 25.2 dB PSNR on ImageNet-1K at 256 × 256 with accelerated convergence.

  • Takeaways & Limitations

    Respecting latent-space geometry across representation learning and generation provides a structurally coherent and efficient visual-synthesis paradigm.

  • Takeaways & Limitations

    The framework assumes standard ViT patch embedding is a primary bottleneck and models each patch-level latent on a fixed-radius sphere.

Abstract

from arXiv · show

Recent studies have explored using pretrained Vision Foundation Models (VFMs) such as DINO for generative autoencoders, showing strong generative performance. Unfortunately, existing approaches often suffer from limited reconstruction fidelity due to the loss of high-frequency details. In this work, we present the \textbf{\em Hyperspherical Autoencoder (HAE)}, a framework that bridges semantic representation and pixel-level reconstruction. Our key insight is that while semantic information in contrastive representations is primarily directional, enforcing strict magnitude matching hinders the preservation of fine-grained details. To address this, we introduce a {\em Directional Feature Alignment} objective that enforces semantic consistency while allowing flexible feature magnitudes for detail retention, alongside a {\em Hierarchical Convolutional Patch Embedding} module to enhance local structure preservation. Furthermore, observing that SSL-based representations intrinsically lie on a hypersphere, we employ {\em Riemannian Flow Matching} to train a Diffusion Transformer (DiT) directly on this spherical latent manifold. Notably, our manifold-aware DiT exhibits highly efficient convergence, achieving an exceptional gFID of \textbf{1.96} alongside a reconstruction rFID of \textbf{0.78} and a PSNR of \textbf{25.2} dB, validating the advantages of our manifold-aware approach.

1 Introduction

VFM-based latent autoencoders enrich semantic representations but can lose pixel-level detail. HAE addresses this trade-off with detail-preserving patch embedding, directional alignment, and hyperspherical generation.

  • Motivation: VFM-based autoencoders capture richer semantic information but often have lower PSNR than standard VAEs.The central challenge is balancing semantic expressiveness with pixel-level reconstruction fidelity.
  • Motivation: Aggressive ViT patch downsampling and rigid MSE alignment jointly hinder learning the high-frequency features required for faithful reconstruction.MSE constrains both feature magnitude and direction, creating a gradient conflict between semantic preservation and pixel reconstruction.
  • Hyperspherical Autoencoder: HAE combines Hierarchical Convolutional Patch Embedding with Directional Feature Alignment to preserve local textures while maintaining semantic consistency.Cosine-similarity alignment relaxes magnitude matching and leaves the encoder flexibility to reduce reconstruction error.
  • Hyperspherical Generation: SSL representations naturally lie on a hyperspherical manifold, motivating Riemannian Flow Matching instead of standard Euclidean diffusion.The manifold-aware formulation is reported to accelerate convergence and enhance generation quality.
  • Hyperspherical Autoencoder: 25.2 dB PSNR is achieved with Directional Feature Alignment while retaining high semantic alignment.The objective alleviates the optimization conflict inherent in MSE-based distillation.

2 Related Work

Related work improves diffusion efficiency through representation alignment and improves latent autoencoders through semantic VFM alignment. Riemannian Flow Matching extends flow matching to manifold-valued latent spaces using tangent vector fields and geodesic trajectories.

  • Representation Alignment: Representation-alignment methods improve diffusion training efficiency by matching DiT features with pretrained vision representations or enforcing intermediate-feature separability.REPA, REG, SRA, LSEP, and REPA-E explore different alignment targets, layers, and training settings.
  • Semantic Latent Autoencoders: Semantic latent autoencoders align or replace VAE encoders with pretrained vision foundation models to improve semantic latent quality and diffusion convergence.VAVAE and MAETok use explicit alignment, while RAE directly replaces the VAE encoder.
  • Riemannian Flow Matching: Riemannian Flow Matching trains a time-dependent tangent vector field on a Riemannian manifold to transport samples between base and target distributions.The objective uses the Riemannian metric and a marginal target flow defining a probability path.
  • Riemannian Flow Matching: The RFM objective is formulated as a regression loss against the marginal target flow, with conditional flows obtained by marginalizing conditional vector fields.The construction defines manifold-aware flow matching rather than relying on Euclidean trajectories.
  • Riemannian Flow Matching: Using geodesic distance yields constant-speed geodesic trajectories between endpoints, expressed analytically through exponential and logarithmic maps.This makes RFM simulation-free on simple manifolds while remaining scalable to general geometries.

3 Hyperspherical Autoencoder (HAE)

HAE combines hierarchical patch embedding and directional feature alignment to preserve local detail while retaining semantic structure. Its progressive training and hyperspherical generative formulation support robust reconstruction and efficient generation.

  • Architecture: HAE replaces standard single-layer ViT patch embedding with a hierarchical convolutional stem that captures fine-grained edges and textures.The four-stage CNN progressively downsamples inputs while preserving local features that large-kernel embedding can discard.
  • Architecture: Cosine-based directional alignment preserves semantic consistency while leaving feature magnitudes flexible for pixel-level reconstruction.The objective aligns angular components instead of enforcing exact student-teacher magnitude matching.
  • Semantic preservation: 87% Top-1 and 97% Top-5 ImageNet-1K accuracy show that HAE retains competitive semantic information against DINOv3’s 89% and 98%.PCA visualizations also show faithful reproduction of the teacher’s fine-grained semantic geometry.
  • Progressive training: The initial training stage combines directional alignment, pixel-wise reconstruction, and perceptual losses to establish a stable latent space.This configuration maximizes pure reconstruction capacity before generative regularization, yielding high PSNR and low rFID.
  • Progressive training: Latent smoothing trades reconstruction fidelity for generative quality, changing rFID from 0.37 to 0.78 and PSNR from 26.2 dB to 25.2 dB while improving gFID from above 32.64 to 2.65.The encoder is frozen from this stage onward, while adversarial training and stochastic noise injection bridge reconstruction and sampling.
  • Hyperspherical latent space: Reconstructions remain visually faithful when latent magnitudes are fixed at 8, 10, 16, or 20, but minimal angular noise severely degrades image quality.These experiments indicate that decoder outputs are more sensitive to latent direction than to magnitude.

4 Experiments

Experiments evaluate HAE’s reconstruction components and hyperspherical generative dynamics through benchmark comparisons, convergence analysis, and ablations. The results associate manifold-aware generation with stronger quality and faster convergence, while patch embedding and radial-penalty choices materially affect performance.

  • Image Generation: 2.65 gFID is achieved without guidance, outperforming VAVAE at 4.29 and RAE at 4.28 under the same 80-epoch, 676M-parameter setup.The comparison uses LightningDiT-XL and the same training budget and architecture for the reported baselines.
  • Image Generation: 4× faster convergence is reported for the identical DiT architecture trained on HAE latents than for the RAE-latent baseline.The authors attribute this to Riemannian Flow Matching avoiding arbitrary radial variations on the hyperspherical manifold.
  • Patch-Embedding Ablation: 19.55 dB PSNR is reported for the standard single-layer ViT patchification baseline, motivating the hierarchical CNN patch embedding for high-frequency detail preservation.The accompanying caption states that the table isolates pure architectural impact without latent smoothing.
  • Riemannian vs. Euclidean Flow Matching: 2.65 gFID versus 15.95 for Euclidean Flow Matching demonstrates the reported benefit of constraining diffusion dynamics to spherical geodesics.Both models use LightningDiT-XL and are evaluated at 80 epochs.

5 Conclusion

HAE addresses the tension between semantic abstraction and pixel-level fidelity by preserving local detail while aligning representations directionally. It then models generation on the hyperspherical latent geometry, achieving strong reconstruction and generation results.

  • Conclusion: HAE identifies aggressive single-layer ViT patch embedding as a bottleneck causing irreversible high-frequency detail loss.The conclusion frames this bottleneck as central to the semantic–pixel-fidelity tension.
  • Conclusion: Hierarchical Convolutional Patch Embedding retains local textures, while Directional Feature Alignment uses cosine similarity to preserve angular semantics and flexible magnitudes.The magnitude flexibility is described as supporting intricate structural-detail reconstruction.
  • Conclusion: Riemannian Flow Matching models generative dynamics on the hyperspherical manifold and uses tangent-space projection with penalties on redundant radial variations.The method is presented as matching the geometry of the magnitude-robust latent space.
  • Conclusion: 0.78 rFID, 25.2 dB PSNR, and 1.96 gFID are reported on ImageNet-1K at 256 × 256.The reported reconstruction and generation metrics summarize HAE’s evaluation outcome.

A.1 Training Algorithm

The training algorithm projects encoder features and noise onto a fixed-radius hypersphere, applies dimension-adjusted timing, and trains velocity fields using tangent-space flow matching with radial suppression. Inference follows spherical trajectories through tangent projection, exponential-map updates, and decoding.

  • Training: R = 32.0 fixes the hypersphere radius, while C = 1024 and α = 8.0 define feature dimensionality and the time-shift schedule.The radius is selected for the DINOv3-Large feature space, and the shift calibrates noise for high channel capacity.
  • Training: The algorithm projects both encoded images and sampled Gaussian noise onto the radius-R sphere before computing spherical interpolation quantities.The training procedure uses z0 = R · E(x) / ∥E(x)∥2 and an analogous projection for noise.
  • Training: The loss matches tangent velocity components and adds λrad · Mean(∥vrad∥2^2) to penalize radial motion.Predicted and target velocities are decomposed into radial and tangential components before optimization.
  • Inference: Inference applies classifier-free auto-guidance, removes radial velocity, and advances samples with the spherical exponential map before renormalization and decoding.The Riemannian Euler sampler uses Rodrigues’ rotation formula and the same α = 8.0 schedule as training.

B Autoencoder Architecture and Training Details

HAE training uses staged architectural and optimization choices to balance semantic alignment with reconstruction fidelity. Its hierarchical convolutional stem replaces standard ViT patch embedding and progressively maps RGB inputs into the DINOv3-Large latent space.

  • Training Setup: Autoencoder training is divided into multiple stages to balance semantic alignment and reconstruction fidelity.The reported setup uses 8 NVIDIA B200 GPUs and approximately 1 hour 20 minutes per epoch.
  • Architecture: The Hierarchical Convolutional Stem progressively downsamples RGB inputs by 16× through four convolutional stages into a 1024-dimensional latent space.The design replaces standard single-layer ViT downsampling to mitigate irreversible high-frequency detail loss.
  • Architecture: Spatial Root Mean Square Normalization and SiLU activations are used in intermediate layers for stable gradient flow and early feature extraction.These implementation choices are specified as part of the intermediate-layer configuration.

B.2 Decoder Architecture

HAE uses a progressive decoder and staged training pipeline for hyperspherical latents, combining reconstruction, adversarial adaptation, and latent smoothing.

  • Decoder: The decoder progressively upsamples hyperspherical latent representations back into RGB pixel space using InterpolateConv blocks and duplicating shortcut connections.TRMS2d and SiLU are used throughout the decoder blocks.
  • Stage 1: Stage 1 jointly trains hierarchical convolutional patch embedding and the decoder to establish a stable semantic latent space.Training uses AdamW, EMA, and BF16 mixed precision over 20 epochs.
  • Stage 2: Stage 2 freezes the DINOv3 encoder and trains the decoder with a DINO-feature discriminator for logically consistent textures.Adversarial training is introduced after a two-epoch discriminator warmup, with cosine loss disabled because the encoder is frozen.
  • Stage 2: Latent smoothing uses a spherical noise scale that increases from 0.8 to 1.2 over 20 epochs and then remains fixed through epoch 40.The schedule thickens the latent manifold gradually rather than abruptly degrading reconstruction.

C Ablation on Decoder Latent Scaling Factor

The decoder is highly sensitive to the scale applied to Riemannian Flow Matching latents, with svae = 14.0 providing the best balance between latent variance and decoder statistics.

  • Motivation: RFM outputs latents on a hypersphere of radius R = 32.0, so generated vectors require rescaling before decoding.The decoder is robust to latent magnitude shifts during training, but inference-scale calibration strongly affects visual quality.
  • Low scales: 7.36 gFID at svae = 8.0 shows that excessive compression prevents the decoder from synthesizing intricate high-frequency details.Low magnitude fails to sufficiently excite decoder feature maps.
  • Optimal scale: 2.65 gFID at svae = 14.0 marks the optimal scaling sweet spot, while larger scales increasingly over-smooth outputs.The study sweeps svae from 8.0 to 20.0; Figure 8 visualizes the corresponding detail changes across scales.

D Limitations

HAE’s reported limitations concern training scope, numerical precision, and performance relative to stronger specialized transformer baselines.

  • Training scope: The generative model has been trained for only 550 epochs, leaving longer training regimes insufficiently explored.At this intermediate budget, LightningDiT-XL is reported to perform comparably to RAE with LightningDiT-XL trained for 800 epochs.
  • Numerical precision: Riemannian Flow Matching currently requires full-precision computation because preliminary BF16 training caused numerical instability.This sacrifices some memory and throughput benefits associated with mixed-precision training.
  • Baseline scope: HAE does not yet outperform the strongest specialized baseline, RAE with DiTDH.The paper attributes DiTDH’s strength partly to its heavier architecture and proposes stronger heads or architectural scaling as future directions.

E Uncurated Generated Samples

This section presents uncurated 256 × 256 image samples generated by LightningDiT-XL across multiple ImageNet classes under a fixed sampling configuration.

  • Sampling setup: The samples are generated at 256 × 256 resolution using the Riemannian Euler sampler for 50 steps with CFG scale 1.1.The configuration is shared across the presented uncurated samples.
  • Classes: The sample grid includes natural, object, and scene classes such as Goldfinch, Macaw, Golden Retriever, Polar bear, and Cliff.Additional classes include Jellyfish, Monarch butterfly, Balloon, Daisy, Volcano, Hamburger, and Icecream.
  • Classes: The remaining examples cover Balloon, Cliff, Daisy, Volcano, Hamburger, and Icecream under the same 256 × 256 sampling setup.These captions identify class-specific uncurated sample sets rather than quantitative comparisons.
Loading 2601.22904v2…