TBSM: ‘삼체 산란’으로 원스텝 생성 모델을 학습하다

도입

현대 생성 모델은 대체로 몇 가지 큰 흐름 위에 있습니다. GAN은 적대적 판별기를 통해 학습 신호를 얻고, 확산 모델은 미리 정해진 노이즈-데이터 경로를 따라 학습합니다. 빠른 소수 스텝 생성 모델은 강한 교사 모델의 증류에 기대는 경우도 많습니다. 논문 “Three-Body Scattering for Generative Modeling”이 제안하는 TBSM은 다른 질문에서 출발합니다. 생성 샘플이 실제 데이터 분포에 가까워지려면 어느 방향으로 움직여야 하는가를 직접 학습하자는 것입니다.

핵심 아이디어

TBSM, 즉 Three-Body Scattering Modeling은 생성 과정을 분포 간 수송장을 배우는 문제로 봅니다. 에너지 거리 기반의 분포 에너지를 이용해 샘플 수준의 움직임을 만들고, 이를 원스텝 생성기의 직접적인 회귀 감독 신호로 사용합니다. 하나의 projectile, 즉 이동 대상이 되는 생성 샘플은 두 참조와 상호작용합니다. 하나는 실제 데이터에서 온 소스로, 샘플을 실제 분포 쪽으로 끌어당깁니다. 다른 하나는 독립적으로 생성된 소스로, 현재 생성 분포로부터 밀어내는 역할을 합니다.

이 삼체 구성은 미니배치 전체의 모든 쌍 상호작용을 계산하는 방식보다 단순합니다. 논문에 따르면 projectile과 조건을 고정했을 때, 이 산란 방향의 기대값은 생성 분포와 실제 분포 사이 에너지 거리 제곱의 절반에 대한 2-Wasserstein 그래디언트 플로우 속도와 일치합니다. 결과적으로 생성기는 간접적인 판별 점수나 고정된 노이즈 스케줄이 아니라, 샘플을 수정해야 할 방향 자체를 학습합니다.

주요 포인트

• 원스텝·소수 스텝 생성 지향: 긴 반복 샘플링이 아니라 한 번 또는 몇 번의 함수 평가로 고품질 생성을 목표로 합니다.
• 끌림과 밀어냄의 조합: 실제 샘플을 향한 인력과 다른 생성 샘플로부터의 척력을 함께 사용합니다.
• O(B) 샘플 손실: B개의 고정 타깃 이벤트에서 O(B)개의 샘플 단위 손실을 만들 수 있어, 전쌍 필드를 쓰는 일부 Drift 계열 방법보다 간결합니다.
• 온라인 추적: 조건부 기대값을 온라인으로 추적해 필드 추정의 노이즈를 줄이는 것을 목표로 합니다.
• 특징 공간 산란: 이미지 실험에서는 동결된 이미지 특징 안에서 산란 신호를 사용합니다.

결과와 의미

논문은 ImageNet-256에서 NFE=1의 원스텝 생성 결과를 보고합니다. 픽셀 공간 PixelDiT-XL은 FID 2.23, 잠재 공간 DiT-XL은 FID 1.63을 달성했다고 합니다. 또한 프로젝트 설명에는 최대 20B 파라미터 규모의 텍스트-이미지 모델에서도 추가 보조 손실 없이 원스텝 및 소수 스텝 생성을 안정적으로 학습한 사례가 언급됩니다.

이 결과가 중요한 이유는 단순히 샘플링 속도 때문만은 아닙니다. 원스텝 생성은 의미 일관성, 세부 묘사, 분포 커버리지를 단 한 번의 전방 계산에 담아야 하므로 매우 어렵습니다. TBSM은 미리 정한 궤적을 따라가게 하는 대신, 생성 분포를 실제 분포 쪽으로 밀어주는 수송 신호를 직접 회귀하도록 설계했습니다.

영향과 남은 질문

TBSM은 확산 모델의 감독 방식, Drift류 동역학, GAN류 목적함수를 분포 수송장이라는 공통 언어로 연결하려는 설계 지도도 제시합니다. 이는 서로 다른 생성 모델 패러다임을 비교하고 조합하는 데 유용한 관점입니다.

다만 학습 비용, 특징 공간 선택의 민감도, 다양한 도메인에서의 안정성, 강한 기준 모델과의 공정한 비교는 전체 논문과 재현 실험을 통해 더 확인해야 합니다. 그럼에도 TBSM은 효율적인 생성 모델을 위한 흥미로운 방향을 제안합니다. 긴 생성 경로를 모사하기보다, 샘플을 올바른 방향으로 움직이는 힘을 직접 배우는 접근이라는 점에서 주목할 만합니다.

출처: Hugging Face Daily Papers