Yan Shuicheng/Cheng Mingming의 신작! Sora의 핵심 구성 요소인 DiT 교육은 10배 가속화되었으며 Masked Diffusion Transformer V2는 오픈 소스입니다.

Sora의 강력한 핵심 기술 중 하나인 DiT는 Diffusion Transformer를 활용하여 생성 모델을 더 큰 규모로 확장하여 뛰어난 이미지 생성 효과를 얻습니다.

그러나 모델 크기가 커지면 훈련 비용이 급등하게 됩니다.

난카이대학교 Sea AI Lab의 Yan Shuicheng, Cheng Mingming 연구팀과 Kunlun Wanwei 2050 연구소는 ICCV 2023 컨퍼런스에서 Masked Diffusion Transformer라는 새로운 모델을 제안했습니다. 이 모델은 마스크 모델링 기술을 사용하여 의미 표현 정보를 학습함으로써 Diffusion Transformer의 학습 속도를 높이고 이미지 생성 분야에서 SoTA 효과를 달성합니다. 이러한 혁신은 이미지 생성 모델 개발에 새로운 혁신을 가져오고 연구자에게 보다 효율적인 훈련 방법을 제공합니다. 연구팀은 다양한 분야의 전문 지식과 기술을 결합해 훈련 속도를 높이고 생성 결과를 향상시키는 솔루션을 성공적으로 제안했습니다. 그들의 작업은 인공 지능 분야의 발전에 중요한 혁신적인 아이디어를 제공했으며 향후 연구 및 실습에 유용한 영감을 제공했습니다. 2303.14389

GitHub 주소: https://github.com/sail-sg/MDT

최근 Masked Diffusion Transformer V2는 다시 한 번 SoTA를 갱신하여 DiT에 비해 훈련 속도를 10배 이상 높이고 ImageNet 벤치마크 점수 1.58을 달성했습니다.

최신 버전의 논문과 코드는 오픈 소스입니다.

Background

DiT로 대표되는 확산 모델은 이미지 생성 분야에서 큰 성공을 거두었지만, 연구원들은 확산 모델이 이미지에서 개체 부분 간의 의미 관계를 효율적으로 학습하기 어려운 경우가 많다는 사실을 발견했습니다. 제한으로 인해 훈련 ​​프로세스의 수렴 효율성이 낮아집니다.

Pictures

예를 들어 위 그림과 같이 DiT는 50k번째 훈련 단계에서 강아지의 털 질감을 생성하도록 학습했고, 이후 200k번째 훈련 단계에서 강아지의 눈 중 하나를 생성하도록 학습했습니다. 훈련 단계와 입이 있었지만 또 다른 눈이 누락되었습니다.

300k 훈련 단계에서도 DiT에 의해 생성된 강아지 두 귀의 상대적인 위치는 그다지 정확하지 않습니다.

이 훈련 및 학습 과정은 확산 모델이 이미지에서 객체의 다양한 부분 간의 의미 관계를 효율적으로 학습하지 못하고 각 객체의 의미 정보만 독립적으로 학습한다는 것을 보여줍니다.

이 현상의 원인은 확산 모델이 각 픽셀의 예측 손실을 최소화하여 실제 이미지 데이터의 분포를 학습하기 때문이라고 연구원들은 추측합니다. 이 프로세스는 객체의 다양한 부분 간의 의미론적 상대적 관계를 무시합니다. 이미지로 인해 모델이 천천히 수렴됩니다.

방법: Masked Diffusion Transformer

위의 관찰에서 영감을 받아 연구원은 확산 모델의 훈련 효율성과 생성 품질을 향상시키기 위해 Masked Diffusion Transformer(MDT)를 제안했습니다.

MDT는 Diffusion Transformer의 상황별 의미 정보 학습 능력을 명시적으로 강화하고 이미지 내 객체 간 의미 정보의 연관 학습을 강화하기 위해 Diffusion Transformer용으로 설계된 마스크 모델링 표현 학습 전략을 제안합니다.

Picture

