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

🔼 이 그림은 잠재 확산 모델의 재구성-생성 경계면을 보여줍니다. 고차원 잠재 공간에서 VA-VAE(Vision Foundation model Aligned Variational AutoEncoder)는 특징 분포를 개선합니다. 비전 기반 모델과의 정렬을 통해 재구성과 생성 간의 경계면을 확장하여, 고차원 잠재 공간에서도 양쪽 성능을 동시에 향상시킬 수 있음을 시각적으로 보여줍니다. 즉, 기존 모델들은 재구성 성능을 높이면 생성 성능이 떨어지는 trade-off를 보였으나, VA-VAE는 이러한 제약을 극복하고 양쪽 성능을 모두 향상시킬 수 있음을 의미합니다.

read the caption

Figure 2: Reconstruction-generation frontier of latent diffusion models. VA-VAE improves the feature distribution of high-dimensional latent. Through alignment with vision foundation models, we expand the frontier between reconstruction and generation in latent diffusion models.

🔼 그림 3은 제안된 비전 기반 모델 정렬 VAE(VA-VAE)의 개념도를 보여줍니다. 고차원 시각 토큰화기를 학습시킬 때, 미리 학습된 비전 기반 모델(예: DINOv2, MAE)을 활용하여 잠재 공간을 제어함으로써 재구성과 생성 성능 간의 최적화 딜레마를 완화하고 생성 성능을 향상시키는 방법을 나타냅니다. 고차원 시각 토큰화기는 높은 해상도 이미지를 효과적으로 압축하고 재구성할 수 있지만, 학습 과정에서 정보 손실이 발생하여 생성 성능이 저하되는 문제점이 있습니다. VA-VAE는 사전 학습된 비전 기반 모델의 지식을 활용하여, 고차원 잠재 공간을 효과적으로 학습하고 생성 성능을 향상시켜 이러한 문제를 해결합니다. 이 그림은 VA-VAE의 아키텍처와 비전 기반 모델을 활용한 학습 과정을 시각적으로 보여주며, 각 구성 요소의 역할과 상호 작용을 명확히 설명합니다.

read the caption

Figure 3: The proposed Vision foundation model Aligned VAE (VA-VAE). Vision foundation models are used to guide the training of high-dimensional visual tokenizers, effectively mitigating the optimization dilemma and improve generation performance.

🔼 그림 4(a)는 서로 다른 토크나이저(f16d32, f16d64)를 사용하여 LightningDiT-B 모델을 ImageNet 256 해상도에서 160 에폭 동안 학습시킨 결과를 보여줍니다. VF 손실(Vision Foundation model alignment Loss)을 적용했을 때, 적용하지 않았을 때에 비해 FID(Fréchet Inception Distance) 값이 훨씬 빠르게 감소하는 것을 확인할 수 있습니다. 이는 VF 손실이 수렴 속도를 최대 2.7배까지 향상시킨다는 것을 의미합니다. f16d32 토크나이저의 경우 VF loss (DINOv2)를 사용했을 때 FID가 약 2.54배 빠르게 감소했고, f16d64 토크나이저의 경우 약 2.76배 빠르게 감소했습니다.

read the caption

(a)

🔼 그림 (b)는 다양한 토크나이저(f16d32, f16d64)를 사용하여 LightningDiT-B 모델을 160 에폭 동안 ImageNet 256 해상도에서 학습시킨 결과를 보여줍니다. VF 손실(VF Loss)을 적용했을 때, VF Loss를 적용하지 않았을 때보다 훨씬 빠르게 수렴하는 것을 보여줍니다. f16d64 토크나이저의 경우 VF Loss 적용 시 최대 2.7배의 속도 향상을 보였습니다. 이는 고차원 토크나이저를 사용하는 경우 VF Loss가 수렴 속도를 크게 개선한다는 것을 의미합니다.

read the caption

(b)

