1. Vision foundation model Aligned Variational AutoEncode (VA-VAE)

As visual tokenizers become more sophisticated with higher-dimensional latent spaces to improve reconstruction quality, they paradoxically become harder for diffusion models to work with, leading to poor generation performance.

To resolve this dilemma, the research introduces a Vision Foundation Model Aligned Variational AutoEncoder (VA-VAE) that uses a novel Vision Foundation Model Alignment Loss (VF Loss). This approach leverages the structured, semantically meaningful representations learned by pre-trained vision foundation models to guide the tokenizer's latent space toward being more "generation-friendly."

Relationship to REPA: REPA aims to employ vision foundation models to constrain DiT, thereby enhancing the convergence speed of generative models. In contrast, our work takes into account both the reconstruction and generative capabilities within the latent diffusion model, with the objective of leveraging foundation models to regulate the highdimensional latent space of the tokenizer, thereby resolving the optimization conflict between the tokenizer and the generative model.

1.1. Align VAE with Vision Foundation Models

Vision Foundation model alignment loss (VF loss) consists of two components: marginal cosine similarity loss and marginal distance matrix similarity loss.

1.1.1. Marginal Cosine Similarity Loss

We project the image latents $𝑍$ to match the dimensionality of foundational visual representations $𝐹$ using a linear transformation $𝑊$, producing ${𝑍}^{\prime }=𝑊𝑍$.

The Marginal Cosine Similarity Loss enforces element-wise alignment between the VAE's latent features and foundation model features, focusing alignment on less similar pairs:

$${ℒ}_{\text{mcos }}=\frac{1}{ℎ\times 𝑤}\sum _{𝑖=1}^{ℎ}\sum _{𝑗=1}^{𝑤}\max \left(0,{𝑚}_1-\cos \left({𝑧}_{𝑖𝑗}^{\prime },{𝑓}_{𝑖𝑗}\right)\right)$$

1.1.2. Marginal Distance Matrix Similarity Loss

Complementary to ${𝐿}_{\text{mcos}}$, which enforces point-to-point absolute alignment, we also aim for the relative distribution distance matrices within the features to be as similar as possible. The Marginal Distance Matrix Similarity Loss aligns the internal structure and relationships within the latent space:

$${ℒ}_{\text{mdms }}=\frac{1}{{𝑁}^2}\sum _{𝑖𝑗}\max \left(0,\left|\cos \left({𝑧}_{𝑖}^{\prime },{𝑧}_{𝑗}^{\prime }\right)-\cos \left({𝑓}_{𝑖},{𝑓}_{𝑗}\right)\right|={𝑚}_2\right)$$

Here, $𝑁=ℎ\times 𝑤$ represents the total number of elements in each flattened feature map.

1.1.3. Adaptive Weighting

$${𝐿}_{\text{vf }}={𝑤}_{\text{ hyper }}\cdot {𝑤}_{\text{ adaptive }}\left({ℒ}_{\text{mcos }}+{ℒ}_{\text{ mdms}}\right)$$

The adaptive weighting function ${𝑤}_{\text{adaptive}}$ is defined as $\frac{\Vert \nabla {ℒ}_{\text{ rec}}\Vert }{\Vert \nabla {ℒ}_{\text{ vf}}\Vert }$ to ensure ${ℒ}_{\text{vf}}$ and ${ℒ}_{\text{rec}}$ have similar impacts on model optimization.

References

1. Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models