MicroGrad: 簡単な自動微分フレームワークを実装する

Micrograd の紹介 (Introduction)#

Micrograd はAndrej Karpathyが作成した学習プロジェクトで、オートマチック微分、バックプロパゲーションなどの基本的な知識を理解するためのものです。

動画のリンク：https://www.bilibili.com/video/BV1De4y1p7Z8

コードのリンク：https://github.com/karpathy/micrograd

バックプロパゲーション (Backpropagation) は機械学習と深層学習で最も一般的に使用されるアルゴリズムの一つで、神経ネットワークにおける損失関数の重みパラメータに対する勾配を計算し、損失を最小化することでモデルの重み更新を指導します。

導関数 (Derivative)#

導関数の数学的定義#

数学における導関数の定義：関数のある点における導関数は、その点の近くでの関数の変化率、すなわちその点における関数の接線の傾きを指します。公式は次のように表されます：

$f'(x) = \lim_{h \to 0} \frac{f(x+h) - f(x)}{h}$

x 点において、h の小さな値（例えば 0.0001）を取り、上記の公式を用いて関数 $f(x)$ の $x$ における導関数を近似計算します。

例えば：

import math
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

# 関数を定義
def f(x):
    return 3*x**2 - 4*x + 5
# matplotlibを使用して関数のグラフを描く
xs =np.arange(-5, 5, 0.25)
ys = f(xs)
plt.plot(xs, ys)

h = 0.00000001
a = -3.0
(f(a + h) - f(a))/h

出力：-22.00000039920269 これは関数 f の - 3 における導関数の近似値です。

実際には導関数の法則を用いて $f$ の導関数を直接求めることができ、 $6x - 4$ となり、-3 における導関数は - 22 です。
しかし、実際の現代の神経ネットワークライブラリでは、すべての関数に対してこのような表現を求めることはなく、巨大で複雑な関数に遭遇した場合、このような導関数の表現を求めることは困難で非効率的です。

導関数のさらなる理解#

上記では関数の未知数に対する導関数を求めましたが、実際の神経ネットワークでは、特定のスカラーに対する関数の導関数を求める必要があります。この場合、導関数の定義から始めると理解しやすくなります。

h = 0.0001
a = 2.0
b = -3.0
c = 10.0

d1 = a*b+c
c += h
d2 = a*b+c

print('slope', (d2 - d1)/h)

定義に従い、ここでは c に対して導関数を求めています（c が少し変化したときの関数 d の変化率を求めています）。結果は：
slope 0.9999999999976694

これは、c が 0.0001 の微小な変化を増加させると、d の変化量は c の変化量の約 1 倍（正の増加、負の減少）であることを意味します。

導関数の定義を用いて d の変化量と c の変化量の関係を求めることもできます：
$\frac {d_2-d_1} {h} = 1$

同様に、関数の a および b に対する導関数の近似値も求めることができます。

神経ネットワークの原理#

単純な多層パーセプトロン MLP を例にとると、神経ネットワークは実際には数学的な公式の集合です。データは入力として、隠れ層を経由して神経細胞内で数学的な演算（通常は入力データの加重和にバイアスを加え、その後非線形活性化関数を適用）を通じて次の層への活性値を生成し、最終的に予測結果などの他のデータを出力します。主な演算は一連の加算と乗算（線形代数では行列演算として表現されます）および非線形活性化です。

前向き伝播はネットワークの出力を計算し、損失値を得る役割を担います；前向き伝播が完了した後、バックプロパゲーションを実行し、その核心的なタスクは損失値と前向き伝播中に保存された中間結果に基づいて、チェーンルールを用いて損失関数の各層のパラメータ（重みやバイアス）に対する勾配を計算することです。これらの計算から得られた勾配は、隠れ層の神経細胞の重みパラメータの更新を指導し、トレーニング誤差を減少させます（実際にはトレーニング誤差を減少させるだけでなく、テスト誤差を減少させ、モデルの一般化能力を向上させることが最終目的です）。

micrograd では、オートマチック微分の原理を示すために単純なスカラー演算を使用していますが、実際の神経ネットワークの実装では、効率のために大量の行列演算が行われます。

オートマチック微分フレームワークの実装#

神経ネットワークには多くのパラメータを維持する必要があるため、データの維持を支援するデータ構造を設計する必要があります。

カスタムクラス#

class Value:
    def __init__(self, data):
        self.data = data

    def __repr__(self):
        return f"Value(data={self.data})"
    
    def __add__(self, other):
        out = Value(self.data + other.data)
        return out
    
    def __mul__(self, other):
        out = Value(self.data * other.data)
        return out

