다양체 위에서의 학습: 표현 인코더로 표준 확산 트랜스포머의 성능 향상하기

초록

표현 인코더를 생성 모델링에 활용하면 효율적이고 높은 충실도의 합성 경로를 제공합니다. 그러나 표준 확산 트랜스포머는 이러한 표현을 직접적으로 수렴하는 데 실패합니다. 최근 연구에서는 이를 확산 트랜스포머의 계산 비용이 많이 드는 폭 확장을 제안하는 용량 병목 현상으로 돌리지만, 우리는 이 실패가 근본적으로 기하학적임을 입증합니다. 우리는 기하학적 간섭을 근본 원인으로 규정합니다: 표준 유클리드 흐름 정합은 매니폴드 표면을 따르기보다는 표현 인코더의 초구형 특징 공간의 저밀도 내부를 통해 확률 경로를 강제합니다. 이를 해결하기 위해 우리는 야코비 정규화를 적용한 리만 흐름 정합(RJF)을 제안합니다. 생성 과정을 매니폴드 측지선에 제한하고 곡률로 인한 오차 전파를 보정함으로써, RJF는 표준 확산 트랜스포머 아키텍처가 폭 확장 없이도 수렴할 수 있게 합니다. 우리의 방법인 RJF는 표준 DiT-B 아키텍처(1억 3천만 개 매개변수)가 효과적으로 수렴하여 기존 방법들이 수렴하지 못했던 3.37의 FID를 달성합니다. 코드: https://github.com/amandpkr/RJF

English

Leveraging representation encoders for generative modeling offers a path for efficient, high-fidelity synthesis. However, standard diffusion transformers fail to converge on these representations directly. While recent work attributes this to a capacity bottleneck proposing computationally expensive width scaling of diffusion transformers we demonstrate that the failure is fundamentally geometric. We identify Geometric Interference as the root cause: standard Euclidean flow matching forces probability paths through the low-density interior of the hyperspherical feature space of representation encoders, rather than following the manifold surface. To resolve this, we propose Riemannian Flow Matching with Jacobi Regularization (RJF). By constraining the generative process to the manifold geodesics and correcting for curvature-induced error propagation, RJF enables standard Diffusion Transformer architectures to converge without width scaling. Our method RJF enables the standard DiT-B architecture (131M parameters) to converge effectively, achieving an FID of 3.37 where prior methods fail to converge. Code: https://github.com/amandpkr/RJF

다양체 위에서의 학습: 표현 인코더로 표준 확산 트랜스포머의 성능 향상하기

Learning on the Manifold: Unlocking Standard Diffusion Transformers with Representation Encoders

초록

Support