MOSS-Image-heal: fp16을 안전하게 만드는 전체 파인튜닝

공식 모델 카드에서 가져온, 동일한 프롬프트, 시드, 16 스텝, CFG 3, shift 5로 생성한 베이스와 heal 결과입니다.

학습으로 없앤 문제

MOSS-Image는 5.08B 규모의 단일 스트림 diffusion transformer이며, 잔차 스트림에 몇 개의 매우 큰 채널이 존재합니다. 작성자의 6개 샘플 프로브에서 베이스 모델의 최대 스트림 값은 58,411이고 fp16 최댓값은 65,504이므로 여유는 1.12배에 불과합니다. 활성값을 조금만 더 밀어붙이는 요소나 중간값을 조금 더 넓게 계산하는 커널은 inf로 오버플로되어 이미지가 검정색 또는 깨진 상태로 돌아옵니다. Comfy-Org의 리팩에 fp16 대신 bf16과 int8 ConvRot 파일이 포함된 이유입니다.

1024² 영역에서 6개의 텍스트-이미지 샘플에 대한 블록별 최대 절대 스트림 값입니다. 모델 카드의 차트입니다.

여유는 같은 변화에서 나옵니다. 모델이 기존에 32개 채널에 몰려 있던 값을 더 많은 채널로 분산시키는 법을 학습한 것입니다. 에너지 집중도가 0.837에서 0.708로 낮아진 것이 그 가시적인 결과이며, 이것이 fp16 공간을 확보해 줍니다.

학습 방식

이번 학습은 동결된 텍스트 리파이너를 제외한 모든 transformer 가중치를 전체 파인튜닝한 것입니다. 텍스트 인코더, tokenizer, VAE, scheduler, pipeline 코드는 공식 MOSS-Image-Base의 것을 사용합니다. 손실 함수는 flow matching에 동결된 베이스로 향하는 증류 항, 어텐션과 FFN 쓰기가 소수 채널에 몰리는 것을 억제하는 스트림 쓰기 집중도 손실, 그리고 스트림 RMS 손실을 더한 구성입니다. 가중치는 시작 RMS의 상한 배수를 초과할 때마다 별도로 감쇠되며, 상한은 처음 10,000 스텝 동안 64에서 16까지 조여집니다.

2차원 텐서별 가중치 이상치: 베이스 transformer에서는 349개 요소가 64배 RMS를 넘었지만 heal에서는 하나도 없습니다. 모델 카드의 차트입니다.

출력은 베이스 모델 수준을 유지

스트림을 길들였다는 것은 모델이 여전히 베이스처럼 그릴 때만 의미가 있으며, 카드에는 두 수치가 모두 보고되어 있습니다. flow 손실은 0.4280에서 0.4289로 이동하고, 프로브 행에서 베이스 대비 velocity 코사인은 0.9983입니다. 이는 스타일 재학습이 아니라 수치적 수리에 가깝습니다.

카드의 두 번째 그리드로, 세로 비율이며 캡션 프롬프트와 GBC 레이아웃 프롬프트 모두로 출력했습니다.

ComfyUI용 파일

세 파일 모두 Comfy-Org 키 레이아웃의 단일 파일 체크포인트이므로, 공식 MOSS-Image 리팩과 동일한 확산 모델 로드 노드로 불러올 수 있습니다. 그 외에 바뀌는 것은 없습니다. 기존에 사용하던 qwen3.5_9b 텍스트 인코더와 MOSS VAE를 그대로 유지하면 됩니다.

int8 파일은 Comfy-Org의 moss_image_base_int8_convrot 레이아웃과 정확히 일치하도록 작성되었습니다. 동일한 757개 키, 양자화되지 않은 모든 텐서는 동일하며, 양자화 코드는 5e-8만큼만 1씩 차이가 납니다. 따라서 절반의 디스크 용량으로 공식 int8 체크포인트와 같은 자리에 그대로 넣을 수 있습니다.

같은 프로브 행에서 fp16 heal은 코사인 0.999996으로 velocity 오류 2.1e-3에 도달하며, 스트림 최댓값은 3,278로 fp16 최댓값보다 20배 낮습니다. 해당 행에서 베이스 모델은 최댓값 61,696으로 그보다 6% 낮습니다.

제공 위치

• 가중치 및 보고서: KBlueLeaf/MOSS-Image-heal
• 트레이너 및 문서: KohakuBlueleaf/KohakuMossTrainer
• 베이스 모델: OpenMOSS-Team/MOSS-Image-Base