__repr__は開発者に優しい文字列表現を生成するために使用されます。
__str__はユーザーに優しい文字列表現を生成するために使用され、通常は print () や str () によって呼び出されます。
__str__が存在しない場合、print () は__repr__を使用します。
__repr__が存在しない場合、repr () およびインタラクティブ環境の直接表示はデフォルトの <...> 表現を使用します。

上記のコードでは、Valueクラスを簡単に定義し、属性 data を含み、Value クラスのメソッド add と mul を定義しています。

a = Value(2.0)
b = Value(-3.0)
c = Value(10.0)
d = a * b + c
d

Value(data=4.0)

次に、Value に他の属性を追加して、より多くの操作の詳細を保持できるようにします。ここでは、集合_childrenを使用して現在の data の子ノードを表し、opは操作演算子を記録し、label は現在のデータの別名を示します。

class Value:
    def __init__(self, data, _children=(), _op='', label=''):
        self.data = data
        self._prev = set(_children)
        self._op = _op
        self.label = label

    def __repr__(self):
        return f"Value(data={self.data})"
    
    def __add__(self, other):
        out = Value(self.data + other.data, (self, other), '+')
        return out
    
    def __mul__(self, other):
        out = Value(self.data * other.data, (self, other), '*')
        return out

計算グラフを表示するための可視化 API の使用#

計算プロセスをより直感的に表示するために、以下のコードを使用して可視化操作を行います（重要ではないのでスキップ可能です）。

from graphviz import Digraph

def trace(root):
    nodes, edges = set(), set()
    def build(v):
        if v not in nodes:
            nodes.add(v)
            for child in v._prev:
                edges.add((child, v))
                build(child)
    build(root)
    return nodes, edges

def draw_dot(root):
    dot = Digraph(format='svg', graph_attr={'rankdir':'LR'})

    nodes, edges = trace(root)
    for n in nodes:
        uid = str(id(n))

        dot.node(name=uid, label='{%s | data %.4f | grad %.4f}' % (n.label, n.data, n.grad), shape='record')

        if n._op:
            dot.node(name=uid+n._op, label=n._op)
            dot.edge(uid + n._op, uid)
    
    for n1, n2 in edges:
        dot.edge(str(id(n1)), str(id(n2)) + n2._op)
    return dot

深い式を定義し、計算グラフを可視化します（前向き伝播）。

a = Value(2.0, label='a')
b = Value(-3.0, label='b')
c = Value(10.0, label='c')
e = a * b; e.label = 'e'
d = e + c; d.label = 'd'
f = Value(-2.0, label='f')
L = d * f; L.label = 'L'
L

draw_dot(L)

output

バックプロパゲーションによる勾配計算#

バックプロパゲーションに入る前に、チェーンルールを復習します。

チェーンルール#

バックプロパゲーションの過程では、最終出力が途中のある入力に対する導関数を求める必要がありますが、最終出力の関数は途中の変数に直接関係していません。

例えば、y が u に関する関数 $(y=f(u))$ であり、u が x に関する関数 $(u=g(x))$ である場合、 $y$ の $x$ に対する導関数は：

$\frac{dy}{dx} = \frac{dy}{du} \cdot \frac{du}{dx}$

もし $h(x)$ が合成関数 $f(g(x))$ であれば、 $h(x)$ の $x$ に対する導関数 $h'(x)$ は外側の関数 $f$ の導関数が内側の関数 $g(x)$ での値に等しく、内側の関数 $g(x)$ の導関数 $g'(x)$ を掛けたものになります：

$h'(x) = f'(g(x)) \cdot g'(x)$

導関数を求める際には、まず外側の関数の導関数を求め（内側の関数はそのままに）、次に内側の関数の導関数を掛けます。

手動での逆伝播#

Value クラスを修正して、手動での逆伝播を実装します。

class Value:
    def __init__(self, data, _children=(), _op='', label=''):
        self.data = data
        self._prev = set(_children)
        self._op = _op
        self.label = label
        self.grad = 0.0
        self._backward = lambda: None

    def __repr__(self):
        return f"Value(data={self.data})"
    
    def __add__(self, other):
        out = Value(self.data + other.data, (self, other), '+')
        def _backward():
            self.grad = 1.0 * out.grad
            other.grad = 1.0 * out.grad
        out._backward = _backward

        return out
    
    def __mul__(self, other):
        out = Value(self.data * other.data, (self, other), '*')
        def _backward():
            self.grad = other.data * out.grad
            other.grad = self.data * out.grad
        out._backward = _backward
        return out

L に対して L 自身の導関数を求めると、導関数は 1 になるため、まず手動で L の勾配を 1 に設定します。

L.grad = 1.0

L に対して勾配の逆伝播を行います。計算グラフから、L=f*d であることがわかるため、L の逆伝播を行うことで f と d の勾配をそれぞれ求めることができます。このように、手動で L に対して逆伝播を行います：

