JA EN
体系数値計算
·無料·14分で読めます

自動微分の仕組み — PyTorchの魔法を1から

loss.backward() と書くだけで何百万個ものパラメータの微分が出てくるのはなぜか。計算グラフ・連鎖律・前進モードと後退モードを前提知識ゼロから解き、40行のミニautogradまで自作します。

対象textタスクalgorithm

「微分してくれる」という不思議

深層学習のコードで、いちばん不思議な行はここです。

loss = F.mse_loss(model(x), t)
loss.backward()          # これだけで全パラメータの勾配が出る

model の中には行列積も活性化関数も正規化も何十層と積み重なり、パラメータは軽く数百万個。それなのに、微分の式はどこにも書かれていません。誰も d(loss)/d(w_483921) を導出していないのに、実行すればすべての w.grad に正しい数値が入ります。

結論を先に言うと、そこに魔法はありません。あるのは連鎖律を、機械が忘れずに最後まで適用しているだけです。ただし「機械が」の部分に工夫が詰まっています。

比喩: 領収書を残しながら料理する

レストランで一皿の原価を出す場面を想像してください。完成した皿を見て推測もできますが、精密にやるなら調理の途中で工程ごとに領収書を残すほうが確実です。「玉ねぎを2個使い、うち1個は下処理で半分捨てた」という記録があれば、後から逆にたどって「玉ねぎが10円値上がりしたら原価はいくら上がるか」に正確に答えられます。

自動微分はまさにこれです。順方向に計算しながら、各工程の記録(どの値からどの値を、どんな演算で作ったか)を残す。計算が終わったら、その記録を逆にたどって影響を配り直す。 記録を計算グラフ、逆にたどることを後退(backward)と呼びます。

微分を手に入れる3つの方法

そもそも、なぜ「自動」微分という専用技術が要るのでしょうか。

記号微分は、数式処理システムに f(x)f(x) を渡して f(x)f'(x) の式そのものを出させます。正確ですが、合成が深くなると式が爆発的に膨れ上がり(式膨張)、50層のネットワークでは書き下すだけで破綻します。

数値微分は、定義どおり少しずらして差を取ります。

f(x)f(x+h)f(x)hf'(x) \approx \frac{f(x+h) - f(x)}{h}
(1)

「入力を hh だけ動かしたとき出力がどれだけ動いたか」を hh で割る、というだけの式です。実装は3行で済むのに、2つの理由で実用になりません。精度: hh を小さくすると式(1)の近似は良くなりますが、f(x+h)f(x+h)f(x)f(x) はほぼ同じ値なので、引き算で有効数字が消し飛びます(桁落ち)。この綱引きは数値の落とし穴 — 桁落ち・丸め・logsumexpで扱っています。回数: パラメータが nn 個なら n+1n+1 回モデルを走らせる必要があり、nn が7,000万なら勾配1回に7,000万回の順伝播。話になりません。

自動微分は、式を展開もせず差分も取らず、プログラムの実行そのものに連鎖律を機械的に適用します。

連鎖律という、たった1本のルール

自動微分の全部が、この1本の式に乗っています。

dzdx=dzdydydx\frac{dz}{dx} = \frac{dz}{dy} \cdot \frac{dy}{dx}
(2)

xx から yy を作り、yy から zz を作ったとき、xxzz に与える影響は途中の影響のかけ算になる、と言っています。dz/dydz/dy は「yy が1動くと zz がどれだけ動くか」、dy/dxdy/dx は「xx が1動くと yy がどれだけ動くか」。歯車が2枚噛み合っていて、歯車比をかければ全体の比になるのと同じです。

重要なのは、この式が何段でも伸びることです。そして各段の比は、その段でやった演算(足し算、かけ算、exp、行列積…)ごとにあらかじめ分かっています。つまり演算の種類ごとに微分のルールを1回だけ登録しておけば、あとはどんな組み合わせでも機械が処理できる。これが自動微分の核心です。

計算グラフ: プログラムを部品に割る

1変数の線形回帰の損失で見ます。

L=(wx+bt)2L = (wx + b - t)^2

これをプリミティブ(それ以上分けない基本演算)に割ります。

u = w * x        # かけ算
v = u + b        # 足し算
e = v - t        # 引き算
L = e * e        # かけ算

各行がノード、値の受け渡しが辺です。これが計算グラフで、順方向に計算しながら組み立てられます。各ノードは、自分の出力を自分の入力で微分した値(局所微分)を知っています。かけ算 u=wxu = w \cdot x なら u/w=x\partial u/\partial w = x、足し算 v=u+bv = u + b なら v/u=1\partial v/\partial u = 1。隣だけを見た比なので、そのノード単独で求まります。

