U-Net에서 DiT까지: 지배력 확산 모델에 변압기 기술 적용

최근에는 Transformer를 중심으로 기계 학습이 르네상스를 경험하고 있습니다. 지난 5년 동안 자연어 처리, 컴퓨터 비전 및 기타 분야의 신경 아키텍처는 주로 변환기에 의해 지배되었습니다.

그러나 여전히 이러한 추세에 영향을 받지 않는 이미지 수준 생성 모델이 많이 있습니다. 예를 들어 확산 모델은 지난 해 이미지 생성에서 놀라운 결과를 얻었으며 거의 ​​모든 모델이 컨벌루션 U-모델을 사용합니다. Net을 백본으로 사용합니다. 이건 좀 놀랍네요! 지난 몇 년간 딥 러닝 분야의 큰 화두는 여러 분야에서 Transformer가 우위를 점했다는 것입니다. U-Net이나 컨볼루션이 확산 모델에서 그렇게 잘 작동하게 만드는 특별한 점이 있나요?

U-Net 백본 네트워크를 확산 모델에 처음 도입한 연구는 Ho 등으로 거슬러 올라갑니다. 이 디자인 패턴은 약간의 변경만으로 자동 회귀 생성 모델 PixelCNN++을 상속합니다. PixelCNN++는 많은 ResNet 블록을 포함하는 컨볼루셔널 레이어로 구성됩니다. 표준 U-Net과 비교하여 PixelCNN++의 추가 공간 셀프 어텐션 블록은 변환기의 기본 구성 요소가 됩니다. 다른 연구와 달리 Dhariwal 및 Nichol et al.은 적응형 정규화 계층을 사용하여 조건 정보 및 채널 수를 컨볼루셔널 계층에 주입하는 등 U-Net에 대한 여러 아키텍처 선택을 제거했습니다.

이 기사에서 UC Berkeley의 William Peebles와 New York University의 Xie Saining은 "변환기를 사용한 확장 가능한 확산 모델"을 썼습니다. 목표는 확산 모델에서 아키텍처 선택의 중요성을 밝히고 미래 생성을 위한 실증적 기준을 제공하는 것입니다. 모델 연구. 이 연구는 U-Net 유도 바이어스가 확산 모델의 성능에 중요하지 않으며 변압기와 같은 표준 설계로 쉽게 대체될 수 있음을 보여줍니다.

이 결과는 확산 모델이 아키텍처 통합 트렌드의 이점을 누릴 수 있음을 보여줍니다. 예를 들어 확산 모델은 이러한 모델의 확장성, 견고성 및 효율성을 유지하면서 다른 분야의 모범 사례와 교육 방법을 상속받을 수 있습니다. 표준화된 아키텍처는 또한 교차 영역 연구에 대한 새로운 가능성을 열어줄 것입니다.

• 논문 주소: https://arxiv.org/pdf/2212.09748.pdf
• 프로젝트 주소: https://github.com/facebookresearch/DiT
• 논문 홈페이지: https://www.wpeebles.com/DiT

이 연구는 새로운 유형의 변압기 기반 확산 모델인 확산 변압기(DiTs)에 중점을 둡니다. DiT는 ViT(Vision Transformers)의 모범 사례를 따르며 몇 가지 작지만 중요한 조정이 이루어졌습니다. DiT는 ResNet과 같은 기존 컨볼루션 네트워크보다 더 효율적으로 확장되는 것으로 나타났습니다.

특히 이 문서에서는 네트워크 복잡성 및 샘플 품질 측면에서 Transformer의 확장 동작을 연구합니다. 이 연구는 확산 모델이 VAE의 잠재 공간 내에서 학습되는 LDM(잠재 확산 모델) 프레임워크에서 DiT 설계 공간을 구축하고 벤치마킹함으로써 U-Net 백본을 변압기로 성공적으로 대체할 수 있음을 보여줍니다. 이 문서에서는 DiT가 확산 모델을 위한 확장 가능한 아키텍처임을 추가로 보여줍니다. 네트워크 복잡성(Gflops로 측정)과 샘플 품질(FID로 측정) 사이에는 강한 상관 관계가 있습니다. 간단히 DiT를 확장하고 대용량 백본(118.6Gflops)으로 LDM을 교육함으로써 클래스 조건부 256 × 256 ImageNet 생성 벤치마크에서 2.27 FID의 최첨단 결과를 얻을 수 있습니다.

확산 변압기

DiTs는 확장성을 유지하기 위해 표준 변압기 아키텍처에 최대한 충실하는 것을 목표로 하는 확산 모델을 위한 새로운 아키텍처입니다. DiT는 ViT의 많은 모범 사례를 유지하고 있으며 그림 3은 전체 DiT 아키텍처를 보여줍니다.

