stop-gradient(停止勾配)とは?順伝播はそのまま、逆伝播だけを止める仕組みを徹底理解

ディープラーニングの論文を読んでいると、stop_gradient.detach()sg[·] といった表記によく出会います。自己教師あり学習のSimSiam、強化学習のDQN、VQ-VAE、GANの学習——一見バラバラなこれらの手法が、実は「ある経路の勾配だけをわざと止める」という共通のテクニックを使っています。これが stop-gradient(停止勾配) です。

普通に考えると不思議です。学習とは勾配を流してパラメータを更新することなのに、なぜわざわざ勾配を止めるのでしょう? 実は、停止勾配は「片側を動かない基準(目標)として扱う」ための道具で、これがあるおかげで、表現が一点に潰れる崩壊を防いだり、学習を安定させたり、微分できない操作を迂回したりできるのです。

stop-gradientを理解すると、次のような場面で効いてきます。

  • 現代手法のコードと数式が読める — SimSiam/BYOL、DQN、VQ-VAE、各種異常検知の sg[·] の意図がわかる
  • 学習が不安定なときの対処 — 「ここの勾配を止めれば安定するのでは」と設計できる

本記事の内容

  • stop-gradientとは何か(順伝播は保ち、逆伝播だけ止める)
  • 最小例とPyTorch(.detach())での勾配の実測
  • なぜ片側の勾配を止めると崩壊を防げるのか
  • 敵対・交互最適化/ターゲットネットワーク/ストレートスルー/EMAティーチャーへの応用

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

stop-gradientとは ― 順伝播はそのまま、逆伝播だけ止める

ニューラルネットの学習は2つの流れでできています。前向き(順伝播)に値を計算し、後ろ向き(逆伝播)に勾配を流してパラメータを更新する。stop-gradientは、このうち逆伝播の流れだけを、ある地点でせき止める操作です。順伝播の値は一切変えません。

stop-gradient:順伝播はそのまま逆伝播だけを途中で止める

図のように、損失 $L$ が2つの枝(A と B)から計算されるとき、枝Bに stop-gradient をかけると、Bの値はそのまま損失に使われるのに、逆伝播の勾配はBの手前で止まる。結果、パラメータが更新されるのは枝Aだけになります。記号では $\mathrm{sg}[\cdot]$ と書き、PyTorchでは .detach()、TensorFlowでは tf.stop_gradient() がこれにあたります。

ポイントは「$\mathrm{sg}[x]$ は、順伝播では $x$ そのもの、逆伝播では定数(勾配0)として振る舞う」ことです。数式で書けば、任意の $x$ について

$$ \mathrm{sg}[x] = x \quad (\text{順伝播}), \qquad \frac{\partial\, \mathrm{sg}[x]}{\partial x} = 0 \quad (\text{逆伝播}) $$

です。値は変わらないのに勾配だけ消える——この一見奇妙な操作が何をもたらすか、まず最小の例で確かめましょう。

最小例 ― detachで勾配が0になる

いちばん簡単な例として、$z = a \cdot b$ を考えます。通常の勾配は $\partial z/\partial a = b$、$\partial z/\partial b = a$ です。ここで $b$ に stop-gradient をかけて $z = a \cdot \mathrm{sg}[b]$ とすると、$b$ への勾配だけが0になります。PyTorchで実際に確かめます。

import torch

a = torch.tensor(3.0, requires_grad=True)
b = torch.tensor(4.0, requires_grad=True)

z = a * b                # 通常
z.backward()
print(a.grad, b.grad)    # -> 4.0, 3.0

a.grad = None; b.grad = None
z2 = a * b.detach()      # bを停止(detach)
z2.backward()
print(a.grad, b.grad)    # -> 4.0, None(=勾配0)

最小例:detachすると片側の勾配が0になる

実測の通り、通常は $\partial z/\partial a = 4$、$\partial z/\partial b = 3$ ですが、b.detach() にすると $\partial z/\partial b = 0$(PyTorchでは勾配が伝わらず None)になります。重要なのは $\partial z/\partial a$ は4のまま変わらないこと。つまり $z$ の値も、$a$ への勾配も影響を受けず、$b$ への勾配だけがピンポイントで消えたのです。

連鎖律で見ると明快です。本来 $\partial L/\partial b = (\partial L/\partial z)(\partial z/\partial b)$ ですが、$\mathrm{sg}$ が $\partial z/\partial b = 0$ にするので、後段の損失 $L$ が何であろうと $b$ への勾配は0で打ち切られます。逆に $a$ 側の経路は何も変わりません。「計算グラフの特定の枝だけをハサミで切る」イメージです。

