拡散モデルを 1 ビットの限界まで量子化することで、新たな最先端 (SOTA) ソリューションが生まれました。

北京航空航天大学やチューリッヒ工科大学などの研究機関の研究者らは、拡散モデル（DM）の重みと活性化を初めて完全に2値化するBiDMと呼ばれる新しい手法を提案しました。

具体的には、著者らは時間的および空間的観点から拡散モデルを最適化しました。

時間的な観点から、「タイムステップフレンドリーなバイナリ構造」(TBS) が導入され、学習可能なアクティベーションバイナライザーとクロスタイムステップ機能接続を通じて、DM の高度にタイムステップ相関したアクティベーション機能に対処します。

空間的な観点からは、画像生成タスクとノイズ推定ネットワークの空間的局所性に特に注意しながら、2 値化された特徴のマッチングの難しさを解決することを目指して、「空間ブロック蒸留」(SPD) が提案されています。

実験結果によると、W1A1 BiDM は LSUN-Bedrooms 256×256 上の LDM-4 モデルで 22.74 の FID スコアを達成し、現在の最先端の一般的な 2 値化方式のスコア 59.44 を大幅に上回り、最大28 倍のストレージ節約と52.7 倍の運用効率向上を達成しました。

詳しく見てみましょう。

既存の拡散モデルは、主に量子化における「重み」に焦点を当てています。

拡散モデルは、その高品質と多様な生成能力により、画像、音声、動画など様々な分野で大きな注目を集め、応用されてきました。最大1000段階のノイズ除去処理を経て、ランダムノイズからデータを生成できます。

ただし、一部の高速サンプリング方法ではタスクの生成に必要なステップ数を効果的に削減できますが、時間ステップあたりの浮動小数点計算コストが高いため、リソースが制限されたシナリオではこのモデルの広範な適用が依然として制限されます。

そのため、拡散モデルの圧縮は、その幅広い応用に向けた重要なステップとなっています。既存の圧縮手法には、主に量子化、蒸留、そしてプルーニングが含まれます。これらの圧縮手法の目的は、精度を維持しながら、ストレージと計算コストを削減することです。

量子化は非常に効果的なモデル圧縮技術と考えられており、重みやアクティベーションを低ビット整数に量子化したり、2 値化したりすることで、コンパクトなストレージと推論における効率的な計算を実現します。

そのため、いくつかの研究では、生成品質を維持しながらモデルの圧縮と加速を実現するために、拡散モデルに量子化を適用しています。

1ビット量子化（または2値化）は、モデルのストレージスペースを最大限節約でき、畳み込みニューラルネットワーク（CNN）などの識別モデルにおいて優れたパフォーマンスを発揮します。さらに、完全な2値化のように、重みと活性化が同時に1ビットに量子化される場合、XNORやビットカウントなどの効率的なビット演算が行列乗算に代わるため、最も効率的な高速化を実現できます。

いくつかの既存の研究では、拡散モデルを 1 ビットに量子化しようと試みていますが、それらの研究は主に重みに焦点を当てており、完全な 2 値化にはほど遠いものです。

実際、拡散モデルのような生成モデルの場合、重みと活性化を完全に2値化すると、その影響は悲惨なものになります。

a)生成モデルとして、拡散モデルの豊富な中間表現は時間ステップと密接に関連していますが、2 値化された重みと活性化を使用する場合、高度に動的な活性化範囲は大幅に制限されます。

b)拡散モデルのような生成モデルは、通常、完全な画像を出力することを必要としますが、高度に離散的なパラメータと特徴空間のため、学習中にそれらを真の値と一致させることが困難です。離散空間における最適化の難しさや、時間ステップダイナミクスに関連する表現力の不足により、二値化された拡散モデルは収束に失敗し、最適化中に崩壊することさえあります。

BiDM を提案: 「重み」と「活性化」を完全に二値化する

前述の欠点に対処するため、著者らはBiDMを提案した。重みと活性化を完全に二値化することで、拡散モデルの圧縮と加速の限界まで押し上げた。

要約すると、BiDM は、完全な 2 値化によってもたらされる課題を克服しながら、活性化機能、モデル構造、生成タスクの観点から拡散モデルの固有の要件に対処することを目指しています。