L._backward()

output
同様に、d に対して逆伝播を行い、c、e の勾配を得ます。c に対して逆伝播を行うことで、a と b の勾配を得ます。

d._backward()
e._backward()

output

神経ネットワークでは、神経細胞がデータを受け取った後、単純な線形演算だけでなく、線形演算の結果に対して非線形処理を行います。次に、計算グラフを拡張し、活性化関数 tanh を使用します。

tanh 関数を定義します：

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

tanh(x) = (e^2x - 1) / (e^2x + 1)
その導関数は tanh (x)' = 1 - tanh (x)²です。

複雑な神経細胞を定義します。

x1 = Value(2.0, label='x1')
x2 = Value(0.0, label='x2')

w1 = Value(-3.0, label='w1')
w2 = Value(1.0, label='w2')

b = Value(6.8813735870195432, label='b')

x1w1 = x1 * w1; x1w1.label='x1w1'
x2w2 = x2 * w2; x2w2.label='x2w2'
x1w1x2w2 = x1w1 + x2w2; x1w1x2w2.label='x1w1x2w2'

n = x1w1x2w2 + b; n.label='n'
o = n.tanh(); o.label='o'

draw_dot(o)

output

上記のように手動で逆伝播を実行します。

output

計算プロセスがさらに複雑になると、手動での逆伝播は非常に煩雑になります。次に、自動の逆伝播を実装します。

自動逆伝播#

具体的には、前向き伝播の後に逆伝播操作を実行し、トポロジカルソートに従ってノードの逆順で ._backward() メソッドを呼び出します。

自動逆伝播関数を定義します。

def backward(self):
        topo = []
        visited = set()
        def build_topo(v):
            if v not in visited:
                visited.add(v)
                for child in v._prev:
                    build_topo(child)
                topo.append(v)
        build_topo(self)
        self.grad = 1.0
        for node in reversed(topo):
            node._backward()

実行します。

o.backward()

これで計算グラフ全体の勾配を一度に求めることができます。しかし、現在の設計には一定の欠陥があります。以下の例はその欠陥を明確に示しています。

a = Value(3.0, label='a')
b = a + a; b.label='b'
b.backward()
draw_dot(b)

b の a に対する導関数は 2 であるべきですが、1 になっています。これは、逆伝播を定義する際に単純に代入操作を行ったためであり、特定の変数を複数回使用する操作に対して、勾配は累積されるべきです。

この問題を修正すると、完全な Value クラスは次のようになります：

class Value:
    def __init__(self, data, _children=(), _op='', label=''):
        self.data = data
        self._prev = set(_children)
        self._op = _op
        self.label = label
        self.grad = 0.0
        self._backward = lambda: None

    def __repr__(self):
        return f"Value(data={self.data})"
    
    def __add__(self, other):
      # 型検証：Valueクラスと数値の演算を許可し、Valueクラスの堅牢性を向上させます
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), '+')
        def _backward():
            self.grad += 1.0 * out.grad
            other.grad += 1.0 * out.grad
        out._backward = _backward

        return out
    # __radd__は逆の魔法のメソッドで、2 * aのような表現をより自然にします
    def __radd__(self, other):
        return self.__add__(other)
    
    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
            other.grad += self.data * out.grad
        out._backward = _backward
        return out
    
    def tanh(self):
        n = self.data
        t = (math.exp(2*n) - 1) / (math.exp(2*n) + 1)
        out = Value(t, (self, ), 'tanh')
        def _backward():
            self.grad += (1 - t ** 2) * out.grad
        out._backward = _backward
        return out 
    
    def backward(self):
        topo = []
        visited = set()
        def build_topo(v):
            if v not in visited:
                visited.add(v)
                for child in v._prev:
                    build_topo(child)
                topo.append(v)
        build_topo(self)
        self.grad = 1.0
        for node in reversed(topo):
            node._backward()

Value クラスの他の関連操作については、自分で定義を追加することができます。例えば、div 操作や pow など。

動画の中で著者は、Pytorch を使用して簡単な MLP を実装する方法についても説明していますので、興味がある方はさらに学んでみてください。

まとめ#

このブログを書いて、Micrograd の旅を通じて、皆さんがオートマチック微分に対して神秘的な印象を持たなくなり、さらには自分でミニバージョンを実装できるようになればと思います。このブログが、深層学習フレームワークの背後で静かに働く「英雄」たち（例えばオートマチック微分エンジン）についての理解を深める手助けとなるなら、私の目的は達成されたことになります！皆さんがさらに探求し、Value クラスにもっと多くの機能を追加したり、Karpathy のオリジナルの動画やコードを見たりすることを歓迎します。きっと多くの収穫があるでしょう。