Black Forest Labs가 FLUX.2 모델을 새로운 아키텍처와 사전학습으로 완전히 재설계해 단일 텍스트 인코더 사용 및 트랜스포머 블록 최적화 달성

Context

FLUX.1은 두 개의 텍스트 인코더를 사용해 프롬프트 임베딩 계산이 복잡했으며, 메모리 효율성과 추론 속도에서 개선 여지가 있었다. 또한 기존 아키텍처는 더 큰 모델 규모를 효율적으로 지원하기 위해 재설계가 필요했다.

Technical Solution

• 텍스트 인코더를 2개에서 1개로 감소: Mistral Small 3.1 단일 인코더로 통일하고 중간층 출력을 스택하여 프롬프트 임베딩 계산 단순화
• 모든 계층에서 바이어스 파라미터 제거: Attention과 Feedforward 서브블록의 모든 레이어에서 바이어스 제거로 모델 가중치 축소
• DiT 아키텍처 최적화: 타임/가이던스 정보를 블록별 개별 파라미터 대신 전역 AdaLayerNorm-Zero로 공유하도록 변경
• 단일스트림 블록 비율 확대: 더블스트림 블록 8개 + 단일스트림 블록 48개 구성으로 변경 (FLUX.1은 19/38)
• 완전 병렬 트랜스포머 블록 구현: 단일스트림 블록에서 Attention QKV 프로젝션과 Feedforward 입력 프로젝션을 융합해 병렬 처리 활성화
• 새로운 오토인코더 도입: AutoencoderKLFlux2로 교체하고 해상도 의존 타임스텝 스케줄링 개선

Impact

• FLUX.2[dev]-32B 모델에서 단일스트림 블록의 파라미터 비율을 73%로 증가 (FLUX.1[dev]-12B는 54% 더블스트림)
• CPU 오프로딩 활성화 시 H100 GPU에서 62GB VRAM 사용으로 추론 가능 (오프로딩 없이는 80GB 초과 필요)
• 4비트 양자화로 24GB GPU에서 추론 가능 (약 20GB VRAM 여유)
• Flash Attention 3 활용 시 Hopper 시리즈 GPU에서 추론 속도 증가

Key Takeaway

FLUX.2는 텍스트 인코더 통합과 파라미터 공유를 통해 복잡도를 줄이면서도 더 큰 모델 규모(32B)를 지원하도록 설계했으며, 바이어스 제거와 완전 병렬 블록 구조는 양자화 및 메모리 최적화와 상호작용하여 제약 환경에서의 배포 가능성을 확대했다.