Repulsor: 메모리 뱅크 기반 대비학습으로 가속화된 생성 모델

초록

Repulsor는 외부 사전학습 인코더 없이 메모리 뱅크와 저차원 프로젝터를 활용해 대규모 음성 샘플을 제공함으로써, 디퓨전·플로우‑매칭 기반 생성 모델의 학습을 크게 가속화한다. 이미지넷‑256에서 400k 스텝만에 FID 2.40을 달성하며, 파라미터·연산량은 기존 모델과 동일하게 유지한다.

상세 분석

본 논문은 현재 시각 합성 분야에서 가장 널리 쓰이는 디퓨전 및 플로우‑매칭 모델이 겪는 두 가지 근본적인 문제—높은 학습 비용과 비효율적인 표현 학습—에 주목한다. 기존 접근법은 (1) 외부 사전학습 비전 인코더를 이용해 생성 모델의 중간 특징을 정렬(alignment)하거나, (2) ‘disperse’라 불리는 양성‑없는 대비학습을 적용해 배치 내 특징 간의 강한 분리를 유도한다. 전자는 인코더 전파 비용과 도메인 쉬프트 위험을 내포하고, 후자는 배치 크기에 의존해 음성 샘플 수가 제한된다.

Repulsor는 이 두 한계를 동시에 해소한다. 핵심 아이디어는 대규모 메모리 뱅크를 도입해 현재 배치와 무관하게 수천~수만 개의 음성 샘플을 유지·갱신함으로써, 배치 크기와 무관하게 풍부한 음성 쌍을 제공한다는 점이다. 메모리 뱅크는 MoCo와 동일하게 FIFO 방식으로 업데이트되며, stop‑gradient 연산을 통해 과거 샘플이 현재 파라미터에 영향을 주지 않도록 설계한다.

또한, 고차원 특징을 그대로 메모리 뱅크에 저장하면 메모리·대역폭 부담이 급증한다. 이를 해결하기 위해 저차원 프로젝션 헤드(단일 선형 레이어)를 삽입해 특징을 D‑차원(예: 128)으로 압축한다. L2 정규화 후 저장된 벡터는 InfoNCE 형태의 대비 손실에 바로 활용될 수 있다. 이 설계는 메모리 사용량을 크게 줄이면서도 충분한 구분력을 유지한다는 실험적 증거를 제공한다.

대조 손실은 기존 ‘disperse’ 손실을 확장한 형태로, 양성 쌍을 완전히 배제하고 오직 음성 쌍 간의 거리(exp(−‖z_i−m_k‖²/τ))를 최소화한다. 여기서 τ는 온도 파라미터이며, γ는 디퓨전 손실과 대비 손실 간의 가중치이다. 손실 함수는

L = L_diff + γ·L_disp

으로 정의되며, L_disp는 메모리 뱅크 전체 K개의 음성 샘플에 대해 평균을 취한다. 이때 K는 배치 크기 B와 독립적으로 설정 가능하므로, 음성 샘플 수를 기하급수적으로 확대할 수 있다.

실험에서는 SiT‑XL/2 아키텍처를 기반으로 ImageNet‑256을 32×32×4 라텐트 공간에 VAE로 인코딩한 뒤, 8개의 A100 GPU에서 400k 스텝까지 학습하였다. 주요 하이퍼파라미터는 배치 크기 256, 메모리 뱅크 크기 65,536, 프로젝션 차원 128, 온도 τ=0.07, γ=0.1 등이다. 결과는 기존 ‘disperse’(FID 5.09)와 SiT 기본 모델(FID 6.02)을 크게 앞서며, 400k 스텝에 FID 2.40을 기록한다. 이는 동일한 연산량·파라미터 수 하에서 가장 높은 성능이며, 특히 외부 인코더를 사용하지 않아 추론 시 추가 비용이 전혀 없다는 점이 큰 장점이다.

또한, 메모리 뱅크와 프로젝터는 학습 단계에서만 활성화되며, 추론 시 완전히 분리된다. 따라서 모델 구조 자체는 기존 디퓨전/플로우‑매칭 파이프라인과 동일하게 유지된다. 이는 실제 서비스 환경에서 레거시 파이프라인을 교체하거나 플러그‑인 형태로 적용하기에 매우 용이함을 의미한다.

한계점으로는 메모리 뱅크 크기와 프로젝션 차원 선택이 데이터셋·모델 규모에 따라 민감하게 작용할 수 있다는 점이다. 또한, 현재는 라텐트 공간에서만 적용했으며, 픽셀‑레벨 디퓨전이나 텍스트‑투‑이미지와 같은 멀티모달 설정에 대한 확장은 추가 연구가 필요하다.

요약하면, Repulsor는 (1) 외부 사전학습 인코더 의존성을 제거, (2) 배치 크기에 구애받지 않는 대규모 음성 샘플 제공, (3) 메모리·연산 효율성을 유지하면서 학습 수렴 속도를 크게 개선한다는 세 가지 핵심 기여를 한다.