MOSS-Image-heal: полный файнтюн, который делает fp16 безопасным

Базовая модель и heal с одинаковым промптом, сидом, 16 шагами, CFG 3 и shift 5, из официальной карточки модели.

Проблема, которую он устраняет обучением

MOSS-Image это диффузионный трансформер с одним потоком на 5,08 млрд параметров, и его остаточный поток несёт в себе несколько очень крупных каналов. На проверке автора из шести образцов наибольшее значение потока у базовой модели составляет 58 411 при максимуме fp16 в 65 504, то есть запас всего 1,12x. Всё, что разгоняет активации чуть сильнее, или любой kernel, вычисляющий чуть более широкую промежуточную величину, приводит к переполнению в inf, и изображение возвращается чёрным или испорченным. Именно поэтому в репаке от Comfy-Org поставляются файлы bf16 и int8 ConvRot, а не fp16.

Наибольшее абсолютное значение потока по блокам на шести образцах текст-изображение при площади 1024². Диаграмма из карточки модели.

Запас появляется из того же изменения: модель научилась распределять то, что раньше сидело в 32 каналах, по большему их числу. Снижение концентрации энергии с 0,837 до 0,708 это видимая сторона этого процесса, и именно оно освобождает место для fp16.

Как его обучали

Это полный файнтюн всех весов трансформера, кроме замороженного text refiner; текстовый энкодер, tokenizer, VAE, scheduler и код pipeline взяты официальные из MOSS-Image-Base. Функции потерь: flow matching плюс дистилляционный член в сторону замороженной базовой модели, потеря концентрации записи в поток (stream write-concentration), которая штрафует записи attention и FFN, сбивающиеся в несколько каналов, и потеря RMS потока. Веса отдельно уменьшаются всякий раз, когда превышают кратный порог от своего начального RMS, причём порог ужесточается с 64 до 16 за первые 10 000 шагов.

Выбросы весов на двумерный тензор: 349 элементов превышали 64x RMS в базовом трансформере, а в heal таких нет. Диаграмма из карточки модели.

Выходы остаются там же, где была базовая модель

Укрощение потока полезно только если модель по-прежнему рисует как базовая, и карточка приводит оба числа. Flow loss меняется с 0,4280 до 0,4289, а косинус скорости относительно базовой модели на строках проверки составляет 0,9983. Это численный ремонт, а не переобучение стиля.

Вторая сетка из карточки, портретное соотношение сторон, напечатана и с промптом из подписи, и с промптом раскладки GBC.

Файлы для ComfyUI

Все три файла это одиночные чекпойнты в раскладке ключей Comfy-Org, поэтому они загружаются через ту же ноду Загрузить модель диффузии, что и официальный репак MOSS-Image. Больше ничего не меняется: вы оставляете текстовый энкодер qwen3.5_9b и VAE MOSS, которые уже используете.

Файл int8 записан так, чтобы точно соответствовать раскладке moss_image_base_int8_convrot от Comfy-Org: те же 757 ключей, каждый неквантованный тензор идентичен, а коды квантования отличаются на единицу в 5e-8 случаев. Так что он встаёт на то же место, что и официальный int8 чекпойнт, занимая вдвое меньше диска.

На тех же строках проверки fp16 heal достигает ошибки скорости 2,1e-3 при косинусе 0,999996 и пика потока 3 278, что в 20 раз ниже максимума fp16. Базовая модель на этих строках даёт пик 61 696, то есть на 6% ниже него.

Доступность

• Веса и отчёты: KBlueLeaf/MOSS-Image-heal
• Тренер и описания: KohakuBlueleaf/KohakuMossTrainer
• Базовая модель: OpenMOSS-Team/MOSS-Image-Base