소식

한어Русский языкEnglishFrançaisIndonesianSanskrit日本語DeutschPortuguêsΕλληνικάespañolItalianoSuomalainenLatina

새로운 지혜 보고서

편집자: 앨런

[새로운 지혜 소개]최근 캘리포니아 대학교 어바인 캠퍼스(University of California, Irvine) 및 기타 기관의 연구자들은 지연 마스킹, MoE 및 계층적 확장과 같은 전략을 사용하여 확산 모델의 훈련 비용을 US$1,890로 줄였습니다.

확산 모델을 훈련하는 데 비용이 얼마나 드나요?

이전의 가장 저렴한 방법(Wuerstchen)의 비용은 $28,400이고 Stable Diffusion과 같은 모델은 훨씬 더 비쌉니다.

대형 모델 시대에 대부분의 사람들은 대형 모델을 가지고 놀 여유가 없습니다. 모든 종류의 Vincentian 여성을 원한다면 제조업체에 의존하여 그들의 무게를 앞으로 옮겨야 합니다.

이 엄청난 비용을 줄이기 위해 연구자들은 다양한 솔루션을 시도했습니다.

예를 들어 원래 확산 모델은 노이즈에서 이미지까지 약 1,000단계가 필요하지만 약 20단계 이하로 줄었습니다.

확산 모델의 기본 모듈이 점차 Unet(CNN)의 DiT(Transformer)로 대체되면서 Transformer의 특성을 기반으로 한 일부 최적화도 이어졌습니다.

예를 들어 파이프라인과 같은 Attention의 일부 중복 계산을 건너뛰는 등의 수량화입니다.

최근 캘리포니아 대학교 어바인 캠퍼스(University of California, Irvine) 및 기타 기관의 연구자들은 "돈 절약"이라는 목표를 크게 발전시켰습니다.

논문 주소: https://arxiv.org/abs/2407.15811

——1,890달러에 11억 6천만 개의 매개변수 확산 모델을 처음부터 훈련합니다!

SOTA에 비해 한 단계 더 개선되어 일반 사람들도 사전 훈련을 맛볼 수 있다는 희망을 볼 수 있게 되었습니다.

더 중요한 것은 비용 절감 기술이 모델의 성능에 영향을 미치지 않는다는 것입니다. 11억 6천만 개의 매개변수는 다음과 같은 매우 좋은 결과를 제공합니다.

모양과 느낌 외에도 모델의 데이터 지표도 우수합니다. 예를 들어 아래 표에 제공된 FID 점수는 Stable Diffusion 1.5 및 DALL·E 2에 매우 가깝습니다.

대조적으로, Wuerstchen의 비용 절감 계획은 이상적인 시험 점수에 미치지 못하는 결과를 가져왔습니다.

돈을 절약하는 팁

연구진은 '각 달러 늘리기'를 목표로 확산 모델의 기본 모듈인 DiT부터 시작했습니다.

우선, 시퀀스 길이는 Transformer의 계산 비용의 적이므로 제거해야 합니다.

이미지의 경우 성능에 영향을 주지 않으면서 계산에 참여하는 패치 수를 최소화하고 메모리 오버헤드도 줄여야 합니다.

이미지 타일 수를 줄이는 방법에는 두 가지가 있습니다. 하나는 각 블록의 크기를 늘리는 것이고, 다른 하나는 패치(마스크)의 일부를 제거하는 것입니다.

전자는 모델 성능을 크게 저하시키므로 마스킹 방법을 고려합니다.

가장 순진한 마스크(순진한 토큰 마스킹)는 컨볼루셔널 UNet의 무작위 자르기 훈련과 유사하지만 이미지의 비연속 영역에 대한 훈련을 허용합니다.

기존의 가장 진보된 방법(MaskDiT)은 출력 전에 복원 및 재구성 구조를 추가하고, 학습을 통해 손실된 정보를 보충하기 위해 추가적인 손실 함수를 통해 학습시킵니다.

두 마스크 모두 계산 비용을 줄이기 위해 처음에 대부분의 패치를 폐기합니다. 정보 손실로 인해 Transformer의 전반적인 성능이 크게 저하됩니다. 이를 보완하려고 해도 별로 개선되지 않습니다.

