PaperlayerView exact revision on arXiv

Source-linked AI summary

Aligning Latent Geometry for Spherical Flow Matching in Image Generation

Tuna Han Salih Meral, Kaan Oktay, Hidir Yesiltepe, Adil Kaan Akan, Pinar Yanardag

arXiv:2605.15193v1cs.CV

TL;DR

Standard latent flow matching uses straight Euclidean paths even though Gaussian noise and VAE latents lie on thin shells, causing motion in a decoder-insensitive radial direction. The paper instead aligns latent tokens to a fixed-radius sphere and trains with spherical paths, improving matched-CFG ImageNet-256 FID across tokenizers and model scales without changing the diffusion architecture.

  • Problem

    Latent flow matching overlooks that Gaussian noise and VAE latents occupy thin spherical shells while straight paths allocate substantial supervision to decoder-insensitive radial motion.

  • Method

    The method projects tokens and Gaussian noise onto a fixed-radius sphere, finetunes only the decoder, and replaces linear interpolation with slerp paths and tangent velocity targets.

  • Results

    The method consistently improves matched-CFG ImageNet-256 FID across three tokenizers and transfers gains from SiT-B to SiT-XL backbones.

  • Takeaways & Limitations

    These results support latent support geometry as a relevant design principle for flow matching while preserving the existing diffusion architecture and avoiding auxiliary encoders.

  • Takeaways & Limitations

    The geometric analysis is defined in flow-training coordinates after tokenizer preprocessing, which need not coincide with raw encoder coordinates.

Abstract

from arXiv · show

Latent flow matching for image generation usually transports Gaussian noise to variational autoencoder latents along linear paths. Both endpoints, however, concentrate in thin spherical shells, and a Euclidean chord leaves those shells even when preprocessing aligns their radii. By decomposing each latent token into radial and angular components, we show through component-swap probes that decoded perceptual and semantic content is carried predominantly by direction, with radius contributing much less. We therefore project data latents onto a fixed token radius, use the radial projection of Gaussian noise as the spherical prior, finetune the decoder with the encoder frozen, and replace linear interpolation with spherical linear interpolation. The resulting geodesic paths stay on the sphere at every timestep, and their velocity targets are purely angular by construction. Under matched training, the method consistently improves class-conditional ImageNet-256 FID across different image tokenizers, leaves the diffusion architecture unchanged, and requires no auxiliary encoder or representation-alignment objective.

1 Introduction

The paper identifies a mismatch between spherical-shell latent geometry and linear flow-matching paths, then aligns the latent space and transport process with fixed-radius spherical geometry. The resulting spherical-slerp method improves ImageNet-256 FID across tokenizers and model scales without changing the diffusion architecture.

  • Motivation: Gaussian noise and VAE latents concentrate in thin spherical shells, while straight-line interpolation passes through unoccupied interior radii.Decoder probes indicate that token direction carries predominantly perceptual and semantic content, whereas length contributes much less.
  • Problem: Standard linear flow matching ignores shell geometry and direction-dominant decoding, spending roughly half or more of endpoint velocity on radial motion.This radial component changes token distance from the origin instead of only changing direction.
  • Results: The spherical-slerp method improves FID across FLUX.2, VA-VAE, and REPA-E FLUX.1 tokenizers under matched guidance, with gains from SiT-B to SiT-XL backbones.A decoder-swap control degrades FID in both directions, tying the learned flow to its latent geometry rather than decoder finetuning alone.
  • Method: The method projects each pretrained-VAE latent token onto a fixed-radius sphere, finetunes only the decoder, and replaces linear interpolation with spherical linear interpolation.The encoder remains frozen, avoiding an auxiliary encoder and representation-alignment objective.
  • Method: Slerp between projected endpoints stays on the sphere at every timestep, producing velocity targets that change token direction but never length.Integration therefore keeps generated samples on the sphere.

2 Related Work