위 그림과 같이 MDT는 확산 훈련 과정을 유지하면서 마스크 모델링 학습 전략을 도입합니다. MDT는 잡음이 있는 이미지 토큰을 마스킹함으로써 비대칭 확산 변환기(Asymmetric Diffusion Transformer) 아키텍처를 사용하여 마스킹되지 않은 잡음이 있는 이미지 토큰에서 마스킹된 이미지 토큰을 예측함으로써 마스크 모델링과 확산 훈련 프로세스를 동시에 달성합니다.

추론 프로세스 중에 MDT는 여전히 표준 확산 생성 프로세스를 유지합니다. MDT의 설계는 Diffusion Transformer가 마스크 모델링 표현 학습을 통해 가져온 의미 정보 표현 능력과 확산 모델의 이미지 세부 정보 생성 능력을 모두 갖도록 돕습니다.

구체적으로 MDT는 VAE 인코더를 통해 이미지를 잠재 공간에 매핑하고 잠재 공간에서 처리하여 컴퓨팅 비용을 절감합니다.

훈련 과정에서 MDT는 먼저 노이즈가 추가된 이미지 토큰 중 일부를 마스크하고 나머지 토큰을 비대칭 확산 변환기로 보내 노이즈 제거 후 모든 이미지 토큰을 예측합니다.

Asymmetric Diffusion Transformer 아키텍처

Picture

위 그림과 같이 Asymmetric Diffusion Transformer 아키텍처에는 인코더, 측면 보간기(보조 보간기) 및 디코더가 포함됩니다.

Pictures

인코더는 추론 프로세스 중에 마스크되지 않은 토큰만 처리합니다. 마스크 단계가 없기 때문에 모든 토큰을 처리합니다.

따라서 훈련 또는 추론 단계에서 디코더가 항상 모든 토큰을 처리할 수 있도록 하기 위해 연구원들은 훈련 과정에서 DiT 블록으로 구성된 보조 보간기를 통해 솔루션을 제안했습니다(그림 참조). 위), 인코더의 출력에서 ​​마스킹된 토큰을 보간 및 예측하고 추론 오버헤드를 추가하지 않고 추론 단계에서 이를 제거합니다.

MDT의 인코더와 디코더는 표준 DiT 블록에 전역 및 로컬 위치 인코딩 정보를 삽입하여 마스크 부분의 토큰 예측을 돕습니다.

Asymmetric Diffusion Transformer V2

Pictures

위 그림에서 볼 수 있듯이 MDTv2는 Masked Diffusion 프로세스를 위해 설계된 보다 효율적인 매크로 네트워크 구조를 도입하여 확산과 마스크를 더욱 최적화합니다. 모델링 과정.

여기에는 U-Net 스타일의 긴 단축키를 인코더에 통합하고 밀도가 높은 입력 단축키를 디코더에 통합하는 것이 포함됩니다.

그 중 Dense input-shortcut은 디코더에 노이즈를 추가한 후 마스크된 토큰을 전송하므로 마스크된 토큰에 해당하는 노이즈 정보를 유지하여 확산 과정의 훈련을 촉진합니다.

또한 MDT는 더 빠른 Adan 최적화 프로그램, 시간 단계 관련 손실 가중치, 확장된 마스크 비율을 포함하여 더 나은 훈련 전략을 도입하여 Masked Diffusion 모델의 훈련 프로세스를 더욱 가속화했습니다.

실험 결과

ImageNet 256 벤치마크 생성 품질 비교

Pictures

위 표는 다양한 모델 크기의 ImageNet 256 벤치마크에서 MDT와 DiT의 성능을 비교한 것입니다.

MDT가 모든 모델 크기에서 더 적은 훈련 비용으로 더 높은 FID 점수를 달성한다는 것은 분명합니다.

MDT의 매개변수와 추론 비용은 기본적으로 DiT와 동일합니다. 왜냐하면 위에서 언급한 것처럼 MDT 추론 과정에서도 DiT와 일치하는 표준 확산 과정이 계속 유지되기 때문입니다.

