기술 요약: Twins: Focal Loss를 이용한 통합 표현 예측 학습

1. 문제 정의

통합 멀티모달 모델은 단일 모델과 공유된 표현 공간을 활용하여 멀티모달 이해와 이미지 생성을 모두 지원하는 것을 목표로 합니다. 기존의 접근 방식들은 이해(Understanding), 재구성(Reconstruction), 생성(Generation) 사이의 소위 "불가능한 삼각형"이라 불리는 지속적인 트레이드오프에 직면해 있습니다:

• 이산적 방법(Discrete Methods): 공유 코드북을 통해 인터페이스를 통일하지만, 재구성 충실도가 떨어지거나 복잡한 토큰화 과정이 필요한 경우가 많습니다.
• 연속적 방법(Continuous Methods): 일반적으로 두 개의 이질적인 표현에 의존합니다. 즉, 이해를 위한 고수준 시맨틱 특징(예: ViT)과 합성을 위한 저수준 잠재 변수(예: VAE)입니다. 이러한 분리는 서로 일치하지 않는 잠재 공간을 생성하며, 시스템이 자신의 생성물을 "이해"하기 위해 추가적인 디코드-인코드 과정을 거치게 만들어 표현의 일관성을 깨뜨립니다.
• 최근의 시도들: UniFlow나 UniLip 같은 방법들은 고차원 CLIP 특징을 압축하여 표현을 통일하려 하지만, 이러한 차원 축소는 이해 능력을 저하시킵니다. 반대로, 고차원 시맨틱 임베딩(예: DINOv2)을 생성하는 RAE와 같은 방식은 고주파 세부 사항을 재구성하는 데 실패하여 이미지 품질이 떨어지는 결과를 초래합니다.

핵심 과제는 계산 오버헤드(예: 어텐션 비용)를 증가시키지 않으면서도 강력한 시맨틱 이해, 고충실도 재구성, 그리고 생성 친화적인 예측력을 동시에 지원하는 단일 연속 표현을 얻는 것입니다.

2. 방법론

2.1. Twins 표현

저자들은 동일한 토큰 그리드 상에서 두 개의 사전 학습된 인코더로부터 추출한 특징을 **채널 단위로 결합(channel-wise concatenating)**하여 형성된 통합 연속 토큰 공간인 Twins를 제안합니다:

1. ViT 구성 요소: 판별적 시맨틱 추론에 최적화된 SigLIP2로부터 추출된 시맨틱 풍부 특징.
2. VAE 구성 요소: 픽셀 수준의 재구성 충실도에 최적화된 Flux.2 VAE로부터 추출된 세부 사항 보존 잠재 변수.

특징을 시퀀스 차원이 아닌 채널 차원을 따라 결합함으로써, Twins는 원래의 시퀀스 길이를 유지합니다. 이는 토큰 수에 따라 이차 함수적 어텐션 비용($O(L^2)$)이 증가하지 않도록 보장하여 계산 효율성을 보존합니다.

2.2. 최적화 불균형 (Optimization Imbalance)

Diffusion Transformer (DiT)가 통합된 Twins 표현을 예측하도록 학습할 때, 저자들은 심각한 최적화 불균형을 관찰했습니다. 모델은 ViT (SigLIP) 구성 요소는 잘 학습하지만, VAE 잠재 분포를 맞추는 데 어려움을 겪어 이미지가 흐려지고 FID 점수가 낮아지는 현상이 발생합니다.

논문은 이 불균형의 원인을 세 가지 이질성으로 설명합니다:

1. 스펙트럼 편향 (Spectral Bias): ViT 특징은 저주파 신호가 지배적인 반면, VAE 특징은 상당한 고주파 에너지(질감, 노이즈)를 포함합니다. 신경망은 자연스럽게 저주파 함수를 우선적으로 학습하는 경향(Spectral Bias 이론)이 있어, DiT가 ViT 특징을 먼저 학습하고 VAE의 세부 사항을 학습하기 어려운 노이즈로 취급하게 됩니다.
2. 고유 차원 (Intrinsic Dimensionality, ID): SigLIP은 물리적 차원(768)이 VAE(128)보다 높지만, 고유 차원은 훨씬 낮습니다 (~15 vs ~35). 저차원 ID를 가진 SigLIP 매니폴드는 학습하기 쉬운 반면, 고차원 ID를 가진 VAE 매니폴드는 더 복잡한 최적화 지형을 제시합니다.
3. 조건 의존성 (Conditional Dependency): SigLIP 특징은 구조화되어 있고 조건과 정렬되어 있는 반면(클래스 레이블이 주어지면 결정론적임), VAE 특징은 조건 독립적인 불확실성을 상당 부분 보유하고 있습니다. 모델은 자연스럽게 예측 가능하고 조건에 정렬된 목적 함수를 우선시합니다.

2.3. Flow Matching을 위한 Focal Loss

이 불균형을 해결하기 위해, 저자들은 Focal Loss 전략을 Flow Matching 목적 함수에 적응시켰습니다. 모든 차원을 동일하게 취급하는 표준 MSE 대신, VAE 채널에 대한 재가중치 스케마를 도입합니다:

• 손실 함수는 VAE 차원의 "어려운" 잔차(큰 오차)에 더 높은 페널티를 부여합니다.
• 가중치 인자는 $w_i = |v_i - v_\theta(z, t)_i|^{2\gamma}$로 정의되며, 여기서 $\gamma$는 0.5로 설정됩니다.
• 이를 통해 모델이 학습하기 어려운 고주파 VAE 특징에 집중하도록 강제하여, 최적화가 쉬운 ViT 특징에 의해 국소 최적점(local optimum)에 정체되는 것을 방지합니다.

3. 주요 기여

• 통합 표현 (Twins): 시맨틱이 풍부한 ViT 특징과 세부 사항을 보존하는 VAE 특징을 채널 단위로 단순하고 효율적으로 결합하여, 시퀀스 길이나 어텐션 비용의 증가 없이 이해와 생성을 모두 지원합니다.
• 불균형 분석: 통합 모델링에서의 최적화 불균형을 체계적으로 식별하고, 이를 스펙트럼 편향, 고유 차원, 조건 의존성으로 규명했습니다.
• Focal Loss 적응: 어려운 VAE 차원의 가중치를 높임으로써 어려운 특징을 효과적으로 학습하게 하는 Focal 재가중치 전략을 Flow Matching에 적용했습니다.
• 성능 향상: 이 접근 방식이 단순한 MSE 손실 대비 생성 성능을 실질적으로 개선하는 동시에 경쟁력 있는 멀티모달 이해 성능을 유지함을 입증했습니다.

4. 실험 결과

4.1. 재구성 (Reconstruction)

ImageNet-1K 검증 세트에서, Twins는 통합 토크나이저 중 최첨단(SOTA) 재구성 지표를 달성했습니다:

• PSNR: 31.46
• SSIM: 0.90
• rFID: 0.11
  Twins는 RAE (PSNR 18.83)를 크게 앞서며, Wan2.2 및 SD-VAE 3와 같은 전용 생성 오토인코더의 재구성 품질과 대등한 수준을 보여줌으로써, 통합 표현이 픽셀 단위의 일관성을 유지함을 증명했습니다.

4.2. 멀티모달 이해 (Multimodal Understanding)

VLM 파이프라인(Qwen2.5-7B 사용)에 통합되었을 때, Twins는 특화된 인코더들과 경쟁력 있는 성능을 보였습니다:

• 강력한 SigLIP2 베이스라인보다 일반적으로 우수한 성능을 보입니다.
• 저수준 VAE 특징의 포함은 GQA (64.93 vs 64.54) 및 TQA (58.89 vs 56.92)와 같은 미세한 작업에서 성능 향상을 가져왔으며, 이는 시맨틱 전용 인코더에 의해 추상화되어 사라지기 쉬운 질감 및 정확한 형태의 세부 사항이 보존되었기 때문입니다.

4.3. 이미지 생성 (Image Generation)

ImageNet@256 및 @512에서:

• 가이던스 미사용 (Without Guidance): MSE 손실을 사용한 기본 Twins 모델은 낮은 FID를 보였습니다 (Flux.2 VAE 베이스라인 대비 저하). Focal Loss를 적용하면 VAE 특징의 예측이 크게 개선되어, 80 에포크 기준 FID가 14.41 (MSE)에서 3.84 (Focal Loss)로 감소했습니다.
• 가이던스 사용 (With Guidance): Twins는 Classifier-Free Guidance를 통해 강력한 gFID (ImageNet@256에서 1.59, ImageNet@512에서 1.79)를 달성했습니다.
• 트레이드오프 해결: RAE가 약간 더 낮은 FID를 달성하기도 하지만, Twins는 재구성 품질(PSNR) 측면에서 RAE를 크게 앞서며, 이해 중심 표현과 생성 중심 표현 사이의 간극을 성공적으로 좁혔습니다.

5. 의의 및 주장

본 논문은 시맨틱이 풍부한 ViT 특징과 세부 사항을 보존하는 VAE 특징을 명시적으로 융합하는 통합 표현을 구축함으로써, 시각적 토큰화의 "불가능한 삼각형"을 성공적으로 돌파했다고 주장합니다.

그 의의는 다음과 같습니다:

1. 단순성: 복잡한 아키텍처 재설계나 추가적인 잔차 인코더 없이, 기존의 오프더쉘프(off-the-shelf) 인코더를 직접 결합하는 방식을 사용합니다.
2. 균형 잡힌 최적화: 이질적인 특징을 공동 모델링할 때 발생하는 특정 실패 모드(DiT가 고주파 세부 사항을 무시하는 현상)를 식별하고 해결하였으며, 이를 위한 일반화 가능한 솔루션(Focal Loss)을 제공합니다.
3. 통합 인터페이스: 동일한 토큰 공간 내에서 모델이 이해와 생성을 모두 수행할 수 있게 하여, 추가적인 디코드-인코드 단계를 제거하고 작업 전반에 걸쳐 표현의 일관성을 보장합니다.

저자들은 이질적인 특징을 공동 모델링하는 것이 쉽지 않지만, 제안된 전략이 통합 표현 생성을 위한 새로운 기준을 세우며, 단일 프레임워크 내에서 시맨틱 및 미세한 잠재 변수 모두를 고품질로 예측할 수 있음을 결론지었습니다.