🔼 그림 (c)는 다양한 토크나이저 크기(차원)에서 VF Loss(Vision Foundation model alignment Loss)의 확장성을 보여줍니다. x축은 DiT(Diffusion Transformer) 모델의 크기(십억 매개변수)를 로그 스케일로 나타내고, y축은 FID(Fréchet Inception Distance) 점수를 나타냅니다. 여러 토크나이저 크기(f16d16, f16d32, f16d64)에 대해, VF Loss를 사용한 경우와 사용하지 않은 경우의 FID 점수를 비교하여, VF Loss가 고차원 토크나이저에서도 더 작은 모델 크기로 높은 성능을 달성할 수 있음을 시각적으로 보여줍니다. 즉, VF Loss를 통해 고차원 토크나이저의 확장성 문제를 효과적으로 해결함으로써, 더 적은 매개변수로도 우수한 이미지 생성 성능을 얻을 수 있음을 보여줍니다. 특히, 모델 크기가 커짐에 따라 VF Loss의 효과가 더욱 두드러지는 것을 확인할 수 있습니다.

read the caption

(c)

🔼 그림 4는 VF Loss(Vision Foundation model alignment Loss)가 Latent Diffusion Model의 학습 속도를 향상시키고 모델의 확장성을 개선하는 효과를 보여줍니다. (a)와 (b)는 서로 다른 토크나이저를 사용하여 ImageNet 256 해상도에서 LightningDiT-B 모델을 160 에폭 동안 학습시킨 결과를 보여줍니다. VF Loss를 적용한 경우, 최대 2.7배까지 학습 속도가 향상되는 것을 확인할 수 있습니다. (c)는 고차원 토크나이저를 사용하는 생성 모델에서 VF Loss가 필요한 파라미터의 수를 줄여 확장성을 개선하는 효과를 보여줍니다.

read the caption

Figure 4: (a)&(b) VF Loss Improves Convergence. We train LightningDiT-B for 160 epochs on ImageNet at 256 resolution using different tokenizers. The VF loss significantly accelerates convergence, with a maximum speedup of up to 2.7 times. (c) VF Loss Improves Scalability. VF loss reduces the need for large parameters in generative models of high-dimensional tokenizer, enabling better scalability.

🔼 이 그림은 논문에서 제안된 VA-VAE와 LightningDiT-XL 모델을 사용하여 ImageNet 256x256 해상도의 이미지를 생성한 결과를 보여줍니다. 다양한 종류의 이미지들이 생성되었으며, 모델의 이미지 생성 능력을 시각적으로 보여주는 역할을 합니다. 각 이미지는 모델이 얼마나 다양하고 사실적인 이미지를 생성할 수 있는지 보여주는 예시입니다.

read the caption

Figure 5: Visualization Results. We visualize our latent diffusion system with proposed VA-VAE together with LightningDiT-XL trained on ImageNet 256×256256256256\times 256256 × 256 resolution.

🔼 표 2는 VF 손실이 생성 성능을 향상시키는 방법을 보여줍니다. f16d16 토크나이저 사양은 널리 사용됩니다 [33, 21]. 차원이 증가함에 따라 재구성 성능은 향상되지만 생성 품질은 저하되는 것을 알 수 있습니다. 이는 잠재 확산 프레임워크 내의 최적화 딜레마를 강조합니다. VF 손실은 재구성 성능에 거의 영향을 미치지 않으면서 고차원 토크나이저의 생성 성능을 크게 향상시킵니다.

read the caption

Table 2: VF loss Improves Generation Performance. The f16d16 tokenizer specification is widely used [33, 21]. As dimensionality increases, we observe that (1) higher dimensions improve reconstruction but reduce generation quality, highlighting an optimization dilemma within the latent diffusion framework; (2) VF Loss significantly enhances generative performance in high-dimensional tokenizers with minimal impact on reconstruction.