2つの革新的なテクノロジーが組み込まれています。

時間的な観点から、著者らは、拡散モデルの活性化特性が時間ステップと高度に相関していることを観察しました。

そこで、 「タイムステップフレンドリーなバイナリ構造」（TBS）が導入されました。これは、学習可能な活性化バイナリ量子化器を介して拡散モデルの動的な活性化範囲に適合し、隣接するタイムステップ間の特徴量の類似性を利用してバイナリモデルの表現力を高めるために、タイムステップ間での特徴量の接続を設計します。

空間的な観点から、著者らは生成タスクにおける空間的局所性と拡散モデルで使用される畳み込みベースの U-Net 構造に注目しました。

そこで、 「空間蒸留」（SPD）法が提案されました。これは、教師として完全精度モデルを導入する手法です。ブロックの自己注意を模倣することで、局所的な特徴に焦点を当て、バイナリ拡散モデルの最適化方向をより適切に導きます。

△時間ステップフレンドリーなバイナリ構造と空間分割蒸留を備えたBiDM

広範囲にわたる実験により、既存の最先端の完全二値化手法と比較して、BiDM は同じ推論効率を維持しながら精度を大幅に向上させ、さまざまな評価基準で既存のすべてのベースライン手法を上回ることが示されています。

具体的には、ピクセル空間拡散モデルでは、BiDM は IS を 5.18 まで向上させる唯一の方法であり、これはフル精度モデルのレベルに近く、最良のベースライン方法よりも 0.95 高くなります。

LDM では、BiDM により、最先端のアプローチによる LSUN ベッドルームの FID が 59.44 から 22.74 まで大幅に削減され、ストレージ スペースは 28.0 倍、OP は 52.7 倍節約されます。

最初の完全に二値化された拡散モデル法として、生成されたサンプルの多数は、BiDM が現在、許容可能な完全に二値化された DM 画像を生成できる唯一の方法であることも証明しており、これにより、DM を低リソースのシナリオに効果的に適用できるようになります。

具体的にはどのように実装されたのでしょうか?

バイナリ拡散モデルのベースライン

拡散モデル: データ分布0~(0)が与えられた場合、順方向プロセスは、遷移カーネル(|-1)（通常はガウス摂動）を通じて生成されるランダム変数のシーケンス∊{1, ...,}を生成します。その表現は次のとおりです。

ここで、∊(0,1)はノイズスケジュールである。ガウス型転送カーネルを用いることで、著者らは結合分布を周辺化することができ、以下のように利用することができる：=1-および

これは、ガウスベクトル ~N(0, I) をサンプリングし、変換を適用することで実現できます。

サンプルは簡単に入手できました。

逆のプロセスは、ノイズを除去してサンプルを生成することを目的としており、学習可能な遷移カーネル（-1|）を使用して利用できない条件付き分布q（-1|）を近似します。これは次のように表されます。

平均

および分散

再パラメータ化技術を使用すると、次のものを取得できます。

ここでは、学習可能なパラメータを持つ関数近似を使用して予測を行います。

拡散モデルのトレーニングでは、サンプルの品質を向上させるために、変分下限値の簡略化された変種が損失関数としてよく使用されます。その表現は次のとおりです。

U-Netは低次元と高次元の特徴を融合できるため、拡散モデルの主流のバックボーンとなっています。U-Netの入力ブロックと出力ブロックは次のように表現できます…

そして

より小さいmはより低レベルのブロックに対応する。スキップ接続は低レベルの情報をDm(·)からUm(·)へ伝播するため、Umが受信する入力は次のように表される。

二値化：量子化は、重みと活性化変数を低ビット幅に離散化することで、ノイズ推定モデルを圧縮・高速化します。二値化拡散モデルのベースラインでは、重みWは1ビットに二値化されます。

符号関数は W を +1 または -1 に制限し、しきい値として 0 を設定します。は浮動小数点スカラーで、次のように初期化されます...

(n は重みの数を表します)、トレーニング プロセス中に学習されます。

一方、活性化は通常、単純な BNN 量子化器を使用して量子化されます。

