MOSS-Image-heal : un finetune complet qui rend la fp16 sûre

Base et heal avec le même prompt, la même graine, 16 étapes, CFG 3 et shift 5, d'après la carte de modèle officielle.

Le problème que l'entraînement élimine

MOSS-Image est un transformer de diffusion single-stream de 5,08B, et son flux résiduel transporte une poignée de canaux très grands. Sur la sonde de six échantillons de l'auteur, la plus grande valeur de flux du modèle de base est de 58 411 contre un maximum fp16 de 65 504, soit 1,12x de marge. Tout ce qui pousse un peu plus fort les activations, ou tout noyau qui calcule un intermédiaire légèrement plus large, déborde vers inf, et l'image revient noire ou corrompue. C'est pourquoi le repack de Comfy-Org livre des fichiers bf16 et int8 ConvRot plutôt qu'un fichier fp16.

Plus grande valeur absolue du flux par bloc sur six échantillons texte-image pour une surface de 1024². Graphique issu de la carte de modèle.

La marge vient du même changement : le modèle a appris à répartir sur davantage de canaux ce qui tenait auparavant dans 32 d'entre eux. La concentration d'énergie qui passe de 0,837 à 0,708 en est la face visible, et c'est elle qui achète l'espace fp16.

Comment il a été entraîné

L'entraînement est un finetune complet de chaque poids du transformer, à l'exception du text refiner qui reste gelé ; l'encodeur de texte, le tokenizer, le VAE, le scheduler et le code du pipeline sont ceux officiels de MOSS-Image-Base. Les pertes sont le flow matching plus un terme de distillation vers la base gelée, une perte de concentration d'écriture du flux qui pénalise l'attention et les écritures FFN s'accumulant dans peu de canaux, et une perte de RMS du flux. Les poids sont décroissés séparément dès qu'ils dépassent un multiple plafond de leur RMS de départ, le plafond se resserrant de 64 à 16 sur les 10 000 premiers pas.

Valeurs aberrantes des poids par tenseur 2-D : 349 éléments se situaient au-dessus de 64x le RMS dans le transformer de base, et aucun dans heal. Graphique issu de la carte de modèle.

Les sorties restent là où était le modèle de base

Dompter le flux n'est utile que si le modèle dessine toujours comme la base, et la carte rapporte les deux chiffres. La perte de flow passe de 0,4280 à 0,4289, et le cosinus des vélocités par rapport à la base sur les lignes de la sonde est de 0,9983. C'est une réparation numérique, pas un réentraînement de style.

La deuxième grille de la carte, en format portrait, imprimée avec le prompt de légende et le prompt de mise en page GBC.

Fichiers pour ComfyUI

Les trois fichiers sont des checkpoints single-file dans la disposition de clés de Comfy-Org, ils se chargent donc via le même nœud Charger Modèle Diffusion que le repack officiel MOSS-Image. Rien d'autre ne change : vous conservez l'encodeur de texte qwen3.5_9b et le VAE MOSS que vous utilisez déjà.

Le fichier int8 est écrit pour correspondre exactement à la disposition moss_image_base_int8_convrot de Comfy-Org : mêmes 757 clés, chaque tenseur non quantifié identique, et 5e-8 des codes de quantification décalés d'un. Il se glisse donc dans le même emplacement que le checkpoint int8 officiel, pour la moitié du disque.

Sur les mêmes lignes de sonde, heal en fp16 atteint une erreur de vélocité de 2,1e-3 avec un cosinus de 0,999996, et culmine à 3 278 dans le flux, 20x sous le maximum fp16. Le modèle de base sur ces lignes culmine à 61 696, soit 6 % en dessous.

Disponibilité

• Poids et rapports : KBlueLeaf/MOSS-Image-heal
• Entraîneur et documentation : KohakuBlueleaf/KohakuMossTrainer
• Modèle de base : OpenMOSS-Team/MOSS-Image-Base