あとは連鎖律で端から端までかけ算するだけ。L/w\partial L/\partial w2e×1×1×x2e \times 1 \times 1 \times x になります。

前進モードと後退モード

端から端まで比をかけ算するとき、どちら側から始めるかという選択肢があります。ここが自動微分でいちばん面白いところです。

前進モードは入力側から始めます。「ww が1動いたら」を出発点に、各ノードで「じゃあ自分はどれだけ動くか」を計算しながら順方向に進みます。値と一緒に微分を運ぶ二重数a+bεa + b\varepsilonε2=0\varepsilon^2 = 0 とする数)を使えば、通常の演算を書き直すだけで済み、記録を溜める必要もありません。

後退モードは出力側から始めます。まず順伝播で計算グラフを作り、それから「LL が1動くには」を出発点に、逆向きに各ノードへ影響を配って回ります。

どちらも答えは同じで、違うのはコストです。入力が nn 個、出力が mm 個の関数を考えてください。

深層学習の損失は、入力(パラメータ)が数百万〜数十億、出力(スカラーの損失)が1個です。nn が巨大で m=1m = 1。後退モードが圧勝します。しかも後退モードで勾配を1本取るコストは、関数を1回評価するコストの定数倍に収まることが知られています(Baur–Strassenの定理、しばしば「安い勾配の原理」と呼ばれます)。パラメータが10倍になっても、勾配計算が10倍高くはなりません。

代償はメモリです。後退で局所微分を求めるには順伝播の途中の値が要るので、中間結果を保持したまま逆走することになります。学習時にGPUメモリを食うのは重みではなくこの中間活性で、バッチサイズを上げるとすぐOOMになる理由もここにあります。逃げ道が勾配チェックポイント(torch.utils.checkpoint)で、中間結果を捨てて後退時に再計算する、メモリと計算のトレードオフです。

とはいえ前進モードも無用ではありません。入力が少なく出力が多い場面(感度解析など)では前進が有利で、PyTorchにも torch.func.jvp / jacfwd があります。

backward() の中で起きていること

ここまで分かれば、PyTorchの挙動はすべて説明がつきます。

1. requires_grad=True に触れた演算は記録を残す。 結果テンソルに grad_fn(どの演算で作られたかを指す関数オブジェクト)がぶら下がります。y = w * x なら y.grad_fnMulBackward0。これが領収書です。

2. .backward() は出力に「1」を置いて始める。 L/L=1\partial L/\partial L = 1 が逆伝播の種火です。損失がスカラーでないと backward() が怒るのは、どこに1を置くか決まらないからです。

3. グラフを逆順にたどる。 各ノードは「上流から届いた勾配 × 自分の局所微分」を下流に渡します。訪問順はトポロジカル順序で、流れ込む勾配が全部揃ってから処理されます。分岐したテンソルに複数経路から勾配が届く場合、それらは足し合わされます

4. 葉テンソルの .grad に加算する。 nn.Parameter などの葉に結果が書かれますが、代入ではなく加算です。毎ステップ optimizer.zero_grad() が要るのはこのため。設計ミスではなく、勾配累積(小さいバッチを複数回回して大きいバッチに見せる技法)を素直に書けるようにするための仕様です。

5. 使い終わったグラフは解放する。 既定では逆伝播後にバッファが捨てられます。同じグラフに2回 backward() して怒られるのはこのためで、意図的に2回流すなら retain_graph=True を渡します。

つまり backward() は、領収書の束を逆順に読んで責任を配って回る集金人です。誤差逆伝播は、この後退モード自動微分をニューラルネットに適用した特別な場合にすぎません。層ごとの具体的な式は誤差逆伝播を1から解説 — 連鎖律だけで全部わかるにあります。

局所微分を体で確かめる

連鎖律でかけ算される「各段の比」は、活性化関数のところでいちばん劇的に効きます。関数を切り替えながら入力をドラッグしてみてください。曲線の傾きがその点の局所微分で、逆伝播ではその値が勾配にかけられます。

FIG 1曲線の傾きがそのまま局所微分。sigmoid を両端に引っ張ると傾きがほぼ0になり、そこを通った勾配が消えることが目で分かる

sigmoid の両端で傾きが潰れるのが見えたはずです。層を重ねるとこの小さい比が何度もかけ算され、勾配が指数的に小さくなる——これが勾配消失です。ReLUが好まれる理由も同じ図で分かります。正の側の傾きがちょうど1なので、何度かけても縮まないのです。

