14 Feb 2025
https://arxiv.org/pdf/2502.09509
https://github.com/zelaki/eqvae.
GitHub - zelaki/eqvae: EQ-VAE: Equivariance Regularized Latent Space for Improved Generative Image Modeling.
EQ-VAE: Equivariance Regularized Latent Space for Improved Generative Image Modeling. - zelaki/eqvae
github.com
EQ-VAE
Equivariance Regularized Latent Space for Improved Generative Image Modeling
eq-vae.github.io
핵심:
- **입력 이미지 x**에 transformation \tau을 적용하고 encoder로 latent 표현을 구함: \(E(\tau \circ x)\)
- **원래 이미지 x**를 encoder에 통과시킨 후 latent 표현에 transformation을 적용함: \(\tau \circ E(x)\)
- 이 둘의 차이를 최소화하도록 학습시킴
0. Abstract
Latent generative models은 고품질 이미지 생성을 위한 주요 접근 방식으로 자리 잡음.
이러한 모델은 이미지들을 latent space로 압축하는 autoencoder와, 이 latent 공간의 분포를 학습하는 generative model로 구성됨.
하지만 기존 autoencoder들은 scaling, rotation처럼 의미 보존적인 변환(semantic-preserving transformation)에 대해 equivariance를 갖추지 못함.
그 결과, latent 공간이 복잡해지고 이는 이미지 생성 성능을 저해하는 요인이 됨.
이를 해결하기 위해, EQ-VAE라는 간단한 **정규화 방식(regularization approach)**을 제안함.
이 방식은 latent space 내에서 equivariance를 강제함으로써, 복잡도는 줄이면서도 복원 품질(reconstruction quality)은 유지하게 만듦.
사전학습된 autoencoder를 EQ-VAE로 fine-tuning함으로써, DiT, SiT, REPA, MaskGIT 등 여러 SOTA generative model들의 성능을 향상시킴.
특히, SD-VAE를 단 5 에폭만 fine-tuning해도 DiT-XL/2에서 7배 빠른 속도 향상을 달성함.
EQ-VAE는 연속형(continuous) autoencoder와 이산형(discrete) autoencoder 모두와 호환되므로, 다양한 latent generative model에 범용적으로 적용 가능함.
프로젝트 페이지와 코드는 다음에서 확인 가능함: https://eq-vae.github.io
1. Introduction
Latent generative models (Rombach et al., 2022)은 고해상도 이미지 생성을 위한 주요 프레임워크로 부상함.
이들은 diffusion models (Rombach et al., 2022; Yao et al., 2024; Ma et al., 2024), masked generative modeling (Chang et al., 2022; Li et al., 2023), autoregressive models (Esser et al., 2021; Li et al., 2024; Tian et al., 2024) 등에서 SOTA 결과를 달성해옴.
이 모델들은 두 단계로 작동함:
- autoencoder가 고차원 이미지를 저차원의 latent space로 압축함.
- 이 latent 공간은 연속적일 수도 있음(e.g., SD-VAE for diffusion)
- 또는 이산적일 수도 있음(e.g., VQ-GAN for autoregressive & masked generative modeling).
- 이 공간은 의미적(semantic)이고 구조적인 정보는 유지하면서도, 고주파(high-frequency) 세부정보는 버리게 됨.
- 그 다음, generative model이 이 latent 표현들의 분포를 학습하게 됨.추론 시에는, generative model이 latent 코드를 샘플링하고, 이를 autoencoder가 다시 이미지로 복원함.
이로써 시각적으로 자연스러운 이미지 생성이 가능해짐.
연속적인(latent space)vs 이산적인(latent space)?
1.연속적인 latent 공간
- 숫자가 부드럽게 바뀔 수 있음.
- 예: 0.123, 0.124, 0.125처럼 아주 미세하게 조절 가능
- 수학적으로 **실수(float)**로 표현됨
- SD-VAE 같은 diffusion model에서 이걸 씀
- 장점: 부드럽게 변화시킬 수 있어서 학습이 유연함
예시: 고양이 → 개로 조금씩 바꿀 때, 중간 단계가 자연스럽게 생성됨
2. 이산적인 latent 공간
- 숫자가 딱 정해진 선택지 중 하나만 가짐, 예: [1, 2, 3, 4] 중 하나만 선택 가능 (중간 없음)
- 수학적으로 정수 또는 토큰 같은 discrete한 값으로 표현됨
- VQ-GAN 같은 모델에서 이걸 씀
- 장점: 구조가 단순해서 압축 잘되고 학습이 안정적임
예시: 고양이 사진이 “토큰 132번, 418번, 5번” 같은 식으로 표현됨
그동안 연구들은 주로 generative phase의 개선에 집중해왔음
(e.g., architecture 발전(Peebles & Xie, 2023), objective 함수(Ma et al., 2024), 최적화 기법(Yao et al., 2024)).
하지만 autoencoder가 latent space를 어떻게 구성하느냐 또한 전체 성능에 중요함.
실제로, latent space의 품질은 다음 두 가지 측면에서 핵심적인 역할을 함:
- 계산 효율성 (차원을 줄이고 학습 수렴 속도를 높임)
- 고품질 이미지 생성 가능성
Diffusion models의 경우, SOTA 방식들(DiT, SiT, PixArt, SD3, Flux 등)은 대부분 Latent Diffusion Model (LDM)에서 제안된 SD-VAE의 구조와 목적함수를 따름.
LDM은 크게 두 가지 정규화 전략을 사용함:
- Continuous variational approach
- KL divergence 항을 통해 latent 분포를 Gaussian prior와 정렬시킴
- → latent space를 부드럽고 구조화된 공간으로 유지함 (Kingma & Welling, 2014)
- Discrete codebook framework
- 학습된 임베딩들의 유한한 집합 내에서 latent 표현을 제한함
- → latent 공간의 복잡도를 제한하고 다른 방식으로 regularization을 제공함 (Esser et al., 2021)
이러한 정규화 전략들은 본질적으로 **trade-off(절충)**을 수반함.
예를 들어, KL divergence 항의 가중치를 높이는 강한 정규화는 latent 공간을 더 부드럽고 학습하기 쉬운 형태로 만들어 generative model이 잘 학습할 수 있도록 도와줌 (Tschannen et al., 2025).
하지만 이와 동시에 latent 표현의 **정보 용량(information capacity)**을 감소시켜, 세부적인 정보 손실이 발생하고 최종적으로는 복원 품질이 저하됨.
실험적 증거에 따르면, 이러한 trade-off는 latent generative model의 전체 성능에 **상한선(upper bound)**을 설정하게 됨 (Rombach et al., 2022).
autoencoder가 세밀한 정보를 보존하는 데 제한이 있다면, 결국 latent generative model이 **고해상도 이미지(high-fidelity image)**를 합성하는 능력 자체가 떨어지게 됨.
이러한 상황은 하나의 근본적인 질문을 제기함:
복원 품질을 훼손하지 않으면서도 생성에 더 적합한 latent space를 만들 수 있을까?
나아가 전체 생성 과정 자체를 개선할 수 있을까?
이 문제를 해결할 수 있는 핵심은 latent space 자체의 구조와 성질에 존재함.
특히 현재의 SOTA autoencoder들이 가지는 중요한 한계를 지적함:
이들의 latent representation은 scaling이나 rotation 같은 **기본적인 공간 변환(spatial transformation)**에 대해 equivariant하지 않음 (그림 2, 3.2절 참고).
이러한 결여는 latent manifold에 불필요한 복잡성을 도입하게 되며,
generative model이 원래 학습하지 않아도 될 **비선형적 관계(nonlinear relationship)**까지 학습하게 만듦.
이 문제를 해결하기 위해, 우리는 autoencoder의 훈련 목적 함수에 간단하지만 효과적인 수정을 제안함.
이 수정은 latent space가 앞서 언급한 equivariance 특성을 갖도록 유도함.
우리가 제안하는 방법은 EQ-VAE로, 변형된 latent representation을 복원한 결과와
입력 이미지 자체를 변형한 결과 간의 **불일치(discrepancy)**를 **패널티(penalty)**로 부과함.
중요한 점은, EQ-VAE는 기존 autoencoder 구조에 어떠한 변경도 필요하지 않음.
처음부터 학습할 필요도 없음.
단지 사전학습된 autoencoder를 **몇 에폭(few epochs)**만 fine-tuning하면,
latent space에 equivariance 특성이 주입되고 공간의 복잡도가 줄어듦 (그림 1-왼쪽, 표 5 참고).
이로써 generative model이 더 쉽게 학습할 수 있게 됨 (그림 1-오른쪽 참고).
이 모든 과정은 autoencoder의 복원 품질을 전혀 해치지 않음.
우리의 방법은 continuous와 discrete autoencoder 모두에 호환됨.
덕분에 다양한 latent generative model에 넓게 적용 가능함.
예를 들어, EQ-VAE를 continuous 모델인 **SD-VAE (Rombach et al., 2022)**에 적용하면,
downstream diffusion model인 DiT (Peebles & Xie, 2023), SiT (Ma et al., 2024), REPA (Yu et al., 2025) 등의 성능이 FID 기준으로 향상됨.
또한, EQ-VAE를 discrete 모델인 **VQ-GAN (Esser et al., 2021)**에 적용하면,
masked generative modeling 방식인 **MaskGIT (Chang et al., 2022)**의 성능도 향상됨.
We make the following contributions:
- 기존 autoencoder의 latent space가 공간 변환에 대해 equivariant하지 않음을 발견하고,이를 기반으로, EQ-VAE라는 간단한 정규화 방식을 제안하여 복원 성능 저하 없이 생성 성능을 향상시킴.
- 이것이 latent generative modeling을 방해함을 밝힘.
- 이 방법은 continuous, discrete autoencoder 모두와 호환되며,
- diffusion이나 masked generative model 같은 대표적인 생성 모델에 plug-and-play 방식으로 적용 가능함.
- 우리의 목적 함수를 이용해 사전학습된 autoencoder를 fine-tuning함으로써,예를 들어, SD-VAE를 단 5 에폭 fine-tuning한 것만으로도
- DiT-XL/2에서 7배, REPA(w/ SiT-XL/2)에서 4배 빠른 속도를 기록함 (그림 1 오른쪽 참고).
- latent generative model의 학습 속도를 크게 향상시킬 수 있음.
2. Related Work
Autoencoders for Latent Generative Models
픽셀 공간에서 직접 diffusion model을 학습하는 것은 계산적으로 매우 비효율적임.
왜냐하면 디지털 이미지의 대부분의 비트는 지각적으로 중요하지 않은 미세한 정보를 담고 있기 때문임.
이 문제를 해결하기 위해 Rombach et al. (2022)은 latent diffusion model을 제안함.
이 모델은 autoencoder가 생성한 압축된 latent 공간에서 작동함.
이때 사용된 KL 정규화 autoencoder인 SD-VAE는 이후 다양한 diffusion 모델에서 광범위하게 사용됨 (Yao et al., 2024; Ma et al., 2024; Chen et al., 2024).
후속 연구들은 주로 재구성 오류(reconstruction error)를 최소화하여 생성 성능의 상한을 끌어올리는 데 집중함.
이를 위해 latent channel 수를 늘리거나 (Esser et al., 2024; Black Forest Labs, 2023; Dai et al., 2023),
작업(task) 특화 priors를 도입하는 등의 방법이 사용됨 (Zhu et al., 2023).
또한, Xie et al. (2025)와 Chen et al. (2025)은 고해상도 이미지 학습의 효율을 높이기 위해
재구성 품질을 유지한 채 압축률(compression ratio)을 극적으로 향상시키는 연구를 수행함.
Hu et al. (2023)은 generative model에 최적화된 latent 공간의 조건을 조사했으며,
상대적으로 약한 decoder가 오히려 생성 성능을 향상시키는 latent 분포를 만든다는 결과를 제시함.
한편, discrete autoencoder는 처음에 VQ-VAE(van den Oord et al., 2017)를 통해 도입되었으며,
이미지 패치를 이산적(discrete) 시각 토큰으로 양자화함.
이후 VQ-GAN(Esser et al., 2021)은 adversarial loss와 perceptual loss를 결합하여
더 정확하고 세밀한 표현을 가능하게 만듦.
후속 연구들은 architecture 개선 (Yu et al., 2022a),
codebook의 크기 확장 및 활용 극대화 전략 (Yu et al., 2024; Zhu et al., 2024a) 등에 집중함.
하지만 위와 같은 기존 접근들과 달리, 본 연구는 공간적 equivariance 특성을 활용하여,
generative modeling에 더 적합한 latent 공간을 구성하는 새로운 관점을 탐색함.
Auxiliary Objectives and Regularization in VAEs
autoencoder는 관측된 데이터에서 의미 있는 특징을 압축하여 표현할 수 있는 latent space를 학습하는 데 목적이 있음.
하지만 아무런 정규화 없이 학습하면, latent code는 구조를 갖지 못함.
이를 해결하기 위해, Variational Autoencoder (VAE)가 Kingma & Welling (2014)에 의해 제안됨.
이 방법은 latent 분포와 Gaussian prior 간의 KL divergence를 최소화함으로써 의미 있는 구조를 부여함.
이후 수많은 연구들이 이 프레임워크를 채택하거나 확장함 (Higgins et al., 2016; Dilokthanakul et al., 2016; Tomczak & Welling, 2018; Takahashi et al., 2019).
다른 연구들은 Wasserstein distance 기반 정규화 (Tolstikhin et al., 2018; Kolouri et al., 2018),
adversarial objective (Zhao et al., 2018; Makhzani et al., 2015),
vector quantization(VQ) (van den Oord et al., 2017) 등을 제안함.
우리 연구와 밀접한 Sinha & Dieng (2021)은 spatial transformation에 대해 불변(invariant)한 latent code를 유지하는 consistency regularization을 제안함.
반면, 우리의 EQ-VAE는 invariance가 아니라 equivariance를 유도함.
즉, transformation에 따라 latent도 같은 방식으로 바뀌도록 하는 특성을 정규화함.
그리고 이러한 equivariance regularization이 latent generative modeling에 미치는 영향을 폭넓게 실험을 통해 검증함.
Equivariance in Computer Vision
**CNN(Convolutional Neural Networks)**이 수많은 비전 과제에서 성공한 핵심은,
convolution 연산 특성 덕분에 translation equivariance를 갖기 때문임.
데이터 내 다른 대칭성(symmetry)을 반영하기 위해, 다양한 group-equivariant convolutional network들이 제안됨.
예를 들어, 2D roto-translation equivariance (Cohen & Welling, 2016; Hoogeboom et al., 2018; Weiler & Cesa, 2019),
3D 확장 (Worrall & Brostow, 2018; Thomas et al., 2018; Kondor, 2018),
scale equivariance (Rahman & Yeh, 2023; Sosnovik et al., 2020) 등이 있음.
이러한 group equivariance 제약을 적용하기 위해 steerable filters가 사용되며, 이는 특정 basis로부터 구성됨.
또한, parameter sharing(Ravanbakhsh et al., 2017), frame averaging(Puny et al., 2022), canonicalization functions(Kaba et al., 2023) 등을 통해서도 equivariance 달성이 가능함.
autoencoder 관련 연구로는, **Winter et al. (2022)**가 latent representation을
group invariant와 equivariant 요소로 나누어 생성하는 방식을 제안함.
그러나 이들은 equivariant 표현이 generative modeling에 미치는 영향은 탐구하지 않음.
3. Method
본 섹션에서는 우리의 방법론을 소개함.
우선 latent generative modeling을 위한 autoencoder 모델의 개요를 다루되, 특히 diffusion model에 사용되는 continuous case에 초점을 맞춤 (3.1절).
다음으로, latent representation이 equivariance를 결여하고 있음을 강조하고 (3.2절), 이를 해결하기 위한 EQ-VAE를 제안함 (3.3절).
3.1. Preliminary: Continuous Autoencoders for Latent Generative Modeling
모델링의 첫 단계는 autoencoder를 사용해 **픽셀 공간(pixel space)**을 continuous(Rombach et al., 2022) 혹은 discrete(Esser et al., 2021) latent 공간으로 압축하는 것임.
이 절에서는 continuous한 경우를 다룸.
입력 이미지 \(x \in \mathbb{R}^{H \times W \times 3}\)가 주어졌을 때,
encoder E는 이를 압축된 표현 z = E(x) \(\in \mathbb{R}^{\frac{H}{f} \times \frac{W}{f} \times c}\)로 변환함.
여기서 f는 compression ratio, c는 latent channel 수를 의미함.
그다음 decoder D는 latent 표현 z를 받아서 이미지 \(\hat{x}\) = D(z)로 복원함.
이미지 x에 대한 학습 목적 함수는 다음과 같음:
\(L_{VAE}(x) = L_{rec}(x, \hat{x}) + \lambda_{gan} L_{gan}(\hat{x}) + \lambda_{reg} L_{reg}\)
- L_{rec}: 픽셀 공간 재구성 손실 + perceptual loss (예: LPIPS, Zhang et al., 2018)
- L_{gan}: patch 기반 adversarial loss (Isola et al., 2017)
- L_{reg}: 일반적으로 Gaussian prior를 따르는 KL 정규화 항 (Kingma & Welling, 2014)
3.2. Lack of Equivariance under Spatial Transformations
이 연구는 다음의 핵심적인 관찰에서 출발함:
SOTA autoencoder인 SD-VAE(Rombach et al., 2022) 등은,
scaling, rotation과 같은 기본적인 공간 변환에 대해 equivariant하지 않음.
즉, 이런 변환에 따라 latent representation이 일관되게 변하지 않음.
이를 공식화하면 다음과 같음:
▫ Spatial Transformation
이미지(또는 latent 표현) \(x(p): \mathbb{R}^2 \to \mathbb{R}^c\)는 2D 좌표 p = [u, v]^T 상에 정의됨.
공간 변환 \(\tau \in \mathbb{R}^{2 \times 2}\)는 좌표 p에 작용하며, 이미지를 다음과 같이 변환함:
\(x_{\tau}(p) = x(\tau^{-1} p)\)
전체 p에 대해 간략히 \tau \circ x로 표현함.
▫ Equivariance
latent 표현 E(x)가 transformation \tau에 대해 equivariant하다는 것은, 변환이 latent 공간에도 그대로 적용될 수 있음을 의미함:
\(\forall x \in \mathcal{X}: \quad E(\tau \circ x) = \tau \circ E(x)\)
이러한 equivariance 존재 여부를 실험하기 위해, 우리는 latent code에 직접 scaling과 rotation 변환 \tau를 적용하고 복원 결과를 비교함.
즉, 다음 두 reconstruction을 비교함:
- \(D(E(\tau \circ x))\): 입력 이미지에 변환을 먼저 적용한 후 인코딩한 것
- \(D(\tau \circ E(x))\) : latent 코드에 변환을 직접 적용한 것
Figure 2, 3에서 질적 및 정량적 결과를 제시함.
실험 결과는 명확한 차이를 보임:
입력 이미지에 변환을 적용한 후 인코딩한 경우 D(E(\tau \circ x))는 정확하게 복원되지만,
latent에 직접 변환을 적용한 경우 D(\tau \circ E(x))는 복원 품질이 크게 저하됨.
이러한 한계는 다음 두 가지 원인에서 비롯됨:
- SD-VAE와 같은 latent generative model의 autoencoder는 convolutional architecture를 기반으로 하며,
- 이 구조는 arbitrary한 spatial transformation (e.g., scaling, rotation)에 대해 equivariant하지 않음.
- 학습 시 사용하는 목적 함수 (reconstruction loss, KL divergence 등)는 equivariance를 명시적으로 혹은 암시적으로 유도하지 않음.
이로 인해, 의미적으로 유사한 입력 x와 \tau \circ x가
각각 E(x), E(\tau \circ x)로 인코딩되더라도 이 둘은
E(\tau \circ x) \neq \tau \circ E(x)로 일치하지 않으며,
이는 latent 공간의 구조를 불필요하게 복잡하게 만듦.
3.3. EQ-VAE: Regularization via Equivariance Constraints
이러한 한계를 해결하기 위해, 우리는 EQ-VAE를 제안함.
이 방법은 latent 표현이 공간 변환에 대해 equivariant해지도록 **정규화(regularization)**함.
그림 1 (왼쪽)에서 보듯, 이는 보다 부드럽고 구조화된 latent representation을 만들어,
생성 모델 학습을 더 효율적으로 만듦.
Explicit Regularization
equivariance를 강제하는 가장 직접적인 방법은, 앞서 정의한 식 (3)의 equivariance 조건을
훈련 시 loss 항으로 추가하는 것임:
\(L_{\text{explicit}}(x) = \| \tau \circ E(x) - E(\tau \circ x) \|_2^2\)
여기서 \tau는 정의된 공간 변환 집합에서 샘플링됨.
하지만 이 loss만 최소화하면 trivial solution으로 수렴할 수 있음.
예를 들어, 모든 x에 대해 **E(x) = c (상수)**가 되도록 표현이 붕괴됨.
실제로 우리 실험에서도 이러한 현상이 관찰되며 (표 7 참고),
따라서 explicit regularization만으로는 효과적이지 않음.
3.3. Implicit Regularization
explicit regularization의 한계를 극복하기 위해, 우리는 implicit 방식을 채택함.
Figure 2의 실험 결과에 영감을 받아, 이 방식은 변형된 latent representation의 복원 결과 \(D(\tau \circ E(x))\)가
변형된 입력 이미지 \(\tau \circ x\)의 복원 결과와 정렬되도록 만듦.
구체적으로는, 기존 training objective 함수 (식 1)를 다음과 같이 수정함:
\(\(L_{\text{EQ-VAE}}(x, \tau) = L_{\text{rec}}(\tau \circ x, D(\tau \circ E(x))) + \lambda_{\text{gan}} L_{\text{gan}}(D(\tau \circ E(x))) + \lambda_{\text{reg}} L_{\text{reg}}\)\)
식 (1)과 비교했을 때 달라진 부분을 강조함.
여기서 \(\tau\)가 항등 변환(identity transform)이면, 위 식은 기존 목적 함수 식 (1)과 동일해짐.
이 방식은 reconstruction과 adversarial objective 양쪽에서 나오는 감독 신호를 활용하여,
encoder가 equivariant한 latent 표현을 생성하도록 유도하면서 동시에 mode collapse를 방지함 (A.1절 참고).
Transformation Design
우리는 두 가지 유형의 공간 변환에 집중함:
- anisotropic scaling
- rotation
Isotropic scaling?
isotropic(등방성) scaling은 말 그대로, 가로(x)와 세로(y)를 똑같이 확대/축소
예시:
- 256x256 이미지를 0.8배로 줄이면 → 204x204가 됨
- 이때 x, y 방향 모두 동일하게 줄었으니까 → isotropic scaling
반대로, anisotropic scaling은 x, y를 다르게 줄이거나 늘리는 거
예시:
- x방향만 0.8배, y방향은 그대로 → 204x256
- x, y 비율이 달라짐 → anisotropic
이들은 다음과 같이 파라미터화됨:
\(S(s_x, s_y) = \begin{bmatrix} s_x & 0 \\ 0 & s_y \end{bmatrix}, \quad R(\theta) = \begin{bmatrix} \cos \theta & -\sin \theta \\ \sin \theta & \cos \theta \end{bmatrix}\)
최종 변환은 scaling과 rotation의 조합이며, \(\tau = S(s_x, s_y) \cdot R(\theta)\)로 정의됨.
변환 파라미터는 다음 범위에서 균일하게 샘플링됨:
- 0.25 < s_x, s_y < 1
- \(\theta \in (\frac{\pi}{2}, \pi, \frac{3\pi}{2})\)
90°의 배수만 선택함으로써 corner artifact를 피함.
다운샘플링에는 bicubic interpolation을 사용함.
실험적으로, scaling equivariance가 rotation equivariance보다 생성 성능에 더 유리함이 확인됨 (표 5 참고).
autoencoder의 원래 재구성 성능을 보존하기 위해, 우리는 식 (5)에서 확률 p_\alpha로 **항등 변환 \tau = I**을 적용하여
기존 목적 함수 (식 1)로 돌아가는 방식으로 구성함.
최종 전체 목적 함수는 다음과 같이 정의됨:
\(L_{\text{total}}(x) = \begin{cases} L_{\text{VAE}}(x) & \text{if } p < p_\alpha \\ L_{\text{EQ-VAE}}(x, \tau) & \text{if } p \geq p_\alpha \end{cases}\)
여기서 p는 [0, 1] 구간에서 균일하게 샘플링됨.
이 값은 정규화 강도를 조절함.
기본값으로 p_\alpha = 0.5로 설정함 (A절에서 정규화 강도에 대한 ablation 수행함).
우리가 제안한 방식은 데이터 증강에 의존하는 기존 방법들과 달리,
변환을 입력이 아닌 latent space에 직접 적용함으로써 equivariance를 유도함 (Brehmer et al., 2024).
Extending EQ-VAE to Discrete Autoencoders
지금까지는 continuous autoencoder를 기준으로 EQ-VAE를 설명함.
하지만 discrete autoencoder (예: VQ-GAN, Esser et al., 2021)에도 적용 가능함.
VQ-GAN에서는 encoder가 생성한 연속적 feature E(x)가
학습된 codebook 내 가장 가까운 항목으로 매핑되어 **양자화(quantization)**된 latent 공간을 구성함.
이 경우에도 EQ-VAE는 쉽게 확장 가능함.
3.3절에서 설명한 equivariance regularization loss를 그대로 사용하되,
**양자화 전에 E(x)에 변환 \tau**를 적용하는 방식으로 구현함.
4. Experiments
4.1. Setup
Implementation Details
모든 autoencoder는 LDM (Rombach et al., 2022) 프레임워크에 맞추어 OpenImages 데이터셋에서 fine-tuning함.
모든 autoencoder는 batch size 10, 총 5 에폭 동안 fine-tuning함.
각 autoencoder의 공간 압축 비율, latent 채널 수 등 세부 스펙은 Appendix E에 정리되어 있음.
DiT (Peebles & Xie, 2023), SiT (Ma et al., 2024), REPA (Yu et al., 2025)는
기본 설정을 그대로 따르며 ImageNet (Deng et al., 2009)에서 학습함.
이미지는 256 \times 256으로 리사이즈되며, batch size는 256임.
모델은 대부분 patch size 2를 사용하는 B/2, XL/2 구조를 사용하며, 단 Table 1의 SD-VAE-16 실험에서는 해상도가 낮아 B/1 구조를 사용함.
이 모델들은 SD-VAE-FT-EMA의 latent 분포 상에서 학습됨.
이는 원래 SD-VAE를 **LAION-Aesthetics (Schuhmann et al., 2022)**에서 EMA 기반으로 추가 fine-tuning한 버전임
(Table 6 및 Peebles & Xie, 2023에서 성능 차이 비교 가능).
MaskGIT은 (Besnier & Chen, 2023)의 설정을 따르며, ImageNet에서 300 에폭, batch size 256으로 학습함.
모든 데이터 전처리는 ADM (Dhariwal & Nichol, 2021) 프로토콜을 따름.
Evaluation
생성 성능 측정은 각 autoencoder의 latent 분포에서 latent generative model을 학습한 뒤,
다음 지표를 사용해 평가함:
- FID (Frechet Inception Distance, Heusel et al., 2017)
- sFID (Nash et al., 2021)
- IS (Inception Score, Salimans et al., 2016)
- Precision, Recall (Kynkäänniemi et al., 2019)
→ 모두 50,000개 샘플, ADM 평가 프로토콜에 따라 측정함.
복원 성능 평가에는 다음 지표를 사용함:
- FID
- PSNR (Peak Signal-to-Noise Ratio)
- SSIM (Structural Similarity Index, Wang et al., 2004)
- LPIPS (Perceptual Similarity, Zhang et al., 2018)
→ ImageNet validation set을 기준으로 측정함.
복원과 생성을 구분하기 위해 RFID (Reconstruction FID), **GFID (Generation FID)**로 구분 표기함.
EQ-VAE의 효과를 정량화하기 위해, equivariance error도 측정함 (Appendix C 참고).
4.2. Equivariance-Regularized VAEs
실험 분석의 시작으로, EQ-VAE가 continuous와 discrete autoencoder 모두에 자연스럽게 적용 가능함을 보여주며
**높은 범용성(versatility)**을 입증함.
Continuous Autoencoders
우리는 EQ-VAE 정규화를 다양한 latent 차원을 갖는 기존의 continuous autoencoder에 통합함.
적용한 모델은 SD-VAE, SD-VAE-16 (Rombach et al., 2022), SDXL-VAE (Podell et al., 2024), SD3-VAE (Esser et al., 2024)임.
정규화가 generative 성능에 미치는 영향을 평가하기 위해, 우리는 정규화 전후의 latent 코드에 대해 DiT-B 모델을 학습하고 비교함.
결과는 Table 1에 제시됨.
간단한 목적 함수임에도 불구하고, EQ-VAE는 모든 autoencoder에 대해 equivariance error를 효과적으로 줄임.
또한, 복원 성능(reconstruction fidelity)을 유지하면서도 생성 성능을 일관되게 크게 향상시킴.
실험 결과는 GFID와 equivariance error 감소 간에 상관관계가 존재함을 시사함.
특히 SD-VAE, SDXL-VAE, SD-VAE-16의 경우, EQ-VAE 정규화가 생성 성능을 크게 향상시킴.
SD3-VAE의 경우에는 equivariance error 감소는 비교적 작았지만, GFID는 여전히 개선됨.
Discrete Autoencoders
EQ-VAE가 discrete autoencoder에도 적용 가능한지 확인하기 위해,
우리는 **VQ-GAN (Esser et al., 2021)**에 대해 실험하고,
**masked image modeling framework인 MaskGIT (Chang et al., 2022)**에서의 효과를 검증함.
Table 1에서 보이듯, EQ-VAE는 discrete 환경에서도 효과적이며,
equivariance error를 줄이고 GFID를 6.8에서 5.9로 개선함.
4.3. Boosting Generative Image Models
EQ-VAE를 continuous 및 discrete autoencoder에 모두 적용함으로써, 다양한 SOTA generative models의 성능을 향상시킴.
포함된 모델은 다음과 같음:
- DiT: 순수 Transformer 기반 diffusion model
- SiT: continuous flow 기반 모델
- REPA: Transformer representation을 self-supervised feature와 정렬시키는 최신 방식
- MaskGIT: 널리 쓰이는 masked generative model
DiT & SiT
Table 2에서 보이듯, EQ-VAE 정규화는 DiT-B, DiT-XL 모델 모두에 걸쳐 의미 있는 성능 향상을 보여줌.
예를 들어, DiT-XL/2를 정규화된 latent 분포에서 학습한 경우,
400K iteration에서 GFID가 19.5 → 14.5로 향상됨.
1.5M iteration에서는 GFID 8.8까지 도달하며,
이는 SD-VAE-FT-EMA 기반 DiT-XL/2 모델의 7M iteration 성능을 초과함.
EQ-VAE가 제공하는 학습 속도 향상 효과는 Figure 4에서 직관적으로 확인 가능함.
또한 Table 2에서는, SiT 모델 역시 EQ-VAE의 정규화 latent 분포로부터 GFID 17.2 → 16.1로 개선됨을 보임 (400K step 기준).
REPA
EQ-VAE 정규화는 latent generative modeling의 첫 단계에서 적용되며,
이는 REPA 방식과 보완적 관계를 가짐.
그 결과 수렴 속도와 생성 성능 모두 개선됨.
예를 들어, **REPA (SiT-XL-2)**를 EQ-VAE로 학습하면
GFID 5.9에 도달하는 데 필요한 iteration 수가 4M → 1M으로 4배 단축됨.
Figure 1에서 이 속도 차이를 시각적으로 확인할 수 있음.
이는 의미 있는 결과임. 왜냐하면, REPA 자체도 원래 diffusion model의 학습 속도를 상당히 가속시키는 방식이기 때문임.
MaskGIT
Table 3에 따르면, EQ-VAE로 학습한 MaskGIT은
130 에폭만에 GFID 6.80에 도달하며, 기존 방식이 300 에폭 소요되던 것에 비해 2배 빠르게 수렴함.
뿐만 아니라, 300 에폭 기준으로 GFID 5.91에 도달하며,
이는 (Besnier & Chen, 2023), (Chang et al., 2022) 양쪽 결과를 모두 초과함.
SOTA Generative Model과의 비교
EQ-VAE가 학습 과정을 어떻게 가속시키는지를 추가적으로 보여주기 위해,
우리는 classifier-free guidance 기반 diffusion 방식들과 EQ-VAE를 비교함.
Table 4에서 보이듯,
EQ-VAE를 적용한 DiT-XL/2는 단 300 에폭만에 GFID 2.37에 도달,
이는 SD-VAE 또는 SD-VAE-FT-MAE 기반 DiT-XL/2와 동일한 성능임.
또한 EQ-VAE를 SOTA 방식인 REPA와 결합했을 때, REPA 표준 방식과 비슷한 성능을 4배 적은 연산량으로 달성함 (200 vs 800 에폭 기준).
4.4. Analysis
Spatial Transformations Ablation
우리는 EQ-VAE의 equivariance 정규화가 각 공간 변환별로 생성 성능에 미치는 영향을 ablation 실험을 통해 분석함.
적용한 변환은 다음과 같음:
- isotropic scaling: S(s, s)
- anisotropic scaling: S(s_x, s_y)
- rotation: R(\theta)
- 결합 변환(combined): scaling + rotation
각 latent 분포에 대해 DiT-B/2 모델을 학습하고 결과를 Table 5에 제시함.
분석 결과:
- scale equivariance를 유도하면 생성 성능에 가장 큰 영향을 줌
- rotation equivariance도 도움이 됨
- 두 변환을 결합하면 성능이 더욱 향상됨, 이는 두 효과가 상호 보완적임을 보여줌
특히, anisotropic scaling은 보다 공격적인 정규화로 인해 생성 성능은 향상되지만,
복원 품질은 저하됨.
따라서 EQ-VAE의 기본 설정은 isotropic scaling과 rotation의 결합을 사용함.
Latent Space Complexity and Generative Performance
EQ-VAE 정규화가 latent manifold의 **복잡성(complexity)**에 미치는 영향을 더 잘 이해하기 위해,
우리는 **Intrinsic Dimension (ID)**을 측정함.
ID란, 어떤 데이터 분포를 설명하는 데 필요한 최소 변수 개수를 의미함 (Bennett, 1969).
Table 5에서는 latent manifold의 intrinsic dimension과 생성 성능(GFID) 간의 상관관계를 확인할 수 있음.
이는 정규화된 latent 분포가 모델링하기 더 간단해졌음을 시사하며,
우리 방식의 효과성을 다시 한 번 뒷받침함.
latent 표현의 복잡성 감소는 **Figure 1 (왼쪽)**에서도 직관적으로 확인 가능함.
ID에 대한 자세한 내용은 Appendix B를 참고하면 됨.
향상된 성능은 단순한 학습 반복 때문이 아님
생성 성능의 향상이 단순히 추가 학습(추가 에폭) 때문이 아니라, 실제로 equivariance 정규화 (식 5)의 효과인지를 검증하기 위해
우리는 EQ-VAE와 **SD-VAE†**을 비교함 (Table 6).
여기서 SD-VAE†는 기존 SD-VAE를 단지 식 (1)의 원래 목적 함수로 5 에폭 추가 학습한 버전임.
결과적으로, SD-VAE†는 생성 성능(GFID)에 거의 영향을 주지 않음.
반면, EQ-VAE는 명확한 성능 향상을 가져옴.
유사하게, SD-VAE에서 파생된 SD-VAE-EMA-FT 또한 GFID 점수에 미미한 영향만을 미침.
이는 EQ-VAE의 정규화가 가진 실질적인 효과를 더욱 부각시킴.
5. Conclusion
본 연구는 autoencoder가 만들어내는 latent 표현의 구조가,
latent generative model의 수렴 속도와 성능에 결정적인 역할을 한다는 점을 강조함.
우리는 기존 autoencoder의 latent representation이 기본적인 공간 변환 (scaling, rotation 등)에 대해 equivariant하지 않음을 관찰함.
이에 대응하기 위해, 우리는 autoencoder의 훈련 목적 함수에 대한 간단한 수정인 EQ-VAE를 제안함.
실험을 통해, 사전학습된 autoencoder를 단 몇 에폭만 fine-tuning하더라도
equivariance error를 줄이고,
복원 성능을 유지하면서도 latent generative model의 성능을 유의미하게 향상시킬 수 있음을 보여줌.
또한, 본 연구는 latent 분포의 **기하학적 구조(geometry)**와 모델 성능 간의 이론적/실증적 관계에 대한
향후 연구 방향의 가능성을 제시함.
Impact Statement
이 논문은 기계학습 전반, 특히 이미지 생성 분야의 발전을 목표로 하는 연구임.
이 연구가 사회에 미칠 수 있는 다양한 영향이 있을 수 있으나, 이 중 특정 항목을 별도로 강조할 필요는 없다고 판단함.
Acknowledgements
본 연구는 그리스 회복 및 회복탄력성 계획 Greece 2.0의 프로젝트 MIS 5154714에 따라,
유럽연합의 NextGenerationEU 프로그램으로부터 지원을 받음.
또한 GRNET의 지원을 통해 하드웨어 리소스를 제공받았으며,
**GENCI-IDRIS의 HPC 리소스(Grants 2024-AD011012884R3)**를 활용하여 수행되었음.
EQ-VAE에서 “regularization”이란?
regularization이란 쉽게 말하면:
“원래의 reconstruction 목표(이미지 복원)” +
- “latent가 spatial 변환(회전, 스케일)에도 일관성 있게 변하게 만들자”
- 라는 추가적인 제약(규칙)을 거는 것
즉, latent 공간에서 “transform 잘 따라가도록” 추가로 학습시키는 과정
1. 원래 reconstruction loss
VAE 기본 학습 :
reconstruction_loss = MSE(x, D(E(x)))
- 입력 이미지 x
- 인코딩 → 디코딩 → 복원 이미지 D(E(x))
- 이 둘의 차이를 MSE로 계산
→ 이건 기본 VAE 학습 과정
2. 추가된 EQ regularization
EQ-VAE에서는 이렇게 추가 과정을 넣음:
# 입력 이미지를 spatial transform 해
x_t = random_transform(x) # (회전, 스케일링 등)
# 변형된 이미지도 encode
z_gt = E(x_t)
# 원래 latent를 변형한 결과
z_pred = random_transform(E(x))
# 두 latent가 비슷해지도록 loss 추가
eq_loss = MSE(z_pred, z_gt)즉,
“이미지를 변환한 후 인코딩한 것” (z_gt)
“latent를 변환한 것” (z_pred)
둘이 비슷해야 한다! 라고 학습시키는 거야.
최종 학습 목표?
전체 loss를 이렇게 바꿈:
total_loss = reconstruction_loss + λ * eq_loss- reconstruction_loss: 원래 VAE가 하던 “이미지 복원” 목표
- eq_loss: 새로 추가한 “latent가 transform 따라야 함” 목표
- λ(lambda_eq): 둘의 중요도 조절하는 가중치
| reconstruction loss | 이미지 복원 제대로 하게 만드는 것 |
| EQ regularization loss | latent가 spatial transform에도 자연스럽게 따라가게 만드는 것 |
EQ regularization은…
- VAE 학습 중에
- latent z를 더 좋은 성질(= transform equivariance)을 가지게 만들어주는 추가 학습 규칙임
- diffusion model은 이 과정을 거친 좋은 z를 쓰게 되니까 학습이 더 빠르고 안정적이게 됨
실제 EQ-VAE 깃허브
train_eqvae.py :
outputs = model(batch['pixel_values'])
loss, rec_loss, eq_loss = model.loss_function(batch['pixel_values'], outputs, batch['transformed_pixel_values'])
# loss = rec_loss + lambda_eq * eq_loss
- rec_loss: reconstruction
- eq_loss: equivariance
- 합쳐서 total_loss 만듦
EQ-VAE는 어떤 이미지 데이터셋을 썼나?
EQ-VAE 논문과 깃허브를 보면, OpenImages 데이터셋을 사용했어.
| 사용한 데이터셋 | OpenImages (Large-Scale 데이터셋, 9M장 이상의 라벨링 이미지) |
| 데이터 목적 | VAE를 학습해서 latent space equivariance 성질을 만들기 위해 사용 |
| 경로 설정 | configs/eqvae_config.yaml 파일에서 train_dir, val_dir, dataset_name 등을 OpenImages 경로로 설정해야 함 |
| 준비 방법 | 직접 OpenImages를 다운로드 받아서 폴더 구조에 맞게 정리해야 했음 (자동 다운로드 스크립트 제공 안 함) |
깃허브 README.md에 이렇게 써 있어:
“Currently, we provide experiments for OpenImages.
After downloading modify paths of train_dir, val_dir, dataset_name in the config file.”
(= OpenImages 데이터만 지원하고, 직접 다운로드 받아서 configs/eqvae_config.yaml에 경로를 수정해야 한다는 뜻)
추가 정보
- OpenImages는 기본적으로 ImageNet보다 훨씬 크고, 더 다양한 real-world 이미지가 포함되어 있어.
- EQ-VAE는 이걸 이용해서 “Affine 변환(e.g., scale, rotate)“을 줘도 latent 공간이 자연스럽게 변화하도록 훈련한 거야.
- OpenImages는 다운로드가 엄청 크고 오래 걸려. (수십 GB 이상)
- 그래서 실제 실험할 때는 조그만 subset을 만들어서 사용하는 경우도 많아 (예: 100K장만 샘플링).
EQ-VAE는 OpenImages를 사용해서 VAE를 학습했고, 학습할 때 Affine 변환을 주면서 latent space equivariance를 유도했음.
필요하면,
- eqvae_config.yaml 안에 데이터셋 세팅 어떻게 되어있는지도 바로 열어서 자세히 분석해줄게.
- 또는 내가 OpenImages 없이 작은 dummy dataset으로 돌릴 수 있는 방법도 같이 짜줄게.
EQ-VAE의 Input과 Output
| Input | 원본 이미지 (OpenImages 같은 real-world 이미지) |
| Output | (1) 원본 이미지를 복원한 reconstruction 이미지 (2) Latent 공간의 distribution (posterior) |
즉, VAE처럼:
- 이미지를 encode해서 latent vector z를 만들고,
- 그 z를 다시 decode해서 원래 이미지를 복원하는 걸 목표로 학습
EQ-VAE 학습 과정에서 뭐가 다른가?
일반 VAE랑 다르게, EQ-VAE는 latent space가
- 이미지의 Affine 변환(scale, rotation 등)에 대해
- 일관성 있게 변해야 한다는 걸 추가로 학습
- 이미지를 회전/스케일 조정해서 넣으면, 그에 맞춰 latent vector도 “일정한 방법”으로 변해야
그래서 loss 함수에 equivariance regularization을 추가 (기존 reconstruction loss + KL divergence + EQ regularization)
EQ-VAE 학습이 잘 됐는지 어떻게 판단해?
EQ-VAE 자체만 보고 학습 평가하는 방법은 두 가지야:
- Reconstruction Loss 감소
- input 이미지를 잘 복원하는지 본다.
- 예를 들면 L2 loss, perceptual loss 등을 모니터링.
- Equivariance Regularization Loss 감소
- 변형된 input을 encode 했을 때,
- latent space 안에서도 예상 가능한 변형이 일어나는지를 본다
(즉, latent space가 스케일, 로테이션 같은 변형을 제대로 따라가는지를 loss로 측정.)
꼭 diffusion에 붙여야 성능을 알 수 있을까?
아니야. 꼭 diffusion까지 갈 필요 없어.
- EQ-VAE만 학습시켜서 reconstruction loss + EQ loss가 충분히 낮아지면,
- “latent representation이 잘 정돈되었다”고 볼 수 있어.
- Diffusion 모델은 “EQ-VAE를 쓰면 downstream task에서 더 빠르게 학습이 된다”를 보여주기 위한 추가적인 실험일 뿐이야.
즉, EQ-VAE는 독립적으로 학습 완료하고 성능 평가할 수 있어.
다만,
진짜 EQ-VAE가 diffusion 같은 downstream model에서 학습을 더 빠르게 한다는 걸 입증하려면, diffusion에 붙여서 비교 실험하는 게 필요해.
항목정리
| Input | 원본 이미지 |
| Output | 복원 이미지 + latent posterior |
| 학습 성공 여부 | (1) Reconstruction Loss 감소 (2) EQ Regularization Loss 감소 |
| 꼭 diffusion 붙여야 하나? | ❌ 아니야, EQ-VAE 자체로도 학습 상태 판단 가능함.✔️ 다만 diffusion 붙이면 “EQ-VAE 덕분에 학습이 빨라진다”를 입증할 수 있음. |
Prior preservation은 왜 필요한가?
이건 특히 diffusion이나 VAE에서 자주 쓰이는 모델 안정화 기법
역할:
- EQ loss만 학습하게 되면, 모델은 “변환된 상황만” 보게 되고 “원래 모습”을 잘 재구성 못할 수 있음.
- 그래서 일정 비율(1 - p_prior)로 그냥 원본 이미지를 reconstruction하는 task도 같이 섞어주는 것.
장점:
- 균형 잡힌 학습이 가능: transformation-invariant 표현도 잘 학습하면서, 일반적 재구성 능력도 유지함.
- overfitting 방지, training 안정성 증가