MOSS-Image-heal: fp16推論を安全にするComfyUI向けフルファインチューン

公式モデルカードより、同じプロンプト、シード、16ステップ、CFG 3、shift 5でのベースとheal。

学習によって解消する問題

MOSS-Imageは5.08Bの単一ストリーム拡散transformerで、その残差ストリームには非常に大きなチャネルがいくつか含まれています。作者の6サンプルプローブでは、ベースモデルの最大ストリーム値は58,411で、fp16の最大値65,504に対して1.12倍の余裕しかありません。アクティベーションを少しでも強く押し上げるもの、またはわずかに広い中間値を計算するカーネルはinfにオーバーフローし、画像は黒または破損した状態で返ってきます。これが、Comfy-Orgのリパックがfp16ではなくbf16とint8 ConvRotファイルを提供している理由です。

1024²エリアでの6つのテキスト画像サンプルにおける、ブロックごとの最大絶対ストリーム値。モデルカードのチャート。

この余裕は同じ変化から生まれています。モデルは、以前は32チャネルに集中していたものをより多くのチャネルに分散させることを学習しました。最大32チャネルのエネルギー占有率が0.837から0.708へ下がったことがその目に見える側面であり、これがfp16の余地を生み出しています。

学習方法

この実行は、凍結されたテキストリファイナーを除くすべてのtransformer重みのフルファインチューンです。テキストエンコーダー、トークナイザー、VAE、スケジューラー、パイプラインコードは公式のMOSS-Image-Baseのものです。損失は、フローマッチングに加えて、凍結ベースへの蒸留項、アテンションとFFNの書き込みが少数のチャネルに集中することにペナルティを与えるストリーム書き込み集中損失、そしてストリームRMS損失です。重みは、開始時RMSの上限倍率を超えるたびに個別に減衰され、上限は最初の10,000ステップで64から16へ厳しくなります。

2次元テンソルごとの重み外れ値: ベースtransformerでは349要素が64倍RMSを超え、healではゼロです。モデルカードのチャート。

出力はベースモデルの水準を保つ

ストリームを抑えることが役立つのは、モデルがベースと同じように描ける場合だけであり、カードは両方の数値を報告しています。フロー損失は0.4280から0.4289へ変化し、プローブ行におけるベースに対する速度コサインは0.9983です。これは数値的な修復であり、スタイルの再学習ではありません。

カードの2つ目のグリッド。縦長アスペクトで、キャプションプロンプトとGBCレイアウトプロンプトの両方を用いて生成されています。

ComfyUI向けファイル

3つのファイルはすべてComfy-Orgのキーレイアウトの単一ファイルチェックポイントなので、公式MOSS-Imageリパックと同じ拡散モデルを読み込むノードで読み込めます。ほかに変更はありません。すでに使用しているqwen3.5_9bテキストエンコーダーとMOSS VAEをそのまま使います。

int8ファイルはComfy-Orgのmoss_image_base_int8_convrotレイアウトに正確に一致するよう書き出されています。同じ757キー、すべての非量子化テンソルは同一で、量子化コードの5e-8が1だけずれています。そのため、公式int8チェックポイントと同じスロットに半分のディスク容量で収まります。

同じプローブ行では、fp16のhealは速度誤差2.1e-3、コサイン0.999996に達し、ストリームのピークは3,278で、fp16最大値より20倍低くなります。これらの行におけるベースモデルのピークは61,696で、最大値より6%低い値です。

入手先

• 重みとレポート: KBlueLeaf/MOSS-Image-heal
• トレーナーと解説: KohakuBlueleaf/KohakuMossTrainer
• ベースモデル: OpenMOSS-Team/MOSS-Image-Base