DiT에 대한 입력은 공간 표현 z입니다(256 × 256 × 3 이미지의 경우 z 모양은 32 × 32 × 4입니다). DiT의 첫 번째 레이어는 patchify로, 각 패치를 입력에 선형적으로 삽입하여 공간 입력을 T 토큰 시퀀스로 변환합니다. 패치를 적용한 후 모든 입력 토큰에 표준 ViT 주파수 기반 위치 임베딩을 적용합니다.

patchify에서 생성된 토큰 T의 수는 패치 크기 하이퍼 매개변수 p에 의해 결정됩니다. 그림 4에 표시된 것처럼 p를 절반으로 줄이면 T가 4배가 되고 따라서 변압기 Gflops는 최소한 4배가 됩니다. 이 기사에서는 DiT 디자인 공간에 p = 2,4,8을 추가합니다.

DiT 블록 설계: 패치 후 입력 토큰은 일련의 변환기 블록에 의해 처리됩니다. 노이즈가 있는 이미지 입력 ​​외에도 확산 모델은 노이즈 시간 단계 t, 클래스 라벨 c, 자연어 등과 같은 추가 조건 정보를 처리하는 경우도 있습니다. 이 기사에서는 다양한 방식으로 조건부 입력을 처리하는 네 가지 변환기 블록 변형을 살펴봅니다. 이러한 설계에는 표준 ViT 블록 설계에 대한 사소하지만 중요한 수정 사항이 포함되어 있습니다. 모든 모듈의 설계는 그림 3에 나와 있습니다.

이 기사에서는 모델 깊이와 너비에 따라 DiT-S, DiT-B, DiT-L 및 DiT-XL의 네 가지 구성을 시도했습니다. 이러한 모델 구성의 범위는 33M에서 675M 매개변수이고 Gflops는 0.4에서 119입니다.

실험

연구원들은 가장 높은 Gflops로 4개의 DiT-XL/2 모델을 훈련시켰으며, 각각은 컨텍스트 내(119.4Gflops), 교차 주의(137.6Gflops), 적응형 레이어 표준(adaLN) 등 서로 다른 블록 디자인을 사용했습니다. , 118.6Gflops) 또는 adaLN-zero(118.6Gflops)입니다. 그런 다음 훈련 중에 FID를 측정했으며 그림 5에 결과가 나와 있습니다.

모델 크기 및 패치 크기가 확장되었습니다. 그림 2(왼쪽)는 각 모델의 Gflops 개요와 400K 훈련 반복에서의 FID를 보여줍니다. 모델 크기를 늘리고 패치 크기를 줄이면 확산 모델이 크게 향상되는 것을 볼 수 있습니다.

그림 6(상단)은 모델 크기가 증가하고 패치 크기가 일정하게 유지됨에 따라 FID가 어떻게 변경되는지 보여줍니다. 네 가지 설정 전반에 걸쳐 Transformer를 더 깊고 넓게 만들어 훈련의 모든 단계에서 FID가 크게 향상되었습니다. 마찬가지로 그림 6(하단)은 패치 크기가 줄어들고 모델 크기가 일정하게 유지될 때의 FID를 보여줍니다. 연구원들은 단순히 DiT에서 처리하는 토큰 수를 확장하고 훈련 과정 전반에 걸쳐 매개변수를 대략적으로 고정시켜 FID가 크게 향상되었음을 다시 한 번 관찰했습니다.

그림 8은 400K 훈련 단계에서 FID-50K와 모델 Gflops를 비교한 것입니다:

SOTA 확산 모델 256×256 ImageNet. 확장된 분석 후에 연구원들은 7M의 걸음 수로 가장 높은 Gflop 모델인 DiT-XL/2를 계속해서 훈련했습니다. 그림 1은 이 모델의 샘플을 보여주고 이를 카테고리 조건부 생성 SOTA 모델과 비교한 결과를 표 2에 나타내었습니다.

분류자가 없는 안내를 사용할 때 DiT-XL/2는 이전의 모든 확산 모델보다 성능이 뛰어나 LDM이 달성한 이전 최고의 FID-50K인 3.60을 2.27로 줄였습니다. 그림 2(오른쪽)에서 볼 수 있듯이 LDM-4(103.6Gflops)와 같은 잠재 공간 U-Net 모델과 비교할 때 DiT-XL/2(118.6Gflops)는 ADM(1120Gflops)보다 계산 효율성이 훨씬 높습니다. ADM-U(742 Gflops), 픽셀 공간 U-Net 모델이 훨씬 더 효율적입니다.

표 3은 SOTA 방법과의 비교를 보여줍니다. XL/2는 이 해상도에서 이전의 모든 확산 모델을 다시 능가하여 ADM의 이전 최고 FID인 3.85를 3.04로 향상시켰습니다.

자세한 연구 내용은 원문을 참고해주세요.

위 내용은 U-Net에서 DiT까지: 지배력 확산 모델에 변압기 기술 적용의 상세 내용입니다. 자세한 내용은 PHP 중국어 웹사이트의 기타 관련 기사를 참조하세요!