自動微分を実装する —— forward と reverse はどこで逆転するか

本稿の位置づけ
PyTorch の backward() が何をしているかは、使っているうちに何となく分かった気になる。だが「なぜ forward ではなく reverse なのか」を人に説明しようとすると、途端に曖昧になった。両方を自分で書けば分かるはずなので、書いて測った。
実装は外部ライブラリなしの素の Python である。numpy も使っていない。結論から言うと、逆転が起きるのは私が思っていたよりずっと手前で、入力と出力が同じ次元数のときには既に reverse が勝っていた。

1. 何を実装するのか

微分を機械にやらせる方法は3つある。

数式微分は式を式のまま変形する。人間と同じことをするので正確だが、式が指数的に膨らむ。数値微分は $(f (x + h) - f (x - h)) /2 h$ を計算する。実装は一行で済むが、 $h$ を小さくすると桁落ちし、大きくすると打ち切り誤差が乗る。

自動微分は3つ目で、この二つのどちらでもない。プログラムを式とみなし、連鎖律を演算のたびに適用していく。近似ではないので、丸め誤差を除けば厳密な値が出る。

その自動微分に、伝播の向きが違う2つのモードがある。

forward モードは入力側から出力側へ、値と一緒に微分を運ぶ
reverse モードは一度計算を最後まで走らせて記録を残し、出力側から入力側へ遡る

同じ答えを出す。違うのは計算量で、そこに実務上の判断が挟まる。

2. forward モード：二重数

forward の実装は驚くほど短い。値 $v$ と微分 $d$ を組にした「二重数」を作り、演算子を定義するだけである。

class Dual:
    __slots__ = ("v", "d")

    def __init__(self, v, d=0.0):
        self.v = v
        self.d = d

    def __add__(self, o):
        o = o if isinstance(o, Dual) else Dual(o)
        return Dual(self.v + o.v, self.d + o.d)

    def __mul__(self, o):
        o = o if isinstance(o, Dual) else Dual(o)
        return Dual(self.v * o.v, self.d * o.v + self.v * o.d)

__mul__ の第2成分が積の微分則そのものである。 $sin$ や $exp$ も同じ調子で書ける。

def d_sin(x):
    return Dual(math.sin(x.v), math.cos(x.v) * x.d)

def d_exp(x):
    e = math.exp(x.v)
    return Dual(e, e * x.d)

ここで効いてくるのが、1回の実行で取れるのが「1つの入力に関する偏微分」だけ、という制約である。 $x$ に関する偏微分が欲しければ $x$ の微分成分を 1、他を 0 にして関数を走らせる。 $y$ に関する偏微分が欲しければ、もう一度走らせる。

つまり 入力が $n$ 個あれば、関数を $n$ 回実行する。

def forward_grad(f, xs):
    n = len(xs)
    cols = []
    for i in range(n):
        duals = [Dual(x, 1.0 if j == i else 0.0) for j, x in enumerate(xs)]
        out = f(duals)
        cols.append([o.d for o in out] if isinstance(out, list) else [out.d])
    m = len(cols[0])
    return [[cols[i][k] for i in range(n)] for k in range(m)]

3. reverse モード：テープ

reverse は forward ほど素直には書けない。演算のたびに「どの値から作られたか」と「局所的な微分」を記録しておき、後から逆順に辿る必要がある。この記録をテープと呼ぶ。

class Tape:
    def __init__(self):
        self.nodes = []  # (親のindex, 局所微分) の並び

    def push(self, deps):
        self.nodes.append(deps)
        return len(self.nodes) - 1

各変数はテープ上の位置を持つ。演算子は新しいノードを積みながら、親への局所微分を書き込む。

class Var:
    __slots__ = ("v", "i", "tape")

    def _bin(self, o, val, dself, do):
        if isinstance(o, Var):
            i = self.tape.push([(self.i, dself), (o.i, do)])
        else:
            i = self.tape.push([(self.i, dself)])
        return Var(val, self.tape, i)

    def __mul__(self, o):
        ov = o.v if isinstance(o, Var) else o
        return self._bin(o, self.v * ov, ov, self.v)

逆伝播は、出力の随伴変数を 1 に置いてテープを末尾から舐めるだけである。

