数行のコードでUNetが安定！中山大学などが提案したScaleLong拡散モデル：スケーリングへの疑問からスケーリングへ

標準の UNet 構造では、ロング スキップ接続のスケーリング係数は通常 1 です。

ただし、Imagen、スコアベースの生成モデル、SR3 などのよく知られている拡散モデルの研究では、すべてこれを設定しており、そのような設定によって拡散モデルのトレーニングを効果的に加速できることがわかりました。

スケーリングへの疑問ただし、Imagenや他のモデルの原著論文ではスキップ接続のスケーリング操作についての具体的な分析はありませんが、このような設定は拡散モデルのトレーニングを高速化するのに役立つと言われています。

まず第一に、この実証的なデモンストレーションでは、この設定がどのような役割を果たすのかが不明確になります。

また、設定できるのは だけなのか、それとも使用できる他の定数があるのか​​どうかもわかりません。

異なる場所のスキップ接続は同じ「ステータス」を持ちますか? なぜ同じ定数を使用するのですか?

著者はこれについて多くの疑問を抱いています...

写真

スケーリングを理解する

一般的に言えば、ResNet や Transformer 構造と比較すると、UNet は実際の使用ではそれほど深くなく、他の「ディープ」ニューラル ネットワーク構造でよくある勾配消失などの最適化の問題が発生しにくくなります。

さらに、UNet 構造の特殊性により、浅い特徴は長いスキップ接続を介して深い位置に接続されるため、勾配消失などの問題がさらに回避されます。

では逆に考えてみましょう。この構造に気をつけないと、更新によって勾配が過剰になったり、パラメータ（特徴量）が振動したりといった問題が起きるのではないでしょうか？

写真

トレーニング中の拡散モデルタスクの特徴とパラメータを視覚化すると、不安定性が存在することがわかります。

パラメータ（特徴）の不安定性は勾配に影響し、勾配はパラメータの更新に影響します。最終的に、このプロセスはパフォーマンスに悪影響を及ぼすリスクが高くなります。したがって、この不安定性を制御する方法を見つける必要があります。

さらに、拡散モデルの場合。 UNet の入力はノイズの多い画像です。モデルが追加されたノイズを正確に予測する必要がある場合、モデルは入力への追加の妨害に対して非常に堅牢である必要があります。

論文: https://arxiv.org/abs/2310.13545

コード: https://github.com/sail-sg/ScaleLong

研究者は、ロング スキップ接続でスケーリングを実行することで、これらの問題を均一に軽減できることを発見しました。

定理 3.1 から、中間層の特徴の振動範囲 (上限と下限の幅) は、スケーリング係数の二乗の合計と正の相関関係にあります。適切なスケーリング係数は、機能の不安定性を軽減するのに役立ちます。

ただし、スケーリング係数を直接 0 に設定すると、振動は実際に最も緩和されることに注意してください。 （手動犬頭）

しかし、UNet がスキップフリー状態に退化すると、不安定性の問題は解決されますが、表現能力は失われます。これは、モデルの安定性と表現能力の間のトレードオフです。

写真

同様に、パラメータ勾配の観点からも。定理 3.3 は、スケーリング係数による勾配の大きさの制御も明らかにしています。

写真

さらに、定理 3.4 は、長いスキップ接続のスケーリングが、入力摂動に対するモデルの堅牢な上限にも影響を与え、入力摂動に対する拡散モデルの安定性を向上させる可能性があることも明らかにしています。

スケーリングになる

上記の分析を通じて、安定したモデルトレーニングにはロングスキップ接続のスケーリングが重要であることがわかります。これは上記の分析にも当てはまります。

次に、どのようなスケーリングがより良いパフォーマンスを達成できるかを分析します。結局のところ、上記の分析では、スケーリングが有益であることを示すことしかできず、どのようなスケーリングが最良または優れているかを判断することはできません。

簡単な方法は、ロングスキップ接続に学習可能なモジュールを導入して、スケーリングを適応的に調整することです。この方法は、学習可能スケーリング (LS) メソッドと呼ばれます。以下に示すように、SENet に似た構造を採用します (ここでは、非常に優れたコード構成を備えた U-ViT 構造が考慮されています。素晴らしい!)

写真

この記事の結果から判断すると、LS は確かに拡散モデルのトレーニングを効果的に安定化することができます。さらに、LS で学習した係数を視覚化してみます。

下の図に示すように、これらの係数は指数関数的に減少傾向を示しており (ここでの最初のロング スキップ接続は、UNet の先頭と末尾を接続する接続を指すことに注意してください)、最初の係数はほぼ 1 に近く、これも驚くべきことです。

写真

この一連の観察に基づいて（詳細については論文を参照してください）、学習可能なパラメータを必要としない定数スケーリング（CS）法をさらに提案しました。

CS 戦略では、最初に使用されるスケーリング操作のような追加のパラメータは必要ないため、追加の計算コストは​​ほとんど発生しません。

CS はほとんどの場合、安定したトレーニングでは LS ほど優れていませんが、既存の戦略では試してみる価値はあります。

上記の CS と LS の実装は非常に簡潔で、数行のコードしか必要としません。さまざまな UNet 構造では、フィーチャの寸法を揃える必要がある場合があります。 （手動犬頭+1）

最近では、FreeU や SCEdit などの後続の研究でも、スキップ接続でのスケーリングの重要性が明らかになっています。ぜひ皆さんも試して、推進してください。