🔼 표 3은 ImageNet 256x256 데이터셋에서 다양한 방법들을 사용하여 측정된 시스템 성능을 비교한 표입니다. 본 논문에서 제안하는 잠재 확산 모델(LightningDiT + VA-VAE)은 rFID 0.28, FID 1.35라는 최첨단 성능을 달성했습니다. 특히, 기존의 DiT [29] 와 SiT [26] 모델과 비교했을 때, 단 64번의 학습 에폭만으로도 FID 지표에서 더 나은 성능을 보였으며, 이는 기존 모델보다 21.8배 빠른 수렴 속도를 의미합니다. 표에는 재구성 성능(rFID, PSNR, LPIPS, SSIM)과 생성 성능(gFID, IS, Pre, Rec)을 포함하여 다양한 지표가 제시되어 있습니다. CFG(Classifier-Free Guidance)를 사용한 경우와 사용하지 않은 경우의 성능 차이도 비교 분석되어 있습니다. 각 모델의 매개변수 수(params), 학습에 사용된 에폭 수(Epoches)와 매개변수 수(params)도 함께 제시하여 모델의 복잡도 및 학습 효율성을 비교 분석할 수 있도록 돕고 있습니다.

read the caption

Table 3: System-Level Performance on ImageNet 256×\times×256. Our latent diffusion system achieves state-of-the-art performance with rFID=0.28 and FID=1.35. Besides, our LightningDiT together with VA-VAE surpasses DiT [29] and SiT [26] in FID within only 64 training epochs, demonstrating a 21.8 ×\times× faster convergence.

🔼 이 표는 다양한 Vision Foundation Model(VFM)을 사용하여 VF Loss(Vision Foundation model alignment Loss)의 생성 성능에 미치는 영향을 평가한 결과를 보여줍니다. 실험은 f16d32 토크나이저를 50 에폭 동안 학습시키고, 각각의 VFM을 사용하여 LightningDiT-B를 160 에폭 동안 학습시켜 수행되었습니다. 표에서 확인할 수 있듯이, DINOv2가 가장 높은 생성 성능(가장 낮은 gFID)을 달성했습니다. 이는 DINOv2가 다른 VFM에 비해 더 나은 latent space 정렬을 제공함을 시사합니다.

read the caption

Table 4: Ablation on Foundation Models. We evaluate the impact of different VF losses on generative performance. Our results show that DINOv2 achieves the highest generative performance.

🔼 이 표는 LightningDiT-B 모델을 사용하여 VF 손실 공식의 여러 구성에 대한 생성 성능 지표를 비교 분석한 결과를 보여줍니다. 구체적으로, 마진 코사인 유사도 손실(mcos), 마진 거리 행렬 유사도 손실(mdistmat), 마진 값을 제거한 경우의 재구성 및 생성 성능(rFID, PSNR, LPIPS, SSIM, gFID) 변화를 보여주어 각 구성 요소의 효과를 분석합니다. LightningDiT-B 모델은 본 논문에서 제시된 향상된 DiT 기본 모델입니다.

read the caption

Table 5: Ablation Study of VF Loss Formulations: Comparison of different configurations on generative performance metrics using LightningDiT-B.

🔼 표 6은 특징 분포의 균일성을 평가하고, 그 결과를 바탕으로 특징 분포의 균일성과 생성 성능 간의 상관관계를 분석한 내용을 보여줍니다. 단순히 캡션에 제시된 내용보다 자세히 설명하면, 다양한 토크나이저(f16d32, f16d64)와 VF Loss 적용 여부에 따른 잠재 공간의 분포 균일성을 t-SNE 기법을 통해 시각화하고, KDE(Kernel Density Estimation)을 이용하여 균일성 지표(표준편차, 지니 계수)를 계산했습니다. 이를 통해 생성 성능 지표인 gFID와의 관계를 분석하여, 잠재 공간의 분포가 균일할수록 생성 성능이 향상될 가능성이 있음을 제시합니다.

read the caption

Table 6: Relationship between uniformity and generative performance: We evaluate the uniformity of feature distribution. Results indicate a possible positive correlation between the uniformity of feature distribution and generative performance.