実装メモ ― detach と no_grad の違い

混同しやすいので整理しておきます。x.detach()(stop-gradient)は計算グラフの一部分だけを切り離す操作で、同じ式の中の他の経路(上の例の $a$)には勾配が流れ続けます。一方、torch.no_grad()そのブロック全体で勾配計算を止めるもので、推論時など「どこにも勾配が要らない」場面で使います。stop-gradientの本質は「一部の経路だけを止め、残りは学習させる」点にあり、これは no_grad では表現できません。requires_grad=False(パラメータの凍結)とも別物で、こちらは特定テンソルを最初から学習対象外にする設定です。

この「片側だけ勾配を止める」操作が、なぜ強力なのか。最も有名な応用である「表現の崩壊防止」から見ていきます。

なぜ崩壊を防げるのか ― 自己教師あり学習の崩壊問題

自己教師あり学習では、ラベルなしで良い表現を学ぶために、同じ入力の2つの見え方(データ拡張した2ビュー)の埋め込みを近づけるという目標をよく使います。「同じものは似た表現になるべき」という自然な発想です。ところが、これには落とし穴があります。

表現崩壊:stop-grad無で定数へ潰れる/有で多様性を保つ

「2つの埋め込みを近づける」だけだと、すべての入力を同じ一点(定数)に写すのが最も簡単な答えになってしまいます。どんな入力も同じ点に潰せば、2ビューの距離は常に0で損失も0。これが表現の崩壊(collapse)で、情報を一切持たない使い物にならない表現です。

ここで stop-gradient が効きます。SimSiamやBYOLは、損失の片側の枝に stop-gradient をかけ、それを「動かない目標」として扱うことで崩壊を防ぎます。

なぜ崩壊を防げるか:片側を動かない目標にする

直感はこうです。stop-gradientが無いと、両側の枝が同時に「相手に近づこう」として、いちばんの近道である一緒に0へ潰れる方向に進みます。一方、片側に stop-gradient をかけると、その側は今この瞬間の固定された目標になり、もう片側だけがそれを追いかけます。追いかける先が0でない限り、両者そろって潰れることが起きにくくなる——これが崩壊を防ぐ仕組みです。

別の見方をすると、stop-gradientは最適化問題の構造を変えます。両側に勾配を流すと「2つの表現を互いに近づける」問題ですが、片側を止めると「動く側を固定された目標へ近づける」問題になります。前者は崩壊という自明解への勾配が常に存在しますが、後者ではその経路が断たれます。SimSiamの論文は、stop-gradientを外すと実際に崩壊が起き性能が地に落ちること、そして予測器(predictor)と組み合わせることで崩壊を避けつつ良い表現が学べることを実験で示しています。なお本記事の概念図は仕組みを示すための模式図で、実際の崩壊の有無はモデル構造(予測器やBatchNormの有無)に依存します。

「片側を動かない基準にする」というこの考え方は、崩壊防止以外にも幅広く使われます。

応用① 敵対・交互最適化 ― 片方を固定して更新する

2つのモデルが競い合う、あるいは交互に最適化する場面でも stop-gradient が登場します。代表例がGAN(生成器と識別器の競争)です。片方を更新するとき、もう片方を stop-gradient で固定すると、「相手は今は動かない」と見なした上で自分だけを最適化できます。

交互最適化・敵対学習:片方を固定して更新

このような min-max 最適化や交互最適化では、両者の勾配が混ざると学習が不安定になりがちです。stop-gradientで片方を切り離せば、各ステップを「相手固定の単純な最適化」に分解でき、安定します。時系列異常検知などで使われる minimax 型の学習(たとえばAnomaly Transformerが prior と series を交互に最適化する設計)も、この一般形の一例です。

次は、強化学習でおなじみのターゲットネットワークです。

応用② ターゲットネットワーク ― 目標を固定して発散を防ぐ

強化学習のDQNでは、Q関数を「自分自身を使って作った目標」に近づけます。ここに停止勾配を使わないと、目標が自分を追いかけて逃げ続け、学習が振動・発散してしまいます(自分で動く的を撃つようなものです)。

ターゲットネットワーク(DQN):目標を固定し発散を防ぐ

そこでTD目標(学習が目指す右辺)に stop-gradient をかけ、「動かない教師」として固定します。これにより、各更新は「固定された目標への回帰問題」になり、安定します。多くの強化学習アルゴリズムが採る「セミ勾配(semi-gradient)」は、まさにこの停止勾配の考え方です。

停止勾配は、勾配を「止める」だけでなく、勾配を「迂回させる」のにも使えます。

