[딥러닝] Flux.1 파인튜닝 모델 Chroma의 특징 및 아키텍쳐 설명 가이드

Chroma는 FLUX.1-schnell에서 파생된 8.9B(89억) 파라미터 모델로, 접근성과 자유, 혁신을 최우선으로 하는 강력한 오픈 소스 AI 모델입니다. Apache 2.0 라이선스로 출시되어 누구나 자유롭게 사용할 수 있으며, 수정 및 확장이 가능하고 기업의 제약 없이 활용할 수 있습니다.

현재 오픈소스로 공개가 되어있고 학습 중인 Chroma는 방대한 데이터셋을 활용하여 생성형 AI의 한계를 확장하는 것을 목표로 하고 있다고 합니다. 이번 블로그에서는 Chroma 프로젝트의 목표, 기술적 혁신, 모델 구조, 그리고 오픈 소스 커뮤니티와 협력하는 방법을 자세히 살펴보겠습니다.

Chroma의 목표

Chroma는 2천만 개의 원본 샘플 중 5백만 개의 고품질 데이터를 엄선하여 학습되었습니다. 데이터셋에는 다음과 같은 요소들이 포함됩니다:

• 애니메이션 아트워크
• 퍼리 아트
• 일반 예술 작품
• 사진 이미지

Chroma의 주요 특징은 완전히 검열되지 않은 AI 모델이라는 점입니다. 기존 모델에서 빠진 해부학적 개념을 다시 도입하여, 예술가, 연구자, AI 개발자들에게 제약 없는 강력한 도구를 제공합니다.

Chroma의 Huggingface 사이트

• Chroma Huggingface 사이트 : Chroma Huggingface

Chroma 실행 방법

Chroma는 현재 ComfyUI와 diffusers(개발 중)에서 실행할 수 있습니다. 여기서는 ComfyUI를 사용한 실행 방법을 설명합니다.

필수 요구사항

• ComfyUI 설치
• Chroma 체크포인트(최신 버전 다운로드)
• T5 XXL 또는 T5 XXL fp8 모델
• FLUX VAE
• Chroma Workflow 파일

수동 설치 방법

1. ComfyUI의 커스텀 노드 폴더로 이동:
2. cd ComfyUI/custom_nodes
3. 필요한 저장소를 클론:
4. git clone https://github.com/lodestone-rock/ComfyUI_FluxMod.git
5. ComfyUI를 재시작하고, 브라우저를 새로고침합니다.
6. 필요한 모델 파일을 해당 디렉터리에 배치:
   • T5_xxl → ComfyUI/models/clip
   • FLUX VAE → ComfyUI/models/vae
   • Chroma 체크포인트 → ComfyUI/models/diffusion_models
7. Chroma Workflow를 불러와 실행합니다.

아키텍처 수정 사항

12B → 8.9B 파라미터 축소

Chroma는 기존 FLUX 모델의 12B 파라미터에서 8.9B로 최적화되었습니다. 이 과정에서 단일 입력 벡터를 인코딩하는 3.3B 파라미터의 불필요한 모듈레이션 레이어가 제거되었습니다.

실험 결과, 해당 레이어의 가중치를 0으로 설정했을 때도 출력이 거의 변하지 않았습니다. 그 이유로는 이 레이어에 인코딩을 할 수 있는 유일한 정보는 0-1 범위의 단일 숫자 뿐이기 때문입니다. 8 bytes의 부동 소수점 값을 인코딩하는데 있어서 3.3B가 사용이 되었고 이 레이어를 효과적으로 정리하고 FFN으로 대체하였습니다.

이에 따라 해당 레이어를 제거하고 250M 파라미터를 가진 경량 FFN(Fully Connected Feedforward Network)으로 대체하여 모델 크기를 줄였습니다. 이 최적화 과정은 RTX 3090 한 대로 하루 만에 완료되었습니다.

MMDiT 마스킹으로 안정성 향상

기존 FLUX 모델에서는 T5와 MMDiT의 패딩 토큰 마스킹이 누락되는 문제가 있었습니다. 예를 들어:

"A cat sat on a mat"

이 문장이 실제로 모델에 입력될 때는 다음과 같이 변형됩니다:

<bos> a cat sat on a mat <pad><pad>...<pad><pad>

이 때문에 모델이 패딩 토큰에 불필요하게 집중하게 되었습니다. Chroma에서는 패딩 토큰을 모두 마스킹하되, 하나만 남기는 방식을 적용하여 안정적인 학습이 가능하도록 개선되었습니다.

타임스텝 분포 조정으로 학습 효율 개선

일반적인 확산 모델에서는 랜덤한 타임스텝을 샘플링하지만, 기존 FLUX 모델은 로그 정규 분포(lognormal distribution)를 사용하여 학습했습니다. 이는 중간 단계의 타임스텝에 비중을 두는 방식이었지만, 극단적인 타임스텝에서는 학습이 부족해지는 문제를 초래했습니다.

Chroma에서는 -x² 함수 기반의 타임스텝 샘플링 방식을 적용하여 모든 타임스텝이 고르게 학습될 수 있도록 조정했습니다. 이를 통해 손실(loss) 스파이크를 줄이고, 모델의 수렴 속도를 향상시켰습니다.

미니배치 최적 수송 기법으로 학습 가속화

Chroma는 단순한 "노이즈 제거"가 아니라, 노이즈 분포를 이미지 분포로 매핑하는 벡터 필드를 학습하는 모델입니다. 이를 최적화하기 위해 미니배치 최적 수송(Minibatch Optimal Transport) 기법을 적용하여 데이터 샘플 간의 매칭을 개선했습니다. 이 기술을 활용하면 학습 경로의 모호성이 줄어들어 모델 훈련이 더욱 효율적으로 이루어집니다.

결론

Chroma는 오픈 소스 AI의 새로운 가능성을 여는 모델입니다. 기존 모델을 최적화하고, 학습 안정성을 향상시키며, 검열 없는 자유로운 사용을 보장하는 강력한 생성형 AI입니다. 최근 트랜드가 소규모 모델인 것처럼 Flux.1-dev or Flux.1-schell같은 모델도 현재 용량이 매우 커서 사용하기 힘들지만 Chroma와 같이 최적화 모델이 나와 다양한 파인튜닝 모델이 나오길 기대합니다.