重みとアクティベーションの両方が 1 ビットに量子化されると、ノイズ除去モデルの計算を XNOR とビットカウント演算に置き換えることができるため、大幅な圧縮と高速化が実現します。

時間ステップに優しいバイナリ構造

提案された方法を詳しく説明する前に、著者はまず拡散モデル (DM) の特性に関する観察を要約します。

下の図では、(a) CIFAR-10 上の完全精度 DDIM モデルの 4 番目の畳み込み層の活性化範囲がノイズ除去ステップによって変化しています。(b) LSUN-Bedrooms データセットでは、各ステップでの完全精度 LDM-4 モデルの出力特徴は、前のステップのものと同様です。

観察 1 : 長い時間ステップでは活性化範囲が大きく変化しますが、隣接する短い時間ステップでは活性化特性は同様です。

TDQ や Q-DM などのこれまでの研究では、DM の活性化分布はノイズ除去中の時間ステップに大きく依存しており、図 2(a) に示すように、隣接する時間ステップ間では類似性が見られ、離れた時間ステップ間では大きな差が見られることが示されています。

したがって、すべてのタイムステップにわたって固定のスケーリング係数を適用すると、活性化範囲に深刻な歪みが生じます。Deepcacheは、分布範囲に加えて、図2(b)に示すように、連続するタイムステップにわたる高次元特徴量の時間的一貫性を重視します。

これらの現象により、著者らは既存の2値化構造を再検討することになった。

二値化、特に重みと活性化の完全な二値化は、4ビット量子化などの低ビット量子化と比較して、活性化の範囲と精度の損失が大きくなります。これにより、豊富な活性化特徴量を生成することがはるかに困難になります。活性化の範囲と出力特徴量の不十分さは、DMのような豊富な表現を持つ生成モデルに深刻な悪影響を及ぼします。

したがって、より柔軟な活性化範囲バイナライザーを採用し、その特徴出力を活用してモデルの全体的な表現力を高めることは、完全なバイナライゼーション後の生成能力を向上させるための重要な戦略です。

著者らはまず、長期的な時間ステップ間の相違に焦点を当てています。

BNNやBi-Realなどの既存の活性化量子化器のほとんどは、式(7)に示すように、活性化を{+1, -1}に直接量子化します。この方法は活性化特性に深刻な摂動を与え、生成モデルの表現力に悪影響を及ぼします。XNOR++などの改良された活性化二値量子化器では、学習可能なスケーリング係数kが用いられます。

Kはベクトルまたは複数のベクトルの積ですが、推論中は一定値のままです。この手法は活性化の特徴表現を部分的に復元しますが、相関の高い時間ステップとは一致せず、依然として大幅なパフォーマンスの低下につながる可能性があります。

著者らは、動的に計算された平均値を用いて活性化二値量子化器を構築するオリジナルのXNORに注目する。2次元畳み込みにおけるその演算は次のように表される。

で

、

、

、

。

これは 2 次元フィルターを表します。

そして

これらはそれぞれ、乗算ありと乗算なしの畳み込み演算を表しています。この手法は、活性化特徴量の範囲を自然に保ち、異なる時間ステップで入力範囲を動的に調整します。

しかし、DM 機能の豊富な表現により、モジュールが通過する前と後の範囲にわたってローカル アクティベーションに不一致が見られ、事前に決定された k 値ではアクティベーション表現を効果的に回復できないことがわかります。

そのため、著者らはkを調​​整可能にし、訓練中にkが学習できるようにすることで、訓練前後の活性化範囲の変化に適応的に適応的に対応できるようにした。著者らが提案する学習可能なマイクロ畳み込みkの勾配計算プロセスは以下の通りである。

注目すべきは、kを学習可能にすることで推論の負担が増大しないことです。計算コストは​​変わらないため、効率的な二項演算が可能になります。

一方、著者らは隣接する時間ステップ間の類似性に注目した。

Deepcacheは、高次元の特徴量をキャッシュとして直接抽出し、U-Netにおける膨大な深層計算を省略することで、推論の大幅な高速化を実現します。このプロセスは以下のように表されます。