応用③ ストレートスルー推定 ― 微分できない操作を飛び越す

VQ-VAEのように、途中に量子化(四捨五入や最近傍コードへの割り当て)といった微分できない操作を挟むと、勾配がそこで完全に途切れ、手前のエンコーダが学習できなくなります。これを救うのがストレートスルー推定(straight-through estimator)で、停止勾配を巧みに使います。

ストレートスルー(VQ-VAE):微分不能な操作を勾配が迂回

トリックは次の式です。量子化を $q(z)$ とするとき、

$$ z_q = z + \mathrm{sg}[\,q(z) – z\,] $$

とします。順伝播では $\mathrm{sg}$ は中身そのままなので $z_q = z + (q(z) – z) = q(z)$、ちゃんと量子化された値になります。逆伝播では $\mathrm{sg}[\cdot]$ の勾配が0なので $\partial z_q/\partial z = 1$、つまり量子化を飛び越して勾配がそのままエンコーダへ通ります。実際に、微分できない量子化を挟んだオートエンコーダで効果を確かめます。

import torch, torch.nn as nn

def quantize(z):
    return torch.round(z * 3) / 3          # 微分不能(勾配は至る所0)

# 学習ループ(抜粋): ストレートスルーの有無で比較
z = enc(x)
zq = quantize(z)                            # ST無: 勾配がここで途切れる
# zq = z + (quantize(z) - z).detach()       # ST有: 前向きは量子化、勾配は素通り
rec = dec(zq)
loss = ((rec - x)**2).mean()

ストレートスルー実測:停止勾配で量子化を迂回し学習が通る

実測の学習曲線です。ストレートスルー無し(赤)では、量子化で勾配が途切れてエンコーダにまったく勾配が届かず、再構成損失が早々に頭打ちになります(デコーダだけが固定された乱数コードに当てはめている状態)。ストレートスルー有り(緑)では、停止勾配で勾配を迂回させたおかげでエンコーダまで学習が届き、損失がさらに下がります。「微分できない操作を、停止勾配で勾配だけ通す」——これがストレートスルーの正体です。

最後に、教師モデルを固定するパターンです。

応用④ EMAティーチャー ― 教師側を停止勾配する

BYOLや平均教師(Mean Teacher)などでは、生徒モデル教師モデルを用意し、生徒を教師の出力に近づけます。このとき教師側には stop-gradient をかけ、勾配では学習させません。代わりに教師の重みは、生徒の重みの指数移動平均(EMA)でゆっくり更新します。

EMAティーチャー(BYOL):教師側を停止勾配

教師を停止勾配で固定することで、生徒は「安定したお手本」を追いかけられ、ここでも崩壊や不安定化を避けられます。教師がゆっくり動く(EMA)だけで勾配では動かない、という設計がポイントです。

stop-gradientが効く場面マップ

ここまで見てきた応用を一枚に整理します。

stop-gradが効く場面マップ

バラバラに見える応用の共通点は「片側を動かない基準(目標)にする」ことです。自己教師あり学習では崩壊を防ぐ基準に、敵対学習では固定する相手に、強化学習では追いかける目標に、VQ-VAEでは勾配を通す抜け道に、そして多変量時系列の異常検知(たとえばSARADのように2つの空間を同時に学習し片側を保護する設計)でも、この同じ道具が使われます。停止勾配を知っていると、これら全部が「同じ発想の言い換え」だと見えてきます。

まとめ

本記事では、stop-gradient(停止勾配)の意味と応用を解説しました。

  • stop-gradientは順伝播の値は保ったまま、逆伝播の勾配だけを0にする操作(.detach() / sg[·]
  • 最小例で確認した通り、$z=a\cdot\mathrm{sg}[b]$ は $z$ の値も $a$ への勾配も変えず、$b$ への勾配だけを消す
  • 共通の役割は「片側を動かない基準にする」こと
  • 自己教師あり(SimSiam/BYOL): 表現の崩壊を防ぐ
  • 敵対・交互最適化(GAN/minimax): 片方を固定して安定化
  • 強化学習(DQNターゲット): 目標を固定して発散を防ぐ
  • ストレートスルー(VQ-VAE): 微分できない操作を勾配が迂回
  • EMAティーチャー(BYOL/平均教師): 教師を固定する

「勾配を止める」は一見、学習に逆行する操作に見えますが、何を学習させ、何を固定するかを設計するための強力な道具です。論文で sg[·].detach() を見かけたら、「ここで片側を基準にして、こちらだけを動かしたいのだな」と読めるようになります。

次のステップとして、以下の記事も参考にしてください。