Prior work studies hyperspherical latent spaces, manifold flow matching, representation-space diffusion, and the relationship between latent structure and reconstruction fidelity. These lines motivate spherical generative modeling while distinguishing this work’s setting from prior approaches.

  • Hyperspherical latent spaces: Hyperspherical latent spaces are standard in discriminative learning, while generative approaches use sphere-constrained VAEs, fixed-radius VAE projections, or sphere-mapped encoders.These methods respectively use von Mises–Fisher distributions, stabilize variance under classifier-free guidance, or bypass diffusion by decoding random sphere points.
  • Riemannian and manifold flow matching: Riemannian flow matching extends simulation-free flow matching to manifolds by defining conditional vector fields along geodesic interpolants and projecting velocities onto the tangent space.Earlier manifold generative models include continuous normalizing flows and score-based diffusion with Riemannian drift, noise, and geodesic integration.
  • Representation-space diffusion: Representation-space diffusion trains generators in frozen-encoder feature spaces, including spheres induced by LayerNorm or systems pairing frozen encoders with trained decoders.A related VAE-based variant adds an alignment loss that pulls intermediate generator features toward a frozen representation encoder.
  • Latent-space structure versus reconstruction fidelity: Prior studies find tokenizer reconstruction quality is a weak predictor of downstream diffusion generation quality, whereas latent-distribution structure governs trainability.Related work quantifies this decoupling across autoencoders and identifies sampling-error robustness as a relevant axis.

3 Methodology

The method addresses a geometric mismatch in latent flow matching: Euclidean chords leave concentrated endpoint shells, while decoded content depends more strongly on token direction than radius. It therefore projects latents onto a fixed-radius sphere, finetunes only the decoder, and uses spherical interpolation without changing the diffusion architecture.

  • Geometric mismatch: Latent endpoints concentrate in thin spherical shells, but Euclidean interpolation moves substantially inward and can place supervision far from the training distribution.Linear paths deviate up to 1.4σ (FLUX.2), 1.8σ (VA-VAE), and 2.5σ (REPA-E FLUX.1), whereas Slerp remains on the sphere.
  • Geometric mismatch: Component-swap probes show decoded images remain close when direction is preserved but change nearly as much as a neighbor replacement when direction is swapped.This asymmetry appears in both LPIPS and DINOv2 distances, indicating stronger decoder sensitivity to direction than radius.
  • Geometric mismatch: Linear flow matching assigns about 50% of velocity to radial motion for FLUX.2 and VA-VAE, and about 90% at the noise endpoint for REPA-E FLUX.1.Spherical interpolation makes the radial velocity share identically zero by construction.
  • Spherical latent construction: The method inserts a token-wise L2 projection between the pretrained encoder and decoder, fixing each token radius to the standard Gaussian concentration radius.The projected latent tensor lies on a product of token-wise spheres, aligning latent scale with the Gaussian prior.
  • Spherical latent construction: The encoder is frozen while only the decoder and discriminator are finetuned, retaining the original reconstruction losses and dropping the KL term for deterministic projected latents.The procedure is applied to FLUX.2, VA-VAE, and REPA-E FLUX.1 tokenizers.
  • Spherical latent construction: The projection is confined to the VAE, so SiT processes spherical latents as ordinary vectors in R^d without architectural changes or an auxiliary encoder.This avoids methods that run a frozen DINOv2 encoder on every generated sample.

4 Experiments

Experiments compare spherical-slerp with vanilla-linear under matched training on ImageNet-256, finding gains across tokenizer families, model scales, guidance settings, and extended training. Ablations show that spherical latent support explains most of the improvement, while the geometry-matched slerp path yields the best observed result.

  • Experimental Setup: Experiments hold architecture, data, training budget, and evaluator fixed while testing SiT models across FLUX.2, VA-VAE, and REPA-E FLUX.1 tokenizers.All settings use ImageNet-256 class-conditional generation and 256 tokens per 256 × 256 image, with models trained from scratch for 80 epochs.
  • Transport path comparison: The spherical construction explains most of the gain over vanilla-linear, while adding slerp produces the best observed result in the transport-path ablation.The construction combines fixed-radius projection, a uniform spherical prior, and spherical-decoder finetuning; Spherical-Linear isolates this construction with Euclidean transport.
  • Generation across tokenizers and scales: Spherical-slerp improves matched-CFG FID across FLUX.2 and VA-VAE at both B and XL scales, and across every tested guidance setting on REPA-E FLUX.1.The tokenizer families differ in objective, channel count, and downsampling factor, so the reported gains span distinct tokenizer configurations.
  • Scaling: At CFG=1.0, spherical-slerp improves FLUX.2 XL/2 by 17.6% and VA-VAE XL/1 by 14.8%, remaining ahead on both families at CFG=1.5.These XL results transfer the B-scale recipe without changing the diffusion architecture or adding an auxiliary encoder during diffusion training or inference.
  • Scaling: 8.35 versus 9.15 at CFG=1.0 and 2.91 versus 3.22 at CFG=1.5 show spherical-slerp remains ahead after FLUX.2 B/2 training is extended to 200 epochs.The extension uses the same batch size and otherwise identical settings, while both methods improve substantially over their 80-epoch results.

