【実装】自動微分を自作する — 100行のミニPyTorch

Valueクラス1つから始めて、演算子オーバーロード・トポロジカル順序・勾配の加算までを組み上げ、その上にニューラルネットを載せて学習させます。設計の理由を辿ると、zero_grad() や retain_graph が仕様の暗記ではなく必然に変わります。

読んで分かった気になるものと、動かして分かるもの

loss.backward() の中身は連鎖律の機械的な適用にすぎない——理屈は自動微分の仕組み — PyTorchの魔法を1からで扱いました。ただ、説明を読んで納得するのと、自分の書いたクラスが w.grad に正しい数字を入れるのを画面で見るのとでは、残るものがまるで違います。

ここではスカラー1個だけを扱う自動微分エンジンを、設計の理由を1つずつ確かめながら組み上げます。最終形は100行ちょっと。その上に小さなニューラルネットを載せ、損失が下がるところまで持っていきます。同じ路線の実装ではKarpathyのmicrogradが広く知られていて、本記事の骨格もその系統です。

作りながら分かるのは、実は使う側でずっと引っかかっていた疑問の答えです。なぜ毎ステップ zero_grad() が要るのか。なぜ勾配は代入ではなく加算なのか。なぜ同じグラフに2回 backward() すると怒られるのか。内側に立つと、これらは仕様の暗記ではなく「そう作るしかなかった」に変わります。

ノードは、小さな配達員だと思う

計算に出てくる値ひとつひとつを Value というオブジェクトにします。2.0 も、w * x の結果も、最終的な損失も、全部 Value です。

各 Value が持つのは3つ。自分の値、自分に届いた勾配、そして届いた勾配を親（自分を作った入力）にどう配るかという手続きです。

3つ目が肝です。オブジェクトは「自分がどう作られたか」しか知らず、グラフ全体のことは何も知りません。上流から勾配が届いたら局所微分を掛けて親に渡す、それだけを実行する配達員です。出力側から順に起こしていくと勾配が末端まで行き渡る——誰も全体像を持たないのに全体が回る、というのが自動微分の設計の気持ちよさです。

骨格: 足し算と掛け算だけの Value

class Value:
    def __init__(self, data, _children=(), _op=''):
        self.data = data
        self.grad = 0.0
        self._backward = lambda: None   # 勾配の配り方（既定は何もしない）
        self._prev = set(_children)     # 自分を作った入力
        self._op = _op                  # 表示用のラベル

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), '+')
        def _backward():
            self.grad += out.grad       # d(a+b)/da = 1
            other.grad += out.grad      # d(a+b)/db = 1
        out._backward = _backward
        return out

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), '*')
        def _backward():
            self.grad += other.data * out.grad   # d(ab)/da = b
            other.grad += self.data * out.grad   # d(ab)/db = a
        out._backward = _backward
        return out

__add__ と __mul__ はPythonの演算子オーバーロードです。これがあると a + b と普通に書けて、裏で勝手にノードが増えます。計算グラフを別途「宣言」しなくていいのはここが効いているからで、順方向にコードを実行するだけでグラフが副産物として組み上がります。

各メソッドの1行目にある other if isinstance(other, Value) else Value(other) は、x * 2 のように生の数値が混ざっても動くようにするための変換です。定数もノードとして取り込まれ、勾配も律儀に計算されますが、誰も読まないので害はありません。入口で型を1つにそろえておくと、以降の処理から分岐が消える——実装を短く保つコツはここにあります。

なぜ「関数を持たせる」設計なのか

局所微分は数値で持てば済みそうに見えます。掛け算なら (other.data, self.data) というタプルを添えておけばいい。単純な演算だけならそれで動きます。

破綻するのは、局所微分が入力ではなく出力から決まる演算が出てきたときです。$\tanh$ を足してみます。

「tanhの傾きは、tanhの出力を2乗して1から引いたもの」と読みます。右辺に $x$ が出てこないのが要点で、傾きを知るのに入力は要らず、すでに計算した出力があれば足ります。指数関数も $\frac{d}{dx}e^x = e^x$、つまり出力そのものが傾きです。

    def tanh(self):
        t = math.tanh(self.data)
        out = Value(t, (self,), 'tanh')
        def _backward():
            self.grad += (1 - t * t) * out.grad
        out._backward = _backward
        return out

_backward の中から t にも out にも触れています。関数を作った時点の変数を関数が覚えている、Pythonのクロージャです。この形なら、局所微分が入力から決まる演算も出力から決まる演算も同じ枠組みに何でも足せます。PyTorchが数百種類の演算をそろえられるのも、演算ごとに「前向きの計算」と「勾配の配り方」を1組で登録できる構造だからです。

裏を返せば、この枠組みに新しい演算を足す作業は毎回まったく同じです。前向きの値を計算し、out を作り、局所微分を配る _backward を書いて貼る。ReLUでも指数でも対数でも手順は変わりません。演算を1つ足すコストが一定だから、フレームワークは種類を増やし続けられます。ちなみに _op に入れたラベルは計算には一切使われず、グラフを図に描いてデバッグするときの目印です。

下の図で点をドラッグすると、その点の傾き（＝局所微分）が縦位置だけから読めることが確かめられます。

逆走の順番を間違えると、静かに壊れる