def reverse_grad(f, xs):
    tape = Tape()
    vs = [Var(x, tape) for x in xs]
    outs = f(vs)
    ...
    for k in range(len(outs)):
        adj = [0.0] * len(tape.nodes)
        adj[outs[k].i] = 1.0
        for i in range(len(tape.nodes) - 1, -1, -1):
            a = adj[i]
            if a == 0.0:
                continue
            for (p, dl) in tape.nodes[i]:
                adj[p] += a * dl
        J.append([adj[v.i] for v in vs])

対称性がはっきりする。forward が入力の数だけ実行を繰り返したのに対し、reverse は 出力が $m$ 個あれば、逆伝播を $m$ 回行う。テープを作る前向きの実行は1回で済む。

4. 正しさの確認

速度を測る前に、両方が正しいことを確かめる。 $f (x, y) = sin (x y) + e^{x} / (y + 2)$ を $x = 0.7, y = 1.3$ で微分し、手で出した解析解と比べた。

手法	$\partial f / \partial x$ の誤差	$\partial f / \partial y$ の誤差
forward	$0$	$0$
reverse	$0$	$0$
中心差分（ $h = 1 0^{- 5}$ ）	$7.1 \times 1 0^{- 12}$	$1.6 \times 1 0^{- 12}$

自動微分の2つは解析解とビット単位で一致した。近似ではないので当然ではあるが、数値微分が $1 0^{- 12}$ の誤差を持つのと対照的である。この差は $h$ をどう選んでも消えない。

さらに、forward と reverse でヤコビアン全体を突き合わせた。

次元	両モードの最大差
$n = 3, m = 4$	$2.2 \times 1 0^{- 18}$
$n = 5, m = 5$	$1.3 \times 1 0^{- 18}$
$n = 7, m = 2$	$4.3 \times 1 0^{- 19}$

独立に書いた2つの実装が一致するので、両方が同時に間違っている可能性は低い。

5. 測る：どこで逆転するか

$n$ 個の入力から $m$ 個の出力を作る関数を用意した。各出力は全入力に依存し、その後に深さ20の非線形な連鎖を通す。ヤコビアン全体を求める時間を測った。

$n$	$m$	forward [ms]	reverse [ms]	比	速い方
1	100	1.08	14.78	0.07	forward
2	50	1.10	4.36	0.25	forward
5	20	1.17	1.04	1.12	reverse
10	10	1.31	0.42	3.12	reverse
20	5	1.58	0.21	7.52	reverse
50	2	2.46	0.11	21.84	reverse
100	1	4.04	0.09	46.22	reverse
500	1	78.67	0.37	210.35	reverse

$n$ が大きく $m$ が小さいほど reverse が有利になる。 $n = 500, m = 1$ では 210 倍の差がついた。逆に $m$ が大きい側では forward が勝つ。

理屈どおりではある。だが私が予想を外したのは $n = m$ の場合だった。実行回数は forward が $n$ 回、reverse が $m$ 回なので、 $n = m$ なら互角になるはずだと思っていた。1刻みで測るとこうなる。

$n = m$	forward [ms]	reverse [ms]	比
1	0.012	0.023	0.50
2	0.045	0.049	0.91
3	0.102	0.080	1.28
4	0.185	0.114	1.62
5	0.297	0.155	1.92
6	0.435	0.199	2.19

逆転は $n = m = 3$ で起きていた。互角ではない。

理由は1回あたりの単価にある。forward の1回は関数の完全な再実行で、Dual オブジェクトを作りながら全演算を通る。reverse の1回は、既にあるテープを配列演算で舐めるだけである。前向きの実行は最初の1回しかない。同じ回数でも中身の重さが違う。

$n = m$ で互角だと思っていたのは、回数だけを数えて単価を数えていなかったからだった。

6. 勾配は関数評価の何回分か

機械学習で効くのは $m = 1$ の場合、つまりスカラーの損失に対する勾配である。ここで知りたいのは絶対時間ではなく、関数を1回計算するのに比べて勾配が何倍かかるかである。

$n$	$f (x)$ [ms]	reverse [ms]	倍率	forward [ms]	倍率
10	0.002	0.046	25.5x	0.2	116x
50	0.003	0.084	26.0x	1.7	524x
100	0.005	0.118	24.8x	4.7	986x
500	0.016	0.414	26.0x	77.5	4,859x
1000	0.029	0.730	25.6x	302.1	10,583x