5 Conclusion

The paper identifies a geometric mismatch in latent flow matching and proposes spherical latent flow matching, which keeps tokens on a fixed-radius sphere and trains along geodesic paths with tangent velocities.

  • VAE latents and Gaussian noise concentrate on thin shells, but linear paths leave those shells and supervise radial motion to which the decoder is weakly sensitive.
  • Spherical latent flow matching projects each VAE token to a fixed-radius sphere and uses projected Gaussian directions as a uniform spherical prior.
  • The method trains along slerp geodesics with tangent velocities, addressing the mismatch without changing the di…

Supplementary Material … A.2 Projected Gaussian Noise is Uniform on the Sphere

The supplementary derivations characterize Gaussian latent norms and justify the spherical prior used by spherical flow matching. Gaussian directions are uniform on fixed-radius spheres, while norm statistics quantify the shell around which Gaussian samples concentrate.

  • A.1 Analytical Gaussian Norm Statistics: Independent standard-normal coordinates make the Gaussian norm R = ∥z∥2 chi-distributed with d degrees of freedom.This distribution underlies the analytical norm statistics reported for Gaussian rows.
  • A.1 Analytical Gaussian Norm Statistics: The exact mean norm E[R] is used as the analytical Gaussian mean radius in Tab. 1.The supplementary material distinguishes this exact mean from conventional approximate shell-radius expressions.
  • A.1 Analytical Gaussian Norm Statistics: The coefficient of variation and E[R] provide the Gaussian entries reported in Tab. 1.These quantities summarize the Gaussian norm’s scale and relative variability.
  • A.1 Analytical Gaussian Norm Statistics: Gaussian samples lie in an O(1)-width shell around radius d, so relative shell thickness decreases as dimension grows.The concentration statement is consistent with the analytical norm calculation.
  • A.2 Projected Gaussian Noise is Uniform on the Sphere: Radial projection of isotropic Gaussian noise produces the spherical noise endpoint on the fixed-radius sphere S^(d−1).The projection is well-defined almost surely because P(ϵ = 0) = 0.
  • A.2 Projected Gaussian Noise is Uniform on the Sphere: Polar-coordinate factorization shows that u = ϵ/∥ϵ∥2 is uniform on S^(d−1) and independent of the Gaussian radius.The Gaussian density separates into a radial term and a direction-independent angular term.
  • A.2 Projected Gaussian Noise is Uniform on the Sphere: Token-wise projection preserves independence across N = h · w spatial positions, retaining Gaussian directions while discarding only radii.The resulting full noise tensor is sampled from the product of uniform spherical measures.

A.3 Exponential-Map Integration for Slerp Targets · B Implementation Details · B.1 Slerp Numerical Handling

The paper proves that exponential-map integration exactly follows slerp trajectories with perfect velocity predictions, whereas projected Euler steps undershoot along the same great circle. It also specifies numerically stable per-token slerp branches for ordinary, near-coincident, and near-antipodal endpoints.

  • A.3 Exponential-Map Integration for Slerp Targets: For endpoints z0,z1 on a sphere, slerp traces their great circle at constant angular speed θ0.The path remains on the sphere and connects endpoints with angle θ0.
  • A.3 Exponential-Map Integration for Slerp Targets: With the true velocity, an exponential-map step of size h reaches the exact slerp point at time t+h.Thus, a perfect velocity predictor makes exponential-map sampling exact along the slerp curve.
  • A.3 Exponential-Map Integration for Slerp Targets: A projected Euler step remains on the same great circle but rotates by arctan(hθ0), rather than the intended hθ0.The discrepancy arises because the Euler update is radially projected back onto the sphere.
  • A.3 Exponential-Map Integration for Slerp Targets: The one-step arc-length deficit is R[hθ0 − arctan(hθ0)] between exponential-map and projected Euler integration.Both methods remain on the same great circle, but the projected Euler step advances less distance.
  • B Implementation Details: Per-token slerp treats each spatial position’s d-dimensional vector as a point on the sphere S^(d−1).This establishes the geometric object used by the implementation.
  • B.1 Slerp Numerical Handling: For angles ω ∈[10^-4, π −0.1], standard slerp is used; smaller angles use renormalized linear interpolation, while larger angles use an arbitrary great-circle branch.Cosines are clamped to [−1 + 10^-6, 1 −10^-6] before arccos, and norms are floored at 10^-8.
  • B.1 Slerp Numerical Handling: The small-angle and antipodal branches are correctness insurance rather than hot paths during training for random unit vectors.The implementation therefore treats these special cases as safeguards for numerical correctness.