しかし、このアプローチはバイナリ拡散モデルには適していません。バイナリネットワークの各出力の情報量は非常に限られているためです。バイナリ拡散モデルは本質的に大幅な圧縮と高速化を実現しているものの、表現力には限界があります。著者らは、隣接する時間ステップ間の特徴量の類似性によってバイナリ表現が向上し、表現力の課題を補えると期待しています。

著者らは、現在のタイムステップの表現を前のタイムステップの特徴で豊かにするために、タイムステップをまたがる情報強化接続を構築した。このプロセスは以下のように表される。

で

これは学習可能なスケーリング係数です。図2(b)に示すように、高次元特徴量の類似性は、ブロックや時間ステップによって異なります。

そのため、著者らは、トレーニング中にモデルがより効果的に学習できるように、いくつかの独立した値を設定しました。

要約すると、タイムステップフレンドリーなバイナリ構造（TBS）は、入力の平均化とタイムステップ間の連結後にスケーリング係数に適用される学習可能なマイクロ畳み込みから構成されます。これらの複合効果により、長いタイムステップにわたる拡散モデルの活性化範囲の変化に適応し、隣接するタイムステップ間の高次元特徴の類似性を活用して情報表現を向上させることができます。

エラー削減の観点から、図3にTBSを視覚化したものを示します。

まず、著者らはベースライン法によるバイナリ拡散モデルの出力をベクトル Bt-1 に抽象化します。

スケーリング係数の不一致により、スケーリング係数と完全精度モデルの出力ベクトル Ft-1 との間に大きな長さの差が生じる可能性があります。

提案されたスケーリング係数と学習可能なマイクロ畳み込みを用いることで、Bt-1はLt-1に拡張されます。Lt-1はFt-1に近づきますが、それでも完全精度モデルとは方向性が異なります。

クロスタイムステップ接続は、前のタイムステップの出力Ft、Bt、Ltをさらに結合します。隣接するタイムステップ間の高次元特徴類似性により、Ft-1とFtの差は比較的小さくなり、Lt-1とLtの結合が容易になります。

最後に、著者らはTBSを適用することで、バイナリ拡散モデルの出力を得ました。これはTt-1=(1-)·Lt-1+·Ltと表されます。この出力は、完全精度モデルの出力Ft-1に最も近い値です。

空間蒸留

生成モデルの特性により、拡散モデルの最適化プロセスは、従来の識別モデルの最適化プロセスとは異なります。

観察 2 : 従来の蒸留では、完全に 2 値化された拡散モデルと完全精度モデルの整合を導くことが困難ですが、生成タスクにおける拡散モデルの機能は空間的な局所性を示します。

従来、量子化モデルの学習中に蒸留損失を追加することは一般的な手法でした。しかし、バイナリモデルの数値空間は有限であるため、単純な損失を最適化に直接適用すると、勾配更新方向の調整が困難になり、学習が困難になります。そのため、中間特徴量に蒸留損失を追加することで、モデルの局所的および大域的な最適化プロセスをより適切に導くことができます。

しかし、生成モデルとして、拡散モデルの非常に豊富な特徴表現により、バイナリ モデルが完全精度モデルを細かく模倣することは非常に困難になります。

オリジナルの拡散モデルの学習で使用されたL2ロスは、拡散プロセス中にガウスノイズと整合しますが、中間特徴の蒸留マッチングには適していません。従来の蒸留プロセスでは、一般的に使用されるL2ロスは、差異の大きいピクセルの最適化を優先する傾向があり、より均一で滑らかな最適化結果をもたらします。

このグローバル制約学習プロセスは、画像生成を目的としたバイナリ モデルでは困難です。これは、表現力が限られているため、細かい蒸留模倣を通じてモデルを直接調整して、完全精度モデルの向きに完全に一致させることが困難だからです。

同時に、著者らは、U-Net をバックボーンとして使用する拡散モデルは、畳み込みベースの構造と生成タスクの要件により、自然に空間的局所性を示すと指摘しました。

これは、分類などのタスクでは大域的な特徴抽出のみが必要であり、低レベルの要件を必要としない従来の識別モデルとは異なります。そのため、従来の蒸留手法は空間局所性を伴う生成タスクには適していません。さらに、既存の拡散モデル蒸留手法のほとんどは、時間ステップ数の削減に重点を置いており、画像生成タスクにおける特徴の空間局所性の必要性には対応していません。

