🧠 5. PC-ALM で学習する¶

mujou.pc は PyTorch のモデルをノードに分け、PC-ALM（Augmented Lagrangian Predictive Coding、https://pub.sakana.ai/pc-alm/）で学習させるモジュール。各ノードはモデルの一部を持ち、隣のノードのメッセージだけから学ぶ。平衡点では、その勾配が誤差逆伝播法のものと一致する。このページでは回帰を学習させる。決めるのは時定数、重みの optimizer、各バッチを見せる時間だけ。

プログラム¶

uv run python examples/pc_regression.py

batch   1: loss 0.4157 (the mean of the last 1)
batch  50: loss 0.1686 (the mean of the last 50)
...
batch 300: loss 0.0043 (the mean of the last 50)

鎖¶

SHOWN = 0.032  # seconds of the run's timer each batch is shown
# tau.h sets how fast the credit settles (and the period of the nodes, h / 2),
# and the weights learn once per batch, by SGD at the rate
# SHOWN / optimizer.tau = 0.3.
tau = pc.TimeConstants(h=4e-3, lam=2e-3, r=4e-3)
optimizer = pc.optim.SGD(tau=SHOWN / 0.3)
torch.manual_seed(0)  # the initial weights, before the run
net = pc.Chain(
    nn.Sequential(nn.Linear(2, 64), nn.Tanh()),
    nn.Sequential(nn.Linear(64, 64), nn.Tanh()),
    nn.Linear(64, 1),
    x=TensorSpec((B, 2)),
    y=TensorSpec((B, 1)),
    batch=B,
    loss=pc.loss.MSELoss(B),  # the loss of each sample, (B,)
    tau=tau,
    optimizer=optimizer,
    learn_interval=0.75 * SHOWN,  # learn once, 3/4 of the way through each batch
)

pc.Chain はモジュールごとにノードを 1 つ作り、nn.Sequential が計算する順につなぐ。サンプルを受け取る pc.Input、最後以外の各モジュールの pc.Layer（layer0、layer1）、損失を持つ最後の pc.Output である。鎖の子は input・layer0・layer1・output で、これがそのまま状態のパスになる。層は net.layers で順に取れる（net.layers[0] が layer0）。

x と y は入力と目標の spec。batch はその先頭の shape で、ここでは B 個のサンプル。batch を省くと鎖はサンプルを 1 つずつ受け取り、x は TensorSpec((2,))、y は TensorSpec((1,)) になる
loss はサンプルごとの損失をバッチの shape で返す。pc.loss.MSELoss(B)（サンプルごとの平均二乗誤差。nn.MSELoss と同じくサンプルの次元で平均する）か pc.loss.CrossEntropyLoss(B)、または同じ shape を返す自作の関数。損失には鎖と同じバッチを渡す（バッチ）
tau は活動と双対変数の時定数で、実行のタイマーの秒で表す。h は活動（鎖のどのノードも tau.period、つまり tau.r = tau.h なら tau.h / 2 ごとに更新する）、lam は双対変数、r はずれの減衰の時定数。この組なら credit は 1 周期に 1 ノードずつ戻る。重みの時定数は optimizer が持つ
optimizer はノードが学ぶときに重みを動かす。pc.optim.SGD(tau=...) は勾配降下で、tau は重みの時定数。1 回の学習ではバッチの平均の勾配に（そのノードが前に学んでからの時間）/ tau を掛けた分だけ重みを動かす。ここではノードがバッチごとに 1 回学ぶので、1 バッチあたりの学習率は SHOWN / optimizer.tau = 0.3 になる（SHOWN は各バッチを見せる時間）。もう 1 つの選択肢は pc.optim.AdamW（optimizer を選ぶ）。各ノードは optimizer を自分の重みに結び付けるので、1 つの pc.optim.SGD を鎖全体で使える。optimizer に既定値は無い
learn_interval はノードが学ぶ時刻で、各書き込みから数える。ここではバッチごとに 1 回、見せる時間の 3/4 の時点。既定値は無い。None は毎回の更新で学ぶモードで、フレームごとに入力が変わるストリーム向け。このモードのよい値はまだ測っていない

ノードの内側では BP と同じくモジュールを autograd で微分する。重みはモジュールを作るときに初期化しておく（ここでは実行の前にホストの torch.manual_seed で）。ノードは初期化しない。

見せて学ぶ¶

torch.set_num_threads(1)
g = torch.Generator().manual_seed(1)
every = torch.ones(B, dtype=torch.bool)  # every sample has a target, and is new
losses: list[float] = []
with mujou.run(net) as run:
    for i in range(300):
        x = 2 * torch.rand(B, 2, generator=g) - 1
        y = torch.sin(2 * x[:, :1]) * torch.cos(x[:, 1:])
        sample = pc.msg.Sample(x=x, y=y, mask=every, reset=every)
        p = pc.show(run, net, sample, seconds=SHOWN)
        losses.append(nn.functional.mse_loss(p["y_hat"], y).item())  # before learning
        if i == 0 or (i + 1) % 50 == 0:
            recent = losses[-50:]
            mean = sum(recent) / len(recent)
            print(
                f"batch {i + 1:3d}: loss {mean:.4f} (the mean of the last {len(recent)})"
            )
    # The state of a node, by its path: input, layer0, layer1, output.
    lam = run.state_dict()["nodes"]["layer0"]["state"]["lam"]
print(f"the mean |lam| of layer0 (its credit) at the end: {lam.abs().mean():.2e}")

pc.show(run, net, sample, seconds=SHOWN) はバッチを書き、鎖がそのバッチに対して出す最初の予測を待つ。そしてバッチを実行のタイマーで SHOWN 秒見せ終えたところで、その予測を返す。その間ノードは目標に向かって緩和し、学ぶ。この予測はどのノードもまだそのバッチから学んでいない時点のものなので、その損失は、まだ見ていないデータをどれだけ当てられるかを表す（prequential な損失）。

mask はどのサンプルに目標があるかを表す。mask が False のサンプルは損失に何も足さない。mask がすべて False なら重みは動かない（活動は動く）。走っている同じ鎖を評価するにはこれを使う。切り替える推論モードは無い
reset はどのサンプルが新しく始まるかを表す。ここではバッチを毎回新しく引くので、どのサンプルも新しい。ノードは前のバッチについて持っていたものを捨てる（新しく始めるサンプル）。フレームが続くストリームなら False にする
pc.msg.Prediction は y_hat を、その元になったサンプルの y・mask・time と組にして持つ。予測はいつも自分の目標と対になる
run.state_dict()["nodes"]["layer0"]["state"] は最初の層の状態。活動の予測からのずれ r、双対変数 lam（平衡点では誤差逆伝播法の credit）、モジュール、クロック、optimizer、relaxer を持つ。送るもの（予測 mu）は Port next の状態にある（mujou.pc の仕組み）

次へ¶

部品のつながりと、ノードが 1 回の更新で行うこと: mujou.pc の仕組み
pc.Chain で MNIST を学ぶ 784-256-256-10 のネットワークと BP の比較: examples/mnist/（uv run python examples/mnist/main.py）。説明は PC-ALM
GPU の上の pc.Chain で CIFAR-10 を学ぶ ResNet-20 と BP の比較: examples/cifar10/（uv run python examples/cifar10/main.py）。説明は PC-ALM
時定数・optimizer・learn_interval の決め方、落とし穴、自作のノード: PC-ALM