ついでに、ReLUには自動微分ならではの罠があります。x=0x=0折れているので、数学的にはそこに微分が存在しません。実装は劣勾配の中から1つを選んで返します(PyTorchのReLUは0を返します)。普段は実害がありませんが、torch.sqrt(0)x.norm() を原点で微分すると、局所微分が 1/(2x)1/(2\sqrt{x}) のように発散して、順伝播は正常なのに勾配だけ nan になります。順方向の値を見ていても気づけない種類のバグです。

40行で自動微分を作る

理屈が分かったら、自分で書くのがいちばん早い。スカラー版の後退モード自動微分は、これだけで動きます。

class Var:
    def __init__(self, value, parents=(), local=()):
        self.value = value
        self.parents = parents     # 入力側のVar
        self.local = local         # 各親に対する局所微分
        self.grad = 0.0

    def __add__(self, o):          # d(a+b)/da = 1, /db = 1
        return Var(self.value + o.value, (self, o), (1.0, 1.0))

    def __mul__(self, o):          # d(ab)/da = b, /db = a
        return Var(self.value * o.value, (self, o), (o.value, self.value))

    def backward(self):
        order, seen = [], set()
        def visit(v):              # トポロジカル順に並べる
            if id(v) in seen: return
            seen.add(id(v))
            for p in v.parents: visit(p)
            order.append(v)
        visit(self)

        self.grad = 1.0            # dL/dL = 1 が種火
        for v in reversed(order):  # 出力側から逆走
            for p, d in zip(v.parents, v.local):
                p.grad += v.grad * d   # 連鎖律 + 複数経路は加算

使ってみます。L=(wx+bt)2L = (wx + b - t)^2ww についての勾配です。

w, x, b, t = Var(2.0), Var(3.0), Var(1.0), Var(5.0)
e = w * x + b + Var(-1.0) * t
L = e * e
L.backward()
print(w.grad)      # 2 * e * x = 2 * 2.0 * 3.0 = 12.0

PyTorchのautogradは、これをテンソルに拡張し、演算の種類を数百個そろえ、GPUカーネルとメモリ管理を足したものです。骨格はいま書いた20行と変わりません。

勾配は何のためにあったか

手に入れた勾配は、パラメータを動かすために使います。「勾配の逆向きに学習率の分だけ進む」——勾配降下法です。下の図で学習率を上げると、勾配が正確でも更新が発散します。自動微分が正しいことと、学習がうまくいくことは別問題、という感覚があると現場で迷いません。勾配そのものの意味はAIのための微分 — 勾配は「どちらに動けば良くなるか」の矢印で扱っています。

FIG 2勾配は「どちらへどれだけ」の情報。学習率を上げると、微分が正確でも谷を飛び越えて発散する

現場ではこう使う

自動微分は「フレームワークに任せて終わり」に見えますが、実際にはMLエンジニアと研究者が週に何度も自分でグラフを触る部分です。

場面1: グラフを意図的に切る。 推論や評価は with torch.no_grad(): で囲みます。記録を残さない=中間結果を保持しないので、メモリと速度が大きく違います。学習中でも、GANの識別器を更新するときや疑似ラベルを作るときは .detach() で勾配の逆流を止めます。忘れると、止めたいはずの側までパラメータが動いて学習が静かに壊れます。損失は下がるのに評価が改善しない、という気づきにくい形で出ます。

場面2: メモリを再計算と引き換える。 OOMが出たら torch.utils.checkpoint.checkpoint でブロックを包み、中間活性を捨てて後退時に再計算させます。順伝播が1回分増えて遅くなる代わりに、バッチサイズを維持できます。

場面3: 勾配が nan になったとき。 torch.autograd.set_detect_anomaly(True) を有効にすると、nan を生んだ演算まで遡って教えてくれます(重いのでデバッグ時だけ)。自作の演算(torch.autograd.Function を継承した層)なら torch.autograd.gradcheck で数値微分と突き合わせます。数値微分は勾配を「求める」道具としては遅すぎますが、「検算する」道具としては今も現役です。

落とし穴を3つ。

面接で問われる形にすると「なぜ深層学習は前進モードではなく後退モードを使うのか」です。入力が多く出力が1つだから、そしてその代償としてメモリを払っているからの2点をセットで言えれば十分です。

まとめ

次は、このグラフをGPU上で効率よく回すために、中間結果がメモリ階層のどこに置かれるのかを見ていきます。

コメント

コメントにはログインが必要です