🧮 更新則¶
更新則(学習則)は、ノードが on_update(dt) に書く更新そのもの。mujou はノード・Port・スケジューラに更新則を組み込まないので、ノードごとに選べ、同じネットワークのまま差し替えられる。このページでは、ノードをまたいで学ぶ則の書き方を feedback alignment を例に説明する。1 つのノードの内側の誤差逆伝播法(BP)は 4. ノードの中で学習する、mujou.pc の則は PC-ALM にある。
ノードの内側と、ノードの間¶
ノードの内側では、BP を含め PyTorch で計算できることは何でも使える。ノードの間では計算グラフをつなげない。Port が送るのは切り離したコピーである。そのため、ノードをまたいで学ぶ則は、各ノードが要るものをメッセージで送る。活動を前へ、誤差(やほかの credit)を後ろへ、目標は活動と一緒に。各ノードは自分の状態と隣のメッセージだけから学ぶ。
feedback alignment¶
examples/feedback_alignment.py は \(\hat y = W_2 \tanh(W_1 x)\) を 2 つのノードで学ぶ。hidden が活動を前へ送り、output が誤差 \(e = \hat y - y\) を後ろへ返す。BP なら hidden は credit \(W_2^\top e\) を受け取るが、hidden は W2 を知らない。代わりに自分の固定の乱数行列 \(B\) で \(B e\) を使う。それでも \(W_2\) が \(B^\top\) にそろっていき、学べる(Lillicrap et al., 2016)。
class Hidden(Node):
"""h = tanh(W1 x); learns W1 from the error that comes back through ``B``."""
def __init__(self) -> None:
super().__init__(period=PERIOD, seed=1)
g = self.generator
self.w1 = torch.randn(N_HIDDEN, N_IN, generator=g) / N_IN**0.5
self.b = torch.randn(N_OUT, N_HIDDEN, generator=g) / N_OUT**0.5 # fixed
self.optimizer = OPTIMIZER.bind([self.w1])
self.clock = pc.clock.LearnClock(
period=PERIOD, learn_interval=LEARN, max_dt=2 * PERIOD
)
self.prev = Port(recv=SAMPLE) # the host writes it
self.next = Port(send=ACTIVITY, recv=ERROR) # two-way
@override
def on_update(self, dt: float) -> None:
super().on_update(dt)
sample, t = self.prev.recv_with_time()
if t is None: # nothing written yet
return
x = sample["x"]
h = torch.tanh(x @ self.w1.T)
self.next.send({"h": h, "y": sample["y"], "time": t}) # relays y
error = self.next.recv()
tick = self.clock.tick(dt, t, self.timer.now())
if tick.learn and error["time"] == t: # the error of this sample
delta = (error["e"] @ self.b) * (1 - h**2) # B e, not W2^T e
self.optimizer.step_([(delta.T @ x) / BATCH], tick.elapsed)
class Output(Node):
"""y_hat = W2 h; sends the error back, and the loss on each new sample."""
def __init__(self) -> None:
super().__init__(period=PERIOD, seed=2)
self.w2 = torch.randn(N_OUT, N_HIDDEN, generator=self.generator) / N_HIDDEN
self.optimizer = OPTIMIZER.bind([self.w2])
self.clock = pc.clock.LearnClock(
period=PERIOD, learn_interval=LEARN, max_dt=2 * PERIOD
)
self.prev = Port(send=ERROR, recv=ACTIVITY) # two-way
self.loss = Port(send=TensorSpec(()))
@override
def on_update(self, dt: float) -> None:
super().on_update(dt)
act, t = self.prev.recv_with_time()
if t is None: # hidden has not sent yet
return
h = act["h"]
e = h @ self.w2.T - act["y"] # the prediction against its own target
self.prev.send({"e": e, "time": act["time"]})
tick = self.clock.tick(dt, act["time"], self.timer.now())
if tick.new: # the first update with this sample, before learning from it
self.loss.send(0.5 * e.pow(2).sum(dim=1).mean())
if tick.learn:
self.optimizer.step_([(e.T @ h) / BATCH], tick.elapsed)
プログラムは各バッチから学ぶ前の output の損失を表示する。損失は 6 秒で約 1.5 から約 0.02 まで下がる。hidden を固定すると 0.15 を下回らない。
両方向の Port¶
hidden.next は Port(send=ACTIVITY, recv=ERROR)、output.prev は Port(send=ERROR, recv=ACTIVITY)。1 本の接続で活動を前へ、誤差を後ろへ運ぶ。循環してつないでもよいので、誤差を別の Port で返してもよい。
目標を中継する¶
活動はそのサンプルの目標 y と、元になった書き込みの時刻を運ぶ。2 つの Port の値は別々に届くので、目標を別の Port で送ると、予測を別のサンプルの目標と組んでしまうことがある。1 つのメッセージにすれば、output はいつも予測を自分の目標と比べ、hidden は持っている入力と同じ time の誤差だけから学ぶ。mujou.pc のノードも同じように、y・mask・time を pc.msg.Activity で運ぶ(PC-ALM)。
新しいかどうかを判定する¶
recv() は新しいかどうかによらず最新の値を返す。各サンプルから 1 回だけ学ぶノードは、時刻で新しいサンプルを見分ける。recv_with_time() が返す時刻か、メッセージが運ぶ時刻を使う。誰も送る前の時刻は None(メッセージが運ぶ時刻なら、その初期値。TimestampSpec なら Timestamp(ns=0))になる。
pc.clock.LearnClock¶
どちらのノードも、いつ学ぶかを pc.clock.LearnClock で決める。mujou.pc の部品だが、どの則でも使える。更新ごとに tick(dt, time, now) に dt、入力が書かれた時刻、今の時刻を渡すと、pc.clock.Tick が返る。
new: 入力が新しい(時刻が前回見たものと違う)learn: この更新で学ぶ。learn_intervalを渡したなら、各書き込みからlearn_interval後の最初の更新(入力が変わらない間は以後learn_intervalごと)。ただし入力を初めて見た更新では学ばない。learn_interval=Noneなら毎回elapsed: 前に学んでからの時間。各dtはmax_dtで打ち切る。刻みはこれで重み付けする。\(W \leftarrow W - \frac{\text{elapsed}}{\tau}\,\nabla\) は勾配流 \(\tau \dot W = -\nabla\) をたどるので、どれだけ学ぶかはノードの更新回数ではなく時間で決まるdt:max_dtで打ち切ったdt
クロックは mujou.Stateful なので、属性に置くだけでノードの状態に入る。経路上のノードは実行の 1 本のタイマーを読むので、深さによらず各書き込みからほぼ同じ時刻に学ぶ。