가장 큰 XL 모델의 경우 400,000단계로 훈련된 MDTv2-XL/2는 7000,000단계로 훈련된 DiT-XL/2보다 FID 점수가 1.92 향상되었습니다. 이 설정에서 결과는 MDT가 DiT보다 약 18배 더 빠른 훈련을 제공한다는 것을 보여줍니다.

소형 모델의 경우 MDTv2-S/2는 훨씬 적은 훈련 단계로 DiT-S/2보다 훨씬 더 나은 성능을 달성합니다. 예를 들어, 400,000 단계의 동일한 교육을 통해 MDTv2의 FID 지수는 39.50으로, 이는 DiT의 FID 지수인 68.40보다 훨씬 높습니다.

더 중요한 것은 이 결과가 400,000 훈련 단계(39.50 대 43.47)에서 더 큰 모델 DiT-B/2의 성능을 초과한다는 것입니다.

ImageNet 256 벤치마크 CFG 생성 품질 비교

images

또한 위 표의 분류자 없는 지침에 따라 MDT의 이미지 생성 성능을 기존 방법과 비교했습니다.

MDT는 FID 점수 1.79로 이전 SOTA DiT 및 기타 방법을 능가합니다. MDTv2는 성능을 더욱 향상시켜 더 적은 훈련 단계로 이미지 생성에 대한 SOTA FID 점수를 1.58이라는 새로운 최저 수준으로 끌어올렸습니다.

DiT와 마찬가지로 훈련을 계속하면서 훈련 중에 모델의 FID 점수 포화도를 관찰하지 못했습니다.

MDT는 PaperWithCode 리더보드에서 SoTA를 새로 고쳤습니다

수렴 속도 비교

그림

위 그림은 ImageNet 256 벤치마크에서 8×A100을 비교합니다. GPU 2 FID의 DiT-S/ 다양한 훈련 단계/훈련 시간에 따른 기준선, MDT-S/2 및 MDTv2-S/2의 성능.

더 나은 상황별 학습 기능 덕분에 MDT는 성능과 생성 속도 모두에서 DiT를 능가합니다. MDTv2의 학습 수렴 속도는 DiT보다 10배 이상 빠릅니다.

MDT는 훈련 단계와 훈련 시간 측면에서 DiT보다 약 3배 빠릅니다. MDTv2는 MDT에 비해 훈련 속도를 약 5배 더 향상시킵니다.

예를 들어, MDTv2-S/2는 학습하는 데 약 100시간(1500k 단계)이 걸리는 DiT-S/2보다 단 13시간(15,000단계) 만에 더 나은 성능을 보여줍니다. 이는 상황별 표현 학습이 중요하다는 것을 보여줍니다. 확산 모델의 더 빠른 생성 학습이 중요합니다.

Summary & Discussion

MDT는 확산 훈련 과정에서 MAE와 유사한 마스크 모델링 표현 학습 방식을 도입했습니다. 이는 이미지 객체의 상황별 정보를 사용하여 불완전한 입력 이미지의 완전한 정보를 재구성하여 학습할 수 있습니다. 이미지의 의미 부분 간의 상관관계를 파악하여 이미지 생성 품질과 학습 속도를 향상시킵니다.

연구원들은 시각적 표현 학습을 통해 물리적 세계에 대한 의미론적 이해를 높이면 생성 모델이 물리적 세계에 미치는 시뮬레이션 효과를 향상시킬 수 있다고 믿습니다. 이는 생성 모델을 통해 물리적 세계 시뮬레이터를 구축하려는 Sora의 비전과 일치합니다. 이 작업이 표현 학습과 생성 학습 통합에 대한 더 많은 작업에 영감을 주기를 바랍니다.

참고자료:

https://arxiv.org/abs/2303.14389

위 내용은 Yan Shuicheng/Cheng Mingming의 신작! Sora의 핵심 구성 요소인 DiT 교육은 10배 가속화되었으며 Masked Diffusion Transformer V2는 오픈 소스입니다.의 상세 내용입니다. 자세한 내용은 PHP 중국어 웹사이트의 기타 관련 기사를 참조하세요!