EQ-VAE: Equivariance Regularized Latent Space for Improved Generative Image Modeling
Authors: Theodoros Kouzelis
Ioannis Kakogeorgiou
Spyros Gidaris
Nikos Komodakis
Venue: ICML 2025 poster (2025)
Keywords: autoencoders
latent generative models
regularization
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: Equivariance regularization of autoencoders boosts latent generative modeling

Abstract: Latent generative models have emerged as a leading approach for high-quality image synthesis. These models rely on an autoencoder to compress images into a latent space, followed by a generative model to learn the latent distribution. We identify that existing autoencoders lack equivariance to semantic-preserving transformations like scaling and rotation, resulting in complex latent spaces that hinder generative performance. To address this, we propose EQ-VAE, a simple regularization approach that enforces equivariance in the latent space, reducing its complexity without degrading reconstruction quality. By finetuning pre-trained autoencoders with EQ-VAE, we enhance the performance of several state-of-the-art generative models, including DiT, SiT, REPA and MaskGIT, achieving a ×7 speedup on DiT-XL/2 with only five epochs of SD-VAE fine-tuning. EQ-VAE is compatible with both continuous and discrete autoencoders, thus offering a versatile enhancement for a wide range of latent generative models.

EQ-VAE: Equivariance Regularized Latent Space for Improved Generative Image Modeling
Theodoros Kouzelis 12
Ioannis Kakogeorgion
Spyros Gidaris
Nikos Komodakis 145
Abstract
images into a lower-dimensional latent space, which can be continuous (e.g., SD-VAE for diffusion (Rombach et al., 2022)) or discrete (e.g., VQ-GAN for autoregressive (Esser et al., 2021; Yu et al., 2022b) and masked generative model- ing (Chang et al., 2022)). This latent space retains essential semantic and structural information while discarding high- frequency details. Second, a generative model learns to model the distribution of these latent representations, en- abling the synthesis of visually coherent images. At infer- ence time, the generative model first samples latent code, which is then decoded back into the image space by the autoencoder. While much research has focused on improv- ing the generative (Peebles & Xie, 2023), objectives (Ma al., 2024), and op- timization techniques (Yao et al., 2024) the autoencoder's role in shaping the latent space remains equally critical to overall performance.
Latent generative models have emerged lead- ing approach for high-quality image synthesis. These models rely on an autoencoder to compress images into a latent space, followed by a gener- ative model to learn the latent distribution. We identify that existing autoencoders lack equivari- ance to semantic- preeerringg transformations like scaling and rotation, resulting in complex latent spaces that hinder generative performance. To address this, we propose EQ-VAE, simple regu- larization approach that enforces equivariance in the latent space, reducing its complexity with- out degrading reconstruction quality. By fine- tuning pre-trained autoencoders with EC-VAE, we enhance the performance of several state-of- the-art generative models, including DiT, SiT, REPA and MaskGIT, achieving x7 speedup on DiT-XLL/ XL/2 2 with only five epochs of SD-VAE fine-tuning. EQ-VAE compatible with both continuous and discrete autoencoders, thus offer- ing a versatile enhancement for a wide range of latent generative models. Project page and code: https: /ttl///eqg vae-va...eeeeeeeeeeeeeeeeeeeeeee github.io.
advances in architectures
In fact, the quality of the latent space is pivotal, influencing both computational efficiency (by reducing dimensionality and accelerating convergence in the generative phase) and the model's ability produce high-fidelity outputs (Rom- bach et al., 2022). In diffusion models, most state-of-the- art approaches-Sunn as DiT (Peebles & Xie, 2023), SiT (Ma et al., 2024), P PixArt (Chen et al., 2024), SD3 (Esser et al., 2024), and Flux (Black Forest Labs, 2023)-reh on autoencoders with architectures and training objectives sim- ilar to the SD-VAE introduced in Latent Diffusion Models (LDM) (Rombach al., 2022). LDM explores two widely adopted regularization strategies: continuous variational approach and discrete codebook framework. The varia- tional approach uses a KL divergence term to align the latent distribution with a Gaussian prior, promoting a smooth and structured latent space (Kingma & Welling, 2014). Alter- natively, the discrete codebook framework constrains the latent space a finite set of learned embeddings, limiting its complexity and providing different form of regularization (Esser et al., 2021).
1. Introduction
Latent generative models (Rombach et al., 2022) have be- come a dominant framework for high-fidelity image synthe- sis, achieving otate-of-the-art results across diffusion models (Rombach et al., 2022: Yao et al., 2024; Ma et al., 2024), masked generative modeling (Chang al., 2022; Li al.,, 2023), and autoregressive models (Esser et al., 2021; Li et al., 2024; Tian et al., 2024). These models operate in two phases. First, an autoencoder compresses high-dimensional
Aachimeddes.Ahennt Reaserch Center, Greece National Tech- nical University of Athens, Greece 'valeo.ai, France 'University of Crete, Greece IACM-FFrth Greece, Correspondence to: Theodoros Kouzelis cthoodoros @athenare.r>>.......
These regularization strategies inherently introduce a trade- off. Stronger regularization, such as increasing the weight of the KL divergence term, produces a smoother and more learnable latent space for the generative model in the second phase (Tschannen al., 2025). However, it also reduces the
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
Image
SD-VAE
+Ours
SDXI-VAE
+Ours
DiT-XI/2
DiT-XI./2 Ours
40
REPA
REPA Ours
20
GFID-50K
10
x speed-up
x4 speed-up
5
100K
400K
1M
2.3M
4M
7M
Training Iteration
Figure 1: Latent Space Structure (Left) Top three principal components of SD-VAE and SDXL-VAE, with and without EQ-VAE, demonstrating visually that our regularization produces smoother latent representations without compromising reconstruction (See Table 1). Accelerated Training (Right) Training curves (without classifier-free guidance) for DiT-XL/2 and REPA (w/ SiT-XL/2) 2), showing that our EQ-VAE accelerates convergence by x7 and x4, respectively.
information capacity of the latent representation, leading to loss of fine-gu tne-ggrined details and ultimately degrading reconstruction quality. Empirical evidence suggests that this trade-off can set an upper bound on the overall performance of latent generative models (Rombach al., 2022), as the autoencoder's 's limited capacity to preserve detailed informa- tion restricts the overall ability of latent generative models to synthesize highly-fidelity images. This raises funda- mental question: Can we mitigate this trade-off. creating a latent space that is more optimized for generative model- ing, without compromising reconstruction quality, thereby improving the overall generative modeling process?
facilitating learning for generative models (e.g., Figure/ 1- right). This is achieved without degrading the autoencoder's reconstruction quality.
Our method is compatible with both continuous and dis- crete autoencoders, enabling broad applicabílity across la- tent generative models. For example, applying EQ-VAE to the continuous SD-VAE (Rombach et al., 2022) signifi- cantly improves the performance of downstream diffusion models such as DiT (Peebles & Xie, 2023), SiT (Ma et al., 2024), and REPA (Yu et al., 2025), as measured by FID scores. Similarly, applying EQ-VAE to discrete VQ-GAN (Esser al., 2021) enhances performance in the masked gen- erative modeling approach MaskGIT (Chang al., 2022).
A key aspect that could address this challenge lies in the structure and properties of the latent space itself. In particu- lar, we identify essential limitation of current state-of-the- art autoencoders: their latent representations are not equiv- ariant to basic spatial transformations, such as scaling and rotation (see Figure 2; extended discussion in See. 3.2). This introduces unnecessary complexity into the latent manifold, forcing the generative model learn nonlinear relationships that could otherwise be avoided.
We make the following contributions:
We identify that the latent space of established au- toencoders lacks equivariance under spatial transfor- mations, which impedes latent generative modeling. Building on this observation, we propose EQ-VAE, a simple regularization strategy that improves genera- tive performance without compromising reconstruction quality.
To address this issue, we propose a simple yet effective modi- fication to the training objective of autoencoders that encour- ages latent spaces to exhibit the aforementioned equivari- ance. Our method called EQ penalizes discrepancies between reconstructions of transformed latent representa- tions and the corresponding transformations of input images. Notably, EQ-VAE requires no architectural changes to isting autoencoder models and does not necessitate training from scratch. Instead, fine-tuning pre-trained autoencoders for a few epochs with EQ EQ-VAE suffices to imbue the latent space with equivariance properties, reducing its complex- ity (see Figure l-left; quantitative results in Table 5) and
Our method is compatible with both continuous and discrete autoencoders, enabling a plug-and-play ap- proach for commonly used generative models such as diffusion and masked generative models.
We show that by fine-tuning well-established autoen- coders with our objective, we significantly acceler- ate the training of latent generative models. For in- stance, fine-tuning SD-VAE for just 5 epochs yields a x speedup on DiT-XL/ XL/2 2 and x4 speedup on REPA (w/ SiT-XL/2 2)
(see Figure 1 (right)).
2

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
2. Related work
and we extensively demonstrate the impact of equivariance regularization on latent generative modeling.
Autoencoders for Latent Generative Models Training diffusion models directly in pixel space is computation- ally inefficient, as most of the bits in digital image cor- respond to subtle details with little perceptual significance. To overcome this issue, Rombach et al. (2022) propose la- tent diffusion models that operate in compressed latent space produced in a separate stage by an autoencoder. Their KL-regularized autoencoder, SD-VAE, has been extensively utilized in numerous diffusion models (Yao et al., 2024; Ma et al., 2024; Chen et al., 2024). Subsequent research has primarily focused minimizing the reconstruction error that sets upper bound generative performance, by in- creasing the number of latent channels (Esser et al., 2024; Black Forest Labs, 2023; Dai et al., 2023) and incorporating task specific priors (Zhu et al., 2023). To enable efficient training on high-resolution images Xie et al. (2025) and Chen et al. (2025) extensively increase the compression ra- tio without compromising the reconstruction quality. Hu et al. (2023) investigate the ideal latent space for generative models and find that a relatively weak decoder produces a latent distribution that enhances generative performance. Discrete autoencoders are initially introduced with VQ-VAE (van den Oord et al., 2017) to quantize image patches into discrete visual tokens. VQ-GAN (Esser al., 2021) further refines VQ-VAF. by integrating adversarial and perceptual losses, enabling more accurate and detailed representations. Subsequent works have focused on architectural improve- ments (Yu et al.. 2022a), strategies to increase the code- book size and maximize its utilization (Yu et al., 2024; Zhu et al., 2024a). Unlike these prior approaches, we investigate a novel Srapecttvee-everagiin spatial equivariance to shape latent space better suited for generative modeling.
Equivariance in Computer Vision The success of Con- volutional neural networks (CNN) in numerous computer vision tasks can be largely attributed to their approximate translation equivariance that arises due to the nature of con- volution. To incorporate other symmetries in the data, var- ious group-equivariant convolutional networks have been proposed, including roto-translation equivariance in 2D (Co- hen & Welling, 2016; Marcos et al., 2017; Hoogeboom et al., 201 8; Weiler & Cesa, 2019), extensions in 3D (Worrall & Brostow, 2018; Thomas et al., 2018; Kondor, 2018), and scale equivariance (Rahman & Yeh, 2023; Sosnovik et al., 2020). The derivation of group equivariance constraint typi- cally results in steerable filters constructed from a basis, Be- sides architectural constraints, equívariance can be achieved by parameter sharing (Ravanbakhsh et al., 2017), frame av- eraging (Puny et al., 2022), and canonicalization functions (Kaba et al., 2023). For autoencoder models, Winter et al. (2022) produce latent representations that are separated into group invariant and equivariant part. Closely related to our work Ryu (2024), train autoencoders to be equivariant under horizontal and vertical flips. However, they do not in- vestigate the impact of equivariant representations on latent generative modeling.
3. Method
This section presents our methodology. We first provide an overview of autoencoder models for latent generative modeling (Sec. 3.1), focusing on the continuous case used in diffusion models. We then highlight the lack of equivariance in latent representations (Sec. and introduce EQ-VAR to address it (Sec. 3.3).
Auxiliary Objectives and Regularization in VAEs Au- toencoders are designed to learn latent spaces that com- pactly represent meaningful features of the observed data. However, without any regularization, their latent code lacks meaningful structure. Variational Autoencoders (VAEs) were introduced in Kingma & Welling (2014) to address this by minimizing the KL divergence between the latent distribution and Gaussian prior. Many subsequent works have adopted and extended this framework (Higgins et al., 2016; Dilokthanakul et al., 2016; Tomczak & Welling, 2018; Takahashi et al., 2019). Other works have proposed alter- native regularizations based on the Wasserstein distance (Tolstikhin et al., 2018; Kolouri et al., 2018), adversarial objectives (Zhao et al., 2018; Makhzani et al., 2015) and vec- tor quantization (VQ) (van den Oord et al., 2017). Closely related to our work, Sinha & Dieng (2021) proposes consis- tency regularization enforcing the latent code to be invariant under spatial transformations. Our EC-VAE promotes equiv- ariance rather than invariance under spatial transformations
3.1. Preliminary: Continuous Autoencoders for Latent Generative Modeling
The first modeling stage consists of an autoencoder that compresses the pixel space into a continuous (Rombach et al. (2022)) or discrete (Esser et al. (2021)) latent space. We focus here on the continuous case. Given an input image x E H x3 an encoder & transforms the image into a compressed representation 2 = E(x) E TR #x where f is the compression ratio and e are the latent channels. Then a decoder takes as input the latent representation and reconstructs the image x = D(z). For an input image x the training objective reads follows:
LVVE = Cree x) + A (x) + XregGreg (1)
where Cree consists of pixel space reconstruction objective and a perceptual loss such LPIPS (Zhang et al., 2018), Lyon is a patch- based adversarial loss (Isola et al., 2017) and Freg
3

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
SD-VAE
Ours
DDSD-VAE VAE DD VAE+(ours)
Input Image
DIE(T ox))
(()) D(r E
I SDXL- VAE D SDXL- -ALA--s
200
192
150
130
117
RFID
107
100
83
78
64
52
50
31
34
10
3
12
12
8 = 0.25
0.50
8
0.75
Latent Space Equivariance. Reconstructed images using SD-VAE (Rombach et al., 2022) and our EQ- VAE when applying scaling transformation т, with fac- tor 0.5, the input images (T x)) versus directly to the latent representations D o {x)). Our approach pre- serves reconstruction quality under latent transformations, whereas SD-VAE exhibits significant degradation. See Fig- ure 6 for additional examples.
Figure 2:
Transformations
Figure 3: Enhanced Reconstruction under Latent Trans- formations. Reconstruction RFID measured between o x and (x)+E(x)) for various spatial transformations. We con- sider scaling transforms with factors s = 0.75, 0.50, 0.25 and also measure the average RFID over rotation angles 0 x n, 2 and SDXL-VAE (Podell al., 2024), with and without EQ-VAE, Our approach significantly reduces RFID com- pared to baselines, improving image fidelity under latent transformations. For readability, we show [RFID].
3r Results for SD- (Rombach al., 2022)
is usually Kullback-Leibler regularization with a Gaussian prior (Kingma& & Welling, 2014).
3.2. Lack of Equivarance under Spatial Tansformations
Our work is motivated by a key observation: state-of-the- art autoencoders, such as SD-VAE (Rombach et al., 2022), produce latent representations (x) that are not equivariant under basic spatial transformations like scaling and rotation.
Our findings reveal a clear disparity: while autoencoders reconstruct images accurately when transformations are ap- plied the input (i.e., D(E(T o x))). applying transforma- tions directly the latent representation (i.e., D(r {())) leads to significant degradation in reconstruction quality.
We formalize this as follows:
This limitation arises because (1) convolutional architec- tures commonly used the autoencoders of latent gener- ative models, such as SD-VAE, are not equivariant under arbitrary spatial transformations such scaling and rota- tion, and (2) their standard training objectives (for example, reconstruction loss and KL divergence) do not explicitly or implicitly encourage equivariance. As result, semantically similar inputs, such as an image x and its scaled counterpart T o x, are encoded into latent codes (xx) and E o x) that are not related by the corresponding spatial transformation, i.e. E( x) A o E(x), thus unnecessarily complicating the structure of the latent space.
Spatial Transformation Let x(p) : R² R be im-
age (or latent representation) defined over 2D coordinates p [u, w] A spatial transformation 7 E R222 acts on the coordinates p transforming x as follows:
x, (p) = 'p)-
(2)
denoted compactly for all p as T x.
Equivariance A latent representation E(x) is equivariant with transformation of the input image if the transfor- mation can be transferred to the representation output:
Vx E X: E(T x E
3.3. EQ-VAE: Regularization via equivariance constraints
To test whether the latent representations of autoencoder models are equivariant under spatial transformations, we ap- plied scaling and rotations 7 directly to the latent code and evaluated the corresponding reconstructions. Specifically, we compare decoding transformed latent representations, D(T {x)), to decoding latents of transformed input im- ages, DEE o x)). We present qualitative and quantitative results in Figure 2 and Figure 3 respectively.
To address this limitation, propose EQ-VAE, which reg- ularizes the latent representations to promote equivariance under spatial transformations. As seen in Figure (left) this produces smoother latent representations, enabling more efficient learning.
Explicit Regularization. direct way to enforce equivari-

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
ance is to include the equivariance constraint from Equa- tion (3) as a loss term during training:
We note that our approach enforces equivariance by apply- ing transformations directly the latent space, distinguish- ing it from methods relying on input data augmentation (Brehmer et al., 2024).
(x)) = ||T o {x) - xl)
(4)
where is sampled from set of spatial transformations. However, minimizing this loss alone can lead to trivial so- lutions, such as collapsing the latent representation to a constant value {(x) = c, Vx, which we observe in our experiments (see Table 7), making explicit regularization ineffective.
Extending EQ- to Discrete Autoencoders. So far, we described EQ VAE in the context continuous autoen- coders. In discrete autoencoders e.g.. VQ-GAN (Esser et al., 2021), the encoder outputs continuous features & E(x) that are mapped to the nearest entry in learned codebook, form- ing a discretized latent space via quantization. Adapting our method for discrete autoencoders, such as VQ-GAN, is straightforward. We employ our equivariance regularization loss as described in Sec. 3.3 and apply the transformations T on the latent features {(x) before the quantization.
Implicit Regularization. To overcome this limitation of explicit regularization, we adopt an implicit approach. In- spired by the findings in Figure 2, this approach aligns the reconstructions of transformed latent representations (D(T o E(x)) with the corresponding transformed inputs (T rox x) ). Specifically, we modify the original training objec-
4. Experiments
tive of Equation (1) as follows:
4.1. Setup
LEQ VAE(x, = Crea rox, ((TOE(x)))
(5)
Implementation Details We finetune all autoencoders on OpenImages to adhere to the framework used in LDM (Rom- bach et al., 2022). We finetune for 5 epochs with batch size 10. Detailed specifications of each autoencoder, including spatial compression rates and latent channels, are provided in Appendix E. For DiT (Peebles & Xie, 2023), SiT (Ma et al., 2024) and REPA (Yu et al., 2025), we follow their de- fault settings and train on ImageNet (Deng et al., 2009) with batch size of 256, where each image is resized to 256 x x56.. We use B/2, XL/2 architectures which employ a patch size 2, except for the experiment with SD AADO-VAE-1 1 6 in Table in which we used B/1, due to its lower spatial resolution com- pared other autoencoders. These models are originally trained in the latent distribution of SD-VAE-FT-EMA subsequent version of the original SD-VAE that has been further fine-tuned with an exponential moving average on LAsON-Ass (Schuhmann et al., 2022) (see Table 6 and (Peebles & Xie, 2023) for their performance differ- ences). For MaskGIT, we follow (Besnier Chen, 2023) and train on ImageNet for 300 epochs with batch size of 256. We follow ADM (Dhariwal & Nichol, 2021) for all data pre-processing protocols.
Agan Gyan (((TO((x))) + dreg
where the changes compared to Eq. (1) are highlighted in color. Notice that when T is the identity transformation, this formulation reduces to the original objective in Eq. (1). By leveraging the rich supervision signal from both recon- struction and adversarial objectives, this approach implicitly encourages the encoder to produce equivariant latent repre- sentations while avoiding mode collapse (see Sec. A.I).
Transformation Design. We focus on two types of spatial transformations: anisotropic scaling and rotations. These are parameterized as:
S(22) =
Sz
R()) = cos 0 sin 0
(6)
Syj
sin cos
The final transformation is the composition of scaling and rotation: f = S(2z, 8y) R()). We sample uniformly 0.25 < $z, Sy < 1. and E (,, ,. ). ). We consider these three rotation angles (multiples of 90%) ) to avoid corner artifacts. For downsampling, we use bicubic interpolation. Empirically, we find scaling equivariance is more beneficial for generation than rotation equivariance (see Table 5).
Evaluation For generative performance, we train latent generative models on the latent distribution of each au- toencoder and we report Frechet Inception Distance (FID) (Heusel et al., 2017), sFID (Nash al., 2021), Inception Score (IS) (Salimans et al., 2016), Precision (Pre.) and Recall (Rec.) (Kynkäänniemi et al., 2019) using 50, 000 samples and following ADM evaluation protocol (Dhariwal & Nichol, 2021). To evaluate reconstruction, we report FID, Peak Signal-to-Noon larity (SSIM) (Wang et al., 2004), and Perceptual Similarity (LPIPS) (Zhang et al., 2018) using the ImageNet validation
To preserve the prior reconstruction capabilities of the au- toencoder, we return to the standard objective (Eq. (1)) by sampling the identity transform 7 = I in Eq. (5) with proba- bility Pa- Our total objective can thus be written as:
[VVR(() (ESOAAE(XXT)
P < Pas
(xxtat(x)=
p 2 Pa-
Ratio (PSNR), Structural Simi-
where p is sampled uniformly from J0. 1]. This controls the strength of our regularization. By default we set Pa = 0.5 (we ablate regularization strength in Sec. A.2).
hhtp://htuginngaacccooobiiblllll vae-f---mma
5

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
EQUIV.
ERROR
MODEL
#PARAMS
ITER.
GFID
AUTOENCODER
RFID
GFID
R(0))
S(s)
DiT-B/
130M
400K
43.5
w/ EQ-VAE (ours)
130M
400K
34.1
3D-VAE
0.90
43.8
0.89
0.69
SiT-B/2
130M
400K
33.0
BQ- (ours)
0.82
34.1
0.56
0.43
w/ EQ-VAE (ours)
130M
400K
31.2
SDXL-VAE
0.67
46.0
1.25
0.97
+ BQ-VAB (ours)
0.65
35.9
0.65
0.35
DiT-XL/2
675M
400K
19.5
DIsC CONT.
GD3-VAL VAE
0.20
58.9
0.51
0.16
w/ EQ-VAE (ours)
675M
400K
14.5
VQ-VAE (ours)
0.19
54.0
0.37
0.11
SiT-XL/2
675M
400K
17.2
SD-VAB-16
0.87
64.1
0.95
0.85
w/ EQ-VAE (ours)
675M
400K
16.1
+ EQ-VAE (ours)
0.82
49.7
0.39
0.17
DiT-XL/2
675M
7M
9.6
VQ-GAN
7.94
6.8
1.35
1.22
w/ EQ-VAE (ours)
675M
1.5M
8.8
BQ-VAE (ours)
7.54
5.9
0.64
0.55
Si SiT-XL/2 2
675M
1M
5.9
w/ EQ-VAE (ours)
675M
1M
5.9
Table 1: Comparison of Autoencoders with and without EQ-VAE. We evaluate reconstruction quality, equivariance errors (defined in Appendix C), and generative performance for continuous VVD, EDXL-VAS SD3-VAR and dis- crete (VO-GAN) autoencoders, with and without EQ-VAE. Generative FID (GFID) is measured using Di DiT-B for con- tinuous VAEs and MaskGIT for VQ-GAN. Our approach reduces reconstruction RFID and equivariance errors while enhancing generative performance (GFID). For additional reconstruction metrics see Table 12.
Table 2: GFID Comparisons. GFID scores on ImageNet 256 x 256 for DiT, SiT, and REPA trained with either SD-VAE-FT-ENE or our EC-VAE. No classifier-free guid- ance (CFG) is used. EQ-VAE consistently enhances both generative performance and training efficiency across all generative models.
MODEL
EPOCH
GFID
1St
MaskGIT
300
6.19
182.1
MaskGIT'
300
6,80
214.0
w/ EQ-VAE (ours)
130
6,80
188.1
set. To distinguish reconstruction and generation FID, we write GFID and RFID, respectively. To quantify the effec- tiveness of EQ-VAE we further measure the equivariance error (see Appendix C).
w/ EQ-EE (ours)
300
5.91
228.8
Table 3: Boosting Masked Generative Modeling. Com-
parison of GFID) and IS on ImageNet 256 x 256 for MaskGIT (Chang et al., 2022) and its open-source PyTorch reproduction' (Besnier & Chen, 2023), trained with either VQ-GAN our EQ-VAE. RQ- accelerates training by more than x2 (130 vs. 300 epochs), highlighting EQ-VAE can be effectively applied to vector-quantized autoencoders.
4.2. Equivariiance-reguarizee VAEs
We begin our experimental analysis by demonstrating the versatility of EQ-VAF, showing that it seamlessly adapts to both continuous and discrete autoencoders.
Discrete Autoencoders To investigate if EQ-VAE can be applied to discrete autoencoders, we experiment on VQ-GAN (Esser al., 2021) and validate the effectiveness of regularization the masked image modeling frame- work MaskGIT (Chang et al., 2022). In Table 1, we show that EQ-VAE is effective in the discrete case, reducing the equivariance error as well as improving the generative per- formance from 6.8 to 5.9 in GFID.
Continuous Autoencoders We integrate our EQ-VAE regularization into established continuous autoencoders with varying latent dimensions. Namely, SD-VAE, SD SD-VAE-16, (Rombach et al., 2022), SDXL-VAE (Podell al., 2024), and SD3-VAE (Esser et al., 2024). To evaluate the effect of the regularization generative performance we train DiT-B models on the latent codes before and after our regularization. We present our results in Table 1. We observe that our simple objective effectively reduces the equivariance error for all autoencoders. Further, EQ-VAE maintains the original autoencoders' reconstruction fidelity while consistently delivering significant improvements in generative performance. The results hint that there is a cor- relation between the generative performance (GFID) and the reduction in equivariacne error. Notably, for 5D-VAE, SDXL-VAE and SD-VAE-1- 16 our regularization signifi- cantly boosts generative performance. For SD3-VAE, al- though the reduction in equivariance error relatively mod- est, still results in a GFID improvement.
4.3. Boosting Generative Image Models
By applying EQ-VAR to both continuous and discrete au- toencoders, we enhance the performance of state-of-the-a generative models, including DiT a pure transformer dif- fusion model, SiT that employs continuous flow-based modeling, REPA a recent approach aligning transformer representations with self-supervised features and MaskGIT well-established masked generative model.
DiT & SiT As demonstrated in Table 2, our regular- ization approach yields significant improvements across
6

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
50K
100K
400K
50K
100K
400K
50K
100K
400K
DiT XLXX
Figure EQ-VAE accelerates generative modeling. We compare results from two DiT-XL/2 models 50K, 100K, and 400K iterations, trained with DA-VAE-FT-EERE (top) and with EQ-VAE (bottom). The same noise and number of sampling steps used for both models, without classifier-free guidance. Our approach delivers faster improvements in image quality, demonstrating accelerated convergence.
both DiT-B and DiT-XL models. Specifically, training DiT-XL/2 on the regularized latent distribution achieves GFID 14.5 at 400K iterations, compared to 19.5 without regularization. Notably, by 1.5M iterations, DiT-XI./2 trained with EQ-VAE achieves GFID 8.8, outperforming the DiT-XL/2 model trained with SD-VAE-PT-EMA even at 7M iterations. The speed-up provided by EQ-VAE can be qualitatively observed in Figure 4. Moreover, in Table 2, we show that SiT models can also benefit from the regularized latent distribution of EQ-VAE, improving GFID from 17.2 to 16.1 400K steps.
MODEL
EPOCHS
GFIDJ SFID)
ISt
PRE.
REC.
LDM
200
3.60
247.7
0.87
0.48
MaskDiT
1600
2.28
5.67
276.6
0.80
0.61
SD-DiT
480
3.23
S1T-xi1/2
1400
2.06
4.50
270.3
0.82
0.59
DiT-XI/2
1400
2.27
4.60
278.2
0.83
0.57
DiT-XI/2
1400
2.47
5.18
276.1
0.82
0.57
EQ-VAE (ours)
300
2.37
4.78
277.3
0.82
0.57
REPA# EQ-VAE (ours)
800
1.42
4.70
305.7
0.80
0.65
200
1.70
5.13
283.0
0.79
0.62
Table 4: Comparison ImageNet 256 with CFG. i indicates that the used autoencoder is the original SD-VAE (instead of SD-VAE-FT-EMA) REPA uses SiT-XI/2. * denotes that guidance inter- val (Kynkäänniemi et al., 2024) applied.
REPA We show that our regularization (which is per- formed in the first stage latent generative modeling) is complementary to REPA, thus leading to further improve- ments in convergence and generation performance. Specif- ically, training REPA (SiT-XL-2) with our EQ-VAE reaches 5.9 GFID in 1M instead of 4M iterations. Thus, the regularized latent distribution of EQ-VAE can make the convergence of REPA x4 faster (Figure 1). This is striking because REPA was shown to already significantly speed-up the convergence of diffusion models.
DiT-XL/2 with EQ-VAE reaches 2.37 GFID in just 300 epochs, matching the performance of DiT-XL/2 trained with SD-VAE or SD-VAE-FT-MAL Even when combin- ing EQ- VAE with the state-of-the-ar approach REPA, are able to achieve comparable results with standard REPA A while using x4 less training compute (200 vs 800 epochs).
MaskGIT As shown in Table 3, MaskGIT trained with our EQ- VQ-VAE converges twice as fast reaching 6.80 GFID in 130 epochs, instead of 300. Furthermore, by epoch 300 reaches 5.91 GFID surpassing the performance reported in both (Besnier & Chen, 2023) and (Chang et al., 2022).
4.4. Analysis
Spatial transformations ablation We begin the analysis of our method by ablating the effect of our equivariance regularization on generative performance with each spatial transformation to understand their respective impact. We consider isotropic S(s, s) or anisotropic S(sg, sy) scaling, rotations R(0), and combined transformations. We then
Comparison with state-of-the-art generative models To further demonstrate how EQ-VAE accelerates the learn- ing process, we compare it with recent diffusion methods using classifier-fres guidance. Notably, as shown in Table 4,

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
AUTOENCODER
T
GFID
RFID]
ID
AUTOENCODER
GFID 4 RFID 4
SD-VAE
43.5
0.90
62.2
SD-VAE (Rombach et al., 2022)
43.8
0,90
SD-VAE- FF--EM (Rombach et al.. 2022)
43.5
0.73
EQ-VAE
R(O)
41.2
0.73
57.9
EQ-VAE
S(s,
35.8
0.78
41.0
SDD-VAR
43.5
0.81
EQ-VAE
R(0) S(s,8) 8)
34.1
0.82
39.4
EQ-VAN
34.1
0.82
EQ-VAE
R(0) S(5r, 8y)
33.2
0.92
38.9
Table 6: Additional Training vs. Equivariance Regu- larization. Comparing various fine-tuning strategies for SD-VAE confirms that EQ-VAE 's improvements stem from equivariance regularization. Denotes additional training with the standard objective (Eq. (1)) for epochs.
Table 5: Spatial Transformation Ablation in EQ-VAE. We measure GFID, RFID, and intrinsic dimension (ID) for la- tents regularized via rotations, isotropic scaling, anisotropic scaling, and combinations. Combining transformations lowers ID and enhances generative performance, though anisotropic scaling can slightly degrade reconstruction.
How many epochs does EQ-VAE need to enhance gen- eration? To demonstrate how quickly our objective regu- larizes the latent distribution, we conduct an ablation study by varying the number of fine-tuning epochs. We train a DiT-B/2 model the resulting latent distribution of each epoch and present the results in Figure 5. Notably, even with a single epoch (10K steps) of fine-tuning the GFID drops from 43.5 to 36.7, highlighting the rapid refinement our objective achieves. For context, ---V-AFFFTTT EMMA has been fine-tuned for 300K steps.
44
SD-VAE
GFID-50K
42
40
38
36
+EQ-VAE VAE
(ours)
0
2
3
5
Finetuning Epochs w/ EQ-VAE
The enhancement in generative performance is not re- sult of the additional training To verify that the improve- ment in generative performance stems from our equivariance regularization (Eq. (5)) rather than additional training, we compare EQ-VAE with SD-VAEt in Table 6. SD-VAE is obtained by fine-tuning VAE for five extra epochs using only the original objective (Eq. (1)). The results show that this additional training has negligible effect on genera- tive performance, whereas EQ-VAE leads to a significant improvement. Similarly, -D-VAE-EEN EMA,,T, derived from SD-VAE, has minimal impact on the cFID score, further underscoring the effectiveness of EQ-VAE.
Figure 5: Rapid Improvement via EQ-VAE Fine-tuning. Even a single epoch of EQ-VAE fine-tuning significantly improves generative modeling performance, reducing GFID from 43.5 to 36.7. Generative modeling with DiT-B/2.
train a DiT-B/2 on each latent distribution. In Table 5, we observe that encouraging scale equivariance has a sig- nificant impact on generative performance. Furthermore, rotation equivariance is also beneficial in generation perfor- mance. Combining transformations yields further improve- ment, demonstrating their complementary effects. While anisotropic scaling yields a better generative performance since the regularization is more aggressive, it negatively impacts reconstruction quality. Thus, our EQ-VAE default setting uses combinations of rotations and isotropic scaling.
5. Conclusion
In this work, we argue that the structure of latent repre- sentations produced by the autoencoder is crucial for the convergence speed and performance of latent generative models. We observed that latent representations of estab- lished autoencoders not equivariant under simple spatial transformations. To address this, we introduce BAE,AE, simple modification to the autoencoder's 's training objective. We empirically demonstrated that fine-tuning pre-trained autoencoders with EQ- VAE for just a few epochs, is enough to reduce the equivariance error and significantly boost the performance of latent generative models while maintaining their reconstruction capability. We believe that our work in- troduces several promising future directions, particularly in exploring the theoretical and empirical relationship between the geometry of the latent distribution and the performance of latent generative models.
Latent space complexity and generative performance To better understand the impact of our regularization on the complexity of the latent manifold, we measure its Intrinsic Dimension (ID). The ID represents the minimum number of variables needed to describe data distribution (Bennett, 1969). Notably, in Table 5, we observe correlation be- tween the intrinsic dimension of the latent manifold and the resulting generative performance. This suggests that the regularized latent distribution becomes simpler to model, further validating the effectiveness of our approach. This reduction in the complexity of latent representations can also be qualitatively observed in Figure (left). For further details on ID, see Appendix
8

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
Impact Statement
Wen, Y. Emu: Enhancing image generation models us- ing photogenic needles in haystack. arXiv preprint 23X1::3309.18000 2023.
This paper presents work whose goal is to advance the field of machine learning in general and image synthesis in par- ticular. There are many potential societal consequences of work, none of which we feel must be specifically highlighted here.
Deng, J., Dong, W. Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In CVPR, pp. 248 255, 2009.
Acknowledgements This work has been partially sup- ported by project MIS 5154714 of the National Recovery and Resilience Plan Greece 2.0 funded by the European Union under the NextGeneration Program.
Dhariwal, P. and Nichol, A. Q. Diffusion models beat GANs on image synthesis. In NeurIPS, 2021.
Dilokthanakul, N., Mediano, P. A., Garnelo, M., Lee, M. C.,
Salimbeni, H., Arulkumaran, K., and Shanahan, M. Deep unsupervised clustering with gaussian mixture variational autoencoders. arXiv preprint arXiv:1611.02444 2016.
Hardware resources were granted with the support of GR- NET. Also, this work was performed using HPC resources from GENCI-IDRIS (Grants R02-AD0010012884B5
Esser, P., Rombach, R., and Ommer, B. Taming transformers for high-resolution image synthesis. In CVPR, pp. 4195- 4205, 2021.
References
Bennett, R. The intrinsic dimensionality of signal collec- tions. IEEE Transactions on Information Theory, 15(5): 517-525, 1969. doi: 10. 10.1199/TT.1199.1054444
Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J.. Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F, Podell, D., Dockhorn, T., English, Z., and Rombach, R. Scaling rectified flow transformers for high-resolution image synthesis. In ICML, pp. 22606-12633 2024.
Besnier, V. and Chen, M. A pytorch reproduction of masked generative image transformer. arXiv preprint arXiv:2310. 14400, 2023.
Black Forest Labs. Flux, 2023.
Facco, E., d' A'Erico, M., Rodriguez, A., and Laio, A. Esti- mating the intrinsic dimension of datasets by a minimal neighborhood information. Scientific reports, 7(1):12140, 2017.
Brehmer, J., Behrends, S., de Haan, P., and Cohen, T. Does equivariance matter at scale? arXiv preprint arXiv:2410.23119, 2024.
Glielmo, A., Macocco, I., Doimo, D., Carli, M., Zeni, C., Wild, R., d' d'Errico, M., Rodriguez, A., and Laio, A. Dadapy: Distance-based analysis of data-manifolds in python. Patterns, 3(10), 2022.
Chang, H., Zhang, H., Jiang, L., Liu, C., and Freeman, W. T. Maskgit: Masked generative image transformer. In CVPR, pp. 11315-11325, 2022.
Chen, J. YU, J., GE, C., Yao, L., Xie, E., Wang, Z., Kwok, J. Luo, P., Lu, H., and Li, Z. Pixart-S'aalphaS: Fast training of diffusion transformer for photorealistic text- to-image synthesis. In ICLR, 2024.
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Chen, J., Cai, H., Chen, J., Xie, E., Yang, S., Tang, H., Li, M., Lu, Y., and Han, Deep compression autoencoder for efficient high-resolution diffusion models. In ICLR, 2025.
Higgins, I., Matthey, L., Pal, A., Burgess, C. P., Glorot, X., Botvinick, M. M., Mohamed, S., and Lerchner, A. beta- vae: Learning basic visual concepts with a constrained variational framework. In ICLR, 2016.
Cheng, E., Kervadec, C., and Baroni, M. Bridging coommation-theoretiin and geometric compression in lan- guage models. arXiv preprint arXiv:2310. 13620, 2023.
Hoogeboom, E., Peters, W., Cohen, S., and Welling, M. Hcxaconv. arXiv preprint arXiv:1803.00 2018.
Cohen, T. and Welling, M. Group equivariant convolutional networks. In ICLR, pp. 2990-2999, PMLR, 2016.
Hu, T., Chen, F., Wang, H., Li, J., Wang, W., Sun, J., and Li, Z. Complexity matters: Rethinking the latent space for generative modeling. NeurIPS, 2023.
Dai, X., Hou, J., Ma, C.-Y., Tsai, S., Wang, J., Wang, R., Zhang, P., Vandenhende, S., Wang, X., Dubey, A., Yu, M., Kadian, A., Radenovic, F., Mahajan, D., Li, K., Zhao, Y., Petrovie, V., Singh, M. K., Motwani, S., and
Isola, P., Zhu, J.-Y,, Zhou, T., and Efros, A. A. Image-to- image translation with conditional adversarial networks. CVPR, pp. 1125-1134, 2017.

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
Kaba, S.-O., Mondal, A. K., Zhang, Y., Bengio, Y., and Ravanbakhsh, S. Equivariance with learned canonicaliza- tion functions. In International Conference Machine Learning, Pp. 15546-15555
Nichol, A. Q. and Dhariwal, Improved denoising diffusion probabilistic models. In ICML,, volume 139, pp. 8162- 8171, 18-24 Jul 2021.
PMLR, 2023.
Peebles, W. and Xie, Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF Interna- tional Conference on Computer Vision, pp. 4194-419000 2023.
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. In ICLR, 2014.
Kolouri, S., Pope, P. E., Martin, C. E., and Rohde, G. K. Sliced wasserstein auto-encoders. In ICLR, 2018.
Podell, D., English, Z., Lacey, K., Blattmann, A., Dockhorn, T., Müller, J., Penna, J., and Rombach, SDXL: Im- proving latent diffusion models for high-resolution image synthesis. In ICLR, 2024.
Kondor, R. N-body networks: covariant hierarchical neu- ral network architecture for learning atomic potentials. arXiv preprint arXiv: arXiv:/803.01e: 1803.01588, 2018.
Pope, P., Zhu, C., Abdelkader, A., Goldblum, M., and Gold- stein, T. The intrinsic dimension of images and its impact on learning, In JCLR, 2021.
Kvinge, H., Brown, D., and Godfrey, C. Exploring the rep- resentation manifolds of stable diffusion through the lens of intrinsic dimension. arXiv preprint 1:X1v:2300.0931, 2023.
Puny, O., Atzmon, M., Smith, E. J., Misra, I., Grover, A., Ben-Hame H., and Lipman, Y. Frame averaging for invariant and equivariant network design. In ICLR, 2022.
Kynkäänniemi, T., Karras, T., Laine, S., Lehtinen, J., and Aila, T. Improved precision and recall metric for assess- ing generative models. Advances in neural information processing systems, 32, 2019.
Rahman, M. A. and Yeh, R. A. Truly scale-equivariant deep nets with fourier layers. Advances in Neural Information Processing Systems, 36:6092-611 6104, 2023.
Kynkäänniemi, T., Aittala, M., Karras, T., Laine, S., Aila, T., and Lehtinen, J. Applying guidance in a limited interval improves sample and distribution quality in diffusion models. In NeurIPS, 2024.
Ravanbakhsh, S., Schneider, J., and Poczos, B. Equivariance through parameter-sharing. In International conference on machine learning, pp. 2892-2901. PMLR, 2017.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and
Li, T., Chang, H., Mishra, S., Zhang, H., Katabi, D., and Krishnan, Mage: Masked generative encoder unify representation learning and image synthesis. In CVPR, pp. 2142-2152, 2023.
Ommer, B. High-resolution image synthesis with latent diffusion models. In CVPR, pp. 10684-10695, 2022.
Ryu, S. Training vggan and vae, with detailed expla- nation. attps://oge //tpt/gghtbh////// com/ f vggan- trganntraining, 2024. GitHub repository.
Li, T., Tian, Y. Li, H., Deng, M., and He, K. Autoregres- sive image generation without vector quantization. arXiv preprint arXiv:2406. 11338, 2024.
Salimans, T., Goodfellow, L., Zaremba, W., Cheung, V., Radford, A., and Chen, X. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
Ma, N., Goldstein, M., Albergo, M. S., Boffi, N. M., Vanden- Eijnden, E., and Xie, S. Sit: Exploring flow and diffusion- based generative models with scalable interpolant trans- formers. In ECCV, pp. 23-40, 2024.
Schuhmann, C., Beaumont, R., Vencu, R., Gordon, C., Wightman, R., Cherti, M., Coombes, T., Katta, A., Mullis, C., Wortsman, M., Schramowski. P., Kundurthy, S., Crow- son, K., Schmidt, L., Kaczmarczyk, R., and Jitsev, J. Laion-5b: An open large-scale dataset for training next generation image-text models. In NeurIPS, volume 35, pp. 25278-25294 2022.
Makhzani, A., Shlens, J.. Jaitly, N., Goodfellow, I., and Frey, B. Adversarial autoencoders. arXiv preprint arXiv: 10l::1511..544 2015.
Marcos, D., Volpi, M., Komodakis, N., and Tuia, D. Rota- tion equivariant vector field networks. In Proceedings of the IEEE International Conference on Computer Vision, pp. 5048-5057 2017.
Simonyan, K. Very deep convolutional networks for large- scale image recognition. arXiv preprint arXiv:1409.1550 2014.
Nash, C., Menick, J., Dieleman, S., and Battaglia, P. W. Generating images with sparse representations. arXiv preprint arXiv:2 103. 003.03841, 2021.
Sinha, S. and Dieng, A. B. Consistency regularization for variational auto-encoders, In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, W. (eds.),
10

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
Advances in Neural Information Processing Systems, vol- ume 34, pp. 12943-12954 Curran Associates, Inc., 2021.
Worrall, D. and Brostow, G. Cubenet: Equivariance to 3d rotation and translation. In ECCV, pp. 567-584, 2018.
Sosnovik, L, Szmaja, M., and Smeulders, A. Scale- equivariant steerable networks. In ICLR, 2020.
Xie, E., Chen, J., Chen, J., Cai, H., Tang, H., Lin, Y., Zhang, Z., Li, M., Zhu, I.., Lu, Y., and Han, S. SANA: Effi- cient high-resolution text-to-image synthesis with linear diffusion transformers. In ICLR, 2025.
Szegedy, C., Vanhoucke, V., loffe, S., Shlens, J., and Wojna, Z. Rethinking the inception architecture for computer vision. In CVPR, pp. 2818 2-266, 2016.
Yao, J., Wang, C., Liu, W., and Wang. X. Fasterdit: Towards faster diffusion transformers training without architecture modification. In NeurIPS, 2024.
Takahashi, H., Iwata, T., Yamanaka, Y., Yamada, M., and Yagi. S. Variational autoencoder with implicit optimal pri- ors. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 5066-5066660 0066-577, 2019.
Yu, J., Li, X., Koh, J. Y., Zhang, H.,, Pang, R., Qin, J., Ku, A., Xu, Y., Baldridge, J., and Wu, Y. Vector-quantized image modeling with improved VQGAN. In ICLR, 2022a.
Thomas, N., Smidt, T., Kearnes, S., Yang. L., Li, L., Kohlhoff, K., and Riley, P. Tensor field networks: Rotation-and translation-equiia 3d point clouds. arXiv preprint neural networks for 2018.
Yu, J., Xu, Y., Koh, J. Y., Luong, T., Baid, G., Wang, Z., Vasudevan, V., Ku, A., Yang, Y., Ayan, B. K., et al. Scal- ing autoregressive models for content-rich text-to-image generation. arXiv preprint arXiv:2206. 10789, 2022b.
Yu, L., Lezama, J., Gundavarapu, N. B., Versari, L., Sohn, K., Minnen, D., Cheng, Y., Gupta, A., Gu, X., Haupt- mann, A. G., Gong, B., Yang, M.-H., Essa, I., Ross, D. A., and Jiang, L. Language model beats diffusion tokenizer is key to visual generation. In The Twelfth International Conference on Learning Representations, 2024.
Tian, K., Jiang, Y., Yuan, Z., Peng, B., and Wang. L. Visual autoregressive modeling: Scalable image generation via next-scale prediction, arXiv preprint arXiv:2404.02905. 2024.
Tolstikhin, I., Bousquet, O., Gelly, S., and Schoelkopf, B. Wasserstein auto-encoders, In ICLR, 2018
Yu, S., Kwak, S., Jang, H., Jeong, J., Huang. J.. Shin, J., and Xie, S. Representation alignment for generation: Training diffusion transformers is easier than you think. In ICLR, 2025.
Tomezak, J. and Welling, M. Vae with a vampprior. In International conference on artificial intelligence and statistics, pp. 1214-1223. PMLR, 2018.
Tschannen, M., Eastwood, C., and Mentzer, F. Givt: Gen- crative daiinte-vocabulary transformers. In ECCV, pp. 292-309, 2025.
Zhang, R., Isola, P., Efros, A. A., Shechtman, E., and Wang, O. The unreasonable effectiveness of deep features as a perceptual metric. In CVPR, pp. 186-595,40 2018.
Valeriani, L., Doimo, D., Cuturello, F., Laio, A., Ansuini, A., and Cazzaniga, A. The geometry of hidden represen- tations of large transformer models. Advances in Neural Information Processing Systems, 36:51133451122 2023.
Zhao, J., Kim, Y., Zhang, K., Rush, A., and LeCun, Y. Adversarially regularized autoeneoders. In ICML, vol- ume 80, pp. 5902-5911, 2018
Zheng, H., Nie, W., Vahdat, A., and Anandkumar, A. Fast training of diffusion models with masked transformers. arXiv preprint arXiv:2.306.0980
van den Oord, A., Vinyals, O., and kavukcuoglu, k. Neural discrete representation learning. In NeurlPS, volume 30, 2017.
2023.
Zhu, L., Wei, F., Lu, Y., and Chen, D. Scaling the codebook size of VQ-GAN to 100,000 with a utilization rate of 99%. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024a.
Wang, Z., Bovik, A., Sheikh, H., and Simoncelli, E. Image quality assessment: from error visibility to structural similarity. IEEE Transactions Image Processing, 13 (4):600-66 11:600-612,,0 2004. doi: 10. 10.1109/1IP:20038881988
Zhu, R., Pan, Li, Y., Yao, T., Sun, Z.. Mei. T., and Chen, C. W. Sd-dit: Unleashing the power of self-supervised discrimination in diffusion transformer. In CVPR, pp. 8435 8445, 2024b.
Weiler, M. and Cesa, G. General (2)-equivariant steerable cnns. Advances in neural information processing systems, 32, 2019.
Zhu, Z., Feng, X., Chen, D., Bao, J., Wang, L., Chen, Y., Yuan, L., and Hua, G. Designing a better asymmetric vq- gan for stablediffusion. arXiv preprint 3rXiv22366.04322 2023.
Winter, R., Bertolini, M., Le, T., Noe, F., and Clevert, D.-A. Unsupervised learning of group invariant and equivariant representations. NeurlPS, 2022.

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
Contents
A. Additional Ablations
A.1. Implicit vs Explicit Equivariance Regularization
Here, we provide an analysis of the design choice of our objective. We aim to design an objective that reduces the equivariance error of the encoder while avoiding mode collapse and preserving reconstruction performance. For each objective investigated, we finetune SD-VAE and evaluate the effect on generative performance by training a DiT-B/2 on the resulting latent distribution. Initially, we perform fine-tuning with the standard objective along with the explicit loss in (Eq. (4)): LV AF Aeexplici and set 0.1. We further experiment with adding stop-gradient (sg) in the 8(T o x) term in Expplicit In Table 7, observe that using Lesplic successfully reduces the equivariance enror for both rotation and scaling transformations. However, both reconstruction and generative performance degrade severely, indicating mode collapse in the latent space.
EQUIVARIANCE ERROR
LOSS
GFID
RFID
R(0) 4
S(s) 4
SD-VAE
43.5
0.90
0.93
0.80
w/ explicit
141.3
117.93
0.32
0.11
w/ explicit + sg
134.7
109.25
0.35
0.13
w/ implicit (ours)
34.1
0.82
0.49
0.15
Table 7: Implicit vs. Explicit Equivariance Regularization. Comparing SD-VAE along with explicit implicit regularization objectives shows that explicit regularization drastically lowers equivariance errors but triggers mode collapse, while implicit regularization enhances significantly the generative performance.
A.2. Regularization Strength
We evaluate the impact hyperparameter Pa which controls the strength of our regularization in Table 8.
AUTOENCODER Pa
GFID + RFID 4
EQ-VAE
0.3
35.4
0.78
EQ-VAE
0.7
34.4
0.88
EQ-VAE
0.5
34.1
0.82
Table 8: Ablation on regularization strength. We perform two experiments, with lower (pa = 0.7) and higher (pa = 0.3)
regularization strength. We observe that our method is relatively robust to choices of Pa- We highlight the setting used throughout all our experiments.
B. Details on the Intrinsic Dimension Estimation.
Several recent works (Valeriani al., 2023; Kvinge al., 2023: Cheng et al., 2023) have utilized ID to measure the complexity of latent representations in deep learning modeling. Further Pope al. (2021) has demonstrated strong comelation between dataset's relative difficulty and its ID. We compute the ID of the latent representations using the TwoNN estimator (Facco et al., 2017), which relies solely on the distances between each point and its two nearest neighbors. In practice, the TwoNN estimator can be affected by noise, which typically leads to an overestimation of the ID. Nevertheless, is a robust tool to evaluate relative complexity and has been used effectively to analyze representations in deep neural networks (Valeriani al., 2023). We adopt the TwoNN implementation of DADAPY (Glielmo et al., 2022).
12

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
C. Details on Evaluation Metrics
C.1. Generation Metrics
We follow the setup and use the same reference batches of ADM (Nichol Dhariwal, 2021) for evaluation, utilizing their official implementation'. We use NVIDIA A100 GPUs for our evaluation. We briefly explain each metric used for the evaluation.
FID (Heusel et al., 2017) quantifies the feature distance between the distributions of two image datasets by leveraging the Inception-v3 network (Szegedy et al., 2016). The distance is calculated based on the assumption that both feature distributions follow multivariate Gaussian distributions.
sFID (Nash et al., 2021) computes FID with intermediate spatial features of the Inception-v3 network, capture spatial distribution of the generated images
(Salimans et al., 2016) measures a KL-divergence between the original label distribution and the distribution of Inception-v3 network's logits after the softmax normalization.
Precision and Recall (Kynkäänniemi et al., 2019) are the fraction of realistic images and the fraction of training data covered by generated data respectivly.
http:s////gthhuu..c/////eeeesst httts//o///eeee/mi////eette iaati ons
13

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
C.2. Reconstruction Metrics
We evaluate reconstruction on the validation of Imagenet which contains 50K images. We provide description of each metric used for the reconstruction evaluation.
PSNR measures the quality of reconstructed images by comparing the maximum possible signal power to the level of noise introduced during reconstruction. Expressed in decibels (dB).
SSIM (Wang et al., 2004) assesses the similarity between two images by evaluating their structural information, luminance, and contrast.
LPIPS (Zhang al., 2018) evaluates the perceptual similarity between two images by comparing their deep feature representations using VGG (Simonyan, 2014)
C.3. Equivariance Error
To quantify the effectiivnness of EQ- VVE constraining the latent representations of the autoencoders to equivary under scale and rotation transformation we measure the equivariance error. Similar to (Sosnovik et al., 2020) we define the equivariance error as follows: A 77 ΣΣIT
||T o E(x) - E(r o x) x)11 / [|ε(T x) 1 where N = 50K in the number of samples in
I
ImageNet validation and T is the set of transformations considered. We conduct our evaluation with T, = { 2: n, 3m for
rotations and Tg = {0.25, 0.50, ] for scale.
D. Detailed Benchmarks
D.1. Detailed generative performance
We provide detailed evaluation of all the generative models presented in the main paper, including additional metrics and training iterations. Specifically, Table 9 details the performance of the DiT-XL/2 and SiT-XL/S models, while Table 10 presents results for the REPA (SiT-XL/2) 2) models trained with both SD-VAE-FT-ENH (as reported in the respective papers) and EQ-VAE. Additionally, Table 11 provides results for MaskGIT models trained using VQ-GAN and EQ-VAE. For all models, we use the evaluation metrics originally reported original publications.
MODEL
#ITERS.
FID
SFID
IS
PREC.1
REC.1
DiT-XL/2 (Peebles & Xie, 2023)
400K
19.5
6.5
77.5
0.60
0.60
VAE
50K
73.6
13.1
34.5
0.50
0.37
w/ EQ-VAE
100K
39.9
6.8
62.2
0.60
0.53
w/ EQ-VAE
200K
22.8
5.9
73.6
0.61
0.62
w/ EQ-VAE
400K
14.5
5.6
81.5
0.63
0.66
SiT-XI/2 (Ma et al., 2024)
400K
17.2
5.1
76.5
0.64
0.63
w/ EQ-VAE
50K
76.1
38.4
15.2
0.50
0.37
w/ EQ-VAE
100K
41.3
10.9
30.9
0.60
0.53
w/ EQ-VAE
200K
24.9
6.4
54.6
0.61
0.62
w/ EQ-VAE
400K
16.1
4.2
79.7
0.64
0.66
Table 9: Detailed evaluation for DiT-XL/2 and SiT-XL/2 models. All results are reported without classifier-free
guidance (CFG 1.0).
14

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
MODEL
#lTERS.
FIDJ
SFID
1St
PREC.1
REC.1
REPA (Yu et al., 2025)
50K
52.3
31.2
24.3
0.45
0.53
w/ EQ-VAB
50K
48.7
26.3
27.6
0.44
0.53
REPA (Yu et al., 2025)
100K
19.4
6.1
67.4
0.64
0.610
w/ EQ-VAE
100K
18.7
5.4
67.8
0.65
0.59
REPA (Yu et al., 2025)
200K
11.1
5.1
100.4
0.69
0.64
w/ EQ-VAE
200K
10.7
5.1
103.5
0.70
0.62
REPA ( Yu et al., 2025)
400K
7.9
5.1
122.6
0.70
0.65
w/ EQ-VAE
400K
7.5
5.0
128.8
0.71
0.63
Table 10: Detailed evaluation on REPA (SIT-XL/2) models. All results are reported without classifier-free guidance
(CFG 1.0)
MODEL
EPOCHS
FID
1St
PREC.
REC.1
MaskGIT (Chang et al., 2022)
300
6.2
182.1
0.80
0.51
MaskGIT* (Besnier & Chen, 2023)
300
6.8
214.0
0.82
0.51
w/ EQ-VAE
50
11.1
116.1
0.73
0.52
w/ EQ VAE
100
7.6
167.5
0.78
0.53
w/ EQ-VAE
200
6.0
211.8
0.79
0.55
w/ FQ-VAE
300
5.9
228.7
0.80
0.55
Table 11: Detailed evaluation MaskGIT models. We the open-source PYTORCH reproduction for experiments. All results are reported without classifier-free guidance (CFG 3.0)
D.2. Detailed reconstruction performance
We detail Table 12 the reconstruction evaluation metrics of each autoencoder with and without EQ-VAE regularization.
AUTOENCODER
RFID, PSNR1
LPIPS
SSIM
VD-VAE
0.90
25.82
0.146
0.71
w/ BQ-VAE (ours)
0.82
25.95
0.141
0.72
SDXL-VAE
0.67
27.36
0.121
0.76
w/ EQ-VAE (ours)
0.65
27.48
0.118
0.76
CONT.
3D3-VAE
0.20
31.27
0.060
0.87
w/ EQ-VAE (ours)
0.19
31.06
0.061
0.87
SD-VAE-16
0.87
24.67
0.161
0.61
w/ BQ-VAE (ours)
0.82
25.21
0.152
0.69
DISC.
VQ-GAN
7.94
19.41
0.540
0.54
w/ EQ-VAE (ours)
7.54
19.61
0.510
0.56
Table 12: Comparison of Autoencoders with and without EQ-VAE. Additional reconstruction evaluation metrics.
15

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
E. Specifications of Autoencoder Models
AUTOENCODER
Creg
DATASET
C
(Rombach et al., 2022)
KL
OpenImages
4
SD-VAB-FT-EMA. (Rombach et al.. 2022)
KL
Openlmages + Laion Aesthetics
4
SD-VAE-15 (Rombach et al., 2022)
KI.
OpenImages
16
16
SDXL SDXL-VAE VAE (Podell al., 2024)
KL
4
VQ-GAN (Esser al., 2021)
VQ
ImageNet
256
16
Table 13: Specifications of Autoencoders. We provide additional information for the autoencoders used in experiments regarding their original training dataset, latent channels and compression rate f.
F. Latent Generative Models
Here we provide brief description of the latent generative models, mentioned in the main paper:
MaskGIT (Chang et al., 2022) utilizes a bidirectional transformer decoder synthesize images by iteratively predicting
masked visual tokens produced by VQ-GAN (Esser al., 2021).
LDM (Rombach et al., 2022) proposes latent diffusion models, modeling the image distribution in a compressed latent
space produced by KL- or VQ-regularized autoencoder.
DiT (Yao et al., 2024) proposes a pure transformer backbone for training diffusion models and incorporates AdaIN-zen -ero modules.
MaskDiT (Zheng al., 2023) trains diffusion transformers with an auxiliary mask reconstruction task.
SD-DiT (Zhu al., 2024b) extends the MaskDiT architecture by incorporating a discrimination objective using momentum encoder.
SiT (Ma et al., 2024) improves diffusion transformer training by moving from discrete diffusion to continuous flow-based modeling.
REPA (Yu et al., 2025) aligns the representations of diffusion transformer models to the representations of self- supervised models.
16

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
G. Additional Qualitative Results
SD-VAE
Ours
Input Image x
D(E(T o x))
D(r E(x)
D(r o E(x))
Scaling Transformation (s 0.5)
(b) Rotation Transformation (0 180%)
Figure 6: Latent Space Equivariance. Reconstructed images using SD-VAS (Rombach al., 2022) and EQ-VAE when applying transformations T to the input images (D(E(T e x))) versus directly to the latent representations (D(r o {x))). We present results for scaling and rotation transformations Our approach preserves reconstruction quality under latent transformations, whereas SD-VDD VVE exhibits significant degradation.
17

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
Image
SD-VAE
+Ours
SDXL-VAE
+Ours
R
R
Figure 7: Additional comparisons of latent representations across different VAE models. EQ-VAE (+Ours) produces smoother latent representations for both SD-VAE and SDXL-VAE.
18

---

Equivariance Regularized Latent Space for Improved Generative Image Modeling
Class label "panda" (388)
Class label = "golden retriever" (207)
Class label "macaw" (88)
Figure 8: Uncurated samples 256 x 256 DiT-XL/2 /w EQ-VAE, Classifier-frre guidance scale = 4.0.
19


Reviews:
[{"review_id": "IhQXcpVyuX", "reviewer": "Reviewer_bi5M", "summary": "This paper observes that existing autoencoders lack equivariance to semantic-preserving transformations like scaling and rotation, resulting in complex latent spaces that hinder generative performance. Based on this observation, the authors propose to regularize the latent by enforcing equivariance in the latent space, reducing its complexity without degrading reconstruction quality. Experiments on different generative models demonstrate the effectiveness of the proposed method.\n\n## update after rebuttal\nMost of my concerns have been addressed and I lean to keep the positive rating.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your insightful comments and efforts in reviewing our manuscript. Below, we provide our responses to each of your comments:\n\n---\n\n**W1. EQ-VAE for T2I generation.**\n\nWe appreciate the reviewer’s concern regarding the applicability of our method to text-to-image (T2I) generation. To address this, we conducted an additional T2I experiment using the MS-COCO dataset [1]. While this serves as a preliminary result in a small-scale setting—given that large-scale T2I experiments exceed our computational resources—it provides valuable empirical validation.\n\nFor this experiment, we employed U-ViT-S/2 and followed their experimental setup [2]. We used SD-VAE to extract image latents in the baseline setting. During sampling, we use CFG with w=2.0. The table below reports gFID at every 50K iterations. We observe that EQ-VAE demonstrates improvements in T2I generation, highlighting the significance of equivariance regularization. These findings suggest that incorporating EQ-VAE into large-scale T2I models is a promising direction for future research.\n\n\n| **Iter.** | **U-ViT-S/2 w/ SD-VAE** | **U-ViT-S/2 w/ EQ-VAE** |\n|:--------:|:-------------------:|:-------------------:|\n| **50K**  |         15.5        |         12.4        |\n| **100K** |         8.6         |         7.6         |\n| **150K** |         7.7         |         7.1         |\n| **200K** |         7.5         |         6.9         |\n| **250K** |         7.3         |         6.8         |\n| **300K** |         7.2         |         6.7         |\n| **350K** |         7.1         |         6.6         |\n| **400K** |         7.0         |         6.6         |\n| **450K** |         7.0         |         6.5         |\n\n---\n\n\n**W2.  Performance with EQ-VAE under same training resources as the baselines.**\n \nWe appreciate the reviewer’s question about whether EQ-VAE solely accelerates convergence or also leads to better overall performance given the same training resources as the baselines. Due to our limited computational resources and short rebuttal timeline, we were unable to train XL models to full convergence (e.g. 1400 epochs for DiT/SiT). However, we recognize the importance of this evaluation and will include experiments with fully converged models in the final version of our paper to provide a more thorough assessment of convergence speed and final performance.\n\nThat said, we emphasize that DiT-XL/2 with EQ-VAΕ, trained for only 300 epochs, already outperforms DiT-XL/2† (which uses SD-VAE before our EQ-VAE fine-tuning), even though the latter is considered fully converged after 1400 epochs (Table 4). This suggests that EQ-VAE is not only accelerating convergence but also contributing to improved performance. We will further investigate this in our updated experiments.\n\n---\n\n**W3. Detailed Reconstruction Metrics.**\n\nWe present detailed evaluation metrics for three models: SD-VAE, SD-VAE† (SD-VAE finetuned for 5 epochs with the original objective), and EQ-VAE (SD-VAE finetuned for 5 epochs with our objective). EQ-VAE significantly boosts the gFID performance of DiT-B/2. Importantly, this gain is achieved without sacrificing reconstruction quality as both EQ-VAE and SD-VAE† show similar improvements over the baseline SD-VAE across all the reconstruction metrics.\n\n\n| **Model**   | **gFID↓** | **rFID↓** | **PSNR↑** | **LPIPS↓** | **SSIM↑** |\n|:-----------:|:---------:|:---------:|:---------:|:----------:|:---------:|\n| **SDVAE**   |   43.5    |   0.90    |   25.82   |   0.146    |   0.71    |\n| **SD-VAE†** |   43.5    |   0.81    |   25.98   |   0.139    |   0.72    |\n| **EQ-VAE**  |   34.1    |   0.82    |   25.95   |   0.141    |   0.72    |\n\n---\n\n**Q1. Autoencoder Training Steps.**\n\nWe thank the reviewer for this question. We provide a detailed breakdown of the original training epochs for SD-VAE and SD-VAE-16, calculated based on the training steps reported in the [official repository](https://github.com/CompVis/latent-diffusion?tab=readme-ov-file#model-zoo) and the batch size specified in the [hugging face repository](https://huggingface.co/stabilityai/sd-vae-ft-ema-original). The epochs are calculated as follows epochs = (steps * batch_size) / dataset_size, where the dataset size for OpenImages is ~1.74M.\nFor VQ-GAN, SD3-VAE, and SD-XL-VAE, to the best of our knowledge, their training iterations are not explicitly stated in their respective papers or official repositories.\n\n\n| **Model**       | **OpenImages Epochs** |\n|:-----------:|:-------------------:|\n| **SD-VAE**      | 27.1          |\n| **SD-VAE-16**   | 48.8           |\n\nWe note that while our primary benchmark experiments use 5 fine-tuning epochs, our ablation study in Fig. 5 shows that even 1 epoch leads to notable performance improvements, particularly in terms of gFID.\n\n---\n\n[1] Microsoft COCO: Common Objects in Context, In ECCV 2014\n\n[2] All are Worth Words: A ViT Backbone for Diffusion Models, In CVPR 2023", "claims_and_evidence": "Yes, the claims made in the submission are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "Yes, the proposed methods and evaluation criteria make sense for the problem or application at hand.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Yes, I have checked the soundness of the experimental designs and analyses. Seems fine to me.", "relation_to_broader_scientific_literature": "This paper shares similar motivation with prior work like REPA which aims to enhance the training efficiency of latent generative models.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n1. The idea of exploring the property of latent space is applaudable, which not only brings practical speedup of latent generative models, but provide insights for the design of VAEs.\n2. The proposed EQ-VAE can generalize to both continuous and discrete autoencoders with only a few epochs finetuning, and experiments on multiple generative models demonstrate the effectiveness of the proposed method.\n3. The paper is well-organized and the writing is clear.\n\nWeaknesses:\n1. While the effectiveness of the proposed method has been validated with image generation on ImageNet, its efficacy on text-to-image, which requires larger training sets, remains unverified. The reviewer understands that it would take much larger training resources for T2I experiments, but the translation of the effectiveness to T2I cannot be guaranteed without empirical validation.\n2. Although EQ-VAE has shown advantage in reducing the training costs compared to the baseline methods, whether further training could bring additional performance improvement is unclear, that is, whether EQ-VAE just makes the latent generative models converge faster or it could result in better performance of latent generative models with the same training resources as the baselines.\n3. While EQ-VAE has compared with baseline VAE in rFID, this metric may be aligned with the training iterations and it is suggested to provide comprehensive comparisons with baseline VAE, finetuned VAE with old objectives and the proposed method with more metrics like PSNR, SSIM and LPIPS.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Although the authors mentioned that they only finetune the VAE for 5 epochs on OpenImages, it would be better to have a sense of comparisons between the finetuning costs and the original pre-training costs if applicable. The reason is that different VAEs may train their models with different iterations and batch sizes.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UTH5lw5zZk", "reviewer": "Reviewer_JW8k", "summary": "The paper proposes EQ-VAE, a framework that introduces equivariance regularization into the training of autoencoders. By incorporating 2D transformations such as rotation and scaling, the method improves the structure and representation ability of the latent space. As a result, it accelerates the training of generative models and enhances generation quality. The authors demonstrate the effectiveness of EQ-VAE through extensive experiments on both discrete and continuous autoencoders, showing consistent improvements in performance across various generative modeling tasks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your insightful comments and efforts in reviewing our manuscript. Below, we provide our responses to each of your comments:\n\n---\n\n**Theoretical validation**\n\nWhile we focus on empirical evidence in this work,  we believe that mathematically formalizing the underlying mechanisms behind the success of equivariance regularization in generative models will be an interesting future direction to explore. Our empirical observation that equivariance regularization reduces the intrinsic dimension of the latent manifold—correlating with improvements in generative performance (as shown in Table 5) is an interesting starting point for future theoretical research.\n\nThe reviewer’s suggested directions, such as investigating how equivariance regularization influences the geometry of the latent space and its impact on the optimization landscape of generative models, are particularly compelling. We will incorporate these insights into the Future Work section of our paper.", "claims_and_evidence": "Yes, the claims made in the submission are supported by clear and convincing evidence. The authors demonstrate that their proposed training strategy effectively improves the reconstruction capability of the VAE itself. Furthermore, they provide experimental results showing that, by offering a more expressive latent space through equivariance regularization, the subsequent training of generative models is also improved. This leads to better generation quality and faster convergence. The experiments cover both discrete and continuous autoencoders, and the evaluations are consistent with the claims presented in the paper.", "methods_and_evaluation_criteria": "Yes, the proposed methods and evaluation criteria make sense for the problem at hand. The authors adopt standard datasets commonly used in image generation tasks, such as OpenImages and ImageNet. They also evaluate their method using widely accepted metrics, including FID and sFID, which are appropriate for assessing the quality and diversity of generated images. The experimental setup aligns well with the goals of the paper and provides a fair basis for comparison.", "theoretical_claims": "Why does the proposed fine-tuning strategy effectively enhance the representational capacity of VAEs and improve the convergence behavior of downstream generative models? While the paper provides intuitive explanations and empirical results to justify these benefits, a more rigorous theoretical analysis or mathematical justification would significantly strengthen the work. For instance, formalizing how equivariant regularization influences the geometry of the latent space, or how it impacts the optimization landscape of generative models, would offer deeper insights beyond empirical validation.", "experimental_designs_or_analyses": "Yes, I checked the soundness and validity of the experimental designs and analyses. The authors conduct thorough comparative experiments on various VAE models, as well as a series of generative models.", "relation_to_broader_scientific_literature": "The key contributions of the paper are closely related to the broader literature on improving latent space representations in generative models, particularly Variational Autoencoders (VAEs). Prior works have demonstrated that the structure and consistency of the latent space play a critical role in the performance of generative models, both in terms of reconstruction quality and sample generation. This paper builds on these findings by introducing equivariance regularization, ensuring that the latent space behaves consistently under geometric transformations such as rotation and scaling.\n\nUnlike previous approaches that focus on improving VAE expressiveness through architecture changes or better priors, EQ-VAE emphasizes geometric consistency, which has been underexplored in this context. The work is also related to recent efforts in incorporating equivariance into deep learning models, but it uniquely applies this concept to enhance the quality and robustness of latent representations in VAEs, which in turn benefits downstream generative tasks.", "essential_references_not_discussed": "No, the paper covers the most relevant prior work related to equivariant representation learning and variational autoencoders. The cited literature provides sufficient context for understanding the key contributions of the paper. I did not identify any essential references that are missing or overlooked.", "other_strengths_and_weaknesses": "Please refer to other sections.", "other_comments_or_suggestions": "No", "questions_for_authors": "No.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "b0QVv2JyWG", "reviewer": "Reviewer_28Pq", "summary": "The paper introduces EQ-VAE, a novel variant of autoencoders designed to enhance the performance of latent generative models. The authors first identify that commonly used autoencoders in modern generative models are not equivariant to spatial transformations of the input, such as scaling and rotation. They argue that enforcing this property can lead to improved generation quality. To address this, they propose a new implicit regularization loss that encourages the encoder network to be equivariant with respect to scaling and rotation. Experimental results demonstrate that incorporating this regularization not only accelerates the training of generative models but also improves their performance in certain cases.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your insightful comments and efforts in reviewing our manuscript. Below, we provide our responses to each of your comments:\n\n---\n\n**W1. Results with CFG for converged models.**\n\nWe appreciate the reviewer’s concern regarding potential performance degradation when integrating EQ-VAE into generative models with classifier-free guidance (CFG). However, we respectfully clarify that EQ-VAE does not inherently degrade performance.\n\nAs shown in Table 4, REPA (trained for 800 epochs) and REPA with EQ-VAE (trained for only 200 epochs) cannot yet be directly compared due to the significant difference in training duration. Notably, DiT-XL/2 with EQ-VAE—trained for just 300 epochs—already outperforms DiT-XL/2† (which uses SD-VAE before EQ-VAE fine-tuning), even though the latter was trained to full convergence (1400 epochs). This suggests that EQ-VAE not only accelerates convergence but may also enhance performance under CFG.\n\nDue to our limited computational resources and short rebuttal timeline, we were unable to train XL models to full convergence. However, we understand the importance of assessing convergence speed in the CFG setting and will include fully converged experiments in the final version of our paper to provide a clearer evaluation.\n\n---\n\n**S1. EQ-VAE with efficient autoencoders.**\n\nWe appreciate the reviewer’s suggestion. Investigating whether our equivariance regularization can be applied to efficient autoencoders such as [1] and [2] is an interesting topic for further exploration. We will include this direction in the Future Work section of our paper.\n\n**Q1. Reasoning behind selecting scaling and rotation.** \n\nThe development of scale and rotation equivariant networks [3],[4] has been extensively studied for various image understanding problems. \nThis motivated us to explore whether well-established autoencoders utilized in latent generative modeling are equivariant under those basic semantic preserving transformations.  \nOur findings, showcased in Figures 2 and 6, revealed that existing architectures lack this property, directly motivating our equivariance regularization approach.\n\n---\n\n**Q2. VQ-VAE training instability.**\n\nWe did not encounter any instability during VQ-VAE fine-tuning.\n\n---\n\n **Q3. Performance gap widens after 400k iterations?**\n\nWe observe that the performance gap remains consistent across all training iterations. As an example, we present a gFID comparison of REPA without and with EQ-VAE at different training iterations. \n\n| **Iter.** | **REPA (gFID)** | **REPA w/ EQ-VAE (gFID)** |\n|:---------------:|:----------------:|:--------------------------:|\n|          50K   |            52.3 |                     48.7 |\n|         100K   |            19.4 |                     18.7 |\n|         200K   |            11.1 |                     10.7 |\n|         400K   |             7.9 |                      7.5 |\n|           1M   |             6.4 |                      5.9 |\n\n---\n\n[3] Group equivariant convolutional networks. In ICLR, 2016.\n\n[4] Scaleequivariant steerable networks. In ICLR, 2020.", "claims_and_evidence": "The claims are sufficiently supported with various experiments.", "methods_and_evaluation_criteria": "The method is evaluated using common metrics such as recnostruction FID, generation FID, Inception Score, and LPIPS.", "theoretical_claims": "There is no theoretical claim in the paper.", "experimental_designs_or_analyses": "The experiments are well-designed, and the evaluation setup is fair across various models and architectures.", "relation_to_broader_scientific_literature": "While regularization techniques for enforcing equivariance in deep learning models have been explored in recent work, none have directly addressed this issue in the context of generative models. Moreover, no prior study has examined the impact of an equivariant latent space on generation quality. Accordingly, the paper's contributions are well-positioned within the broader literature.", "essential_references_not_discussed": "The paper includes all essential references.", "other_strengths_and_weaknesses": "### **Strengths**\n- The paper is well-presented, making it clear and enjoyable to read.  \n- Given that latent generative models are the predominant approach for high-resolution image generation, the contributions of this work are likely to have a significant impact on the field.  \n- While most recent studies focus primarily on the diffusion or generative components of such systems, this paper explores a relatively underexplored area by improving the latent space of the autoencoder.\n\n### **Weaknesses**\n- The main weakness of the work is that integrating EQ-VAE into generative models can lead to a performance drop for some models when using classifier-free guidance (e.g., REPA in Table 4).  It would be valuable to assess how much additional training is required to recover the original model's performance. This provides a clearer evaluation of the convergence speed, especially since such models are rarely used without classifier-free guidance.", "other_comments_or_suggestions": "It would be interesting to explore how EQ-VAE impacts other autoencoders designed to improve the efficiency of SD-VAE, such as Cosmos-Tokenizer [1] and LiteVAE [2].\n\n[1] Agarwal N, Ali A, Bala M, Balaji Y, Barker E, Cai T, Chattopadhyay P, Chen Y, Cui Y, Ding Y, Dworakowski D. Cosmos world foundation model platform for physical ai. arXiv preprint arXiv:2501.03575. 2025 Jan 7.\n\n[2] Sadat S, Buhmann J, Bradley D, Hilliges O, Weber RM. Litevae: Lightweight and efficient variational autoencoders for latent diffusion models. arXiv preprint arXiv:2405.14477. 2024 May 23.", "questions_for_authors": "1) What is the reasoning behind selecting only scaling and rotation as the equivariant operations?  \n\n2) Does adding the EQ-VAE loss in the context of VQ-VAE introduce any instability or impact due to the additional quantization step?  \n\n3) Do you have any intuition about why the performance gap widens after 400k iterations? Specifically, why do different VAEs perform similarly in the early training phase, but diverge later?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "W3z0E2Em4v", "reviewer": "Reviewer_k5zo", "summary": "This paper addresses the limitations of existing latent generative models, which often lack equivariance to semantic-preserving transformations like scaling and rotation. To overcome this challenge, the authors propose EQ-VAE, a regularization technique that enforces equivariance in the latent space, simplifying it while preserving reconstruction quality. EQ-VAE enhances the performance of various state-of-the-art generative models and is compatible with both continuous and discrete autoencoders, providing a versatile improvement for a range of latent generative frameworks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your insightful comments and efforts in reviewing our manuscript. Below, we provide our responses to each of your comments:\n\n---\n\n**W1. Minor improvements in Table 2 compared to REPA.**\n\nWe respectfully emphasize that in Table 2 EQ-VAE demonstrates substantial improvements across all models (DiT, SiT, and REPA) in both B and XL configurations. While the improvement of EQ-VAE for SiT-B/2 may appear relatively small compared to REPA, this comparison requires important contextual considerations:\n    \n- REPA is a distillation strategy applied directly in the generative diffusion stage, which inherently leads to improved performance, particularly with the powerful pre-trained visual encoder (DINOv2r) it leverages.\n- In contrast, our method, EQ-VAE, is applied in the autoencoding stage and does not rely on any external model. Instead, it regularizes the latent space, resulting in enhanced performance for generative modeling.\n\nMoreover, as shown in Table 2 and Figure 1 (right), *EQ-VAE is orthogonal to REPA*, accelerating its convergence by a factor of four. This highlights EQ-VAE’s ability to improve efficiency without sacrificing performance.\n\nLastly, we note that in our reproduction environment, SiT-B/2 achieved a higher 34.7 FID compared to the 33.0 FID reported in the paper, making the improvement attributed to EQ-VAE slightly larger (34.7 → 31.2 FID).", "claims_and_evidence": "The authors provide empirical results demonstrating the effectiveness of their method, EQ-VAE, in improving the performance of both continuous and discrete autoencoders. Specifically, they mention significant speedups in training times and improvements in downstream model performance, as measured by FID scores.", "methods_and_evaluation_criteria": "The use of benchmark datasets and FID scores as evaluation criteria is appropriate for assessing the performance of image synthesis methods. These metrics provide a clear way to quantify improvements in image quality and training efficiency, making the evaluation relevant and meaningful for the context of latent generative modeling.", "theoretical_claims": "This paper does not involve theoretical claims.", "experimental_designs_or_analyses": "By comparing EQ-VAE against baseline models like SD-VAE and SD-VAE-EMA-FT, the authors convincingly demonstrate that improvements in generative performance are attributed to their proposed method rather than just additional training. No significant issues were noted in the designs or analyses.", "relation_to_broader_scientific_literature": "The paper builds upon the foundational work of latent variable models, particularly in the context of variational autoencoders (VAEs). By applying their method to enhance models like DiT and SiT, the authors demonstrate a significant improvement in generative performance.", "essential_references_not_discussed": "There is no significant related work that is not discussed.", "other_strengths_and_weaknesses": "Strengths\n\nThe paper addresses critical challenges in latent generative modeling, specifically the trade-offs between reconstruction quality and generative performance. By proposing a solution that enhances both aspects, the work has significant implications for advancing state-of-the-art generative models.\n\nWeaknesses\n\nThe experimental results in Table 2 show that the improvement of the baseline by using EQ-VAE is relatively small. For example, SiT-B/2, 400K. The experimental result of REPA under the same configuration is about 25. Does this mean that the improvement of the generated results by this method is relatively small?", "other_comments_or_suggestions": "No other suggestions.", "questions_for_authors": "No other questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper receives unanimously positive reviews for its solid contributions in introducing equivariance regularization into latent generative models. Although there are some relatively minor concerns around its evaluation and theoretical contributions, the work has made substantial contributions to generative modeling.

Therefore, I recommend acceptance of this work.
Decision: Accept (poster)