——정보를 잃는 것은 바람직하지 않은데, 어떻게 정보를 잃지 않고 입력을 줄일 수 있을까요?

지연 마스크

본 논문에서는 마스크 이전에 전처리를 위해 패치 믹서를 사용하고, 폐기된 패치의 정보를 살아남은 패치에 삽입함으로써 높은 마스크 성능 저하의 영향을 크게 줄이는 지연 마스킹 전략을 제안합니다.

이 아키텍처에서는 Attention 레이어와 Feed-forward 레이어의 조합을 통해 패치 믹서가 구현됩니다. 전체 모델의 손실 함수는 다음과 같습니다.

MaskDiT에 비해 여기에는 추가적인 손실 기능이 필요하지 않으며 전반적인 디자인과 교육이 더 간단합니다.

믹서 자체는 매우 가벼운 구조로 되어 있어 비용 절감의 기준을 충족합니다.

미세 조정

마스킹 비율이 매우 높으면 확산 모델이 이미지의 전체 구조를 학습하고 훈련에서 테스트로 분포 이동을 도입하는 능력이 크게 감소하므로 저자는 사전 훈련 후에 소량의 미세 조정(마스크 해제)을 수행했습니다. (마스크) ).

또한 미세 조정을 통해 마스크 사용으로 인해 발생하는 바람직하지 않은 생성 아티팩트를 완화할 수 있습니다.

MoE 및 계층화된 확장

MoE는 훈련 비용을 크게 늘리지 않고도 모델의 매개변수와 표현 능력을 높일 수 있습니다.

저자는 전문가가 선택한 라우팅을 기반으로 단순화된 MoE 레이어를 사용하며, 각 전문가는 전문가 간의 부하 균형을 맞추기 위해 추가 보조 손실 기능 없이 토큰에 대한 경로를 결정합니다.

또한 저자는 Transformer 블록의 너비(즉, Attention 레이어와 Feed-forward 레이어의 숨겨진 레이어 크기)를 선형적으로 증가시키는 계층적 스케일링 방법도 고려했습니다.

비전 모델의 더 깊은 계층은 더 복잡한 기능을 학습하는 경향이 있으므로 더 깊은 계층에 더 많은 매개변수를 사용하면 성능이 향상됩니다.

실험 설정

저자는 패치 크기가 2인 DiT-Tiny/2와 DiT-Xl/2라는 두 가지 DiT 변형을 사용합니다.

모든 모델은 코사인 학습률 감소 및 높은 가중치 감소 기능을 갖춘 AdamW 최적화 프로그램을 사용하여 훈련되었습니다.

모델의 프런트 엔드는 Stable-Diffusion-XL 모델의 4채널 VAE(Variational Autoencoder)를 사용하여 이미지 특징을 추출합니다. 또한 대규모 훈련(저장 버전)에서 최신 16채널 VAE의 성능을 제공합니다. )도 테스트되었습니다.

저자는 EDM 프레임워크를 모든 확산 모델에 대한 통합 교육 설정으로 사용하고 FID 및 CLIP 점수를 사용하여 이미지 생성 모델의 성능을 측정합니다.

텍스트 인코더에는 가장 일반적으로 사용되는 CLIP 모델이 선택되었습니다. T5-xxl과 같은 대형 모델은 텍스트 합성과 같은 까다로운 작업에서 더 나은 성능을 발휘하지만 여기서는 비용 절약을 위해 사용되지 않습니다.

훈련 데이터 세트

2,200만 개의 이미지-텍스트 쌍을 포함하는 세 가지 실제 이미지 데이터세트(Conceptual Captions, Segment Anything, TextCaps)가 사용됩니다.

SA1B는 실제 자막을 제공하지 않기 때문에 여기서는 LLaVA 모델에서 생성된 합성 자막을 사용합니다. 저자는 또한 1,500만 개의 이미지-텍스트 쌍을 포함하는 두 개의 합성 이미지 데이터세트인 JourneyDB와 DiffusionDB를 대규모 훈련에 추가했습니다.

소규모 절제를 위해 연구원들은 더 큰 COYO-700M 데이터세트에서 10개 CIFAR-10 클래스의 이미지를 서브샘플링하여 cifar-captions라는 텍스트-이미지 데이터세트를 구축했습니다.