reverse の倍率は $n$ を100倍にしてもほとんど動かない。24 倍から 26 倍の間に収まっている。一方 forward は 116 倍から 10,583 倍へ、 $n$ に比例して増えていく。

これが reverse モードが機械学習を支えている理由である。パラメータが何億個あろうと、勾配1本の値段は関数1回の定数倍で済む。パラメータを増やしても勾配計算の相対コストが増えない、というのは自明ではない。

ただし倍率そのものは真に受けないでほしい。教科書には「2〜4倍」と書いてあり、実際 C++ や JAX ではその程度になる。ここで 20〜30 倍も出ているのは、Python のオブジェクト生成とリスト操作が支配的だからである。主張できるのは定数倍であることであって、定数の値ではない。

7. reverse の代償

reverse がいつでも良いわけではない。テープを持つ必要がある。

連鎖の深さ	テープのノード数
10	52
100	502
1,000	5,002

深さに比例して増える。前向きの計算が終わるまで、中間結果を捨てられない。深いネットワークで学習時のメモリが問題になるのはこれが理由で、勾配チェックポイント（一部を捨てて再計算する）という手法が要るのもここから来ている。

forward にはこの問題がない。値を運びながら進むだけで、記録を残さない。メモリは入力サイズだけで決まる。

8. 実務でどう選ぶか

測った結果をまとめると、判断は次のようになる。

状況	選ぶモード	根拠
損失の勾配（ $n$ 大・ $m = 1$ ）	reverse	倍率が $n$ によらず一定
ヤコビアンベクトル積、感度解析（ $m$ 大・ $n$ 小）	forward	実行回数が入力次元で決まる
$n \approx m$	reverse	単価が安い（ $n = m = 3$ で逆転）
メモリが厳しい・深い連鎖	forward、または reverse＋チェックポイント	テープが深さに比例
高階微分	forward を重ねる	テープを二重に持つより素直

「入力が多ければ reverse」という覚え方で実務はほぼ回る。ただし $n = m$ が互角ではないことは、測るまで私は間違えていた。

9. 書いてみて分かったこと

forward の実装は演算子オーバーロードだけで完結する。状態を持たないので、書いた端から動く。reverse はテープという可変状態が入る分だけ手間が増える。

reverse で踏みやすい罠が一つある。逆伝播で随伴変数を 代入ではなく加算 しなければならない点である。同じ変数が式の中で複数回使われると、そこへ複数の経路から寄与が集まる。代入で書くと最後の経路だけが残り、他が消える。厄介なのは、各変数が一度しか使われない単純な式では代入でも正しい答えが出てしまうことで、小さい例のテストは通ってしまう。ヤコビアン全体を forward と突き合わせて初めて差が出る種類の誤りである（第4節の照合はこれを検出するために入れた）。

実装の手間がこれだけ違うのに、主要なフレームワークが reverse を採用しているのは、性能差がその手間を正当化するからである。第6節の表のとおり、 $n = 1000$ では 400 倍以上の開きがある。

次にやるとすれば、テープを持たない reverse（クロージャで連鎖を組む方式）と、二重数を入れ子にした高階微分を試したい。

10. 再現方法

コードは GitHub に置いてある。外部ライブラリは使っていない。

https://github.com/makoto-developer/blog-examples/tree/main/automatic-differentiation

python3 bench.py   # 正しさの確認（第4節）と n×m のベンチ（第5節）
python3 bench2.py  # 勾配の倍率（第6節）とテープの大きさ（第7節）

ad.py が forward と reverse の実装本体で、上の2本はそれを呼ぶだけである。本文の表はこのスクリプトの出力をそのまま貼っている。測定は Python 3 の素の実行で、各測定は3〜5回の最小値を取った。

ただし時間は環境で変わる。環境によらないのは、モード間の比と、倍率が $n$ に依存するかどうかである。

参考文献

A. Griewank, A. Walther, Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, 2nd ed., SIAM (2008)：自動微分の標準的な教科書。勾配コストが関数評価の定数倍で抑えられることの議論を含む。
A. G. Baydin, B. A. Pearlmutter, A. A. Radul, J. M. Siskind, Automatic differentiation in machine learning: a survey, JMLR 18 (2018) 1：機械学習分野からの概観。数式微分・数値微分との区別が整理されている。