したがって、既存の損失関数を使用してバイナリ拡散モデルを最適化することの難しさや、拡散モデルの空間的局所性を考慮して、著者らは空間ブロック蒸留 (SPD) を提案します。

具体的には、蒸留前に特徴量を小さなブロックに分割し、ブロックごとに空間自己注意を計算する新しい損失関数を設計しました。従来のL2損失関数では、バイナリ拡散モデルへの直接マッチングが困難で最適化の課題につながりますが、注意メカニズムにより、蒸留最適化においてより重要な部分に焦点を当てることができます。

しかし、これは完全に 2 値化された拡散モデルにとっては依然として課題であり、非常に離散的な 2 値出力には限られた情報しかなく、モデルがグローバル情報を取得することが困難になります。

そのため、著者らは中間特徴を複数の小さなブロックに分割し、各ブロックに対して独立して空間自己注意を計算し、バイナリモデルが最適化プロセス中にローカル情報をより有効に活用できるようにしました。

SPDはまず、バイナリ拡散モデルと完全精度拡散モデルの特定のブロックの中間特徴biとfp∊bxcxwxh出力をp2個の小さなブロックに分割します。

次に、各小さなブロックの自己注意を個別に計算します。

正規化後、対応する位置での損失を計算して合計します。

ここで、||·||² はL2関数を表します。最終的に、総トレーニング損失は次のように計算されます。

ここで、d はアップサンプリングまたはダウンサンプリング プロセス中のブロックの数を表し、最終的には中間ブロックを含む 2d+1 個の中間特徴が生成されます。は損失項のバランスをとるために使用されるハイパーパラメータ係数であり、デフォルト値は 4 です。

下の図は、LSUN-bedroomデータセットにおけるLDMモデルの最後のTimeStepBlockの出力を視覚化したものです。FP32はフル精度モデルの出力fpを表します。Diffはフル精度モデルの出力と2値化モデルの出力||fp-bi||の差を表します。これは自己注意メカニズムを表しています。

著者らは、前述の中間特徴と自己注意を視覚化しました。図に示すように、空間ブロック蒸留（SPD）により、モデルは各ブロック内のローカル情報にさらに注意を払うようになります。

実験

著者らは、CIFAR-10 32×32、L​​SUN-Bedrooms 256×256、LSUN-Churches 256×256、FFHQ 256×256 などのさまざまなデータセットで実験を行い、ピクセル空間拡散モデルと潜在空間拡散モデルでテストしました。

使用された評価指標には、インセプションスコア（IS）、フレシェ・インセプション距離（FID）、スライディング・フレシェ・インセプション距離（sFID）、適合率、再現率が含まれていました。現在までに、拡散モデルをこれほど極端に圧縮した研究はありません。

そのため、著者らは、古典的な二値化アルゴリズムと最新の最先端の汎用二値化アルゴリズムをベースラインとして使用しました。

彼らは拡散モデルからTimestepEmbedBlocksの出力を抽出し、TBSとSPDの演算対象とした。また、畳み込み層にはReActNetと同じショートカット接続を採用した。詳細な実験設定は付録Aに示されている。

主な結果

ピクセル空間拡散モデル：著者らはまずCIFAR-10 32×32データセットを用いて実験を行いました。表1に示すように、ベースライン手法を用いた拡散モデルのW1A1二値化では、大幅な性能低下が見られました。しかし、BiDMは全ての指標において大幅な改善を示し、これまでにない画質回復を実現しました。

具体的には、BiDM は IS 指数を 4.23 から 5.18 に改善しましたが、FID 指数は 27.9% 減少しました。

以下は、CIFAR-10 データセットで 100 ステップを使用して DDIM 2 値化を行った結果です。

潜在的な空間拡散モデル:著者らの LDM 実験には、LSUN-Bedrooms 256×256 および FFHQ 256×256 データセットでの LDM-4 の評価と、LSUN-Churches 256×256 データセットでの LDM-8 の評価が含まれます。

実験では200ステップDDIMサンプラーを使用し、詳細な結果を表2に示します。

