🧮 Update Rules

An update rule (a learning rule) is what a node writes in on_update(dt). Mujou builds no rule into nodes, ports or the scheduler: each node chooses its own, and the same network can swap them. This page shows how to write a rule that learns across nodes, with feedback alignment as the example; 4. Learn in a Node has backpropagation (BP) inside one node, and PC-ALM the rule of mujou.pc.

Within a node, across nodes

Inside a node, anything PyTorch computes is allowed, BP included. Across nodes, no computation graph can be joined: a port sends a detached copy. A rule across nodes therefore sends what each node needs as messages: activities forward, errors (or other credits) back, and the targets along with the activities. Each node learns from its own state and its neighbors' messages only.

Feedback alignment

examples/feedback_alignment.py learns \(\hat y = W_2 \tanh(W_1 x)\) with two nodes. hidden sends its activity forward, and output sends the error \(e = \hat y - y\) back. BP would give hidden the credit \(W_2^\top e\), but hidden never sees W2: it uses \(B e\) with a fixed random matrix \(B\) of its own, and learns anyway, as \(W_2\) comes to align with \(B^\top\) (Lillicrap et al., 2016).

uv run python examples/feedback_alignment.py

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)

The program prints the loss of output on each batch before learning from it, which falls from about 1.5 to about 0.02 in 6 s; with hidden frozen, it stays above 0.15.

Two-way ports

hidden.next is Port(send=ACTIVITY, recv=ERROR) and output.prev is Port(send=ERROR, recv=ACTIVITY): one connection carries the activity forward and the error back. Connections may form cycles, so the error could as well come back on ports of its own.

Relaying the target

The activity carries the target y of its sample and the time of the write it came from. Values on two ports travel independently, so a target sent on a port of its own could pair a prediction with the target of another sample. In one message, output always compares a prediction with its own target, and hidden learns only from the error whose time is that of the input it holds. The nodes of mujou.pc carry y, mask and time in their pc.msg.Activity the same way (PC-ALM).

Telling what is new

recv() returns the latest value whether or not it is new. A node that learns once from each sample tells a new one by its time: the time recv_with_time() returns, or a time the message carries. Before anyone sends, the time is None (or, for a time in the message, its initial value: Timestamp(ns=0) of a TimestampSpec).

pc.clock.LearnClock

Both nodes decide when to learn with a pc.clock.LearnClock, a part of mujou.pc that any rule can use. In each update, tick(dt, time, now) takes the dt, the time the input was written at and the time now, and returns a pc.clock.Tick:

• new: the input is new (its time is not the one seen last)
• learn: learn in this update. With learn_interval, that is the first update learn_interval after each write (and every learn_interval while the input stays), but never the update that first sees an input; with learn_interval=None, every update
• elapsed: the time since the clock last learned, each dt capped at max_dt. Weigh the step with it: \(W \leftarrow W - \frac{\text{elapsed}}{\tau}\,\nabla\) follows the gradient flow \(\tau \dot W = -\nabla\), so how far a node learns depends on time, not on how often it updates
• dt: the dt capped at max_dt

The clock is a mujou.Stateful, so an attribute that holds it is in the node's state by itself. The nodes on a path read the one timer of the run, so they learn at about the same time after each write, however deep they are.