평가하다

모든 평가 실험은 DiT-Tiny/2 모델과 cifar-captions 데이터세트(256×256 해상도)를 사용하여 수행되었습니다.

각 모델은 AdamW 최적화 도구와 지수 이동 평균(마지막 10K 단계에 대해 평활 계수 0.995)을 사용하여 60K 최적화 단계에 대해 훈련되었습니다.

지연 마스크

실험의 기준선에서는 위에서 언급한 Naive 마스킹을 선택했으며, 이 기사의 지연 마스킹에는 백본 네트워크의 10% 미만의 매개변수 수를 사용하여 경량 패치 믹서를 추가했습니다.

일반적으로 패치가 많이 손실될수록(높은 마스킹 비율) 모델의 성능이 저하됩니다. 예를 들어 MaskDiT의 성능은 50%를 초과한 후 크게 떨어집니다.

여기의 비교 실험에서는 기본 하이퍼파라미터(학습률 1.6×10e-4, 가중치 감소 0.01, 코사인 학습률)를 사용하여 두 모델을 학습합니다.

위 그림의 결과는 FID, Clip-FID, Clip Score의 세 가지 지표에서 지연 마스킹 방식이 개선되었음을 보여줍니다.

또한 마스킹 비율이 증가함에 따라 기준선과의 성능 격차가 넓어집니다. 마스킹 비율이 75%일 때 순진한 마스킹은 FID 점수를 16.5로 감소시키는 반면, 우리의 방법은 5.03에 도달하여 마스킹이 없는 FID 점수(3.79)에 더 가깝습니다.

초매개변수

LLM 교육에 대한 일반적인 아이디어에 따라 여기서는 두 작업의 하이퍼파라미터 선택을 비교합니다.

첫째, 피드포워드 계층에서는 SwiGLU 활성화 함수가 GELU보다 우수합니다. 둘째, 가중치 감쇠가 높을수록 이미지 생성 성능이 향상됩니다.

또한 LLM 훈련과 달리 이 문서의 확산 모델은 AdamW 2차 모멘트(β)에 대해 더 높은 실행 평균 계수를 사용할 때 더 나은 성능을 달성할 수 있습니다.

마지막으로 저자는 학습 속도를 가능한 최대 값으로 높이면서(훈련이 불안정해질 때까지) 적은 수의 훈련 단계를 사용하면 이미지 생성 성능이 크게 향상된다는 것을 발견했습니다.

믹서 디자인

일반적으로 기적을 이루기 위해 열심히 일하는 것이 옳으며, 저자는 더 큰 패치 믹서를 사용한 후에도 모델 성능이 계속 향상되는 것을 관찰했습니다.

그러나 비용을 절약하기 위해 여기에서는 여전히 소형 믹서를 선택합니다.

저자는 노이즈 분포를 (-0.6, 1.2)로 수정하여 자막과 생성된 이미지 간의 정렬을 개선했습니다.

아래 그림에서 볼 수 있듯이 75% 마스킹 비율에서 저자는 다양한 패치 크기를 사용하는 경우의 영향도 연구했습니다.

연속된 영역의 수가 커지면(패치가 커짐) 모델의 성능이 저하되므로 각 패치를 무작위로 마스킹하는 원래 전략이 유지됩니다.

계층적 스케일링

이 실험에서는 DiT-Tiny 아키텍처의 두 가지 변형을 교육했습니다. 하나는 일정한 너비를 갖고 다른 하나는 계층적으로 확장된 구조를 사용했습니다.

두 방법 모두 Naive 마스킹을 사용하고 Transformer의 크기를 조정하여 두 경우의 모델 컴퓨팅 능력이 동일하도록 보장하면서 동일한 훈련 단계와 훈련 시간을 수행합니다.

위 표의 결과를 보면 세 가지 성능 지표 모두에서 Hierarchical Scaling 방식이 Baseline Constant Width 방식보다 우수하다는 것을 알 수 있으며, 이는 DiT의 마스킹 훈련에 Hierarchical Scaling 방식이 더 적합하다는 것을 의미합니다.

참고자료:

https://arxiv.org/abs/2407.15811