これら3つのデータセットにおいて、著者らの手法は最良のベースライン手法と比較して大幅な改善を達成しました。他の二値化アルゴリズムと比較して、BiDMはすべての指標で優れた性能を示しました。

LSUN-Bedrooms、LSUN-Churches、FFHQデータセットにおいて、BiDMはベースライン手法と比較して、それぞれFID指標を61.7%、30.7%、51.4%削減しました。XNOR++と比較すると、XNOR++はノイズ除去時に固定の活性化スケーリング係数を使用するため、活性化のダイナミックレンジが非常に限られており、拡散モデルの非常に柔軟な生成表現に一致させることが困難です。

BiDMは、マイクロ畳み込みkを学習可能にし、動的に計算されるスケーリング係数を適用することで、この課題に対処します。この最適化により、様々な指標において1桁以上の大幅な改善が実現します。

LSUN-Bedrooms データセットと LSUN-Churches データセットでは、BiDM の FID インデックスはそれぞれ 319.66 から 22.74 に、292.48 から 29.70 に減少しました。

さらに、BiDM は最先端の二値化手法 ReSTE と比較して複数の指標で大幅な改善を実現し、特に LSUN-Bedrooms データセットで顕著な改善を示しています。

アブレーション実験

著者らは、BiDM で提案された各コンポーネントの有効性を評価するために、LSUN-Bedrooms 256×256 データセットで LDM-4 の包括的なアブレーション研究を実施しました。

提案されたSPD法とTBS法の有効性を評価し、その結果を表3に示します。

LDMにSPD法またはTBS法を適用したところ、元の性能と比較して大幅な改善が見られました。TBS法を導入すると、FIDとsFIDはそれぞれ106.62と56.61から35.23と25.13に大幅に減少しました。

同様に、SPD法を追加すると、FIDとsFIDはそれぞれ40.62と31.61に大幅に減少します。他の指標も大幅な改善を示しています。これは、学習可能な因子を導入し、時間ステップ間の接続を追加することで、学習中に二値化モデルの特徴量をフル精度の特徴量に継続的に近似するという、著者らの手法の有効性を示しています。

さらに、これら2つの手法を組み合わせてLDMに適用すると、それぞれの手法を単独で適用した場合と比較して、さらなる改善が見られます。これは、パッチレベルで完全精度モデルとバイナリモデル間の蒸留を実行することで、バイナリモデルのパフォーマンスを大幅に向上できることをさらに示しています。

著者らは追加のアブレーション実験も実施しており、その結果は付録 B に示されています。

効率分析

著者らは、完全な2値化の下でのモデルの推論効率を分析しました。

表 4 の結果は、BiDM が XNOR ベースラインと同じ 28.0 倍のメモリ効率と 52.7 倍の計算節約を実現しながら、イメージ生成においてベースライン モデルを大幅に上回り、FID が 106.62 から 22.74 に減少したことを示しています。

要約

本論文では、拡散モデルの圧縮を限界まで押し上げる新しい完全二値化手法である BiDM を提案します。

彼らは、異なる時間ステップでの活性化特性と画像生成タスクの特性という 2 つの観察に基づいて、時間と空間の観点からそれぞれ時間ステップフレンドリーなバイナリ構造 (TBS) と空間ブロック蒸留 (SPD) 手法を提案しました。

これらの方法は、完全な二値化における表現能力の厳しい制限と、高度に離散的な空間最適化の課題に対処します。

最初の完全に二値化された拡散モデルとして、BiDM は複数のモデルとデータセットにおいて、現在の最先端の一般的な二値化手法よりも大幅に優れた生成パフォーマンスを実証しました。

LSUN-Bedroomsデータセットにおいて、BiDMはFID22.74を達成しました。これは、最先端手法のFID59.44を大きく上回る数値であり、視覚的に許容できるサンプルを生成できる唯一のバイナリ手法となっています。また、ストレージ容量は最大28.0倍、計算効率は最大52.7倍向上しています。

詳細は原論文をご参照ください。

論文リンク:
https://arxiv.org/abs/2412.05926
プロジェクトホームページ:
https://github.com/Xingyu-Zhe...

- 以上-