B.2 Training Hyperparameters … C.3 Decoder-Flow Coupling

The paper uses a shared training recipe across tokenizers and model scales, while isolating tokenizer-specific timestep shifts and geometry-specific path and latent choices. Ablations show that improvements arise from spherical geometry and require decoder-flow co-adaptation rather than finetuning or compute alone.

  • B.2 Training Hyperparameters: Hyperparameters are grid-searched on SiT-B with FLUX.2 and transferred unchanged to SiT-XL, VA-VAE, and REPA-E FLUX.1, except for tokenizer-specific timestep shifts.Linear and spherical recipes share hyperparameters; they differ in path, latent projection, and slerp-specific tangent projection and exponential-map sampling.
  • B.3 Compute: All settings produce 256 latent tokens per 256×256 image by pairing model patch size with tokenizer VAE downsample factor.Experiments use NVIDIA H200 GPUs, with 2 GPUs for SiT-B and 4 GPUs for SiT-XL.
  • B.3 Compute: An SiT-B run takes approximately 6.5 hours on 2×H200, while an SiT-XL run takes approximately 16.8 hours on 4×H200.These correspond to approximately 13 and 67 H200-hours, respectively, for 80 epochs.
  • B.4 Code and Checkpoint Release: The authors will release code, tokenizer-specific spherical-decoder finetuning configurations, and trained SiT-B and SiT-XL flow-matching checkpoints.The release covers FLUX.2, VA-VAE, and REPA-E FLUX.1 spherical decoders and checkpoints reported in Table 5.
  • C.1 Core Recipe Checks: Recipe ablations use FLUX.2 as the canonical tokenizer, with the resulting B-scale recipe transferred across FLUX.2, VA-VAE, and REPA-E FLUX.1.The matched-recipe decoder finetune alone changes vanilla-linear FID from 26.35 to 26.95, so the remaining improvement is attributed to latent geometry.
  • C.3 Decoder-Flow Coupling: Spherical projection, not decoder finetune compute, provides the gain, and decoder-flow substitutions degrade both vanilla and spherical pairings.The flow and decoder co-adapt to a specific latent geometry, so neither component can be independently substituted without loss.

C.4 Scaling with Number of Function Evaluations (NFE)

Spherical-slerp outperforms the baseline at every tested number of function evaluations, with its largest FID advantage at NFE = 20 and a smaller but persistent advantage at NFE = 500.

  • C.4 Scaling with Number of Function Evaluations (NFE): Spherical-slerp wins at every tested NFE, with the largest gap at NFE = 20 (−19.84 FID) and a narrower gap of −5.30 at NFE = 500.The advantage holds across sampling budgets rather than only at low step counts.

C.5 Direction-vs-Radius Component Ablation · C.6 Scope of Held-Out Comparisons · D Qualitative Results

