[논문 리뷰] Scalable Adaptive Computation for Iterative Generation

Allan Jabri, David J. Fleet외 1명|arXiv (Cornell University)|2022. 12. 22.

Stochastic Gradient Optimization Techniques인용 수 13

한 줄 요약

이 논문은 고차원적이고 높은 레이어의 자연 데이터(예: 이미지 및 영상)에 대해 균일한 계산이 비효율적인 문제를 해결하기 위해 도메인에 관계없이 적용 가능한 Recurrent Interface Networks(RINs)를 소개한다. RINs는 잠재 토큰을 사용해 전역 자기주의를 집중시키고, 계산을 입력 데이터의 차원에서 분리함으로써 가변적이고 확장 가능한 고차원 데이터 생성을 가능하게 한다. RINs는 1024×1024 해상도에서 캐스케이드나 가이던스 없이도 이미지 및 영상 확산 모델링에서 최신 기준 성능(SOTA)을 달성하며, 잠재 자기조건화를 통해 재귀를 활용함으로써 2D/3D U-Net 대비 최대 10배의 FLOPs 절감을 이룬다.

ABSTRACT

Natural data is redundant yet predominant architectures tile computation uniformly across their input and output space. We propose the Recurrent Interface Networks (RINs), an attention-based architecture that decouples its core computation from the dimensionality of the data, enabling adaptive computation for more scalable generation of high-dimensional data. RINs focus the bulk of computation (i.e. global self-attention) on a set of latent tokens, using cross-attention to read and write (i.e. route) information between latent and data tokens. Stacking RIN blocks allows bottom-up (data to latent) and top-down (latent to data) feedback, leading to deeper and more expressive routing. While this routing introduces challenges, this is less problematic in recurrent computation settings where the task (and routing problem) changes gradually, such as iterative generation with diffusion models. We show how to leverage recurrence by conditioning the latent tokens at each forward pass of the reverse diffusion process with those from prior computation, i.e. latent self-conditioning. RINs yield state-of-the-art pixel diffusion models for image and video generation, scaling to 1024X1024 images without cascades or guidance, while being domain-agnostic and up to 10X more efficient than 2D and 3D U-Nets.

연구 동기 및 목표

• 이미지 및 영상과 같은 고차원적이고 부족한 자연 데이터에서 딥러닝 모델의 균일한 계산의 비효율성 문제를 해결한다.
• 입력 크기의 제곱 복잡도를 피하면서 데이터 콘텐츠에 따라 계산을 동적으로 할당하는 확장 가능한 아키텍처를 개발한다.
• 재귀를 활용해 반복적 생성 작업에서 깊이 있는 표현력 있는 라우팅을 가능하게 하여, 시간에 따른 역전파 없이도 맥락 전파를 약간씩 감소시킨다.
• 픽셀 수준의 확산 모델에서 캐스케이드나 가이던스에 의존하는 것을 없애기 위해 더 효율적이고 도메인에 관계없이 적용 가능한 아키텍처를 설계한다.

제안 방법

• 숨겨진 유닛을 입력 크기와 함께 선형적으로 증가하는 인터페이스 토큰과 고정된 크기의 컴act한 잠재 토큰으로 분해함으로써, 계산을 입력 레이아웃에서 분리한다.
• 교차 어텐션을 사용해 인터페이스 토큰에서 잠재 토큰으로, 그리고 그 반대로 정보를 라우팅함으로써 双방향 하향 및 상향 피드백을 가능하게 한다.
• 입력 크기의 제곱 복잡도를 피하기 위해 전역 자기주의를 작은 잠재 토큰 집합에만 적용한다.
• 잠재 자기조건화 도입: 이전 추론 단계의 잠재 활성화를 이후 단계에서 맥락으로 재사용함으로써, 시간에 따른 역전파 없이도 재귀를 모방한다.
• RIN 블록을 스택하여 반복 과정에서 맥락을 누적하는 더 깊은 네트워크를 구축함으로써 라우팅 표현력을 향상시킨다.
• 고정된 계산 그래프를 갖는 표준 미분 가능 네트워크로 학습 및 배포함으로써 현대 하드웨어에서 효율적인 추론을 가능하게 한다.

실험 결과

연구 질문

• RQ1 입력 크기의 제곱 복잡도를 유발하지 않으면서도 고차원 데이터에 대해 계산을 적응적으로 할당할 수 있는 신경망 아키텍처를 설계할 수 있는가?
• RQ2 반복적 생성 작업에서 시간에 따른 역전파 없이도 재귀를 활용해 라우팅의 깊이를 증가시킬 수 있는가?
• RQ3 RINs와 같은 도메인에 관계없이 적용 가능한 아키텍처가 고해상도 이미지 및 영상 생성에서 샘플 품질과 FLOP 효율성 측면에서 U-Net 기반 확산 모델을 능가할 수 있는가?
• RQ4 잠재 자기조건화가 반복적 확산 환경에서 라우팅 깊이와 모델 표현력에 얼마나 기여하는가?
• RQ5 더 효율적인 아키텍처를 사용해 캐스케이드나 가이던스 없이도 픽셀 수준의 확산 모델을 1024×1024 해상도로 확장할 수 있는가?

주요 결과

• RINs는 가이던스나 캐스케이드 없이도 64×64에서 1024×1024까지 다양한 해상도에서 ImageNet에서 최신 기준(FID) 성능을 달성한다.
• 1024×1024 이미지 생성에서 RINs는 2D 및 3D U-Net 기준선 대비 추론 단계당 FLOPs를 최대 10배까지 감소시킨다.
• Kinetics600 영상 생성 벤치마크에서 RINs는 선구적 접근 방식을 능가하며, 단계당 FLOPs를 10배 감소시킨다.
• 잠재 자기조건화 덕분에 반복적 노이즈 제거 단계 동안 라우팅 맥락을 유지함으로써 더 깊은 효과적인 계산 그래프를 가능하게 하여 모델 표현력을 향상시킨다.
• RINs는 다양한 모odal리티에서 우수한 성능을 유지하며 도메인에 관계없이 적용 가능하다. 어텐션과 MLP만을 기반으로 하므로 광범위하게 적용 가능하다.
• 모델는 최소한의 아키텍처 수정으로도 강력한 성능을 달성하였으며, 이는 적응형 계산이 표준 확산 학습 파이프라인에 효율적으로 통합될 수 있음을 보여준다.