The component ablations show that decoded content is carried predominantly by latent direction rather than radius, motivating spherical paths without radial motion. The held-out comparisons clarify that this study instead fixes the tokenizer and modifies latent geometry only.

  • C.5 Direction-vs-Radius Component Ablation: The experiments compare original decodes with keep-direction, keep-radius, and full-substitute conditions using LPIPS and DINOv2 cosine similarity.Same-class partners provide direction or radius substitutions, while a stronger baseline uses population-mean direction and radius.
  • C.5 Direction-vs-Radius Component Ablation: Replacing radius with a population-wide constant leaves decoded images nearly unchanged, whereas replacing direction makes them nearly unrecognizable across tokenizers.Keep-direction reaches LPIPS ≈0.04–0.08 and DINOv2 cosine ≈0.94–0.97 for FLUX.2 and REPA-E FLUX.1, while keep-radius reaches LPIPS ≈0.88–0.93 and DINOv2 cosine ≈0.01–0.02 across all three.
  • C.5 Direction-vs-Radius Component Ablation: The same direction-versus-radius pattern holds for both same-class and population-mean substitutions, with radius sensitivity spreading slightly further for same-class VA-VAE substitutions.Every per-sample point lies above y=x for all three tokenizers in both substitute regimes.
  • C.5 Direction-vs-Radius Component Ablation: The ablation concludes that decoded content is carried almost entirely by angular direction and that the decoder is largely insensitive to radius.Linear flow paths therefore spend substantial velocity on radial motion, whereas spherical interpolation has no radial component by construction.
  • C.6 Scope of Held-Out Comparisons: REPA-E is outside the matched comparison scope because reproducing it requires retraining both the VAE and diffusion model under a representation-alignment objective.This study holds the tokenizer fixed and modifies only latent geometry.
  • C.6 Scope of Held-Out Comparisons: Kumar and Patel’s approach is also outside scope because it performs flow matching in frozen DINOv2 feature space and requires encoder execution at inference.The supplied passage distinguishes that tokenizer regime from the geometric constraint studied here.

D.1 Class-Conditional Samples · D.2 Same-Class Component Swaps · D.3 Reconstructions

The appendix visualizes class-conditional generation, same-class radial–angular swaps, and tokenizer reconstruction comparisons. Samples show class-organized outputs, direction-preserving swaps closely match anchors, and the three reconstruction variants are visually near-identical at 256×256.

  • D.1 Class-Conditional Samples: Each ImageNet-1k class is shown as a 4×4 grid of 16 class-conditional samples from spherical-slerp SiT-XL/2 with the FLUX.2 tokenizer.Figure 10 presents one 16-sample panel per class.
  • D.2 Same-Class Component Swaps: Across all three tokenizers, the keep-direction hybrid is visually nearly indistinguishable from the anchor.The comparison averages LPIPS and DINOv2 distance over 1024 anchor-neighbor pairs and includes representative pairs for each tokenizer.
  • D.3 Reconstructions: Figures 14 to 16 compare reconstructions from Vanilla, Vanilla FT, and Spherical FT on eight ImageNet classes.The reported comparison concerns original, matched-compute vanilla-finetuned, and spherical-finetuned decoders.
  • D.2 Same-Class Component Swaps: The swaps compare anchor, keep-direction, keep-radius, and neighbor columns for each same-class anchor-neighbor pair.The keep-direction hybrid combines anchor direction with neighbor radius, while the keep-radius hybrid combines anchor radius with neighbor direction.
  • D.3 Reconstructions: At 256×256 per tile, the three reconstruction variants are visually near-identical, as expected for high-fidelity tokenizers.The comparison includes the original tokenizer, matched-compute vanilla decoder finetune, and spherical decoder finetune.

E Scope and Future Directions

The study isolates latent geometry in a controlled class-conditional ImageNet-256 setting, while identifying solver efficiency and broader generation settings as directions for future work. These extensions preserve the central fixed-radius spherical projection and geodesic-flow formulation.

  • Controlled scope: Experiments hold the diffusion backbone, training budget, evaluator, and tokenizer-specific preprocessing fixed across three pretrained VAE/tokenizer families.The native token dimensions are d = 16 and d = 32, enabling attribution of observed gains to spherical latent support and geodesic transport.
  • Solver efficiency: Matched sampling budgets and guidance settings isolate latent-geometry effects, leaving solver-efficiency studies for future work.Potential directions include adaptive ODE solvers, lower-NFE samplers, and distillation-based acceleration.
  • Solver efficiency: Because spherical-slerp remains on the fixed-radius manifold, specialized integrators or training-time acceleration could exploit this structure.This is proposed as a future direction rather than evaluated in the present study.
  • Generality: Future evaluations should extend beyond ImageNet-256 to text-conditioned generation, higher resolutions, and non-square aspect ratios.ImageNet-256 provides a clean, widely used benchmark for class-conditional generation.
  • Generality: Jointly trained encoder-decoder tokenizers could further test generality without changing the central fixed-radius projection and spherical-geodesic formulation.The formulation assumes encoder outputs are projected token-wise onto a fixed-radius sphere before training the flow model along corresponding spherical geodesics.
Loading 2605.15193v1…