← Course home All tutorials

Week 06 — A robot’s movements form a tiny language

Can we predict its next command?

Imagine watching a robot move UP, wander around, and then receive RECALL. Its next movement must repeat the first one. You can answer if you remember the beginning. Someone who sees only the last command cannot tell which direction to choose.

Recall here means remembering the initial movement and predicting it again after the RECALL token. For example, BOS UP … RECALL requires the next prediction UP. Recall accuracy is the fraction of these requests answered correctly; recall loss is cross-entropy measured only at those requests. This task-specific use of “recall” is distinct from the classification metric that measures how many actual positives a classifier finds.

Now suppose the robot is still wandering. Its next movement is chosen randomly, independently of its history. Even perfect memory cannot reveal that choice. Our experiment contains both kinds of prediction: a dependency that can be learned, and uncertainty that remains.

Follow one question: how can a model use earlier evidence to predict what comes next?

Route	Scenes	Main question
Build the task	1–3	What can history make predictable?
Open and train a recurrent model	4–8	Can a state retain the cue? Scene 6 is optional.
Add gates	9	Does controlled retention help?
Calculate and train attention	10–11	Can the model retrieve earlier evidence directly?
Compare held-out predictions	12	Which findings survive new episodes and initializations?

The loss epilogue and wave mini-lab are optional extensions.

Read and run

Each scene develops one part of the experiment. Pause at a question before running its calculation; open Show explanation when you want to compare your reasoning. Figures and their captions help you interpret the outputs. The setup and figure helpers below can stay folded while you read.

Open in Colab, or upload this notebook to Colab. Run cells in order, or use Run all to reproduce the chapter. No dataset downloads are required. The notebook uses NumPy, Matplotlib, PyTorch, and IPython.

Allow 150–180 minutes across two teaching blocks for the twelve scenes. The 27 small CPU training runs take about 4–6 minutes on the development laptop; the last cell measures the current run. Reading and discussion take most of the teaching time.

By the end, you should be able to construct next-token targets, explain what a recurrent state carries, interpret gates and causal attention, and assess a held-out comparison using recall-specific metrics. Week 04 introduced learned representations; Week 05 shared parameters across locations. Here we share computations across time.

The optional wave forecasting mini-lab compares equal observation windows and contrasts next-sample prediction with recursive rollouts. It can run separately from the command experiment. Its diagnostic extension asks why recurrent forecasts drift; full retraining is optional and takes about 13.7 minutes on this laptop.

Notation and terminology contract

Start with three symbols: \(s_t\) is the command we observe, \(y_t=s_{t+1}\) is the next command to predict, and \(h_t\) is the model’s internal state. Each scene defines its additional symbols locally.

Full notation reference — symbols, shapes, and conventions

Shapes describe a single episode unless a batch axis is shown. Indices start at zero.

Commands, positions, and shapes

Symbol or term	Meaning	Shape or code connection
\(s_t\)	Command token observed at position \(t\)	An integer ID in TOKENS; IDs are category labels
\(t\), \(j\)	Current/query position and a candidate source position	Time indices; \(s_{\leq t}\) means the prefix through \(t\)
\(y_t=s_{t+1}\)	Target for the prediction made after reading \(s_t\)	Shifted target ID; targets
\(d\)	Number of wandering commands between the initial cue and RECALL	DELAY; experiments use 8, 16, and 32
\(T=d+4\)	Number of input/target positions in one episode	Complete episode length is \(T+1\)
\(K=7\)	Number of command types in the vocabulary	len(TOKENS)
\(B\), \(\mathcal B\)	Batch size and the set of examples in a batch	Input/target ID arrays have shape \((B,T)\)
\(D\), \(H\)	Embedding/representation width and recurrent hidden width	An embedded batch is \((B,T,D)\); recurrent states are \((B,T,H)\)
BOS, RECALL, EOS	Begin an episode, request the initial movement, end an episode	Control tokens; they do not change the displayed coordinates
Recall / recall accuracy / recall CE	Repeat the initial cue / fraction correct at that request / cross-entropy at that request	Score only the input position containing RECALL; not classification sensitivity
Episode / prefix	A complete command sequence / the portion observed so far	States reset at episode boundaries

Representations, parameters, and predictions

Symbol or term	Meaning	Shape or code connection
\(x_t\)	Numerical representation of the current token	\((K,)\) one-hot vector in Scene 4; \((D,)\) embedding in trained recurrent models
Embedding	A learned vector selected from a table by a token ID	nn.Embedding; categorical IDs have no numerical magnitude
\(h_t\), \(z_t\)	Recurrent hidden state and its preactivation	\((H,)\); \(z_t\) is the affine sum before the activation
\(h^{(\ell)}\), \(z^{(\ell)}\)	Representation and preactivation at layer \(\ell\)	Superscripts label layers; subscripts label time
\(W^{(\ell)}\), \(b^{(\ell)}\), \(\theta\)	Layer weights, biases, and their full collection	Course notation retained from Week 04
\(W_x\), \(W_h\), \(b_h\)	Input-to-state weights, state-to-state weights, and state bias	\((H,K)\), \((H,H)\), \((H,)\) in the numerical RNN
\(W_o\), \(b_o\)	State-to-output weights and output bias	\((K,H)\), \((K,)\) in the numerical RNN
\(\phi\), \(\sigma\), \(\odot\)	Activation, logistic sigmoid, and elementwise multiplication	Numerical RNN uses \(\phi=\tanh\); gates use \(\sigma\)
\(a_t\) (logits)	One unnormalized score for each possible next token	\((K,)\); batch logits have shape \((B,T,K)\)
\(p_t\), \(q(y_t\mid s_{\leq t})\)	Predicted distribution and probability assigned to the observed target	Softmax converts logits to probabilities summing to one
\(\ell_t\), \(\mathcal L\), \(\mathcal L_{\mathcal B}\)	Loss at a position and a mean loss over the stated positions/batch	\(\ell_t=-\ln p_t[y_t]\); cross-entropy (CE) is measured in nats
\(g^a_t\), \(g^h_t\), \(g^z_t\)	Loss derivatives with respect to logits, hidden state, and preactivation	Same shapes as the quantities differentiated; used in Scene 6
\(N_{ab}\), \(\alpha\)	Training count of transition \(a\to b\) and additive smoothing amount	Bigram counts; \(\alpha=1\)
\(\ln\), \(\log_2\), \(\mathbb E\)	Natural log, base-two log, and expectation over the stated generator	Natural log is used in code for cross-entropy
\(\operatorname{onehot}(y_t)\)	Vector with one at the target ID and zero elsewhere	\((K,)\); used in the output-gradient calculation
\(T_\mu\)	Number of positions in the small manual-gradient example	\(T_\mu=3\); distinct from the full episode length
Teacher forcing	Predict every next token from the observed prefix during training/evaluation	The input supplies actual earlier tokens, rather than generated guesses
Accuracy	Fraction of targets whose token matches the largest-probability prediction	Equal scores use the lowest token ID
Nat / bit	Loss units using natural / base-two logarithms	Divide nats by \(\ln 2\) to obtain bits

Gates and attention

Symbol or term	Meaning	Shape or code connection
\(u_t\), \(r_t\), \(\widetilde h_t\)	GRU retention/update gate, reset gate, and candidate state	\((H,)\); \(u_t\) near one retains the previous coordinate
\(W_{xu},W_{hu}\) (and \(r,n\) variants), \(b_{xu},b_{hu}\)	Input/recurrent affine parameters for GRU update, reset, and candidate calculations	Input weights \((H,D)\), recurrent weights \((H,H)\), biases \((H,)\); suffixes \(u,r,n\) name roles
\(c_t\), \(\widetilde c_t\), \(f_t\), \(i_t\), \(o_t\)	LSTM cell state, candidate, forget, input, and output gates	State-width vectors; Scene 9 also uses a scalar example
\(Q_{\rm att}\), \(K_{\rm att}\), \(V_{\rm att}\)	Query, key, and value matrices for one attention head	\((T,d_k)\), \((T,d_k)\), \((T,d_v)\); keys are distinct from vocabulary size \(K\)
\(d_k\), \(d_v\), \(N_{\rm heads}\)	Query/key width, value width, and number of attention heads	Two heads of width six in the trained Transformer
\(m\)	Source-position index in an attention normalization sum	Runs from \(0\) to \(T-1\) in Scene 10
\(S\), \(M\)	Scaled query–key scores and the causal mask	\((T,T)\); rows are queries, columns are sources
\(A\), \(O\)	Attention weights and their weighted value output	\((T,T)\) and \((T,d_v)\); each row of \(A\) sums to one
Causal mask	Restriction to current and earlier input positions	\(M_{tj}=0\) for \(j\leq t\) and \(-\infty\) for future positions
\(\operatorname{LN}\), \(\epsilon\)	Feature-wise layer normalization and its numerical stabilizer	Normalization acts within each token; \(\epsilon=10^{-5}\)
MLP / ReLU / residual	Multilayer perceptron / rectified linear unit activation \(\max(0,z)\) / addition of a branch to its input	Transformer MLP acts separately at each token
Train / validation / test	Fit parameters / select checkpoints / evaluate frozen choices	Whole episodes are separated before constructing shifted arrays

Operation convention. Vectors are rows; stored weights are output-by-input. Thus \(x_tW_x^\top\) maps inputs to state coordinates. An attention head uses \(Q_{\rm att}K_{\rm att}^\top\) to match positions. Coordinates in robot figures illustrate motion; the models receive command tokens.

# @title Setup: run once before the scenes
import time
import re
from pathlib import Path
import torch
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from IPython.display import HTML, display

started = time.perf_counter()
TOKENS = ("UP", "DOWN", "LEFT", "RIGHT", "RECALL", "BOS", "EOS")
UP, DOWN, LEFT, RIGHT, RECALL, BOS, EOS = range(7)
K = len(TOKENS)
DELAY = 8
SEEDS = {"train": 6101, "validation": 6102, "test": 6103}
SIZES = {"train": 4096, "validation": 512, "test": 1024}
NAVY, BLUE, GRAY = "#172033", "#2457d6", "#5f6878"
plt.rcParams.update({"svg.hashsalt": "week06", "font.size": 11, "axes.spines.top": False,
                     "axes.spines.right": False, "text.color": NAVY})
# Export into the repository when available; standalone Colab uses ./week06-figures.
root = next((p for p in (Path.cwd(), *Path.cwd().parents)
             if (p / "_quarto.yml").exists()), None)
FIGURES = root / "assets/figures/shared/week06" if root else Path("week06-figures")
FIGURES.mkdir(parents=True, exist_ok=True)

def save_figure(fig, name):
    for suffix in ("svg", "png"):
        metadata = {"Date": None} if suffix == "svg" else {}
        fig.savefig(FIGURES / f"{name}.{suffix}", bbox_inches="tight", dpi=140, metadata=metadata)

Auxiliary visual helpers

Run the setup cell and the figure-helper cells once, then continue to Scene 1. They configure the notebook and define figures; model training begins in Scene 7. The plotting functions are collected here so the scenes can place calculations beside their illustrations.

Expand a helper only when you want to edit its presentation. Sequence construction, model operations, training, and evaluation remain in their scenes. Playback uses locally generated HTML with text controls; the adjacent static figure and caption carry the same teaching content.

# @title Auxiliary: visualization and export helpers
DELTAS = np.array([[0, 1], [0, -1], [-1, 0], [1, 0], [0, 0], [0, 0], [0, 0]])
def positions(sequence):
    return np.cumsum(DELTAS[np.asarray(sequence)], axis=0)

def trajectory_panels(prefixes, stop, reveal=False):
    fig, axes = plt.subplots(1, 2, figsize=(10, 4.1))
    for ax, seq in zip(axes, prefixes):
        xy = positions(seq[:stop])
        ax.plot(xy[:, 0], xy[:, 1], "o-", color=BLUE, lw=2)
        ax.scatter(*xy[-1], s=140, marker="s", color=NAVY, zorder=3)
        ax.scatter(0, 0, s=90, marker="x", color=GRAY, label="start")
        ax.set(xlim=(-2, 3), ylim=(-3, 3), xlabel="horizontal position",
               ylabel="vertical position", title=f"Initial command: {TOKENS[seq[1]]}")
        ax.set_aspect("equal"); ax.grid(alpha=.25); ax.legend(loc="upper left")
        label = " ".join(TOKENS[i] for i in seq[:stop])
        ax.text(.5, -.25, label, transform=ax.transAxes, ha="center", fontsize=9)
        if reveal:
            ax.text(.5, 1.02, f"After RECALL: {TOKENS[seq[1]]}",
                    transform=ax.transAxes, ha="center", color=BLUE)
    fig.subplots_adjust(bottom=.27, wspace=.35)
    return fig

def prefix_animation(prefixes):
    fig, axes = plt.subplots(1, 2, figsize=(10, 4.2))
    def draw(frame):
        for ax, seq in zip(axes, prefixes):
            ax.clear(); xy = positions(seq[:frame + 1])
            ax.plot(xy[:, 0], xy[:, 1], "o-", color=BLUE)
            ax.scatter(*xy[-1], s=140, marker="s", color=NAVY)
            ax.set(xlim=(-2, 3), ylim=(-3, 3), xlabel="horizontal", ylabel="vertical",
                   title=f"Cue {TOKENS[seq[1]]} | observed: {TOKENS[seq[frame]]}")
            ax.set_aspect("equal"); ax.grid(alpha=.25)
        fig.suptitle("At RECALL, pause: which command comes next?", color=NAVY)
    animation = FuncAnimation(fig, draw, frames=len(prefixes[0]), interval=800, repeat=False)
    html = animation.to_jshtml(default_mode="once")
    # Matplotlib includes an optional icon-font CDN; replace it with text controls.
    html = re.sub(r'<link[^>]*>', '', html)
    # Text controls need flexible widths; the animation must fit the page column.
    html = html.replace('width: 36px;', 'width: auto; padding: 4px 7px; margin: 2px;')
    html = html.replace('display: inline-block;', 'display: block; max-width: 100%;')
    html = html.replace('</style>', '.animation img {max-width: 100%; height: auto;}\n</style>')
    labels = {"fast-backward": "First", "backward": "Previous", "step-backward": "Previous",
              "step-forward": "Next", "play": "Play",
              "pause": "Pause", "stop": "Pause", "forward": "Next",
              "fast-forward": "Last", "minus": "Slower", "plus": "Faster"}
    html = html.replace('<i class="fa fa-play fa-flip-horizontal"></i>', "Reverse")
    html = re.sub(r'<i class="fa fa-([\w-]+)"></i>',
                  lambda m: labels.get(m.group(1), m.group(1).replace("-", " ")), html)
    assert "https://" not in html and "http://" not in html
    # JS initializes direction=0 and timer=null; playback begins only on interaction.
    plt.close(fig)
    return HTML(html)


def plot_frequencies(frequencies):
    fig, axes = plt.subplots(1, 2, figsize=(11, 3.5))
    for offset, name in zip((-.18, .18), ("train", "validation")):
        axes[0].bar(np.arange(K) + offset, frequencies[name][0],
                    width=.36, label=name, color=BLUE if name == "train" else GRAY)
        axes[1].bar(np.arange(4) + offset, frequencies[name][1],
                    width=.36, label=name, color=BLUE if name == "train" else GRAY)
    axes[0].set(xticks=np.arange(K), xticklabels=TOKENS, ylabel="fraction", title="All target tokens")
    axes[1].set(xticks=np.arange(4), xticklabels=TOKENS[:4], ylabel="fraction", title="Recall targets only")
    axes[1].axhline(.25, color=GRAY, ls="--", label="iid population: 0.25")
    for ax in axes:
        ax.tick_params(axis="x", labelrotation=30)
    handles, labels = axes[1].get_legend_handles_labels()
    fig.legend(handles, labels, loc="upper center", ncol=3, fontsize=9)
    fig.tight_layout(rect=(0, 0, 1, .87))
    save_figure(fig, "02-target-frequencies"); plt.show()

# @title Auxiliary: recurrent numerical figures
def plot_recurrent_step(row):
    names = ("previous", "input", "memory", "z", "h")
    titles = (r"Previous $h_0$", r"Input $x_1W_x^T$", r"Memory $h_0W_h^T$",
              r"Sum + bias: $z_1$", r"Tanh: $h_1$")
    fig, axes = plt.subplots(1, 5, figsize=(12, 3.1))
    for ax, name, title in zip(axes, names, titles):
        values = row[name].reshape(2, 1)
        ax.imshow(values, cmap="RdBu_r", vmin=-.35, vmax=.35, aspect="auto")
        for j, value in enumerate(row[name]):
            ax.text(0, j, f"{value:+.4f}", ha="center", va="center",
                    color="white" if abs(value) > .23 else NAVY)
        ax.set(xticks=[], yticks=[0, 1], yticklabels=["unit 0", "unit 1"], title=title)
    fig.suptitle("One recurrent step: current input and previous state meet", color=NAVY)
    fig.tight_layout(rect=(0, 0, 1, .88)); save_figure(fig, "03-rnn-step-ledger"); plt.show()

def plot_unrolled_rnn(rows, ids, targets):
    fig, ax = plt.subplots(figsize=(11, 4.7))
    ax.set(xlim=(-.7, 3.0), ylim=(-1.5, 2.1)); ax.axis("off")
    for t, row in enumerate(rows):
        ax.text(t, 1.65, f"input {TOKENS[ids[t]]}\n$x_{t}$", ha="center", color=NAVY)
        ax.text(t, .45, f"$h_{t}$ = [{row['h'][0]:+.3f}, {row['h'][1]:+.3f}]\n"
                r"same $W_x,W_h,b_h$", ha="center", va="center",
                bbox=dict(boxstyle="round,pad=.5", fc="#eef3ff", ec=BLUE))
        ax.annotate("", xy=(t, .8), xytext=(t, 1.35), arrowprops=dict(arrowstyle="->", color=BLUE))
        ax.annotate("", xy=(t, -.5), xytext=(t, .1), arrowprops=dict(arrowstyle="->", color=BLUE))
        ax.text(t, -.85, f"$p_{t}$ -> target {TOKENS[targets[t]]}\n"
                f"loss $\\ell_{t}$ = {row['loss']:.4f}\n" + r"same $W_o,b_o$",
                ha="center", va="center", color=NAVY)
        if t:
            ax.annotate("", xy=(t-.32, .45), xytext=(t-1+.32, .45),
                        arrowprops=dict(arrowstyle="->", color=BLUE, lw=2))
    ax.text(-.65, .45, r"$h_{-1}=0$", ha="center", va="center", color=GRAY)
    ax.annotate("", xy=(-.32, .45), xytext=(-.52, .45), arrowprops=dict(arrowstyle="->", color=GRAY))
    ax.set_title("Three uses of one recurrent cell; three uses of one output head", color=NAVY)
    fig.tight_layout(); save_figure(fig, "04-rnn-unrolled"); plt.show()

def plot_shared_gradients(contributions):
    matrices = [g["W_h"] for g in contributions]
    matrices.append(sum(matrices))
    limit = max(np.abs(a).max() for a in matrices)
    fig, axes = plt.subplots(1, 4, figsize=(11, 3.2))
    for ax, values, title in zip(axes, matrices, ["use t=0", "use t=1", "use t=2", "sum: stored gradient"]):
        ax.imshow(values, cmap="RdBu_r", vmin=-limit, vmax=limit)
        for j in range(H):
            for k in range(H):
                ax.text(k, j, f"{values[j,k]:+.5f}", ha="center", va="center",
                        color="white" if abs(values[j,k]) > .65*limit else NAVY, fontsize=10)
        ax.set(xticks=range(H), yticks=range(H), xlabel="previous hidden unit",
               ylabel="current hidden unit", title=title)
    fig.suptitle(r"$\partial\mathcal{L}/\partial W_h$: sum contributions from every shared use", color=NAVY)
    fig.tight_layout(rect=(0, 0, 1, .87)); save_figure(fig, "05-rnn-shared-gradients"); plt.show()

def plot_rnn_anatomy():
    from matplotlib.patches import FancyBboxPatch, Circle
    fig, axes = plt.subplots(1, 2, figsize=(12, 4.6))
    for ax in axes:
        ax.set(xlim=(0, 10), ylim=(0, 7)); ax.axis("off")
    def box(ax, x, y, w, h, text):
        ax.add_patch(FancyBboxPatch((x,y),w,h,boxstyle="round,pad=.08",
                                   facecolor="#eef3ff",edgecolor=BLUE))
        ax.text(x+w/2,y+h/2,text,ha="center",va="center",color=NAVY,fontsize=11)
    def arrow(ax, start, end):
        ax.annotate("",xy=end,xytext=start,
                    arrowprops=dict(arrowstyle="->",color=GRAY,lw=1.5))
    ax=axes[0]; ax.set_title("One cell carries a vector between steps",color=NAVY)
    box(ax,.1,4.4,2.2,1.0,r"$h_{t-1}$" + "\ntwo numbers")
    box(ax,.1,1.8,2.2,1.0,r"$x_t$" + "\nseven numbers")
    box(ax,3.4,3.0,2.4,1.4,"weighted sum\nthen tanh")
    arrow(ax,(2.4,4.9),(3.3,4.0)); arrow(ax,(2.4,2.3),(3.3,3.4))
    box(ax,6.7,3.2,2.5,1.0,r"$h_t$" + "\ntwo numbers")
    arrow(ax,(5.9,3.7),(6.6,3.7))
    ax.text(7.9,2.35,"to the next step",ha="center",color=NAVY)
    arrow(ax,(7.9,3.1),(7.9,2.65))
    ax.text(5, .75,"An output head also reads this state\nto predict the next command.",ha="center",color=NAVY)
    ax=axes[1]; ax.set_title("Inside: two ordinary tanh neurons",color=NAVY)
    box(ax,.1,4.7,2.3,1.0,r"all inputs $x_{t,k}$")
    for j,y in enumerate([3.5,2.0]):
        ax.add_patch(Circle((1.25,y),.38,facecolor="#eef3ff",edgecolor=BLUE))
        ax.text(1.25,y,rf"$h_{{t-1,{j+1}}}$",ha="center",va="center",fontsize=10)
    for j,y in enumerate([4.6,2.6]):
        ax.add_patch(Circle((6.0,y),.65,facecolor="#eef3ff",edgecolor=BLUE))
        ax.text(6,y,"sum\ntanh",ha="center",va="center",fontsize=10)
        arrow(ax,(2.5,5.2),(5.4,y+.25))
        for previous_y in [3.5,2.0]: arrow(ax,(1.7,previous_y),(5.4,y-.2))
        arrow(ax,(6.7,y),(8.1,y))
        ax.text(8.3,y,rf"$h_{{t,{j+1}}}$",ha="left",va="center",color=NAVY)
    ax.text(5,.65,r"Each neuron has its own weights and bias.",ha="center",color=NAVY)
    fig.tight_layout(); save_figure(fig,"20-rnn-cell-anatomy"); plt.show()

# @title Auxiliary: recurrent experiment figures
def plot_recurrent_curves(run, stem="06-rnn-learning", model_label="RNN"):
    history = run["history"]
    epochs = [r["epoch"] for r in history]
    fig, axes = plt.subplots(1, 3, figsize=(12, 3.6))
    axes[0].plot(epochs, [r["train_all_ce"] for r in history], color=GRAY, label="train, online")
    axes[0].plot(epochs, [r["all_ce"] for r in history], color=BLUE, label="validation")
    axes[0].set(title="All-token cross-entropy", ylabel="nats / token")
    axes[0].legend(fontsize=9)
    axes[1].plot(epochs, [r["recall_ce"] for r in history], color=BLUE)
    axes[1].axhline(np.log(4), color=GRAY, ls="--", label="uniform over movements")
    axes[1].set(title="Validation recall cross-entropy", ylabel="nats / recall")
    axes[1].legend(fontsize=8)
    axes[2].plot(epochs, [r["recall_accuracy"] for r in history], color=BLUE)
    axes[2].axhline(.25, color=GRAY, ls="--")
    axes[2].set(title="Validation recall accuracy", ylim=(0, 1.04), ylabel="fraction correct")
    for ax in axes:
        ax.axvline(run["best_epoch"], color=NAVY, ls=":"); ax.set_xlabel("epoch")
    fig.suptitle(f"Delay 8 {model_label}: dotted line is the validation-selected checkpoint", color=NAVY)
    fig.tight_layout(rect=(0, 0, 1, .88)); save_figure(fig, stem); plt.show()

def plot_recall_predictions(model, ids, targets, delay, stem):
    model.eval()
    with torch.no_grad():
        p = model(ids[:4])[:, delay+2].softmax(-1).numpy()
    fig, axes = plt.subplots(2, 4, figsize=(12, 6.4))
    for j in range(4):
        prefix = ids[j, :delay+3].numpy()
        xy_path = positions(prefix)
        ax = axes[0, j]
        ax.plot(xy_path[:, 0], xy_path[:, 1], "o-", color=GRAY, markersize=3)
        current = xy_path[-1]
        truth, prediction = int(targets[j, delay+2]), int(p[j].argmax())
        ax.scatter(*current, marker="s", s=55, color=NAVY)
        for token, color, label, shift in ((truth, BLUE, "actual", 0), (prediction, "#d1495b", "predicted", .08)):
            delta = DELTAS[token]
            ax.arrow(current[0]+shift, current[1]+shift, delta[0]*.8, delta[1]*.8,
                     head_width=.18, color=color, length_includes_head=True, label=label)
        ax.set_aspect("equal", adjustable="datalim"); ax.margins(.3); ax.grid(alpha=.2)
        ax.set_title(f"Validation episode {j}\ncue {TOKENS[int(prefix[1])]} | guess {TOKENS[prediction]}", fontsize=10)
        ax.set(xlabel="horizontal", ylabel="vertical")
        axes[1, j].bar(np.arange(K), p[j], color=[BLUE if k == truth else GRAY for k in range(K)])
        axes[1, j].set(xticks=np.arange(K), xticklabels=TOKENS, ylim=(0, 1), ylabel="probability")
        axes[1, j].tick_params(axis="x", labelrotation=60, labelsize=8)
    handles, labels = axes[0, 0].get_legend_handles_labels()
    fig.legend(handles, labels, loc="upper center", ncol=2)
    fig.tight_layout(rect=(0, 0, 1, .93)); save_figure(fig, stem); plt.show()

def plot_memory_comparison(runs, kinds, stem):
    delays = (8, 16, 32)
    fig, axes = plt.subplots(1, 2, figsize=(10, 3.7))
    for kind, color in zip(kinds, (BLUE, "#d1495b", "#008577")):
        axes[0].plot(delays, [runs[(kind,d)]["validation"]["recall_accuracy"] for d in delays], "o-", label=kind, color=color)
        axes[1].plot(delays, [runs[(kind,d)]["validation"]["recall_ce"] for d in delays], "o-", label=kind, color=color)
    axes[0].axhline(.25, color=GRAY, ls="--", label="population fixed-guess reference")
    axes[1].axhline(np.log(4), color=GRAY, ls="--", label="uniform movement CE")
    axes[0].set(ylabel="fraction correct", title="Validation recall accuracy", ylim=(0, 1.04))
    axes[1].set(ylabel="nats / recall", title="Validation recall cross-entropy", ylim=(0, 1.65))
    for ax in axes: ax.set_xlabel("number of intervening commands d"); ax.set_xticks(delays); ax.legend(fontsize=8)
    fig.tight_layout(); save_figure(fig, stem); plt.show()

def plot_gradient_reach(diagnostics):
    fig, ax = plt.subplots(figsize=(8, 3.8))
    for delay, values in diagnostics.items():
        distances = np.arange(delay+3)[::-1]
        ax.semilogy(distances, np.maximum(values[:delay+3], 1e-12), "o-", markersize=3, label=f"delay {delay}")
    ax.set(xlabel="positions before the RECALL input", ylabel="mean gradient norm at embedding output",
           title="RNN selected checkpoints: recall-loss sensitivity along the prefix")
    ax.legend(); ax.grid(alpha=.2)
    fig.tight_layout(); save_figure(fig, "09-rnn-gradient-reach"); plt.show()

def plot_lstm_cell():
    fig, ax = plt.subplots(figsize=(10, 3.3)); ax.axis("off")
    ax.set(xlim=(-.2, 4.3), ylim=(-.4, 2))
    for x, text in [(0, r"$c_{t-1}$"), (1, r"retain: $f_t\odot c_{t-1}$"),
                    (2.5, r"add: $c_t$"), (4, r"read: $h_t=o_t\odot\tanh(c_t)$")]:
        ax.text(x, 1, text, ha="center", va="center", fontsize=12,
                bbox=dict(boxstyle="round,pad=.5", fc="#eef3ff", ec=BLUE))
    for left, right in [(0.25,.55),(1.5,2.15),(2.8,3.35)]:
        ax.annotate("", xy=(right, 1), xytext=(left, 1), arrowprops=dict(arrowstyle="->", color=BLUE))
    ax.text(2.5, 0, r"write: $i_t\odot\widetilde c_t$", ha="center", color=NAVY)
    ax.annotate("", xy=(2.5,.75), xytext=(2.5,.2), arrowprops=dict(arrowstyle="->", color=BLUE))
    ax.text(1.6, 1.75, r"Gates and candidate depend on current input and previous $h_{t-1}$", ha="center", color=GRAY)
    fig.tight_layout(); save_figure(fig, "11-lstm-cell-path"); plt.show()

# @title Auxiliary: attention figures
def attention_grid(ax, values, row_labels, column_labels, title, blocked=None, probabilities=False):
    shown = np.where(np.isfinite(values), values, 0)
    ax.imshow(shown, cmap="Blues" if probabilities else "RdBu_r",
              vmin=0 if probabilities else -max(abs(shown).max(), 1),
              vmax=1 if probabilities else max(abs(shown).max(), 1), aspect="auto")
    for t in range(values.shape[0]):
        for j in range(values.shape[1]):
            text = "blocked" if blocked is not None and blocked[t,j] else f"{values[t,j]:.2f}"
            is_blocked = blocked is not None and blocked[t,j]
            dark = values[t,j] > .6 if probabilities else abs(values[t,j]) > .65*max(abs(shown).max(), 1)
            color = "white" if dark and not is_blocked else NAVY
            ax.text(j, t, text, ha="center", va="center", fontsize=9, color=color)
    ax.set(xticks=range(len(column_labels)), xticklabels=column_labels,
           yticks=range(len(row_labels)), yticklabels=row_labels, title=title)
    ax.tick_params(axis="x", labelrotation=45)

def plot_attention_calculation():
    labels = [f"{t}: {TOKENS[token]}" for t,token in enumerate(attention_tokens)]
    fig, axes = plt.subplots(2, 2, figsize=(10, 8))
    attention_grid(axes[0,0], masked_scores_demo, labels, labels, "Scaled scores + causal mask", future_mask_demo)
    attention_grid(axes[0,1], weights_demo, labels, labels, "Weights: softmax over source columns", probabilities=True)
    attention_grid(axes[1,0], V_demo, labels, ["value 0", "value 1"], "Source values (hand-chosen)")
    attention_grid(axes[1,1], output_demo, labels, ["output 0", "output 1"], "Weighted outputs: A @ V")
    for ax in axes[0]: ax.set_ylabel("query position"); ax.set_xlabel("source position")
    fig.tight_layout(); save_figure(fig, "12-attention-calculation"); plt.show()

def plot_mask_comparison():
    labels = [f"{t}: {TOKENS[token]}" for t,token in enumerate(attention_tokens)]
    fig, axes = plt.subplots(1, 2, figsize=(10, 4.4))
    for ax, allowed, title in zip(axes, (~future_mask_demo, np.ones((5,5), dtype=bool)),
                                  ("Causal: earlier and current input", "Bidirectional: all supplied input")):
        ax.imshow(allowed, cmap="Blues", vmin=0, vmax=1)
        ax.set(xticks=range(5), xticklabels=labels, yticks=range(5), yticklabels=labels,
               xlabel="source position", ylabel="query position", title=title)
        ax.tick_params(axis="x", labelrotation=45)
        ax.add_patch(plt.Rectangle((3.5, 2.5), 1, 1, fill=False, ec="#d1495b", lw=3))
    fig.suptitle("Red outline: RECALL query reading its future UP target", color=NAVY)
    fig.tight_layout(rect=(0, 0, 1, .88)); save_figure(fig, "13-causal-bidirectional-masks"); plt.show()

def plot_trained_attention(model, ids):
    model.eval()
    with torch.no_grad():
        _, weights = model(ids[:1], return_attention=True)
    weights = weights[0].numpy()
    labels = [f"{t}:{TOKENS[int(token)]}" for t,token in enumerate(ids[0])]
    fig, axes = plt.subplots(1, 2, figsize=(12, 5.5))
    for head, ax in enumerate(axes):
        image = ax.imshow(weights[head], cmap="Blues", vmin=0, vmax=1)
        fig.colorbar(image, ax=ax, fraction=.046, pad=.04, label="attention weight")
        ax.set(xticks=range(len(labels)), xticklabels=labels, yticks=range(len(labels)), yticklabels=labels,
               xlabel="source position", ylabel="query position", title=f"Head {head}: validation episode 0")
        ax.tick_params(axis="x", labelrotation=90, labelsize=8); ax.tick_params(axis="y", labelsize=8)
        ax.add_patch(plt.Rectangle((.5, 9.5), 1, 1, fill=False, ec="#d1495b", lw=2))
    fig.suptitle("Attention weights; red outline marks RECALL-to-initial-cue edge", color=NAVY)
    fig.tight_layout(rect=(0, 0, 1, .90)); save_figure(fig, "16-transformer-attention"); plt.show()

def teaching_box(ax, x, y, w, h, text):
    from matplotlib.patches import FancyBboxPatch
    ax.add_patch(FancyBboxPatch((x,y),w,h,boxstyle="round,pad=.07",
                               facecolor="#eef3ff",edgecolor=BLUE))
    ax.text(x+w/2,y+h/2,text,ha="center",va="center",color=NAVY,fontsize=12)

def teaching_arrow(ax, start, end):
    ax.annotate("",xy=end,xytext=start,
                arrowprops=dict(arrowstyle="->",color=GRAY,lw=1.6))

def plot_gru_carry():
    fig,ax=plt.subplots(figsize=(10,4.1));ax.set(xlim=(0,10),ylim=(0,5));ax.axis("off")
    teaching_box(ax,.1,3.3,2.0,.95,"old state\n0.50")
    teaching_box(ax,.1,1.1,2.0,.95,"candidate\n−0.40")
    teaching_box(ax,3.0,3.3,2.0,.95,"retain × 0.9\n= 0.45")
    teaching_box(ax,3.0,1.1,2.0,.95,"write × 0.1\n= −0.04")
    teaching_box(ax,7.2,2.2,2.5,1.0,"new state\n0.41")
    teaching_arrow(ax,(2.2,3.8),(2.9,3.8));teaching_arrow(ax,(2.2,1.6),(2.9,1.6))
    teaching_arrow(ax,(5.1,3.8),(6.1,2.7));teaching_arrow(ax,(5.1,1.6),(6.1,2.7))
    ax.text(6.4,2.7,"+",fontsize=24,ha="center",va="center",color=NAVY)
    teaching_arrow(ax,(6.7,2.7),(7.1,2.7))
    ax.text(5,.25,r"One coordinate: $h_t=u_t h_{t-1}+(1-u_t)\widetilde h_t$, with $u_t=0.9$",
            ha="center",color=NAVY,fontsize=12)
    ax.set_title("A gate chooses the mixture; it does not choose a command",color=NAVY)
    fig.tight_layout();save_figure(fig,"21-gru-carry-mixture");plt.show()

def plot_attention_route():
    fig,ax=plt.subplots(figsize=(10,5.2));ax.set(xlim=(0,10),ylim=(0,7));ax.axis("off")
    teaching_box(ax,.1,5.4,2.5,.9,"query at RECALL\nwhat to match")
    teaching_box(ax,.1,3.6,2.5,.9,"source keys\nwhat can match")
    teaching_box(ax,3.5,4.5,2.6,.9,"dot products\nand scaling")
    teaching_arrow(ax,(2.7,5.85),(3.4,5.1));teaching_arrow(ax,(2.7,4.05),(3.4,4.8))
    teaching_box(ax,7.1,4.5,2.6,.9,"mask future\nthen softmax")
    teaching_arrow(ax,(6.2,4.95),(7.0,4.95))
    teaching_box(ax,.1,1.5,2.5,.9,"source values\ncontent to combine")
    teaching_box(ax,4.8,1.5,3.0,.9,"weighted sum\nretrieved representation")
    teaching_arrow(ax,(2.7,1.95),(4.7,1.95))
    teaching_arrow(ax,(8.4,4.4),(6.6,2.5))
    ax.text(8.6,3.15,"weights",ha="center",color=NAVY,fontsize=12)
    ax.text(5,.3,"Only permitted sources contribute. The output is a vector, not a token prediction.",
            ha="center",color=NAVY,fontsize=11)
    ax.set_title("Follow one query through causal attention",color=NAVY)
    fig.tight_layout();save_figure(fig,"22-attention-route");plt.show()

def plot_transformer_block():
    fig,ax=plt.subplots(figsize=(10,6.8));ax.set(xlim=(0,10),ylim=(0,9));ax.axis("off")
    teaching_box(ax,3.1,7.6,4.0,.85,"token + position embeddings")
    teaching_box(ax,3.1,5.5,4.0,1.05,"normalize → causal attention\n→ output projection")
    teaching_box(ax,3.1,3.15,4.0,1.05,"normalize → position-wise MLP\n12 → 24 → 12")
    teaching_box(ax,3.1,.6,4.0,1.0,"normalize → output head\nseven next-command logits")
    teaching_arrow(ax,(5.1,7.5),(5.1,6.65))
    teaching_arrow(ax,(5.1,5.4),(5.1,4.8));ax.text(5.1,4.6,"+",ha="center",va="center",fontsize=23,color=NAVY)
    teaching_arrow(ax,(5.1,4.4),(5.1,4.25))
    teaching_arrow(ax,(5.1,3.05),(5.1,2.45));ax.text(5.1,2.25,"+",ha="center",va="center",fontsize=23,color=NAVY)
    teaching_arrow(ax,(5.1,1.95),(5.1,1.7))
    ax.plot([5.1,1.8,1.8,4.6],[7.0,7.0,4.6,4.6],color=GRAY,lw=1.5)
    teaching_arrow(ax,(4.4,4.6),(4.8,4.6))
    ax.plot([5.1,1.8,1.8,4.6],[4.4,4.4,2.25,2.25],color=GRAY,lw=1.5)
    teaching_arrow(ax,(4.4,2.25),(4.8,2.25))
    ax.text(.8,5.7,"carry\ninput",ha="center",color=NAVY,fontsize=11)
    ax.text(.8,3.15,"carry\ninput",ha="center",color=NAVY,fontsize=11)
    ax.text(8.5,6.0,"mixes permitted\npositions",ha="center",color=NAVY,fontsize=11)
    ax.text(8.5,3.65,"acts within\neach position",ha="center",color=NAVY,fontsize=11)
    ax.set_title("One pre-normalization Transformer block",color=NAVY)
    fig.tight_layout();save_figure(fig,"23-transformer-block-route");plt.show()

# @title Auxiliary: final comparison figures
def plot_test_summary():
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.2))
    colors = (BLUE, "#d1495b", "#008577")
    for k, (kind, color) in enumerate(zip(FINAL_KINDS, colors)):
        for ax, metric, index in zip(axes, ("recall_accuracy", "recall_ce"), (3, 2)):
            means = [summary[(kind,d)]["mean"][index] for d in FINAL_DELAYS]
            spread = [summary[(kind,d)]["sd"][index] for d in FINAL_DELAYS]
            ax.errorbar(FINAL_DELAYS, means, yerr=spread, fmt="o-", capsize=4, color=color, label=kind)
            for d in FINAL_DELAYS:
                ax.scatter([d+(s-1)*.6 for s in range(3)],
                           [test_results[(kind,d,seed)][metric] for seed in FINAL_SEEDS],
                           s=18, color=color, alpha=.6, zorder=3)
    axes[0].axhline(.25, color=GRAY, ls="--")
    axes[1].axhline(np.log(4), color=GRAY, ls="--")
    axes[0].set(title="Held-out recall accuracy", ylabel="fraction correct", ylim=(-.03,1.08))
    axes[1].set(title="Held-out recall cross-entropy", ylabel="nats / recall")
    for ax in axes:
        ax.set_xlabel("intervening commands d"); ax.set_xticks(FINAL_DELAYS); ax.legend(fontsize=9)
    fig.suptitle("Three initializations: small dots are runs; bars are sample SD, not confidence intervals", fontsize=11)
    fig.tight_layout(rect=(0,0,1,.92)); save_figure(fig, "18-held-out-seed-comparison"); plt.show()

def plot_paired_test():
    x, y = held_out[EXAMPLE_DELAY]
    fig, axes = plt.subplots(3, 4, figsize=(12, 8), sharey=True)
    for row, kind in enumerate(FINAL_KINDS):
        with torch.no_grad():
            p = final_runs[(kind,EXAMPLE_DELAY,EXAMPLE_SEED)]["model"](x[:4])[:,EXAMPLE_DELAY+2].softmax(-1).numpy()
        for j, ax in enumerate(axes[row]):
            truth = int(y[j,EXAMPLE_DELAY+2]); guess = int(p[j].argmax())
            ax.bar(np.arange(K), p[j], color=[BLUE if k==truth else GRAY for k in range(K)])
            ax.set(xticks=np.arange(K), xticklabels=TOKENS, ylim=(0,1.05))
            ax.tick_params(axis="x", labelrotation=60, labelsize=8)
            ax.set_title(f"{kind}: guess {TOKENS[guess]}" + (f"\nTest episode {j}, cue {TOKENS[truth]}" if row==0 else ""), fontsize=10)
            if j==0: ax.set_ylabel("recall probability")
    fig.suptitle("Same first four delay-16 test episodes; seed 6206; blue bar is the true cue", fontsize=12)
    fig.tight_layout(rect=(0,0,1,.94)); save_figure(fig,"19-paired-test-recall"); plt.show()


def plot_target_surprise():
    fig,ax=plt.subplots(figsize=(8,4.6))
    probability=np.linspace(.001,1,600)
    ax.plot(probability,-np.log(probability),color=BLUE,lw=2)
    for q,offset in [(.01,(36,-4)),(.25,(22,24)),(.8,(-95,34))]:
        loss=-np.log(q)
        ax.scatter([q],[loss],color=BLUE,zorder=3)
        ax.annotate(f"q={q:g}: {loss:.2f} nats",(q,loss),xytext=offset,
                    textcoords="offset points",color=NAVY,
                    arrowprops=dict(arrowstyle="-",color=GRAY))
    ax.set(xlim=(0,1),ylim=(0,7.2),xlabel="Probability assigned to the actual next command",
           ylabel="Loss at this position (nats)",title="A confident mistake is costly")
    right=ax.secondary_yaxis("right",functions=(lambda n:n/np.log(2),lambda b:b*np.log(2)))
    right.set_ylabel("Same cost in bits")
    fig.tight_layout();save_figure(fig,"24-target-surprise");plt.show()

# @title Auxiliary: sliding-window, rollout, and horizon-error figures
def plot_wave_window(origin=96):
    assert WAVE_WINDOW <= origin <= WAVE_LENGTH - WAVE_HORIZON
    signal = wave_splits["train"][0]
    index = np.arange(WAVE_LENGTH)
    fig, ax = plt.subplots(figsize=(9, 3.5))
    ax.plot(index, signal, color=GRAY, lw=1, label="training trajectory")
    ax.axvspan(origin-WAVE_WINDOW, origin-1, color=BLUE, alpha=.12, label="observed input: 32 samples")
    ax.plot(index[origin-WAVE_WINDOW:origin], signal[origin-WAVE_WINDOW:origin], color=BLUE, lw=2)
    ax.axvspan(origin, origin+WAVE_HORIZON-1, color="#c67c19", alpha=.12, label="future labels: 64 samples")
    ax.axvline(origin, color=NAVY, ls="--", lw=1)
    ax.set(xlabel="Sample index", ylabel="Measured amplitude", title="Slide the window; keep the future outside the input")
    ax.legend(loc="upper center", bbox_to_anchor=(.5, -.2), ncol=2, fontsize=9)
    fig.tight_layout(); save_figure(fig, "25-wave-observation-window"); plt.show()

WAVE_COLORS = {"Persistence":GRAY, "Linear AR":"#c67c19", "MLP":"#8055aa",
               "RNN":BLUE, "GRU":"#218367", "1D CNN":"#b14765", "CNN → RNN":"#007f99"}

def save_wave_comparison(fig, name):
    save_figure(fig, name)
    svg_path = FIGURES / f"{name}.svg"
    svg_path.write_text("\n".join(line.rstrip() for line in svg_path.read_text().splitlines()) + "\n")

def plot_wave_rollout():
    trajectory, origin = wave_test_origins[WAVE_EXAMPLE_INDEX]
    signal = wave_splits["test"][trajectory]
    past = np.arange(origin-WAVE_WINDOW, origin)
    future = np.arange(origin, origin+WAVE_HORIZON)
    fig, axes = plt.subplots(2, 2, figsize=(11, 9), sharex=True, sharey=True)
    groups = (("References", ("Persistence", "Linear AR")), ("Dense window", ("MLP",)),
              ("Recurrent windows", ("RNN", "GRU")), ("Temporal features", ("1D CNN", "CNN → RNN")))
    for ax, (title, kinds) in zip(axes.flat, groups):
        ax.plot(past, signal[past], color=GRAY, lw=2, label="observed")
        ax.plot(future, signal[future], color="black", lw=2, label="actual future")
        for kind in kinds:
            seed = WAVE_SEEDS[0] if kind in WAVE_NEURAL else None
            predicted = wave_results[(kind, seed)]["predictions"][WAVE_EXAMPLE_INDEX]
            ax.plot(future, predicted, color=WAVE_COLORS[kind], lw=1.5,
                    ls="--" if seed is None else "-", label=kind)
        ax.axvline(origin, color=NAVY, ls=":")
        ax.set(title=title, xlabel="Sample index", ylabel="Amplitude")
        ax.legend(fontsize=8, loc="upper center", bbox_to_anchor=(.5, -.24), ncol=2)
    fig.suptitle("Same held-out origin and axes; seed 6406 for every neural example")
    fig.tight_layout(); save_wave_comparison(fig, "26-wave-recursive-rollout"); plt.show()

def plot_wave_horizon_errors():
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.8), sharex=True, sharey=True)
    horizons = np.arange(1, WAVE_HORIZON+1)
    groups = (("Original predictors", ("Persistence", "Linear AR", "MLP", "RNN", "GRU")),
              ("Convolutional extension", ("Linear AR", "MLP", "1D CNN", "CNN → RNN")))
    for ax, (title, kinds) in zip(axes, groups):
        for kind in kinds:
            errors = np.stack([v["horizon_mse"] for (name, seed), v in wave_results.items() if name == kind])
            mean = errors.mean(axis=0)
            ax.plot(horizons, mean, color=WAVE_COLORS[kind], label=kind)
            if len(errors) > 1:
                sd = errors.std(axis=0, ddof=1)
                ax.fill_between(horizons, np.maximum(mean-sd, 1e-8), mean+sd,
                                color=WAVE_COLORS[kind], alpha=.12)
        ax.axhline(WAVE_NOISE**2, color=NAVY, ls=":", label="noise reference")
        ax.set(yscale="log", xlabel="Samples ahead", ylabel="MSE (original units)", title=title)
        ax.legend(loc="upper center", bbox_to_anchor=(.5, -.18), ncol=2, fontsize=8)
    axes[0].set_ylim(bottom=WAVE_NOISE**2/3)  # clip bands at a stated positive display boundary
    fig.tight_layout(); save_wave_comparison(fig, "27-wave-horizon-errors"); plt.show()

def plot_wave_architectures():
    """Implemented 32-sample routes; schematic features, no model predictions."""
    from matplotlib.patches import FancyBboxPatch
    fig, ax = plt.subplots(figsize=(8.2, 8.8))
    ax.set(xlim=(0, 100), ylim=(0, 112)); ax.axis('off')
    def label(x, y, words, size=12, color=NAVY, **kwargs):
        ax.text(x, y, words, ha='center', va='center', fontsize=size, color=color, **kwargs)
    def box(x, y, width, words, color=BLUE, size=12):
        ax.add_patch(FancyBboxPatch((x, y-4.5), width, 9,
                     boxstyle='round,pad=0.3,rounding_size=1',
                     facecolor='white', edgecolor=color, linewidth=1.5))
        label(x+width/2, y, words, size, color)
    def arrow(x1, y1, x2, y2, color=GRAY):
        ax.annotate('', xy=(x2,y2), xytext=(x1,y1),
                    arrowprops=dict(arrowstyle='->', color=color, lw=1.5))
    label(50, 109, 'One observed window → one next-sample prediction', 15, weight='bold')
    label(50, 104, 'Single-example shapes below; batching adds an axis B.', 11, GRAY)
    rows = [('MLP',93), ('RNN',75), ('GRU',57), ('1D CNN',39), ('CNN → RNN',21)]
    for kind, y in rows:
        color = WAVE_COLORS[kind]
        ax.axhspan(y-7.8, y+8, color=color, alpha=.045)
        ax.text(1, y+6.5, kind, fontsize=13, color=color, weight='bold', va='center')
        box(1, y-1, 16, '32 samples\n32 × 1', color)
        arrow(17.5,y-1,21,y-1)
        box(82, y-1, 17, 'Linear head\n1 scalar', color)
        arrow(78.5,y-1,81.5,y-1)
        if kind == 'MLP':
            box(22,y-1,22,'Flatten\n32 values',color)
            arrow(44.5,y-1,48,y-1)
            box(49,y-1,29,'Dense + ReLU\n16 features',color)
            label(51,y-7,'Combine all positions at once',11,GRAY)
        elif kind in ('RNN','GRU'):
            width = 20 if kind == 'RNN' else 12
            for x, words in [(22,'Step 1'),(42,'Step 2'),(62,'Step 32')]:
                box(x,y-1,16,words,color)
            arrow(38.5,y-1,41.5,y-1);arrow(58.5,y-1,61.5,y-1)
            label(60,y-1,'…',12,GRAY)
            gate = 'tanh update' if kind == 'RNN' else 'gated retain / write'
            label(50,y-7,f'Shared {gate}; last state: {width} values',11,GRAY)
            label(56,y+6.5,'Zero state → read samples in order',11,GRAY)
        elif kind == '1D CNN':
            box(22,y-1,17,'Conv5\n28 × 8',color,11)
            arrow(39.5,y-1,41,y-1)
            box(42,y-1,17,'Conv3\n26 × 8',color,11)
            arrow(59.5,y-1,61,y-1)
            box(62,y-1,16,'Flatten\n208',color,11)
            label(51,y-7,'ReLU after each convolution; head reads all positions',11,GRAY)
        else:
            box(22,y-1,20,'Conv5 + ReLU\n28 × 8',color,11)
            arrow(42.5,y-1,45,y-1)
            box(46,y-1,32,'RNN: 28 feature steps\nLast state: 16 values',color,11)
            label(51,y-7,'Shared local filters → shared recurrent update',11,GRAY)
            label(58,y+6.5,'RNN starts from zero state',11,GRAY)
    label(50, 10, 'ConvK: valid stride-1 filter of length K. Shapes: positions × features.', 11, GRAY)
    label(50, 5, 'Neural inputs use training-only scaling; output is converted back to amplitude.', 11, GRAY)
    fig.tight_layout(pad=.5)
    save_wave_comparison(fig, '30-wave-architecture-map'); plt.show()


def plot_wave_feedback():
    """Feedback protocol shared by all predictors; no persistent RNN state."""
    from matplotlib.patches import FancyBboxPatch
    fig, ax = plt.subplots(figsize=(8.2, 2.25))
    ax.set(xlim=(0,100),ylim=(0,25));ax.axis('off')
    boxes = [(1,25,'Window\n32 samples'),(35,28,'Frozen predictor\nOne of the routes above'),(73,26,'Next amplitude\n1 prediction')]
    for x,w,words in boxes:
        ax.add_patch(FancyBboxPatch((x,12),w,9,boxstyle='round,pad=.3',
                                   facecolor='white',edgecolor=BLUE,lw=1.5))
        ax.text(x+w/2,16.5,words,ha='center',va='center',fontsize=12)
    for a,b in [(26.5,34.5),(63.5,72.5)]:
        ax.annotate('',xy=(b,16.5),xytext=(a,16.5),arrowprops=dict(arrowstyle='->',color=GRAY,lw=1.5))
    ax.plot([86,86,13.5],[11.5,7,7],color=BLUE,lw=1.5)
    ax.annotate('',xy=(13.5,11.5),xytext=(13.5,7),arrowprops=dict(arrowstyle='->',color=BLUE,lw=1.5))
    ax.text(50,3,'Drop oldest sample; append prediction; repeat 64 times.',ha='center',va='center',fontsize=12)
    ax.text(50,24,'Recursive rollout: feedback changes the next input',ha='center',va='center',fontsize=14,weight='bold')
    fig.tight_layout(pad=.5)
    save_wave_comparison(fig,'31-wave-feedback-loop');plt.show()

Scene 1 — Watch, pause, predict

The robot moves one grid unit for UP, DOWN, LEFT, or RIGHT. BOS begins an episode, RECALL asks for the initial command again, and EOS ends it; these control tokens do not move the robot.

Every episode has the form

BOS → initial command → d random movement commands → RECALL → initial command → EOS

The initial command and all \(d\) intervening movements are independently uniform over four directions in the underlying generator.

Checkpoint. These two short prefixes have the same recent commands. What should each robot do after RECALL? Could the most recent token alone tell them apart? Run the next cell only after predicting.

Show explanation — which command follows RECALL?

The first robot must repeat UP, and the second must repeat DOWN. Their recent suffix ends in the same commands, so a predictor using only the current token produces the same distribution for both. Useful evidence is the initial cue.

prefixes = [np.array([BOS, cue, RIGHT, DOWN, LEFT, RECALL]) for cue in (UP, DOWN)]
assert np.array_equal(prefixes[0][2:], prefixes[1][2:])
fig = trajectory_panels(prefixes, stop=len(prefixes[0]))
save_figure(fig, "01-recall-prefixes")
plt.show()
display(prefix_animation(prefixes))

Once Loop Reflect

Figure 1. Two histories, one shared suffix. The same recent commands follow different initial cues. At RECALL, the earlier cue determines which command to repeat. The animation steps through the same prefixes as the static figure.

The square marks the robot; the cross marks the origin. Pause or step through the observed prefix using the playback controls. The static figure above shows both completed prefixes.

The toy trajectories have delay 3 for readability. The experiment below uses delay 8. Position is not the prediction input: the task is to retain an earlier categorical command.

Scene 2 — A movement becomes a token

Checkpoint. Write the input row and its next-token target row for the first complete episode. Which input position is scored as recall?

Remove the last token to obtain inputs; remove the first to obtain targets:

input:  BOS    UP     RIGHT  DOWN   LEFT   RECALL UP
target: UP     RIGHT  DOWN   LEFT   RECALL UP     EOS

Each column is one prediction. When the input is RECALL, its target is the earlier cue UP. Each row remains a complete episode; context ends at its boundary.

Show explanation — where does the model predict recall?

For BOS UP RIGHT DOWN LEFT RECALL UP EOS, inputs are BOS UP RIGHT DOWN LEFT RECALL UP; targets are UP RIGHT DOWN LEFT RECALL UP EOS. At zero-based input position 5, the model reads RECALL and must predict UP. For general delay \(d\), that position is \(d+2\).

example = np.array([BOS, UP, RIGHT, DOWN, LEFT, RECALL, UP, EOS])
example_x, example_y = example[:-1], example[1:]
print("position input   target")
for t, (x, y) in enumerate(zip(example_x, example_y)):
    print(f"{t:8d} {TOKENS[x]:7s} {TOKENS[y]}")
one_hot = np.eye(K)[example_x]
assert np.array_equal(one_hot.argmax(axis=-1), example_x)
assert np.all(one_hot.sum(axis=-1) == 1)
# A transparent toy embedding lookup: rows name tokens, columns name features.
embedding_table = np.arange(K * 3, dtype=float).reshape(K, 3) / 10
embedded = one_hot @ embedding_table
assert np.allclose(embedded, embedding_table[example_x])
print("UP ID:", UP, "one-hot:", one_hot[1], "toy embedding:", embedded[1])
print("One episode:", example_x.shape, "one-hot:", one_hot.shape,
      "embedding:", embedded.shape)

position input   target
       0 BOS     UP
       1 UP      RIGHT
       2 RIGHT   DOWN
       3 DOWN    LEFT
       4 LEFT    RECALL
       5 RECALL  UP
       6 UP      EOS
UP ID: 0 one-hot: [1. 0. 0. 0. 0. 0. 0.] toy embedding: [0.  0.1 0.2]
One episode: (7,) one-hot: (7, 7) embedding: (7, 3)

The toy embedding demonstrates lookup only. Its coordinates are untrained. Later, training will learn the embedding table as part of \(\theta\).

Construct reproducible, separate episodes

Generate whole episodes for training, validation, and test, then form input/target pairs. Reject duplicate episodes so none crosses a split boundary. Reserve test performance for Scene 12; the code below checks shapes, overlap, and reproducibility.

Sampling details — duplicate rejection and the population reference

Reject any whole episode already accepted, including within the same split. This gives exact requested sizes, but samples without replacement: the curated splits are not strictly iid or statistically independent. The original iid generator remains our population reference. At delay 8 it has \(4^9=262,144\) distinct episodes; the code reports rejected draws.

The classroom illustration is excluded from the splits. If a future experiment adds counterfactual families, group all members before splitting. Split first; construct shifted arrays afterward.

def draw_episode(rng, delay):
    moves = rng.integers(0, 4, size=delay + 1)
    return np.concatenate(([BOS], moves, [RECALL, moves[0], EOS]))

def make_splits(delay, sizes=SIZES, seeds=SEEDS):
    if sum(sizes.values()) > 4 ** (delay + 1):
        raise ValueError("Requested more unique episodes than the generator can supply")
    seen, result, rejected = set(), {}, {}
    for name, size in sizes.items():
        rng = np.random.default_rng(seeds[name])
        rows, rejected[name] = [], 0
        while len(rows) < size:
            row = draw_episode(rng, delay)
            key = tuple(row)
            if key in seen:
                rejected[name] += 1
                continue
            seen.add(key); rows.append(row)
        result[name] = np.stack(rows)
    return result, rejected

splits, rejected = make_splits(DELAY)
# Zero-based recall prediction position: the input token is RECALL.
recall_t = DELAY + 2
xy = {name: (rows[:, :-1], rows[:, 1:]) for name, rows in splits.items()}
for name, rows in splits.items():
    x, y = xy[name]
    assert rows.shape == (SIZES[name], DELAY + 5)
    assert x.shape == y.shape == (SIZES[name], DELAY + 4)
    assert np.all(rows[:, 0] == BOS) and np.all(rows[:, -1] == EOS)
    assert np.all(x[:, recall_t] == RECALL)
    assert np.array_equal(y[:, recall_t], rows[:, 1])
    assert np.all((rows[:, 1:DELAY + 2] >= 0) & (rows[:, 1:DELAY + 2] < 4))
    print(f"{name:10s}: {len(rows):4d} episodes | seed {SEEDS[name]} | duplicate draws rejected {rejected[name]}")
keys = {name: set(map(tuple, rows)) for name, rows in splits.items()}
assert not (keys["train"] & keys["validation"] or keys["train"] & keys["test"]
            or keys["validation"] & keys["test"])
repeat, _ = make_splits(DELAY)
assert all(np.array_equal(splits[name], repeat[name]) for name in splits)
print("Shapes, recall alignment, split disjointness, and reproducibility: passed")

train     : 4096 episodes | seed 6101 | duplicate draws rejected 34
validation:  512 episodes | seed 6102 | duplicate draws rejected 9
test      : 1024 episodes | seed 6103 | duplicate draws rejected 20
Shapes, recall alignment, split disjointness, and reproducibility: passed

Scene 3 — Establish what is learnable

Checkpoint. If you know the entire past, can you predict a fresh random wandering command? If you remember the cue, can you predict the command after RECALL?

For the original iid generator, a fixed movement guess or a previous-token-only predictor has expected recall accuracy 25%. Retaining the cue can make recall certain. Fresh random movements remain unpredictable: their optimal cross-entropy is \(\ln4\) and expected accuracy is 25%.

Baseline scope — which predictors have the 25% reference?

The recall reference describes fixed train-fitted frequency and previous-token-only bigram predictions. Counts from the whole prefix include cue information and need not obey that limit. A seven-token uniform baseline also differs: it assigns \(1/7\) to every token, including impossible targets.

Inspect training and validation frequencies before fitting. Small imbalances are sampling variation, not a rule of the language.

Show explanation — what can memory make predictable?

The initial cue and wandering moves are independent uniform draws. Observing their past cannot reveal the next fresh draw. At RECALL, the target equals a command already observed, so remembering that cue can make the prediction certain. This is why we measure recall separately from random positions.

frequencies = {}
for name in ("train", "validation"):
    x, y = xy[name]
    frequencies[name] = (np.bincount(y.ravel(), minlength=K) / y.size,
                         np.bincount(y[:, recall_t], minlength=K)[:4] / len(y))
plot_frequencies(frequencies)

Figure 2. Count the targets before fitting. All-token counts and recall-only counts describe different prediction problems. Small direction imbalances vary across splits; the dashed line is the generator’s 25% reference.

Text alternative. Movement commands account for about 21% each of all targets; RECALL and EOS each account for 1/12, and BOS is never a target. Recall targets span the four directions, near the iid reference of 25%, with visible finite-sample imbalance. The numerical evaluation below quantifies how this affects each fixed predictor.

Fit only on training data

Compare uniform probabilities, empirical target frequencies, and a smoothed bigram predictor. A bigram uses the current token \(a\) to predict the next token \(b\):

\[p(b\mid a)=\frac{N_{ab}+\alpha}{\sum_cN_{ac}+K\alpha},\qquad \alpha=1.\]

It cannot distinguish different prefixes ending in RECALL. No sequence position or earlier cue enters these three baselines. The frequency baseline needs no smoothing here because every possible target type appears in training; BOS is never a target. The bigram smoothing avoids zeros for unobserved transitions.

train_x, train_y = xy["train"]
counts = np.bincount(train_y.ravel(), minlength=K)
frequency = counts / counts.sum()
assert np.all(counts[[UP, DOWN, LEFT, RIGHT, RECALL, EOS]] > 0)
transition_counts = np.zeros((K, K), dtype=int)
np.add.at(transition_counts, (train_x.ravel(), train_y.ravel()), 1)
alpha = 1.0
bigram = (transition_counts + alpha) / (transition_counts.sum(axis=1, keepdims=True) + K * alpha)
assert np.allclose(bigram.sum(axis=-1), 1)
print("Fitted P(next | RECALL):")
print(dict(zip(TOKENS, np.round(bigram[RECALL], 4))))

Fitted P(next | RECALL):
{'UP': 0.2445, 'DOWN': 0.2554, 'LEFT': 0.2481, 'RIGHT': 0.2513, 'RECALL': 0.0002, 'BOS': 0.0002, 'EOS': 0.0002}

An oracle that knows the generator

The oracle is a reference algorithm, not a trained model. At positions predicting the initial cue or wandering moves it assigns probability \(1/4\) to each direction. With known fixed delay it predicts RECALL and EOS exactly; at recall it retrieves the observed initial command.

There are \(d+1\) unpredictable targets and three deterministic targets among \(T=d+4\) predictions. Thus the original iid population optimum is

\[\mathbb E[\mathcal L]=\frac{(d+1)\ln4}{d+4},\qquad \mathbb E[\mathrm{accuracy}]=\frac{(d+1)/4+3}{d+4}.\]

This is not a proof of the optimal predictor for the finite, duplicate-filtered sampling scheme. Its oracle cross-entropy still equals the displayed expression on each valid episode because each unpredictable target receives exactly \(1/4\). Accuracy fluctuates with the sampled direction counts.

def oracle_probabilities(x, delay):
    b, t = x.shape
    assert t == delay + 4
    p = np.zeros((b, t, K))
    p[:, :delay + 1, :4] = .25
    p[:, delay + 1, RECALL] = 1
    # Only already observed cue x[:, 1] is used at recall.
    p[np.arange(b), delay + 2, x[:, 1]] = 1
    p[:, delay + 3, EOS] = 1
    return p

def metrics(p, y, mask):
    assert p.shape == (*y.shape, K) and np.allclose(p.sum(axis=-1), 1)
    true_p = np.take_along_axis(p, y[..., None], axis=-1)[..., 0]
    assert np.all(true_p[mask] > 0)
    return (-np.log(true_p[mask]).mean(), (p.argmax(axis=-1)[mask] == y[mask]).mean())

val_x, val_y = xy["validation"]
probabilities = {
    "Uniform (7 tokens)": np.full((*val_x.shape, K), 1 / K),
    "Training frequency": np.broadcast_to(frequency, (*val_x.shape, K)),
    "Smoothed bigram": bigram[val_x],
    "Generator oracle": oracle_probabilities(val_x, DELAY),
}
all_mask = np.ones_like(val_y, dtype=bool)
recall_mask = val_x == RECALL
random_mask = np.zeros_like(val_y, dtype=bool)
random_mask[:, :DELAY + 1] = True  # initial cue plus all wandering moves
print("Validation only; loss in nats, accuracy as a fraction")
print(f"{'Predictor':23s} {'all CE':>8s} {'all acc':>8s} {'recall CE':>10s} {'recall acc':>10s} {'random CE':>10s} {'random acc':>10s}")
results = {}
for name, p in probabilities.items():
    values = [v for mask in (all_mask, recall_mask, random_mask) for v in metrics(p, val_y, mask)]
    results[name] = values
    print(f"{name:23s}" + " ".join(f"{v:9.4f}" for v in values))
expected_ce = (DELAY + 1) * np.log(4) / (DELAY + 4)
assert np.isclose(results["Generator oracle"][0], expected_ce)
assert results["Generator oracle"][2:4] == [0.0, 1.0]
assert np.isclose(results["Uniform (7 tokens)"][0], np.log(K))
# Counterfactual check: change only the cue; oracle recall must follow it.
changed = val_x.copy(); changed[:, 1] = (changed[:, 1] + 1) % 4
assert np.array_equal(oracle_probabilities(changed, DELAY)[:, recall_t].argmax(-1), changed[:, 1])
print("Analytic loss, perfect oracle recall, and cue intervention: passed")

Validation only; loss in nats, accuracy as a fraction
Predictor                 all CE  all acc  recall CE recall acc  random CE random acc
Uniform (7 tokens)        1.9459    0.2059    1.9459    0.2188    1.9459    0.2502
Training frequency        1.7214    0.2051    1.5685    0.2441    1.5687    0.2463
Smoothed bigram           1.6881    0.2096    1.3863    0.2461    1.5849    0.2522
Generator oracle          1.0397    0.4377   -0.0000    1.0000    1.3863    0.2502
Analytic loss, perfect oracle recall, and cue intervention: passed

Interpretation checkpoint

Conclusion. The bigram recalls 24.6% of validation cues; the oracle recalls 100%. Access to the earlier cue solves the dependency, while both remain uncertain about fresh random movements. Finite-sample accuracy need not equal the 25% population reference.

All-token loss mixes distinct prediction problems. As delay grows, random wandering occupies more of the sequence, so even a perfect recall mechanism can barely change the overall average. Recall cross-entropy and accuracy are the primary memory measurements. The random-position column includes the unpredictable initial cue as well as wandering.

TODO(student). Explain why a falling all-token loss alone would not demonstrate memory. Identify the input token whose output distribution must be inspected. Then predict how the iid oracle’s all-token loss changes as \(d\) grows, without executing the test split.

No trained representation or sequence model has been evaluated yet. These results establish the experiment’s contract and a benchmark, not that any neural architecture succeeds or fails.

Scene 4 — One recurrent step, with actual numbers

Question. How can the same command have a different effect after a different history?

An ordinary feedforward layer sees the numbers we give it now. To use earlier commands, we must give it a record of them. A recurrent neural network (RNN) builds that record one step at a time: it reads the current command and the state left by the previous step, then writes a new state. The next step receives that new state.

The state is a vector of real numbers, called a hidden representation. “Hidden” means it is internal to the model, rather than an observed command or the final prediction. It need not resemble a list of past tokens. Training must discover what to keep; having a state does not guarantee useful memory.

Open the recurrent cell

A cell is a convenient box around familiar operations: weighted sums, biases, and an activation. Our microscope contains two ordinary tanh neurons. Each receives all seven input coordinates and both previous hidden coordinates. The box outputs two new hidden coordinates, not two command probabilities.

plot_rnn_anatomy()

Figure 3. The recurrent cell is a small feedforward calculation with an extra input. Both neurons read the previous state. The resulting vector goes to the next step and to a separate output head. Arrows show dependencies; each incoming connection has a weight.

First write the state, then read a prediction

For one episode, \(x_t\) is a seven-coordinate one-hot row vector. \(h_{t-1}\) and \(h_t\) have two coordinates (\(H=2\)). Start a new episode with \(h_{-1}=0\).

\[z_t=x_tW_x^\top+h_{t-1}W_h^\top+b_h,\qquad h_t=\tanh(z_t).\]

The first term reads the current command. Because \(x_t\) is one-hot, it selects that command’s column of \(W_x\). The second term mixes the previous hidden coordinates. Add the bias, then apply tanh to each coordinate; its output lies between \(-1\) and \(1\).

For hidden neuron \(j\), the vector equation is just an ordinary neuron calculation:

\[z_{t,j}=\sum_{k=1}^{K}W_{x,jk}x_{t,k}+\sum_{r=1}^{H}W_{h,jr}h_{t-1,r}+b_{h,j},\qquad h_{t,j}=\tanh(z_{t,j}).\]

These coordinate labels are one-based mathematical indices; Python arrays start at zero. \(W_x:(H,K)\) and \(W_h:(H,H)\) store one row per output neuron, as in Week 04. Every step reuses the same parameters.

A separate linear output head turns the state into seven scores, or logits:

\[a_t=h_tW_o^\top+b_o,\qquad p_t=\operatorname{softmax}(a_t),\qquad \ell_t=-\log p_t[y_t].\]

Here \(W_o:(K,H)\) has one row per possible next command. Softmax converts scores into probabilities that sum to one. The largest probability gives the predicted command; the loss reads the probability assigned to the actual next command \(y_t\). Neither logits nor hidden coordinates are probabilities. We keep \(z_t\) for hidden preactivations and \(a_t\) for output logits.

Checkpoint. Which term vanishes at the first step? At the second step, can you calculate the state from the current token alone?

Show explanation — which term vanishes at the first step?

The previous state is \(h_{-1}=0\), so \(h_{-1}W_h^\top\) is zero. The input and bias can still create a nonzero state. The second step reads that state, so its result generally depends on earlier input as well as the current token.

Choose small numbers we can inspect

The following parameters are hand-picked and untrained. We read BOS, UP, and RIGHT from the earlier handcrafted episode; their next-command targets are UP, RIGHT, and DOWN. This calculation demonstrates the mechanism, not successful recall.

# A three-position microscope from the earlier handcrafted episode, not an evaluation split.
rnn_ids = example_x[:3].copy()       # BOS, UP, RIGHT
rnn_targets = example_y[:3].copy()  # UP, RIGHT, DOWN
rnn_x = np.eye(K, dtype=np.float64)[rnn_ids]
H = 2
W_x = np.array([[ .2, -.1,  .1, -.2,  .0,  .1, .0],
                [-.1,  .2,  .2,  .1, -.1,  .0, .1]], dtype=np.float64)
W_h = np.array([[.3, -.2], [.1, .4]], dtype=np.float64)
b_h = np.array([.05, -.05], dtype=np.float64)
W_o = np.array([[ .3, -.1], [-.2,  .2], [.1, .3], [-.1, -.2],
                [ .2,  .1], [-.3, -.1], [.0, .2]], dtype=np.float64)
b_o = np.zeros(K, dtype=np.float64)
assert W_x.shape == (H, K) and W_h.shape == (H, H) and W_o.shape == (K, H)

Follow one loop iteration

The loop first builds the new state, then asks the output head for a prediction. h_previous = h_t passes the record to the next iteration. Subtracting the largest logit stabilizes softmax without changing its probabilities. The ledger saves intermediate values so we can check each operation.

# Visible recurrent arithmetic: no hidden model class or built-in RNN.
ledger = []
h_previous = np.zeros(H, dtype=np.float64)
for t, (x_t, y_t) in enumerate(zip(rnn_x, rnn_targets)):
    input_part = x_t @ W_x.T
    memory_part = h_previous @ W_h.T
    z_t = input_part + memory_part + b_h
    h_t = np.tanh(z_t)
    a_t = h_t @ W_o.T + b_o
    shifted = a_t - a_t.max()  # stable softmax and log-sum-exp
    log_p_t = shifted - np.log(np.exp(shifted).sum())
    p_t = np.exp(log_p_t)
    ledger.append(dict(x=x_t, previous=h_previous.copy(), input=input_part,
                       memory=memory_part, z=z_t, h=h_t, logits=a_t,
                       probabilities=p_t, loss=-log_p_t[y_t]))
    h_previous = h_t

assert np.allclose(ledger[0]["memory"], 0)
assert all(np.isclose(row["probabilities"].sum(), 1) for row in ledger)
microscope_loss = np.mean([row["loss"] for row in ledger])
print("Mean loss over three prediction positions:", microscope_loss)

Mean loss over three prediction positions: 1.9229781712278502

# @title Full three-step forward ledger
import io
import html
from contextlib import redirect_stdout
forward_text = io.StringIO()
with redirect_stdout(forward_text):
    for t, row in enumerate(ledger):
        print(f"t={t}: {TOKENS[rnn_ids[t]]} -> target {TOKENS[rnn_targets[t]]}")
        for name in ("previous", "input", "memory", "z", "h", "logits", "probabilities"):
            print(f"  {name:13s}", np.round(row[name], 4))
        print(f"  loss          {row['loss']:.6f}")
display(HTML("<details><summary>Full forward ledger — all three positions</summary><pre>"
             + html.escape(forward_text.getvalue()) + "</pre></details>"))

Full forward ledger — all three positions

t=0: BOS -> target UP
  previous      [0. 0.]
  input         [0.1 0. ]
  memory        [0. 0.]
  z             [ 0.15 -0.05]
  h             [ 0.1489 -0.05  ]
  logits        [ 0.0497 -0.0398 -0.0001 -0.0049  0.0248 -0.0397 -0.01  ]
  probabilities [0.1505 0.1376 0.1432 0.1425 0.1468 0.1376 0.1418]
  loss          1.893846
t=1: UP -> target RIGHT
  previous      [ 0.1489 -0.05  ]
  input         [ 0.2 -0.1]
  memory        [ 0.0547 -0.0051]
  z             [ 0.3047 -0.1551]
  h             [ 0.2956 -0.1539]
  logits        [ 0.1041 -0.0899 -0.0166  0.0012  0.0437 -0.0733 -0.0308]
  probabilities [0.1596 0.1315 0.1415 0.144  0.1503 0.1337 0.1395]
  loss          1.937840
t=2: RIGHT -> target DOWN
  previous      [ 0.2956 -0.1539]
  input         [-0.2  0.1]
  memory        [ 0.1194 -0.032 ]
  z             [-0.0306  0.018 ]
  h             [-0.0305  0.018 ]
  logits        [-0.011   0.0097  0.0023 -0.0005 -0.0043  0.0074  0.0036]
  probabilities [0.1412 0.1441 0.143  0.1426 0.1421 0.1438 0.1432]
  loss          1.937249

Keep the numbers attached to the operation

After BOS, the state is approximately \(h_0=(0.1489,-0.0500)\). The next input is UP. Its column of \(W_x\) contributes \((0.2,-0.1)\); mixing the previous state through \(W_h\) contributes approximately \((0.0547,-0.0051)\).

Adding the bias gives \(z_1\approx(0.3047,-0.1551)\). Applying tanh gives \(h_1\approx(0.2956,-0.1539)\). For example, the first neuron’s preactivation is

\[z_{1,1}=0.2+0.3(0.1489)-0.2(-0.0500)+0.05\approx0.3047.\]

The output head now reads those two numbers to produce seven logits. For UP, its row of \(W_o\) gives \(a_1[\mathrm{UP}]\approx0.3(0.2956)-0.1(-0.1539)=0.1041\). Softmax compares this score with all six others. Inspect the ledger: the untrained head’s largest probability need not belong to the actual next command RIGHT.

Check your reading — does a positive hidden coordinate mean “UP”?

No. The hidden coordinates are internal features; their meaning depends on the learned parameters. The UP logit combines both coordinates using its output weights. A positive hidden coordinate alone identifies neither a command nor a probability.

The next figure expands the second step. A two-coordinate state can carry earlier context, but training still has to make that context useful.

plot_recurrent_step(ledger[1])

Figure 4. One recurrent update, term by term. At the second input, UP, both the current input and the state left by BOS contribute. Add those contributions and the bias before applying tanh to each coordinate.

Conclusion. The second step combines the current token with the state left by the first step. Its two hidden coordinates are internal features; a separate head converts them to seven command probabilities. These untrained values demonstrate the mechanism, not learned memory.

Scene 5 — Unroll the same machine

Checkpoint. If we run this cell for three positions, have we tripled the number of trainable parameters?

Unrolling draws separate uses of the same operations. The hidden states differ, but the parameter values are shared. There are \(HK+H^2+H+KH+K=41\) trainable scalars here, independent of sequence length. More positions enlarge the computation graph, not the parameter collection.

Each output sees the current token and previous hidden state. We calculate a mean loss over the three positions, \(\mathcal L=\frac13\sum_{t=0}^{2}\ell_t\). This is a miniature instance of the future training objective, not a new split or a recall evaluation.

Show explanation — does unrolling add parameters?

It adds intermediate states and uses of the operations. Every use refers to the same weights and biases. The numerical RNN still has 41 trainable scalars; a longer sequence creates a larger computation graph.

parameter_count = sum(a.size for a in (W_x, W_h, b_h, W_o, b_o))
assert parameter_count == 41
print("Trainable scalar count:", parameter_count)
plot_unrolled_rnn(ledger, rnn_ids, rnn_targets)

Trainable scalar count: 41

Figure 5. One cell used three times. The hidden state changes with the prefix. The recurrent and output weights are shared across positions; each output is compared with its shifted target.

Conclusion. Three steps still use 41 trainable scalars. History changes the states; unrolling reuses the weights and enlarges the computation graph.

Scene 6 — One shared weight receives several responsibilities

Optional backward-pass excursion. You can continue directly to Scene 7. Its training code does not depend on the calculations in this scene.

The forward pass carries a state to later steps. During training, a later prediction can therefore send a loss gradient back through earlier states. A shared weight receives contributions from every place it was used. PyTorch will perform this calculation for us; the details below connect it to the chain rule from Week 04.

Open the backward equations — current and future responsibilities

Let \(g^a_t=\partial\mathcal L/\partial a_t\). For the mean loss over \(T_\mu=3\) microscope positions,

\[g^a_t=(p_t-\operatorname{onehot}(y_t))/T_\mu.\]

The hidden state has two outgoing branches: the current output head and the next recurrent step. Add the incoming sensitivities before applying the tanh derivative:

\[g^h_t=\underbrace{g^a_tW_o}_{\text{current output}}+ \underbrace{g^z_{t+1}W_h}_{\text{future through recurrence}},\qquad g^z_t=g^h_t\odot(1-h_t^2).\]

At the last position, the future term is zero. Work right to left. For output-by-input parameter storage, each use contributes an outer product:

\[\frac{\partial\mathcal L}{\partial W_x}=\sum_t(g^z_t)^\top x_t,\qquad \frac{\partial\mathcal L}{\partial W_h}=\sum_t(g^z_t)^\top h_{t-1},\qquad \frac{\partial\mathcal L}{\partial W_o}=\sum_t(g^a_t)^\top h_t.\]

Bias gradients sum \(g^z_t\) or \(g^a_t\). In the code, np.outer makes those shapes explicit. This is the same branch-addition rule from Week 04 and shared-gradient sum from Week 05, now applied across time.

Checkpoint. Should we average the parameter contributions again after summing them? Which scaling has already accounted for the mean loss?

Show explanation — should shared gradients be averaged again?

The factor \(1/T_\mu\) already enters the derivative of each output loss. Summing contributions from every parameter use then gives the derivative of that mean loss. Another division would shrink the gradient incorrectly.

# @title Optional: calculate shared gradients by hand
T_mu = len(ledger)
contributions = [None] * T_mu
backward_ledger = [None] * T_mu
future_hidden_gradient = np.zeros(H, dtype=np.float64)
for t in reversed(range(T_mu)):
    row = ledger[t]
    g_logits = row["probabilities"].copy()
    g_logits[rnn_targets[t]] -= 1
    g_logits /= T_mu  # mean loss: apply this scaling once
    local_hidden_gradient = g_logits @ W_o
    g_hidden = local_hidden_gradient + future_hidden_gradient
    g_z = g_hidden * (1 - row["h"] ** 2)
    contributions[t] = {
        "W_x": np.outer(g_z, row["x"]),
        "W_h": np.outer(g_z, row["previous"]),
        "b_h": g_z.copy(),
        "W_o": np.outer(g_logits, row["h"]),
        "b_o": g_logits.copy(),
    }
    backward_ledger[t] = dict(local=local_hidden_gradient.copy(),
                              future=future_hidden_gradient.copy(), total=g_hidden, z=g_z)
    future_hidden_gradient = g_z @ W_h
manual_gradients = {name: sum(g[name] for g in contributions) for name in contributions[0]}
manual_initial_gradient = future_hidden_gradient.copy()
for t, row in enumerate(backward_ledger):
    print(f"t={t}: local {row['local'].round(6)} + future {row['future'].round(6)}"
          f" = total {row['total'].round(6)}; through tanh {row['z'].round(6)}")
assert np.allclose(backward_ledger[-1]["future"], 0)
assert np.linalg.norm(backward_ledger[0]["future"]) > 0
assert np.allclose(contributions[0]["W_h"], 0)  # first recurrent input is zero
print("Mean-loss scaling is already in each contribution; sum without dividing again.")

t=0: local [-0.098079  0.052067] + future [0.019401 0.010886] = total [-0.078679  0.062953]; through tanh [-0.076935  0.062796]
t=1: local [0.037095 0.084509] + future [ 0.015121 -0.032229] = total [0.052216 0.05228 ]; through tanh [0.047655 0.051042]
t=2: local [ 0.066286 -0.047477] + future [0. 0.] = total [ 0.066286 -0.047477]; through tanh [ 0.066224 -0.047461]
Mean-loss scaling is already in each contribution; sum without dividing again.

plot_shared_gradients(contributions)

Figure 6. Shared weights collect shared responsibility. The matrices show contributions from three uses of the recurrent weights and their sum. The first direct contribution is zero because the initial state is zero; later losses can still depend on that first state.

The first use contributes zero directly to \(W_h\) because \(h_{-1}=0\). Its hidden state still affects later losses; zero direct contribution is not a disconnected computation graph. Each displayed matrix uses the total sensitivity, including future losses. These are contributions grouped by parameter use, not gradients from isolated per-position losses.

Ask autograd the identical question

Rebuild the exact three-step calculation in PyTorch with float64 tensors. Keep the hand-picked values, zero initial state, target shift, and mean reduction identical. backward() traverses the unrolled graph and accumulates into each shared parameter’s gradient.

The initial state is fixed in the model. We temporarily request its derivative as a diagnostic; it is not counted as a trainable parameter. Resetting a state to zero between episodes is distinct from detaching a state between positions, which removes future-to-past gradient paths.

# @title Optional: verify the same gradients with autograd
parameters = {name: torch.tensor(value, dtype=torch.float64, requires_grad=True)
              for name, value in dict(W_x=W_x, W_h=W_h, b_h=b_h, W_o=W_o, b_o=b_o).items()}
torch_x = torch.tensor(rnn_x, dtype=torch.float64)
torch_targets = torch.tensor(rnn_targets, dtype=torch.long)
initial_state = torch.zeros(H, dtype=torch.float64, requires_grad=True)
h_t = initial_state
states, logits = [], []
for x_t in torch_x:
    z_t = x_t @ parameters["W_x"].T + h_t @ parameters["W_h"].T + parameters["b_h"]
    h_t = torch.tanh(z_t)
    h_t.retain_grad()
    states.append(h_t)
    logits.append(h_t @ parameters["W_o"].T + parameters["b_o"])
torch_logits = torch.stack(logits)
autograd_loss = F.cross_entropy(torch_logits, torch_targets, reduction="mean")
autograd_loss.backward()
np.testing.assert_allclose(autograd_loss.item(), microscope_loss, rtol=1e-12, atol=1e-12)
np.testing.assert_allclose(torch_logits.detach().numpy(),
                           np.stack([row["logits"] for row in ledger]), rtol=1e-12, atol=1e-12)
print("Parameter  maximum absolute gradient difference")
for name, parameter in parameters.items():
    observed = parameter.grad.numpy()
    np.testing.assert_allclose(observed, manual_gradients[name], rtol=1e-10, atol=1e-12)
    print(f"{name:9s}  {np.max(np.abs(observed - manual_gradients[name])):.3e}")
for t, state in enumerate(states):
    np.testing.assert_allclose(state.grad.numpy(), backward_ledger[t]["total"], rtol=1e-10, atol=1e-12)
np.testing.assert_allclose(initial_state.grad.numpy(), manual_initial_gradient, rtol=1e-10, atol=1e-12)
print("Forward values, mean loss, all shared parameter gradients, hidden-state gradients, and initial-state diagnostic: passed")

Parameter  maximum absolute gradient difference
W_x        6.939e-18
W_h        5.204e-18
b_h        0.000e+00
W_o        6.939e-18
b_o        2.776e-17
Forward values, mean loss, all shared parameter gradients, hidden-state gradients, and initial-state diagnostic: passed

What the backward check establishes

Conclusion. Manual and autograd parameter gradients agree to within about \(3\times10^{-17}\) in this float64 example. This verifies the derivative calculation; successful recall still requires training. The first direct \(W_h\) contribution is zero because its input state is zero, while later losses can depend on the first new state.

If you return to this excursion after the lecture, trace one path from a later loss to an earlier state in the unrolled diagram. Identify where the current and future sensitivities are added. We now let training find useful parameters.

Scene 7 — Train the first remembering robot

Question. Can learning turn the recurrent mechanism into useful memory?

We now use generated episodes rather than the handcrafted microscope. IDs enter a learned embedding with width \(D=16\), then a single forward tanh RNN with \(H=32\), then a linear seven-logit head. Shapes are \((B,T)\to(B,T,D)\to(B,T,H)\to(B,T,K)\). Each episode starts with a zero state; there is no state carried between batches.

The PyTorch RNN module evaluates the recurrence from Scene 4 efficiently. Its two bias vectors add to the earlier single \(b_h\); both are counted as trainable parameters. It is unidirectional, so a prediction cannot inspect later input tokens.

Build the prediction path

Read the model in three stages: embedding looks up a vector for each command, recurrent updates the hidden representation through the sequence, and head reads each state into next-command logits. The normal forward path is those three operations. The optional reset_before branch will let Scene 8 erase memory; return_embeddings will let it inspect sensitivity to inputs.

from torch import nn
import copy

torch.set_num_threads(1)  # small CPU batches; avoids multithreading overhead
TRAINING_SEED = 6206
D_EMBED = 16
EPOCHS, BATCH_SIZE, LEARNING_RATE = 60, 128, .003

class CommandRecurrentModel(nn.Module):
    def __init__(self, kind="RNN"):
        super().__init__()
        self.kind = kind
        self.hidden_width = 32 if kind == "RNN" else 17
        self.embedding = nn.Embedding(K, D_EMBED)
        recurrent_class = nn.RNN if kind == "RNN" else nn.GRU
        self.recurrent = recurrent_class(D_EMBED, self.hidden_width,
                                         batch_first=True, bidirectional=False)
        self.head = nn.Linear(self.hidden_width, K)

    def forward(self, ids, reset_before=None, return_embeddings=False):
        e = self.embedding(ids)
        if reset_before is None:
            h, _ = self.recurrent(e)  # omitted initial state defaults to zeros
        else:
            # Diagnostic intervention: erase the state immediately before this position.
            before, _ = self.recurrent(e[:, :reset_before])
            after, _ = self.recurrent(e[:, reset_before:])  # fresh zero state
            h = torch.cat((before, after), dim=1)
        logits = self.head(h)
        return (logits, e) if return_embeddings else logits

One epoch: predict, measure, update

An epoch visits every training episode once. Each batch makes predictions, computes mean next-command loss, and backpropagates it. Clear old gradients before each batch; clip the new gradient after backward(); then let Adam update the parameters.

Logits have shape \((B,T,K)\). Reshaping to \((BT,K)\) treats each position as one classification example, with its corresponding flattened target. Multiply each batch mean by its episode count when accumulating the epoch mean; all episodes here have the same length.

def train_one_epoch(model, optimizer, ids, targets, order):
    model.train()
    total_loss = 0.0
    for indices in order.split(BATCH_SIZE):
        optimizer.zero_grad(set_to_none=True)
        logits = model(ids[indices])
        loss = F.cross_entropy(logits.reshape(-1, K), targets[indices].reshape(-1))
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        total_loss += loss.item() * len(indices)
    return total_loss / len(ids)

Measure the dependency separately

Validation performs no parameter update (torch.no_grad()). sequence_metrics reports both the complete next-command task and the recall position \(d+2\). The random-position metrics check whether performance is consistent with the generator’s unpredictable movements. This metric helper is also reused by later models.

# @title Metric helper: all positions, recall, and random movements
@torch.no_grad()
def sequence_metrics(model, ids, targets, delay, reset_before=None):
    model.eval()
    logits = model(ids) if reset_before is None else model(ids, reset_before=reset_before)
    recall_index = delay + 2
    return {
        "all_ce": F.cross_entropy(logits.reshape(-1, K), targets.reshape(-1)).item(),
        "all_accuracy": (logits.argmax(-1) == targets).float().mean().item(),
        "recall_ce": F.cross_entropy(logits[:, recall_index], targets[:, recall_index]).item(),
        "recall_accuracy": (logits[:, recall_index].argmax(-1) == targets[:, recall_index]).float().mean().item(),
        "random_ce": F.cross_entropy(logits[:, :delay+1].reshape(-1, K), targets[:, :delay+1].reshape(-1)).item(),
        "random_accuracy": (logits[:, :delay+1].argmax(-1) == targets[:, :delay+1]).float().mean().item(),
    }

Keep the checkpoint that remembers best

Protocol fixed before test evaluation. Later command models reuse this table.

Choice	Setting
Training objective	Mean cross-entropy over all next-token positions
Checkpoint selection	Smallest validation recall cross-entropy
Optimizer	Adam, learning rate 0.003
Budget	60 epochs, batch size 128; 1,920 updates
Gradient clipping	Global norm at 1, after backward and before the update
Data and order	Same episodes and batch orders across architectures
Initialization	Seed 6206 here; three seeds in Scene 12

Clipping bounds the gradient; it cannot restore lost information. Test data never enter the training function.

Development history — why 60 epochs?

A 30-epoch CPU pilot still had improving validation curves, so the shared budget was extended to 60 before test scoring. This is validation-based development, not evidence of convergence. Scene 12 repeats initializations and evaluates held-out episodes.

The outer loop calls one training epoch, measures validation performance, and saves the best parameter snapshot. A deep copy matters: otherwise later updates could overwrite the supposedly saved weights. At the end we restore that snapshot. Fixed batch orders make the later architecture comparison reproducible.

Checkpoint. Why do the training objective and checkpoint criterion differ? What would an all-token metric conceal as the delay grows?

Show explanation — why select using recall loss?

Most targets are random movements. Their uncertainty dominates the all-token mean, especially at longer delays. The training objective fits the entire sequence, while validation recall loss selects the checkpoint for the dependency we are investigating. The criterion is fixed before test evaluation.

# The training function gets only train/validation arrays, never the complete split dictionary.
def train_sequence_model(kind, train_ids, train_targets, validation_ids,
                         validation_targets, delay, model_factory=None,
                         initialization_seed=TRAINING_SEED):
    torch.manual_seed(initialization_seed)
    model = CommandRecurrentModel(kind) if model_factory is None else model_factory()
    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
    orders = [torch.randperm(len(train_ids), generator=torch.Generator().manual_seed(6200+epoch))
              for epoch in range(EPOCHS)]
    history, best_ce, best_epoch, best_state = [], float("inf"), None, None
    start = time.perf_counter()
    for epoch, order in enumerate(orders, start=1):
        train_all_ce = train_one_epoch(model, optimizer, train_ids, train_targets, order)
        validation = sequence_metrics(model, validation_ids, validation_targets, delay)
        history.append(dict(epoch=epoch, train_all_ce=train_all_ce, **validation))
        if validation["recall_ce"] < best_ce:
            best_ce, best_epoch = validation["recall_ce"], epoch
            best_state = copy.deepcopy(model.state_dict())
    model.load_state_dict(best_state)
    model.eval()
    assert best_epoch == 1 + int(np.argmin([r["recall_ce"] for r in history]))
    return dict(model=model, history=history, best_epoch=best_epoch, initialization_seed=initialization_seed,
                seconds=time.perf_counter()-start,
                parameters=sum(p.numel() for p in model.parameters()),
                validation=sequence_metrics(model, validation_ids, validation_targets, delay))

Run the first training experiment

Convert the already separated train and validation arrays into integer tensors, then train the delay-8 RNN. The function receives these two splits explicitly. Its return value holds the selected model, learning history, and validation metrics for the figures below.

train_ids = torch.tensor(xy["train"][0], dtype=torch.long)
train_targets = torch.tensor(xy["train"][1], dtype=torch.long)
validation_ids = torch.tensor(xy["validation"][0], dtype=torch.long)
validation_targets = torch.tensor(xy["validation"][1], dtype=torch.long)
assert train_ids.shape == train_targets.shape == (4096, 12)
trained = {("RNN", 8): train_sequence_model("RNN", train_ids, train_targets,
                                       validation_ids, validation_targets, 8)}
rnn8 = trained[("RNN", 8)]
assert rnn8["parameters"] == 1943
assert rnn8["model"](validation_ids[:4]).shape == (4, 12, K)
print("RNN, delay 8:", rnn8["validation"], "selected epoch", rnn8["best_epoch"],
      "parameters", rnn8["parameters"], "seconds", round(rnn8["seconds"], 2))

RNN, delay 8: {'all_ce': 1.0457680225372314, 'all_accuracy': 0.4348958432674408, 'recall_ce': 0.027155200019478798, 'recall_accuracy': 1.0, 'random_ce': 1.3901737928390503, 'random_accuracy': 0.2465277761220932} selected epoch 60 parameters 1943 seconds 2.5

# @title Verification: the trained module matches the explicit recurrence
# Confirm the efficient module performs the audited row-vector recurrence.
with torch.no_grad():
    model = copy.deepcopy(rnn8["model"]).double()  # float64 audit; trained model stays float32
    e = model.embedding(validation_ids[:2])
    state = torch.zeros(2, model.hidden_width, dtype=torch.float64)
    explicit_states = []
    for t in range(e.shape[1]):
        state = torch.tanh(e[:, t] @ model.recurrent.weight_ih_l0.T
                           + state @ model.recurrent.weight_hh_l0.T
                           + model.recurrent.bias_ih_l0 + model.recurrent.bias_hh_l0)
        explicit_states.append(state)
    explicit_logits = model.head(torch.stack(explicit_states, dim=1))
    torch.testing.assert_close(explicit_logits, model(validation_ids[:2]), rtol=1e-10, atol=1e-12)
print("Trained RNN module agrees with the explicit recurrence")

Trained RNN module agrees with the explicit recurrence

plot_recurrent_curves(rnn8)

Figure 7. Follow learning with the metric that answers the question. The dotted line marks the checkpoint selected by validation recall cross-entropy. Training loss averages snapshots during updates; validation is measured at epoch end.

Conclusion. The selected delay-8 RNN recalls every validation cue, with recall loss about 0.027 nats. Random-position loss stays near \(\ln4\): training learns the dependency while independent movements retain their uncertainty. The selected epoch is 60, so this run does not establish convergence.

Look at the next command

Use validation episodes 0–3, fixed before viewing predictions. Blue marks the actual next move and its probability bar; red marks the argmax prediction. Overlapping arrows mean agreement. The robot’s path is the observed prefix through RECALL; later commands are not fed into the prediction. Control-token predictions have zero displacement and are still shown by name and probability.

Checkpoint. Inspect the cue and path first. Predict the command, then compare it with the model distribution. These are teacher-forced prefix predictions, not generated trajectories or test results.

plot_recall_predictions(rnn8["model"], validation_ids, validation_targets, 8, "07-rnn-recall-predictions")
# Causality: changing only future input tokens cannot change earlier logits.
future_changed = validation_ids[:8].clone()
future_changed[:, 11:] = (future_changed[:, 11:] + 1) % 4
with torch.no_grad():
    original_logits = rnn8["model"](validation_ids[:8])
    changed_logits = rnn8["model"](future_changed)
torch.testing.assert_close(original_logits[:, :11], changed_logits[:, :11], rtol=0, atol=0)
print("Future-token intervention: earlier RNN logits unchanged")

Future-token intervention: earlier RNN logits unchanged

Figure 8. Read a prediction as a distribution. The same four validation prefixes show the actual and predicted recall movement above seven-token probabilities. These predictions use observed prefixes; the displayed paths are not sampled model rollouts.

Scene 8 — How far does memory last?

Train a fresh RNN at each delay, with the same episode counts, seed, optimizer, and update budget. This measures learning at each task length; it does not test extrapolation of one trained model to unseen lengths.

Increasing delay changes more than the gradient path: recall receives fraction \(1/(d+4)\) of the all-token training loss, and the model must also track the later control-token positions. Interpret the curve as this experiment’s combined difficulty. A separately designed recall-weighted experiment would help isolate loss dilution; it is not silently substituted here.

Checkpoint. Will excellent delay-8 recall guarantee delay-32 recall under this budget? Make a prediction before executing the sweep.

Show explanation — what does increasing delay test?

Here we train a separate model at each delay, under the same update budget. The result measures how successfully that protocol learns each task. It also changes recall’s relative weight in the all-token loss. It does not test whether one trained model generalizes to unseen sequence lengths.

# Longer-delay test episodes are generated for split integrity only; not converted or scored.
training_data = {8: (train_ids, train_targets, validation_ids, validation_targets)}
for delay in (16, 32):
    longer_splits, _ = make_splits(delay)
    split_keys = {name: set(map(tuple, rows)) for name, rows in longer_splits.items()}
    assert not (split_keys["train"] & split_keys["validation"] or split_keys["train"] & split_keys["test"]
                or split_keys["validation"] & split_keys["test"])
    tx = torch.tensor(longer_splits["train"][:, :-1], dtype=torch.long)
    ty = torch.tensor(longer_splits["train"][:, 1:], dtype=torch.long)
    vx = torch.tensor(longer_splits["validation"][:, :-1], dtype=torch.long)
    vy = torch.tensor(longer_splits["validation"][:, 1:], dtype=torch.long)
    training_data[delay] = (tx, ty, vx, vy)
    trained[("RNN", delay)] = train_sequence_model("RNN", tx, ty, vx, vy, delay)
    del longer_splits  # no test tensors enter a training/evaluation function
for delay in (8, 16, 32):
    run = trained[("RNN", delay)]
    print(f"RNN d={delay:2d}, selected epoch={run['best_epoch']:2d}, "
          f"recall CE={run['validation']['recall_ce']:.4f}, accuracy={run['validation']['recall_accuracy']:.4f}")

RNN d= 8, selected epoch=60, recall CE=0.0272, accuracy=1.0000
RNN d=16, selected epoch=60, recall CE=1.2383, accuracy=0.3105
RNN d=32, selected epoch=49, recall CE=1.3863, accuracy=0.2754

plot_memory_comparison(trained, ("RNN",), "08-rnn-memory-delay")

Figure 9. Make the dependency more distant. Each point comes from a separately trained RNN. Compare recall performance across delays while remembering that the fixed budget and reduced recall-loss weight also affect learning.

Conclusion. Validation recall accuracy falls from 100% at delay 8 to 31% at delay 16 and 28% at delay 32. Longer dependencies are harder for this RNN training protocol; the sweep also changes recall-loss weight and control-token timing, so it does not isolate gradient decay.

Probe the path, then erase the memory

For the first 64 validation episodes at each delay, differentiate recall-only loss with respect to each position’s embedding output. This is a gradient at the selected checkpoint, not a measurement of how much information a hidden state contains. Plot absolute mean per-position norms on a log scale (display floor \(10^{-12}\)). Future positions must have zero sensitivity in a causal recurrent model.

Two complementary interventions at delay 8 ask whether predictions use the past: erase the state immediately before RECALL, or change only the initial cue and score against that changed cue. The latter uses only the prefix ending at RECALL; no inconsistent future suffix is read. These are validation diagnostics, not additional independent evaluation sets.

gradient_diagnostics = {}
for delay in (8, 16, 32):
    model = trained[("RNN", delay)]["model"]
    vx, vy = training_data[delay][2:]
    model.zero_grad(set_to_none=True)
    logits, embeddings = model(vx[:64], return_embeddings=True)
    embeddings.retain_grad()
    recall_loss = F.cross_entropy(logits[:, delay+2], vy[:64, delay+2])
    recall_loss.backward()
    norms = embeddings.grad.norm(dim=-1).mean(dim=0).detach().numpy()
    assert np.all(norms[delay+3:] == 0)
    gradient_diagnostics[delay] = norms
    model.zero_grad(set_to_none=True)

model = rnn8["model"]
erased = sequence_metrics(model, validation_ids, validation_targets, 8, reset_before=10)
changed_prefix = validation_ids[:128, :11].clone()
changed_prefix[:, 1] = (changed_prefix[:, 1] + 1) % 4
with torch.no_grad():
    counterfactual_predictions = model(changed_prefix)[:, -1].argmax(-1)
changed_cue_accuracy = (counterfactual_predictions == changed_prefix[:, 1]).float().mean().item()
print(f"RNN delay 8: original recall accuracy {rnn8['validation']['recall_accuracy']:.4f}; "
      f"state erased {erased['recall_accuracy']:.4f}; changed cue accuracy {changed_cue_accuracy:.4f}")
print("Recall loss has exactly zero embedding gradient at future positions")

RNN delay 8: original recall accuracy 1.0000; state erased 0.2910; changed cue accuracy 1.0000
Recall loss has exactly zero embedding gradient at future positions

plot_gradient_reach(gradient_diagnostics)

Figure 10. Inspect sensitivity along the recurrent path. Recall-loss gradient norms are measured at embedding outputs of validation prefixes. Small gradients describe local sensitivity, and do not directly measure how much cue information is stored.

Conclusion. At delay 8, erasing state before RECALL lowers accuracy from 100% to 29%; changing the cue makes all 128 altered-prefix predictions follow the new cue. Together, these interventions support use of earlier evidence. Gradient size alone does not measure stored information.

Scene 9 — Learn when to retain information

Question. Can the model leave a useful hidden coordinate almost unchanged while other commands arrive?

The tanh RNN rewrites every coordinate through an affine sum and tanh at every step. A gated recurrent unit (GRU) adds a more direct route for the previous state. It computes a candidate update and learns how much of that candidate to mix with the old coordinate.

A gate is a number, not an instruction

A sigmoid gate lies between zero and one. It multiplies a coordinate; it does not execute an if statement. With update gate \(u_t\) and candidate \(\widetilde h_t\), the new state is

\[h_t=u_t\odot h_{t-1}+(1-u_t)\odot\widetilde h_t.\]

The multiplication is elementwise, so different coordinates can retain different amounts. In our convention, a gate near one favors retaining the old coordinate. Some presentations reverse the coefficients; read the equation before interpreting a gate’s name. PyTorch calls this gate \(z_t\), but we use \(u_t\) because \(z_t\) already denotes a preactivation.

For one illustrative coordinate, let the previous state be \(0.5\), the candidate be \(-0.4\), and the gate be \(0.9\). The retained part is \(0.45\), the new part is \(-0.04\), and their sum is \(0.41\). These values are chosen to explain the mixture, not taken from a trained model.

assert np.isclose(.9*.5 + (1-.9)*(-.4), .41)
plot_gru_carry()

Figure 11. Retain most of the old coordinate and write a little of the candidate. The gate scales both paths through complementary coefficients. This direct mixture is applied separately to every hidden coordinate.

Learn the gate and the candidate

The gate is itself computed from the current input and previous state. Here \(x_t\) is the token’s input representation, an embedding in the trained model. Both gates and the candidate have shape \((H,)\):

\[u_t=\sigma(x_tW_{xu}^\top+b_{xu}+h_{t-1}W_{hu}^\top+b_{hu}),\] \[r_t=\sigma(x_tW_{xr}^\top+b_{xr}+h_{t-1}W_{hr}^\top+b_{hr}).\]

The reset gate \(r_t\) controls how much recurrent information enters the candidate. The candidate still uses a tanh neuron calculation:

\[\widetilde h_t=\tanh\!\left(x_tW_{xn}^\top+b_{xn}+r_t\odot(h_{t-1}W_{hn}^\top+b_{hn})\right).\]

Thus \(u_t\) chooses how much of the old state survives in the final mixture; \(r_t\) controls the old-state contribution to the proposed new content. Their jobs differ. A possible carry path is useful only if training learns when to use it.

Implementation convention — where PyTorch applies reset

The PyTorch GRU definition applies reset after the recurrent affine transformation, including its bias. The original formulation resets hidden coordinates before that transformation. Mixing coordinates and gating them generally do not commute. Our equations and numerical audit follow the implementation we train.

Checkpoint. If \(u_t=(1,0)\), which coordinate retains its old value? Does changing \(r_t\) necessarily change that retained coordinate?

Show explanation — what does the update gate retain?

The first coordinate copies its previous value and the second takes its candidate value. At the ideal endpoint \(u_{t,1}=1\), the first candidate has zero coefficient: changing reset can change that candidate without changing the final retained coordinate. Real finite sigmoid outputs approach, rather than exactly reach, those endpoints. Intermediate values interpolate.

Check the library calculation once

Read the code in the same order: input/recurrent affine sums, reset and update gates, candidate, then the mixture. PyTorch stacks the three parameter groups, so chunk(3) separates them. This two-unit cell is untrained and separate from the models in our experiment.

# One untrained, two-unit GRU cell with a three-feature input representation.
torch.manual_seed(6306)
gru_cell = nn.GRUCell(3, 2, dtype=torch.float64)
gate_input = torch.tensor([.2, -.1, .3], dtype=torch.float64)
gate_previous = torch.tensor([.5, -.25], dtype=torch.float64)
with torch.no_grad():
    input_affine = gate_input @ gru_cell.weight_ih.T + gru_cell.bias_ih
    recurrent_affine = gate_previous @ gru_cell.weight_hh.T + gru_cell.bias_hh
    x_r, x_u, x_n = input_affine.chunk(3)  # PyTorch stores reset, update, candidate
    h_r, h_u, h_n = recurrent_affine.chunk(3)
    reset = torch.sigmoid(x_r + h_r)
    update = torch.sigmoid(x_u + h_u)
    candidate = torch.tanh(x_n + reset * h_n)
    manual_gru_state = (1-update)*candidate + update*gate_previous
    module_gru_state = gru_cell(gate_input, gate_previous)
    torch.testing.assert_close(manual_gru_state, module_gru_state, rtol=1e-12, atol=1e-12)
for name, value in (("reset r", reset), ("update u", update), ("candidate", candidate), ("new state", manual_gru_state)):
    print(name, value.numpy().round(6))
print("Manual GRU update matches GRUCell")

reset r [0.715222 0.574873]
update u [0.417841 0.588137]
candidate [ 0.392204 -0.07202 ]
new state [ 0.437245 -0.176697]
Manual GRU update matches GRUCell

Compare under the same protocol

Use hidden width 17 for the GRU, giving 2,023 parameters versus the RNN’s 1,943 (about 4% more). Both use embedding width 16. This roughly matches parameter count, not hidden-state width or computation cost. The GRU performs more operations per time step; report wall time alongside metrics.

Checkpoint. Predict the comparison without assuming the gate guarantees success. Then inspect the executed evidence, including selected epochs and the fixed budget.

for delay in (8, 16, 32):
    tx, ty, vx, vy = training_data[delay]
    trained[("GRU", delay)] = train_sequence_model("GRU", tx, ty, vx, vy, delay)
assert trained[("GRU", 8)]["parameters"] == 2023
print("Validation-selected models; all losses in nats. Single initialization seed.")
print("model delay params epoch recall_CE recall_acc all_CE random_CE seconds")
for kind in ("RNN", "GRU"):
    for delay in (8, 16, 32):
        run = trained[(kind, delay)]; m = run["validation"]
        print(f"{kind:5s} {delay:5d} {run['parameters']:6d} {run['best_epoch']:5d} "
              f"{m['recall_ce']:9.4f} {m['recall_accuracy']:10.4f} {m['all_ce']:6.4f} "
              f"{m['random_ce']:9.4f} {run['seconds']:7.2f}")
# Check the trained modules remain causal for both architectures.
for kind in ("RNN", "GRU"):
    model = trained[(kind, 8)]["model"]
    with torch.no_grad():
        original = model(validation_ids[:8])
        changed = model(future_changed)
    torch.testing.assert_close(original[:, :11], changed[:, :11], rtol=0, atol=0)
print("Future-token checks passed for both trained architectures")

Validation-selected models; all losses in nats. Single initialization seed.
model delay params epoch recall_CE recall_acc all_CE random_CE seconds
RNN       8   1943    60    0.0272     1.0000 1.0458    1.3902    2.50
RNN      16   1943    60    1.2383     0.3105 1.2433    1.3897    3.56
RNN      32   1943    49    1.3863     0.2754 1.4491    1.3900    5.84
GRU       8   2023    59    0.0418     1.0000 1.0455    1.3891    3.88
GRU      16   2023    60    0.5657     0.7754 1.2093    1.3874    5.76
GRU      32   2023    60    1.3142     0.4336 1.3079    1.3866   10.01
Future-token checks passed for both trained architectures

plot_memory_comparison(trained, ("RNN", "GRU"), "10-rnn-gru-memory")

Figure 12. Does a gate help under this protocol? RNN and GRU checkpoints use the same episodes and update budget, with close parameter counts. These validation curves show one initialization for each delay.

Conclusion. At delay 16, GRU validation recall is 78%, compared with 31% for RNN. At delay 32, GRU reaches only 43% and RNN 28%: gating helps this run but does not solve every delay. These are single-seed results; Scene 12 checks initialization variability.

LSTM — separate carrying from reading

A GRU carries and exposes the same hidden state. A long short-term memory (LSTM) cell separates an internal cell state \(c_t\) from the representation \(h_t\) that the output head reads. It can keep a cell coordinate while exposing only part of it. Input, forget, and output gates are sigmoid functions of the current input and previous hidden state; the candidate \(\widetilde c_t\) uses tanh:

\[c_t=f_t\odot c_{t-1}+i_t\odot\widetilde c_t,\qquad h_t=o_t\odot\tanh(c_t).\]

The forget gate scales carrying, the input gate scales writing, and the output gate controls exposure. Unlike the GRU mixture, carrying and writing coefficients need not sum to one. We illustrate LSTM here without training one.

LSTM gradient qualification and reference

Along the direct cell-state path with gates held fixed, the derivative is \(f_t\). The full derivative also includes gate-computation paths. Gates near one can preserve a direct path, but cannot guarantee unlimited memory or successful training. See the PyTorch LSTM equations.

Checkpoint. For scalar values \(c_{t-1}=0.5\), \(f_t=0.9\), \(i_t=0.2\), \(\widetilde c_t=-0.4\), and \(o_t=0.8\), predict \(c_t\) before calculating. These are illustrative gate values, not a trained LSTM.

Show explanation — carrying and exposing the LSTM state

The retained part is \(0.9(0.5)=0.45\) and the write is \(0.2(-0.4)=-0.08\), so \(c_t=0.37\). The exposed value is \(0.8\tanh(0.37)\approx0.2832\). Changing only the output gate changes what is exposed at this step without changing the already computed cell state. Subsequent steps can still be affected through the changed hidden representation.

cell_previous, forget, input_gate, cell_candidate, output_gate = .5, .9, .2, -.4, .8
cell_next = forget*cell_previous + input_gate*cell_candidate
hidden_next = output_gate*np.tanh(cell_next)
assert np.isclose(cell_next, .37)
print(f"cell state = {cell_next:.4f}; exposed hidden state = {hidden_next:.4f}")
plot_lstm_cell()

cell state = 0.3700; exposed hidden state = 0.2832

Figure 13. Carry, write, and read an LSTM state. The forget gate scales the old cell state, the input gate scales the candidate write, and the output gate controls the exposed hidden state. The diagram illustrates the equations; no LSTM is trained here.

Interpretation checkpoint — what did the experiment establish?

One mechanism question. If a coordinate has a useful value, what does the GRU’s carry path allow the next step to do? Why does that not guarantee successful recall at every delay?

Show explanation — a route for memory, not a guarantee

A gate close to one can keep the coordinate close to its previous value even when the candidate differs. Training still has to learn a useful representation, an appropriate gate, and an output head that reads it. The displayed delay curves show selected models under a finite budget; the presence of a gate alone predicts no particular score.

One evidence question. Suppose recall improves while random-position loss stays near \(\ln4\). Which part of the task has been learned, and how far can you generalize the comparison?

Show explanation — what the recurrent comparison supports

Recall improves on a dependency present in the prefix. Independent wandering commands retain their uncertainty. Cue changes and state erasure help test whether predictions use earlier information, while small gradients alone do not measure stored information. These are fixed datasets and one initialization per model and delay, with close parameter counts and a shared update budget. Several checkpoints reach the budget limit. Scene 12 repeats initializations before making the held-out comparison; even that will not establish a universal architecture ranking.

We have now trained and inspected recurrent predictors. The next scenes replace sequential state propagation with direct, causally restricted attention to earlier representations.

Scene 10 — Retrieve earlier evidence directly

Question. Must information pass through every intervening recurrent state to influence recall?

An RNN has to carry the cue through the wandering commands. Attention offers another route: at RECALL, form a new representation by combining representations at available earlier positions. The cue need not survive in one state passed through every intervening step.

A Transformer is a model built around attention and ordinary neural layers. We will first calculate one attention operation by hand, then assemble a small Transformer. Attention alone produces a representation; an output head still has to turn it into next-command probabilities.

Give a query something to match and something to read

Each position supplies three vectors, formed by learned linear projections of its representation:

Role	What it does in this calculation	What it does not mean
Query	The current position’s vector used to score possible sources	A literal question written in words
Key	A source’s vector compared with the query	A fixed token ID or a dictionary entry
Value	A source’s content vector used in the weighted sum	The matching score itself

Self-attention means the queries, keys, and values come from the same sequence of representations. A query–key dot product determines a source’s score. Softmax turns the permitted scores into weights. Those weights combine the values, not the keys.

plot_attention_route()

Figure 14. Separate matching from retrieved content. Keys help decide the weights; values provide the vectors to combine. The causal mask restricts which sources the query can read.

Follow one query row

Stack the vectors for one head into \(Q_{\rm att},K_{\rm att}:(T,d_k)\) and \(V_{\rm att}:(T,d_v)\). Query position \(t\) scores source position \(j\):

\[S_{tj}=\frac{Q_{\rm att}[t]K_{\rm att}[j]^\top}{\sqrt{d_k}}.\]

A large score gives that source more weight relative to the others. Dividing by \(\sqrt{d_k}\) moderates the dot-product scale; it does not impose a time boundary. We impose that boundary before softmax: \(M_{tj}=0\) for \(j\leq t\) and \(-\infty\) for \(j>t\).

\[A_{tj}=\frac{\exp(S_{tj}+M_{tj})}{\sum_{m=0}^{T-1}\exp(S_{tm}+M_{tm})},\qquad O[t]=\sum_{j=0}^{T-1}A_{tj}V_{\rm att}[j].\]

Here \(m\) runs over source positions in the normalization sum. A blocked source has zero weight. For each query, the remaining weights sum to one. The output \(O[t]\) is a weighted average of the value vectors, with shape \((d_v,)\). It is not yet a distribution over commands.

All query rows can be calculated together:

\[S=Q_{\rm att}K_{\rm att}^\top/\sqrt{d_k},\qquad A=\operatorname{softmax}_{\text{source columns}}(S+M),\qquad O=AV_{\rm att}.\]

Rows are queries and columns are sources, so softmax acts across columns, not down rows. The mechanism follows Attention Is All You Need, Sections 3.2–3.3. The paper stores projection weights feature-by-output; our nn.Linear layers store output-by-input, so the equivalent row-vector projection uses a transpose.

Checkpoint. In BOS UP RIGHT RECALL UP, may RECALL read the final UP? Why may it read its own current input?

Show explanation — may RECALL read its future answer?

The query at position 3 can use positions 0 through 3. Position 4 contains its target UP, so that source must be masked. The diagonal is allowed because the model already observes the current input RECALL; it is predicting the following token.

Choose numbers that make each operation visible

The five positions below have hand-chosen two-coordinate queries, keys, and values. These are untrained teaching arrays, not a model explanation. Follow row 3, the RECALL query: scores first, future masking second, weights third, then the value sum.

import math
attention_tokens = np.array([BOS, UP, RIGHT, RECALL, UP])
Q_demo = np.array([[1., 0.], [0., 1.], [1., 1.], [1., 0.], [0., 1.]])
K_demo = np.array([[.1, .2], [2., 0.], [0., 2.], [-1., 0.], [8., 0.]])
V_demo = np.array([[0., 0.], [1., 0.], [0., 1.], [.2, .2], [0., 10.]])

The matrix product makes every query–source score. Subtracting each row’s largest permitted score stabilizes softmax; it leaves the weights unchanged. axis=-1 is the source axis.

scores_demo = Q_demo @ K_demo.T / math.sqrt(2)
future_mask_demo = np.triu(np.ones((5, 5), dtype=bool), k=1)
masked_scores_demo = np.where(future_mask_demo, -np.inf, scores_demo)
shifted_demo = masked_scores_demo - masked_scores_demo.max(axis=-1, keepdims=True)
weights_demo = np.exp(shifted_demo) / np.exp(shifted_demo).sum(axis=-1, keepdims=True)
output_demo = weights_demo @ V_demo

# @title Attention ledger and checks: inspect the RECALL row
assert np.allclose(weights_demo.sum(axis=-1), 1)
assert np.all(weights_demo[future_mask_demo] == 0)
print("RECALL query:", Q_demo[3])
print("Scaled scores:", scores_demo[3].round(4))
print("After causal masking:", masked_scores_demo[3].round(4))
print("Row-normalized weights:", weights_demo[3].round(4))
print("Weighted value:", output_demo[3].round(4))
print("Each source contribution to the RECALL output:\n", (weights_demo[3, :, None]*V_demo).round(4))
# The diagonal is visible: input at t predicts target at t+1.
assert np.all(np.diag(weights_demo) > 0)

RECALL query: [1. 0.]
Scaled scores: [ 0.0707  1.4142  0.     -0.7071  5.6569]
After causal masking: [ 0.0707  1.4142  0.     -0.7071    -inf]
Row-normalized weights: [0.1607 0.6158 0.1497 0.0738 0.    ]
Weighted value: [0.6306 0.1645]
Each source contribution to the RECALL output:
 [[0.     0.    ]
 [0.6158 0.    ]
 [0.     0.1497]
 [0.0148 0.0148]
 [0.     0.    ]]

At RECALL, the query is \((1,0)\) and the initial UP key is \((2,0)\). Their scaled score is \(2/\sqrt2\approx1.4142\). The future UP key would score \(8/\sqrt2\), but the mask blocks it entirely. Conclusion. The initial cue receives about 62% of the weight; the future answer receives zero. The output combines permitted value vectors, so it is neither a hard token selection nor a command probability.

Check your reading — must attention select exactly one source?

No. Softmax distributes weight across every permitted source with a finite score. A source can dominate without being the only contributor. Each coordinate of the output is the same weighted sum of that coordinate in the values. In this example, changing a key can change the weights; changing only a value changes the content combined under the existing weights.

plot_attention_calculation()

Figure 15. Turn scores into retrieved content. Read one query row across the score and weight matrices, then follow its weighted values into the output. Future sources receive zero attention weight.

Remove the mask and the answer becomes visible

The final UP is a future input token and the correct target for the earlier RECALL query. A bidirectional encoder is useful when the entire supplied sequence is legitimately available, such as encoding a complete sentence. It cannot use that entire shifted training sequence for autoregressive next-command prediction without leaking future answers.

A causal Transformer can compute all training positions in parallel while respecting this restriction. At generation time, predict from the available prefix, append one chosen token, then repeat. Parallel teacher-forced training does not make unknown future commands available during generation.

Checkpoint. Which weight must be zero? What happens to the earlier attention output if we change only the future value vector?

Show explanation — change only a future value

The weight from the RECALL query to the final UP must be zero. Changing that future value leaves the causal output unchanged, because its coefficient is zero. Without the mask, that value can contribute and the output can change. This is a direct leakage check, independent of whether the untrained output happens to resemble the right answer.

plot_mask_comparison()
unmasked_shift = scores_demo - scores_demo.max(axis=-1, keepdims=True)
unmasked_weights = np.exp(unmasked_shift)/np.exp(unmasked_shift).sum(axis=-1, keepdims=True)
changed_values = V_demo.copy(); changed_values[4] += np.array([20., -20.])
assert np.array_equal((weights_demo @ changed_values)[:4], output_demo[:4])
assert not np.allclose((unmasked_weights @ changed_values)[3], (unmasked_weights @ V_demo)[3])
print(f"RECALL weight on its future target: causal={weights_demo[3,4]:.4f}, "
      f"unmasked={unmasked_weights[3,4]:.4f}")
print("Future-value intervention: causal output unchanged; unmasked output changes")

RECALL weight on its future target: causal=0.0000, unmasked=0.9772
Future-value intervention: causal output unchanged; unmasked output changes

Figure 16. Visibility depends on the prediction task. Causal attention permits the current and preceding inputs. Bidirectional attention can read the supplied future answer; the red box identifies the forbidden edge for this next-token task.

Conclusion. Without masking, this query places about 98% of its weight on the future answer. Changing that future value changes the unmasked output, while the causal output remains identical. The mask enforces the prediction boundary before training begins.

Scene 11 — Assemble and train a small Transformer

Question. How do we turn the attention calculation into a trainable next-command model?

We need representations to attend to, learned projections to form queries/keys/values, and an output head to read the result. Our model contains one causal self-attention block followed by seven next-command logits. This small decoder-only model has no encoder or cross-attention.

Read the route before the implementation

An embedding table supplies a vector for each command. A learned position embedding supplies a vector for its index. Adding them gives \(h^{(0)}:(B,T,D)\) with \(D=12\). This tells the model both what token it sees and where it occurs; attention’s score calculation alone supplies no absolute position index.

The attention branch reads a normalized copy of \(h^{(0)}\). Add its output back to the original representation, then do the same with an ordinary MLP:

\[h^{(1)}=h^{(0)}+\operatorname{MultiHeadCausalAttention}(\operatorname{LN}_1(h^{(0)})),\] \[h^{(2)}=h^{(1)}+\operatorname{MLP}(\operatorname{LN}_2(h^{(1)})).\]

These are residual additions: the branch must return the same shape as its input. They are not concatenations. The MLP is \(12\to24\to12\) with ReLU, sharing weights across positions while acting on each position separately. Final normalization and a linear output head produce logits \((B,T,7)\).

plot_transformer_block()

Figure 17. Follow a representation through one Transformer block. Each bypass carries the input to an addition. Attention mixes permitted positions; the MLP transforms features within each position. The diagram shows the pre-normalization variant used here.

1. Normalize features within a token

Layer normalization centers and rescales the feature coordinates of one token, then learns a scale and shift for each feature. It does not average across token positions or examples. In the code, dim=-1 is the feature axis; keepdim=True preserves an axis of width one for broadcasting.

We use population variance, a stabilizer \(\epsilon=10^{-5}\), and an explicit affine transformation, checked against PyTorch LayerNorm. The normalization’s statistics therefore cannot read a future token.

class FeatureLayerNorm(nn.Module):
    def __init__(self, width, eps=1e-5):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(width))
        self.shift = nn.Parameter(torch.zeros(width))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        variance = ((x-mean)**2).mean(dim=-1, keepdim=True)
        return (x-mean)/torch.sqrt(variance+self.eps)*self.scale+self.shift

# @title Verification: explicit feature normalization matches LayerNorm
# Numerical normalization check with nontrivial learned affine values.
norm_probe = torch.arange(24, dtype=torch.float64).reshape(2, 12)/7
manual_norm = FeatureLayerNorm(12).double()
reference_norm = nn.LayerNorm(12, dtype=torch.float64)
with torch.no_grad():
    manual_norm.scale.copy_(torch.linspace(.5, 1.5, 12))
    manual_norm.shift.copy_(torch.linspace(-.2, .2, 12))
    reference_norm.weight.copy_(manual_norm.scale)
    reference_norm.bias.copy_(manual_norm.shift)
torch.testing.assert_close(manual_norm(norm_probe), reference_norm(norm_probe), rtol=1e-10, atol=1e-12)
print("Explicit feature normalization matches LayerNorm")

Explicit feature normalization matches LayerNorm

2. Calculate attention in two heads

A head is one query/key/value calculation. Our learned projections produce 12 coordinates, split into two groups of six. Each group calculates its own weights and retrieved vector. Two heads offer two learned matching routes; they do not come with predefined jobs.

The shape changes from \((B,T,12)\) to \((B,2,T,6)\): reshape separates the feature groups, then transpose(1,2) places the head axis before positions. We use \(N_{\rm heads}=2\) rather than \(H\), which already denotes recurrent hidden width. Scores have shape \((B,2,T,T)\), with query and source as the last two axes. Masking and softmax are exactly the operations from Scene 10, now repeated across batches and heads.

def split_heads(projected, heads, head_width):
    batch, positions, _ = projected.shape
    return projected.reshape(batch, positions, heads, head_width).transpose(1, 2)


def scaled_causal_attention(q, k, v, block_cue=False):
    positions, head_width = q.shape[-2:]
    scores = q @ k.transpose(-2, -1) / math.sqrt(head_width)
    future_mask = torch.ones(positions, positions, dtype=torch.bool, device=q.device).triu(1)
    scores = scores.masked_fill(future_mask, float("-inf"))
    if block_cue and positions > 1:
        # Later diagnostic: block cue source 1; BOS remains available in every row.
        source_mask = torch.arange(positions, device=q.device) == 1
        scores = scores.masked_fill(source_mask, float("-inf"))
    weights = scores.softmax(dim=-1)
    return weights @ v, weights

3. Read both heads back into one representation

attend_from_prefix normalizes the input, forms Q/K/V, and calls the attention calculation. To join the heads, reverse the axis swap and reshape \((B,2,T,6)\) back to \((B,T,12)\). A learned output projection lets the next layer combine the two heads’ features. This function returns the attention branch, which the forward pass will add to its input.

The optional block_cue flag is only for the later information-boundary experiment. Ordinary training leaves it false. return_attention exposes the weight matrices for inspection; the classifier still reads representations through its output head.

def attend_from_prefix(model, h0, block_cue=False):
    batch, positions, _ = h0.shape
    normalized = model.norm_attention(h0)
    q = split_heads(model.query(normalized), model.heads, model.head_width)
    k = split_heads(model.key(normalized), model.heads, model.head_width)
    v = split_heads(model.value(normalized), model.heads, model.head_width)
    head_outputs, weights = scaled_causal_attention(q, k, v, block_cue)
    combined = head_outputs.transpose(1, 2).reshape(batch, positions, model.width)
    return model.attention_output(combined), weights

4. Register the parameters and connect the route

__init__ creates the learned tables and layers. forward follows Figure 17: token plus position, attention plus its residual, MLP plus its residual, then the output head. Each operation has already been defined above. No built-in attention or Transformer block is hidden inside this model.

This is a compact teaching variant of the original Transformer: normalization is before each branch, positions are learned instead of sinusoidal, and there is one block with no dropout, encoder, or cross-attention. Its position table holds at most 36 inputs. We do not claim unseen-length generalization.

class CommandTransformer(nn.Module):
    def __init__(self):
        super().__init__()
        self.width, self.heads, self.head_width = 12, 2, 6
        self.max_positions = 36
        self.token = nn.Embedding(K, self.width)
        self.position = nn.Embedding(self.max_positions, self.width)
        self.norm_attention = FeatureLayerNorm(self.width)
        self.query = nn.Linear(self.width, self.width)
        self.key = nn.Linear(self.width, self.width)
        self.value = nn.Linear(self.width, self.width)
        self.attention_output = nn.Linear(self.width, self.width)
        self.norm_mlp = FeatureLayerNorm(self.width)
        self.mlp_in = nn.Linear(self.width, 24)
        self.mlp_out = nn.Linear(24, self.width)
        self.norm_final = FeatureLayerNorm(self.width)
        self.output_head = nn.Linear(self.width, K)

    def forward(self, ids, return_attention=False, block_cue=False):
        _, positions = ids.shape
        if not 1 <= positions <= self.max_positions:
            raise ValueError("Input length must be between 1 and 36 positions")
        positions_t = torch.arange(positions, device=ids.device)
        h0 = self.token(ids) + self.position(positions_t)
        attention_branch, weights = attend_from_prefix(self, h0, block_cue)
        h1 = h0 + attention_branch
        h2 = h1 + self.mlp_out(F.relu(self.mlp_in(self.norm_mlp(h1))))
        logits = self.output_head(self.norm_final(h2))
        return (logits, weights) if return_attention else logits

Checkpoint. Which part can combine different positions, and which parts transform each position independently? Where is a future source removed?

Show explanation — which operations mix positions?

Attention combines values from permitted source positions. Embedding lookup, feature normalization, the position-wise MLP, and the output head act at each position independently. Residual additions combine matching coordinates at the same position. The mask removes future sources from scores before softmax; a later output head cannot repair leaked attention.

Before training, verify the tensor shapes, 1,915 parameters, normalized weight rows, zero future weights, and invariance of earlier logits when a future input changes.

# @title Verification: Transformer shapes, masking, and future-token invariance
# Shapes, masking, and causality before training.
torch.manual_seed(TRAINING_SEED)
transformer_probe = CommandTransformer()
probe_ids = validation_ids[:4]
with torch.no_grad():
    probe_logits, probe_weights = transformer_probe(probe_ids, return_attention=True)
assert probe_logits.shape == (4, 12, K)
assert probe_weights.shape == (4, 2, 12, 12)
assert sum(p.numel() for p in transformer_probe.parameters()) == 1915
upper_triangle = torch.ones(12,12,dtype=torch.bool).triu(1)
assert torch.all(probe_weights[..., upper_triangle] == 0)
torch.testing.assert_close(probe_weights.sum(-1), torch.ones(4,2,12), rtol=1e-6, atol=1e-6)
changed_probe = probe_ids.clone(); changed_probe[:, 11] = (changed_probe[:,11]+1)%4
with torch.no_grad():
    changed_probe_logits = transformer_probe(changed_probe)
torch.testing.assert_close(probe_logits[:, :11], changed_probe_logits[:, :11], rtol=0, atol=0)
print("Before training: shapes, finite logits, row sums, zero future weights, and future-token invariance passed")
assert torch.isfinite(probe_logits).all()
print("Shape ledger: IDs (B,T); embeddings (B,T,12); Q/K/V (B,2,T,6); scores (B,2,T,T); logits (B,T,7)")

Before training: shapes, finite logits, row sums, zero future weights, and future-token invariance passed
Shape ledger: IDs (B,T); embeddings (B,T,12); Q/K/V (B,2,T,6); scores (B,2,T,T); logits (B,T,7)

Train with the existing protocol

The Transformer has 1,915 parameters, versus 1,943 for the RNN and 2,023 for the GRU. The counts are close, but architecture, widths, normalization, positional parameters, and compute differ. The count includes position-table rows unused at shorter delays, another limit of count matching. Attention forms \(T\times T\) scores per head; this implementation has quadratic score storage as length grows.

Reuse the Scene 7 training protocol and function, training separately at delays 8/16/32. A learned position table always places the initial cue at index 1; this task can be solved by learning to retrieve that fixed position. It does not establish general content-dependent retrieval or length extrapolation.

Checkpoint. Predict whether direct access to the cue will help at longer delays, then inspect the executed curves and runtime. No test data enter this comparison.

for delay in (8, 16, 32):
    tx, ty, vx, vy = training_data[delay]
    trained[("Transformer", delay)] = train_sequence_model(
        "Transformer", tx, ty, vx, vy, delay, model_factory=CommandTransformer)
print("Validation-selected models; single initialization seed; losses in nats")
print("model        delay params epoch recall_CE recall_acc all_CE random_CE seconds")
for kind in ("RNN", "GRU", "Transformer"):
    for delay in (8, 16, 32):
        run = trained[(kind,delay)]; m = run["validation"]
        print(f"{kind:12s} {delay:5d} {run['parameters']:6d} {run['best_epoch']:5d} "
              f"{m['recall_ce']:9.4f} {m['recall_accuracy']:10.4f} {m['all_ce']:6.4f} "
              f"{m['random_ce']:9.4f} {run['seconds']:7.2f}")
transformer8 = trained[("Transformer",8)]

Validation-selected models; single initialization seed; losses in nats
model        delay params epoch recall_CE recall_acc all_CE random_CE seconds
RNN              8   1943    60    0.0272     1.0000 1.0458    1.3902    2.50
RNN             16   1943    60    1.2383     0.3105 1.2433    1.3897    3.56
RNN             32   1943    49    1.3863     0.2754 1.4491    1.3900    5.84
GRU              8   2023    59    0.0418     1.0000 1.0455    1.3891    3.88
GRU             16   2023    60    0.5657     0.7754 1.2093    1.3874    5.76
GRU             32   2023    60    1.3142     0.4336 1.3079    1.3866   10.01
Transformer      8   1915    60    0.0034     1.0000 1.0412    1.3879    9.25
Transformer     16   1915    59    0.0041     1.0000 1.1802    1.3882   11.26
Transformer     32   1915    60    0.0075     1.0000 1.2718    1.3871   20.93

plot_recurrent_curves(transformer8, "14-transformer-learning", model_label="Transformer")

Figure 18. Train direct retrieval. The delay-8 Transformer is evaluated with the same learning and selection protocol as the recurrent models. The dotted line marks the validation-selected checkpoint.

plot_memory_comparison(trained, ("RNN", "GRU", "Transformer"), "15-three-model-memory")

Figure 19. Compare three ways to use history. These validation curves use one initialization per architecture and delay. Close parameter counts leave differences in computation, representation width, and positional information.

Conclusion. The Transformer reaches 100% validation recall at all three delays in this run. Direct retrieval makes this fixed-position dependency learnable under the shared budget; the experiment does not yet establish content-based search or unseen-length generalization.

Check a trained model’s information boundary

Check all three selected Transformer checkpoints. Alter future tokens and compare all earlier logits; check future attention weights and future embedding sensitivities. Then inspect the delay-8 attention maps and intervene on the cue. An attention heatmap displays weights, not a complete causal explanation of a prediction.

At delay 8, changing the cue while keeping the observed suffix fixed tests whether the prediction follows the cue. Blocking source position 1 removes direct retrieval of that cue in our one-block model: values before attention contain only each token’s embedding and position. In a deeper Transformer, other positions might already carry cue information, so removing one edge would not imply erasing every route.

# @title Verification: causal outputs, attention weights, and input gradients
for delay in (8,16,32):
    model = trained[("Transformer",delay)]["model"]
    vx, vy = training_data[delay][2:]
    recall_position = delay+2
    altered = vx[:8].clone()
    altered[:, recall_position+1:] = (altered[:, recall_position+1:]+1)%4
    with torch.no_grad():
        logits, weights = model(vx[:8], return_attention=True)
        altered_logits = model(altered)
    torch.testing.assert_close(logits[:, :recall_position+1], altered_logits[:, :recall_position+1], rtol=0, atol=0)
    upper = torch.ones(vx.shape[1],vx.shape[1],dtype=torch.bool).triu(1)
    assert torch.all(weights[..., upper] == 0)
    # Differentiate recall-only loss with respect to token embedding outputs.
    captured_embeddings = []
    def capture_embedding(module, args, output):
        output.retain_grad(); captured_embeddings.append(output)
    hook = model.token.register_forward_hook(capture_embedding)
    model.zero_grad(set_to_none=True)
    try:
        recall_logits = model(vx[:8])[:, recall_position]
        F.cross_entropy(recall_logits, vy[:8, recall_position]).backward()
        assert torch.all(captured_embeddings[0].grad[:, recall_position+1:] == 0)
    finally:
        hook.remove(); model.zero_grad(set_to_none=True)
print("Trained Transformer: future-token, zero future attention, and zero future embedding-gradient checks passed at every delay")

Trained Transformer: future-token, zero future attention, and zero future embedding-gradient checks passed at every delay

Does recall actually use the cue?

Now keep the wandering suffix fixed and change the initial cue. Does the prediction follow the new cue? Then block source position 1 in attention and measure recall again. Read these two interventions alongside the heatmaps below: one changes the evidence, the other removes its direct route.

model = transformer8["model"]
changed_prefix = validation_ids[:128,:11].clone()
changed_prefix[:,1] = (changed_prefix[:,1]+1)%4
with torch.no_grad():
    changed_prediction = model(changed_prefix)[:, -1].argmax(-1)
    blocked_logits = model(validation_ids, block_cue=True)[:,10]
    blocked_accuracy = (blocked_logits.argmax(-1)==validation_targets[:,10]).float().mean().item()
changed_accuracy = (changed_prediction==changed_prefix[:,1]).float().mean().item()
print(f"Delay-8 Transformer: original recall accuracy={transformer8['validation']['recall_accuracy']:.4f}; "
      f"changed-cue accuracy={changed_accuracy:.4f}; cue-source blocked accuracy={blocked_accuracy:.4f}")

Delay-8 Transformer: original recall accuracy=1.0000; changed-cue accuracy=1.0000; cue-source blocked accuracy=0.2441

plot_trained_attention(transformer8["model"], validation_ids)

Figure 20. Inspect learned attention, then test its role. Each heatmap shows one head on the same validation episode. The outlined edge connects RECALL to the initial cue. The surrounding cue interventions provide evidence beyond the attention weights.

Conclusion. Changing the cue makes every altered-prefix prediction follow it; blocking its source lowers recall from 100% to 24%. This supports a functional role for the cue route in this one-block model, beyond what the heatmaps alone show.

plot_recall_predictions(transformer8["model"], validation_ids, validation_targets, 8, "17-transformer-predictions")

Figure 21. Inspect the Transformer’s recall predictions. Paths and distributions correspond to the same four validation episodes shown for the RNN. Compare the probability on each true cue, alongside the most likely command.

Interpretation checkpoint — what changed from recurrence?

Trace the computation. At the RECALL query, how can changing a key differ from changing its value? Why is a large weight on the cue insufficient by itself to explain the prediction?

Show explanation — weights and content have different roles

A key participates in a query–key score, so changing it can change the whole row of weights. A value is content in the weighted sum. Changing only that value leaves this head’s scores and weights fixed but can change the output. The complete model also has other sources, another head, output projections, residual paths, and an MLP. Cue-change and blocked-source interventions test whether that route matters to the prediction; a heatmap alone does not establish it.

Read the evidence. If this Transformer recalls the cue at long delays, what has this experiment demonstrated, and what has it left open?

Show explanation — the scope of direct retrieval

The model can learn this controlled dependency under the specified training protocol. The cue is always at position 1, so learned positional retrieval is sufficient; general content-based search has not been established. These curves use one initialization and selected validation checkpoints. Widths, normalization, compute, and optimization differ despite close parameter counts. Good recall does not make independent wandering predictable or demonstrate language understanding. Scene 12 repeats initializations and evaluates held-out episodes under frozen choices.

The three architectures are now implemented and trained. Scene 12 fixes the final comparison, evaluates the reserved test splits, and closes with interpretation. Until then, neural results are validation evidence.

Scene 12 — Compare on episodes nobody trained on

Question. Does the memory result survive a new initialization and held-out episodes?

Freeze the comparison before scoring test performance: RNN, GRU, and the one-block Transformer; delays 8/16/32; initialization seeds 6206/6207/6208; unchanged data seeds, batch orders, widths, optimizer, clipping, and 60-epoch budget. Each run selects its own checkpoint by validation recall cross-entropy. Reuse the nine seed-6206 runs already computed; train the other eighteen. Keep every seed, including unsuccessful runs; do not select a best seed on test results.

All 27 checkpoints must exist before the test loop. The test split contains 1,024 complete episodes per delay, with no episode shared with training or validation. We regenerate the same splits and verify the training/validation arrays against those used above. This synthetic episode task needs episode separation; real time series may instead require chronological splits and separation of overlapping windows.

Report all-token, recall, and random-position cross-entropy and accuracy. Show mean ± sample standard deviation over three initializations, plus individual runs. This spread describes optimization variability on the same fixed data; it is neither a confidence interval nor uncertainty over new datasets. The fixed test examples are the first four delay-16 episodes and seed 6206, chosen here before scoring.

Checkpoint. Which evidence tests memory, which tests generalization to new episodes, and which only tests sensitivity to initialization?

Show explanation — what does repeating seeds establish?

Held-out episodes test predictions on examples excluded from fitting and checkpoint selection. Repeating initialization seeds tests optimization variability on these fixed datasets. Cue interventions test whether predictions respond to the relevant evidence. These address different questions; three seeds do not estimate uncertainty across new datasets.

FINAL_KINDS = ("RNN", "GRU", "Transformer")
FINAL_DELAYS = (8, 16, 32)
FINAL_SEEDS = (6206, 6207, 6208)
EXAMPLE_DELAY, EXAMPLE_SEED = 16, 6206
expected_runs = {(kind, d, seed) for kind in FINAL_KINDS for d in FINAL_DELAYS for seed in FINAL_SEEDS}
final_runs = {(kind, d, TRAINING_SEED): trained[(kind, d)]
              for kind in FINAL_KINDS for d in FINAL_DELAYS}
for seed in FINAL_SEEDS[1:]:
    for kind in FINAL_KINDS:
        for delay in FINAL_DELAYS:
            factory = CommandTransformer if kind == "Transformer" else None
            run = train_sequence_model(kind, *training_data[delay], delay,
                                       model_factory=factory, initialization_seed=seed)
            final_runs[(kind, delay, seed)] = run
            print(f"Frozen checkpoint: {kind}, d={delay}, seed={seed}, epoch={run['best_epoch']}", flush=True)
assert set(final_runs) == expected_runs
assert all(run["initialization_seed"] == seed for (kind, d, seed), run in final_runs.items())
print("All 27 validation-selected checkpoints ready; no test performance scored yet.")

Frozen checkpoint: RNN, d=8, seed=6207, epoch=58
Frozen checkpoint: RNN, d=16, seed=6207, epoch=60
Frozen checkpoint: RNN, d=32, seed=6207, epoch=60
Frozen checkpoint: GRU, d=8, seed=6207, epoch=59
Frozen checkpoint: GRU, d=16, seed=6207, epoch=60
Frozen checkpoint: GRU, d=32, seed=6207, epoch=52
Frozen checkpoint: Transformer, d=8, seed=6207, epoch=60
Frozen checkpoint: Transformer, d=16, seed=6207, epoch=60
Frozen checkpoint: Transformer, d=32, seed=6207, epoch=60
Frozen checkpoint: RNN, d=8, seed=6208, epoch=60
Frozen checkpoint: RNN, d=16, seed=6208, epoch=60
Frozen checkpoint: RNN, d=32, seed=6208, epoch=44
Frozen checkpoint: GRU, d=8, seed=6208, epoch=60
Frozen checkpoint: GRU, d=16, seed=6208, epoch=60
Frozen checkpoint: GRU, d=32, seed=6208, epoch=60
Frozen checkpoint: Transformer, d=8, seed=6208, epoch=60
Frozen checkpoint: Transformer, d=16, seed=6208, epoch=60
Frozen checkpoint: Transformer, d=32, seed=6208, epoch=60
All 27 validation-selected checkpoints ready; no test performance scored yet.

held_out = {}
for delay in FINAL_DELAYS:
    rows, _ = make_splits(delay)
    regenerated = tuple(torch.tensor(rows[name][:, offset:stop], dtype=torch.long)
                        for name in ("train", "validation") for offset, stop in ((0, -1), (1, None)))
    assert all(torch.equal(a, b) for a, b in zip(regenerated, training_data[delay]))
    held_out[delay] = (torch.tensor(rows["test"][:, :-1], dtype=torch.long),
                       torch.tensor(rows["test"][:, 1:], dtype=torch.long))
assert set(final_runs) == expected_runs  # no missing run can be silently omitted
# Evaluation starts only after the complete checkpoint set is fixed.
test_results = {}
for key, run in final_runs.items():
    kind, delay, seed = key
    test_results[key] = sequence_metrics(run["model"], *held_out[delay], delay)
assert set(test_results) == expected_runs
METRIC_NAMES = ("all_ce", "all_accuracy", "recall_ce", "recall_accuracy", "random_ce", "random_accuracy")
assert all(np.isfinite([r[m] for m in METRIC_NAMES]).all() for r in test_results.values())

summary = {}
for kind in FINAL_KINDS:
    for delay in FINAL_DELAYS:
        values = np.array([[test_results[(kind, delay, seed)][m] for m in METRIC_NAMES] for seed in FINAL_SEEDS])
        summary[(kind, delay)] = dict(mean=values.mean(0), sd=values.std(0, ddof=1))

Read recall first

The next table answers the chapter’s central question: can the model predict the initial cue on held-out episodes? Each entry is mean ± sample SD over three initializations. Accuracy describes which token wins; recall cross-entropy also describes how much probability the true token receives. Both matter.

Read one delay at a time. At delay 16, compare the spread across RNN and GRU runs with the Transformer results. At delay 32, ask whether a small spread means consistently successful recall or consistently poor recall. The full per-seed ledger below retains every run and all six metrics; the separate compute ledger records parameter counts and runtime.

# @title Display the compact held-out recall comparison
from IPython.display import Markdown
compact_lines = ["| Model | Delay | Recall accuracy | Recall CE (nats) |",
                 "|---|---:|---:|---:|"]
for delay in FINAL_DELAYS:
    for kind in FINAL_KINDS:
        stats = summary[(kind, delay)]
        accuracy = f"{stats['mean'][3]:.4f} ± {stats['sd'][3]:.4f}"
        loss = f"{stats['mean'][2]:.4f} ± {stats['sd'][2]:.4f}"
        compact_lines.append(f"| {kind} | {delay} | {accuracy} | {loss} |")
display(Markdown("\n".join(compact_lines)))

Model	Delay	Recall accuracy	Recall CE (nats)
RNN	8	1.0000 ± 0.0000	0.1194 ± 0.1152
GRU	8	1.0000 ± 0.0000	0.0446 ± 0.0095
Transformer	8	1.0000 ± 0.0000	0.0030 ± 0.0003
RNN	16	0.4404 ± 0.1256	0.9919 ± 0.2216
GRU	16	0.8141 ± 0.1162	0.5479 ± 0.0650
Transformer	16	1.0000 ± 0.0000	0.0044 ± 0.0004
RNN	32	0.2568 ± 0.0000	1.3878 ± 0.0024
GRU	32	0.3223 ± 0.1024	1.3627 ± 0.0429
Transformer	32	1.0000 ± 0.0000	0.0098 ± 0.0015

Result at a glance. At delay 16, GRU recalls about four in five cues, RNN fewer than half, and Transformer every cue across all three initializations. At delay 32, both recurrent models struggle; the Transformer succeeds on the fixed-position dependency. Read the loss column alongside accuracy and retain all seeds.

# @title Display the full neural and compute ledgers
from IPython.display import Markdown
header = "| Model | d | Seed | Epoch | All CE | All acc | Recall CE | Recall acc | Random CE | Random acc |\n"
header += "|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|\n"
rows_text = []
for kind in FINAL_KINDS:
    for delay in FINAL_DELAYS:
        for seed in FINAL_SEEDS:
            key = (kind, delay, seed)
            values = " | ".join(f"{test_results[key][m]:.4f}" for m in METRIC_NAMES)
            rows_text.append(f"| {kind} | {delay} | {seed} | {final_runs[key]['best_epoch']} | {values} |")


display(Markdown("<details>\n<summary>Full neural results — every seed and all six metrics</summary>\n\n"
                 + "Test performance; CE in nats, accuracies as fractions.\n\n"
                 + header + "\n".join(rows_text) + "\n\n</details>"))
compute_lines = ["| Model | Delay | Parameters | Mean training seconds |",
                 "|---|---:|---:|---:|"]
for kind in FINAL_KINDS:
    for delay in FINAL_DELAYS:
        runs = [final_runs[(kind, delay, seed)] for seed in FINAL_SEEDS]
        seconds = np.mean([run["seconds"] for run in runs])
        compute_lines.append(f"| {kind} | {delay} | {runs[0]['parameters']} | {seconds:.1f} |")
display(Markdown("<details>\n<summary>Compute ledger — parameters and observed CPU runtime</summary>\n\n"
                 + "\n".join(compute_lines) + "\n\n</details>"))

Full neural results — every seed and all six metrics

Test performance; CE in nats, accuracies as fractions.

Model	d	Seed	Epoch	All CE	All acc	Recall CE	Recall acc	Random CE	Random acc
RNN	8	6206	60	1.0442	0.4427	0.0278	1.0000	1.3880	0.2569
RNN	8	6207	58	1.0639	0.4396	0.2488	1.0000	1.3887	0.2528
RNN	8	6208	60	1.0479	0.4358	0.0816	1.0000	1.3875	0.2477
RNN	16	6206	60	1.2439	0.3290	1.2463	0.2969	1.3899	0.2520
RNN	16	6207	60	1.2242	0.3336	0.8415	0.4941	1.3899	0.2458
RNN	16	6208	60	1.2281	0.3404	0.8878	0.5303	1.3923	0.2516
RNN	32	6206	49	1.4484	0.2661	1.3881	0.2568	1.3896	0.2522
RNN	32	6207	60	1.3458	0.2824	1.3852	0.2568	1.4114	0.2397
RNN	32	6208	44	1.4340	0.2646	1.3900	0.2568	1.4169	0.2506
GRU	8	6206	59	1.0451	0.4411	0.0412	1.0000	1.3885	0.2548
GRU	8	6207	59	1.0464	0.4357	0.0553	1.0000	1.3885	0.2476
GRU	8	6208	60	1.0431	0.4425	0.0374	1.0000	1.3864	0.2567
GRU	16	6206	60	1.2108	0.3479	0.5888	0.7422	1.3878	0.2480
GRU	16	6207	60	1.2048	0.3608	0.4730	0.9482	1.3871	0.2511
GRU	16	6208	60	1.2094	0.3449	0.5821	0.7520	1.3883	0.2439
GRU	32	6206	60	1.3081	0.2939	1.3132	0.4404	1.3868	0.2467
GRU	32	6207	52	1.3114	0.2936	1.3873	0.2598	1.3884	0.2518
GRU	32	6208	60	1.3099	0.2932	1.3876	0.2666	1.3867	0.2512
Transformer	8	6206	60	1.0401	0.4452	0.0034	1.0000	1.3863	0.2603
Transformer	8	6207	60	1.0402	0.4421	0.0029	1.0000	1.3866	0.2561
Transformer	8	6208	60	1.0399	0.4460	0.0028	1.0000	1.3861	0.2613
Transformer	16	6206	59	1.1799	0.3615	0.0040	1.0000	1.3878	0.2488
Transformer	16	6207	60	1.1807	0.3610	0.0046	1.0000	1.3887	0.2482
Transformer	16	6208	60	1.1806	0.3624	0.0045	1.0000	1.3887	0.2498
Transformer	32	6206	60	1.2723	0.3092	0.0081	1.0000	1.3876	0.2464
Transformer	32	6207	60	1.2726	0.3159	0.0105	1.0000	1.3879	0.2537
Transformer	32	6208	60	1.2731	0.3136	0.0109	1.0000	1.3884	0.2512

Compute ledger — parameters and observed CPU runtime

Model	Delay	Parameters	Mean training seconds
RNN	8	1943	2.5
RNN	16	1943	3.5
RNN	32	1943	5.8
GRU	8	2023	3.9
GRU	16	2023	5.9
GRU	32	2023	10.3
Transformer	8	1915	8.9
Transformer	16	1915	11.3
Transformer	32	1915	20.9

Keep the reference predictors in the comparison

Refit frequency and smoothed bigram tables on each delay’s training targets, never on test counts. The uniform seven-token baseline and the generator oracle use no fitted test information. The oracle’s random-position cost remains two bits; increasing delay raises its all-token cost because unpredictable positions occupy more of the sequence. Finite-sample movement frequencies can make a fixed guess’s test accuracy differ from 25%.

The compact reference table reports recall at delay 16, using the same test episodes as the neural comparison. Open the full ledger for every delay, every position group, and accuracy as well as loss.

# @title Evaluate training-fit references on the reserved episodes
baseline_test = {}
reference_lines = ["| Predictor | d | All CE | All acc | Recall CE | Recall acc | Random CE | Random acc |",
                   "|---|---:|---:|---:|---:|---:|---:|---:|"]
for delay in FINAL_DELAYS:
    tx, ty, _, _ = training_data[delay]
    x, y = (tensor.numpy() for tensor in held_out[delay])
    counts = np.bincount(ty.numpy().ravel(), minlength=K)
    freq = counts / counts.sum()
    transitions = np.zeros((K, K))
    np.add.at(transitions, (tx.numpy().ravel(), ty.numpy().ravel()), 1)
    conditional = (transitions + 1) / (transitions.sum(-1, keepdims=True) + K)
    predictors = {"Uniform (7 tokens)": np.full((*x.shape, K), 1/K),
                  "Training frequency": np.broadcast_to(freq, (*x.shape, K)),
                  "Smoothed bigram": conditional[x], "Generator oracle": oracle_probabilities(x, delay)}
    random = np.zeros_like(y, dtype=bool); random[:, :delay+1] = True
    for label, probabilities_test in predictors.items():
        values = [v for mask in (np.ones_like(y, dtype=bool), x == RECALL, random)
                  for v in metrics(probabilities_test, y, mask)]
        baseline_test[(label, delay)] = values
        reference_lines.append(f"| {label} | {delay} | " + " | ".join(f"{v:.4f}" for v in values) + " |")
    assert np.isclose(baseline_test[("Generator oracle", delay)][0], (delay+1)*np.log(4)/(delay+4))
    assert baseline_test[("Generator oracle", delay)][2:4] == [0., 1.]

# @title Display the compact references and full ledger
reference_compact = ["| Delay 16 reference | Recall CE (nats) | Recall accuracy |",
                     "|---|---:|---:|"]
for label in predictors:
    values = baseline_test[(label, 16)]
    reference_compact.append(f"| {label} | {values[2]:.4f} | {values[3]:.4f} |")
display(Markdown("\n".join(reference_compact)))
display(Markdown("<details>\n<summary>Full reference results — all delays and position groups</summary>\n\n"
                 + "\n".join(reference_lines) + "\n\n</details>"))

Delay 16 reference	Recall CE (nats)	Recall accuracy
Uniform (7 tokens)	1.9459	0.2461
Training frequency	1.4917	0.2461
Smoothed bigram	1.3871	0.2480
Generator oracle	-0.0000	1.0000

Full reference results — all delays and position groups

Predictor	d	All CE	All acc	Recall CE	Recall acc	Random CE	Random acc
Uniform (7 tokens)	8	1.9459	0.2019	1.9459	0.2578	1.9459	0.2406
Training frequency	8	1.7214	0.2038	1.5693	0.2197	1.5686	0.2473
Smoothed bigram	8	1.6880	0.2100	1.3867	0.2646	1.5852	0.2505
Generator oracle	8	1.0397	0.4304	-0.0000	1.0000	1.3863	0.2406
Uniform (7 tokens)	16	1.9459	0.2199	1.9459	0.2461	1.9459	0.2443
Training frequency	16	1.6421	0.2226	1.4917	0.2461	1.4917	0.2474
Smoothed bigram	16	1.6312	0.2212	1.3871	0.2480	1.4975	0.2457
Generator oracle	16	1.1784	0.3576	-0.0000	1.0000	1.3863	0.2443
Uniform (7 tokens)	32	1.9459	0.2339	1.9459	0.2246	1.9459	0.2483
Training frequency	32	1.5623	0.2347	1.4434	0.2520	1.4435	0.2484
Smoothed bigram	32	1.5591	0.2358	1.3867	0.2666	1.4452	0.2492
Generator oracle	32	1.2708	0.3110	-0.0000	1.0000	1.3863	0.2483

plot_test_summary()

Figure 22. Repeat initialization before drawing a conclusion. Small dots show held-out results from all three seeds. Lines join means; error bars show sample standard deviation across initializations on fixed datasets. Dashed lines mark reference accuracy and loss.

plot_paired_test()

Figure 23. Compare on the very same held-out prefixes. Rows show the three models and columns show the first four delay-16 test episodes at seed 6206. Blue marks the true cue, including when another token receives the highest probability.

Interpret the evidence before naming a winner

Conclusion. All three models recall every held-out cue at delay 8. At delay 16, mean accuracy is 44% for RNN, 81% for GRU, and 100% for Transformer. At delay 32, RNN is near the fixed-guess reference (26%), GRU is weak and variable (32%), and Transformer remains at 100% on this fixed-position task.

Random-position cross-entropy remains near \(\ln4\). The learned dependency is recall; independent movements remain unpredictable. Compare recall loss too: perfect accuracy can coexist with different probability quality.

Limits of this architecture comparison

This is a fixed-budget comparison on one synthetic task. Similar parameter counts leave differences in widths, normalization, positional information, and computation; attention stores quadratic score matrices. Some selected checkpoints reach epoch 60, so convergence is unestablished. Increasing delay reduces recall’s share of the training objective. The Transformer can retrieve the cue’s fixed position; these results establish neither arbitrary retrieval nor unseen-length generalization nor general superiority on language or time series.

Three initialization seeds describe optimization variability on fixed datasets, not uncertainty across new datasets. All-token accuracy also mixes random movements, recall, and deterministic control tokens.

Interpretation checkpoint. At delay 32, all three RNN runs have the same recall accuracy. Does that imply identical predictions or reliable recall? Compare their losses with the references before answering.

Show explanation — zero observed spread can describe repeated failure

Accuracy only counts the largest-probability token. Different distributions can make the same number of correct predictions, and the full ledger shows different recall cross-entropies. The RNN delay-32 accuracy is around the fixed-guess reference, with recall loss around \(\ln4\). Zero observed accuracy SD across three runs therefore supports neither useful memory nor identical predictors, and does not establish zero uncertainty.

Experimental checkpoint. What would changing the cue position test that this comparison leaves open?

Show explanation — reserve new evidence for a new question

The current Transformer can retrieve a fixed position. Varying where the cue appears would test a different dependency and could require content-sensitive retrieval. Freeze the new generator, model choices, training/checkpoint criteria, and a fresh test split before comparison. Changing recall-loss weighting is another separate experiment. Neither should be tuned against the test results already shown here.

Optional epilogue — Cross-entropy in bits

Start with the probability of what actually happened

Scene 4 defined the loss at the actual target. Here we examine how its probability changes the penalty:

\[\ell_t=-\ln p_t[y_t].\]

A likely observation receives a small penalty. An observation the model considered nearly impossible receives a large one. The logarithm makes the penalty grow sharply as the true target’s probability approaches zero. A zero probability for an observed target would give infinite loss; trained softmax models use finite logits and stable log-softmax arithmetic rather than taking the log of a rounded display value.

Before running the next cell, predict which probability will cost most: \(0.8\), \(0.25\), or \(0.01\). These are illustrative target probabilities, not another test evaluation.

target_probabilities = np.array([.8, .25, .01])
position_nats = -np.log(target_probabilities)
position_bits = position_nats / np.log(2)
surprise_lines = ["| True-target probability | Loss (nats) | Same loss (bits) |",
                  "|---:|---:|---:|"]
for probability, nats, bits in zip(target_probabilities, position_nats, position_bits):
    surprise_lines.append(f"| {probability:.2f} | {nats:.4f} | {bits:.4f} |")
display(Markdown("\n".join(surprise_lines)))
assert np.all(np.diff(position_nats) > 0)

True-target probability	Loss (nats)	Same loss (bits)
0.80	0.2231	0.3219
0.25	1.3863	2.0000
0.01	4.6052	6.6439

plot_target_surprise()

Figure 24. Measure surprise at the observed target. The horizontal coordinate is the probability assigned to the command that actually occurred. The two vertical scales express the same penalty in nats and bits. This is an illustrative loss curve, not a learning curve.

Accuracy can miss a confident mistake

Suppose the actual next command is DOWN. Both toy predictors below choose UP, so both are wrong by accuracy. The second assigns only \(0.01\) to DOWN; its loss is much larger than the first predictor’s loss. These distributions are restricted to the four movement commands for illustration.

toy_probabilities = np.array([[.60, .30, .05, .05],
                              [.97, .01, .01, .01]])  # UP, DOWN, LEFT, RIGHT
true_movement = DOWN
assert np.allclose(toy_probabilities.sum(axis=1), 1)
assert np.all(toy_probabilities.argmax(axis=1) == UP)
toy_losses = -np.log(toy_probabilities[:, true_movement])
print("Both predict UP; actual target is DOWN.")
print("Target probabilities:", toy_probabilities[:, true_movement])
print("Losses in nats:", toy_losses.round(4))
assert toy_losses[1] > toy_losses[0]

Both predict UP; actual target is DOWN.
Target probabilities: [0.3  0.01]
Losses in nats: [1.204  4.6052]

Check your reading — can two correct predictions have different losses?

Yes. If both predictors put their largest probability on the true token, both count as correct. The one assigning \(0.9\) to that token receives less loss than one assigning \(0.4\). Accuracy keeps only the winner; cross-entropy retains information about the target’s probability. This is why we report both.

Why add logarithms across a sequence?

The probability assigned to a whole observed continuation is the product of the conditional probabilities along it. A logarithm turns that product into a sum:

\[-\ln\prod_{t=0}^{T-1}q(y_t\mid s_{\leq t}) =\sum_{t=0}^{T-1}-\ln q(y_t\mid s_{\leq t}).\]

Our training loss averages these contributions across positions and examples. For a fixed sequence length, dividing by the number of targets changes the scale rather than which parameter values minimize the loss. Across different delays, however, it changes the fraction contributed by recall: there is one recall position among \(d+4\) targets. The objective does not explicitly give remembering the cue extra weight.

The Scene 3 random-movement reference is \(-\ln(1/4)=\ln4\) nats. In bits, the same irreducible uncertainty has a simpler numerical value: two. The iid qualification and finite-sample caveats from Scene 3 still apply.

Bits express the same loss in another unit

Natural logarithms give nats; base-two logarithms give bits:

\[-\log_2 q(y_t\mid s_{\leq t}) =\frac{-\ln q(y_t\mid s_{\leq t})}{\ln2}.\]

An event with probability \(1/4\) costs two bits; one with probability \(1/2\) costs one. Halving probability adds one bit of surprise.

Optional coding interpretation and reference

Under a matched conditional coding scheme, these are ideal code-length contributions. Fractional contributions concern accumulated or average costs, not literal fractional bits in a file. Framing, model information, and coding overhead also matter; this notebook does not implement a compressor. See Shannon, A Mathematical Theory of Communication, Introduction and Sections 9–10.

The generator oracle knows every deterministic target, including recall, and assigns \(1/4\) to each unpredictable movement. There are \(d+1\) unpredictable targets among \(d+4\) total targets, so its mean reference is

\[\frac{2(d+1)}{d+4}\quad\text{bits per target}.\]

At delay 16, this is \(2(17)/20=1.7\) bits per target, even though recall itself costs zero. A model can therefore improve recall greatly while its all-token mean changes only modestly. The following table converts our existing held-out results; it trains and scores no additional model.

bits_lines = ["| Delay 16 predictor | Mean all-token bits | Mean recall bits |",
              "|---|---:|---:|"]
for kind in FINAL_KINDS:
    mean = summary[(kind,16)]["mean"] / np.log(2)
    bits_lines.append(f"| {kind}, three initializations | {mean[0]:.4f} | {mean[2]:.4f} |")
bits_lines.append(f"| Generator oracle reference | {2*17/20:.4f} | 0.0000 |")
display(Markdown("\n".join(bits_lines)))
assert np.isclose(np.log(4)/np.log(2), 2)

Delay 16 predictor	Mean all-token bits	Mean recall bits
RNN, three initializations	1.7775	1.4309
GRU, three initializations	1.7432	0.7905
Transformer, three initializations	1.7030	0.0063
Generator oracle reference	1.7000	0.0000

Final loss checkpoint. A model reaches nearly zero recall loss, but its random-position loss stays near two bits. Has learning failed on the wandering commands?

Show explanation — distinguish uncertainty from a missed dependency

The wandering commands were generated independently and uniformly. Two bits is their ideal generator reference, so that cost is compatible with successful learning of all predictable structure. Recall is different: the cue is present in the prefix, so the oracle assigns its target probability one and loss zero. Evaluate a model against the relevant position group rather than expecting every part of this task to become certain.

Conclusion. Accuracy records which token wins; cross-entropy records the probability of the actual target. Converting nats to bits changes the unit, while the distinction between learnable recall and unpredictable wandering remains the same.

Optional mini-lab — Forecast a continuous wave

The robot task made memory easy to recognize: one command had to survive until RECALL. A continuous signal poses a different question. Can a model infer the motion of an oscillation from a short, noisy window, then continue it when observations stop?

We will compare persistence, linear AR, MLP, RNN, GRU, 1D CNN, and CNN → RNN. Every model receives the same last 32 measurements. The MLP sees those 32 numbers as a vector; the recurrent models read them in order. A 1D CNN applies shared filters along time; the hybrid lets an RNN read the sequence of convolutional features. No model receives the complete trajectory, a frequency label, or the future. Recurrent states reset for each window, including each rollout step, so this comparison isolates representations of the same finite history.

This mini-lab can run separately: execute the chapter’s setup cell, the folded forecast figure helper near the beginning, and then this section in order. It imports its own training utilities and does not need the robot’s 27 models. Allow roughly 20–30 minutes to read and investigate it; the code reports its CPU runtime.

A small notation contract

Symbol	Meaning in this mini-lab
\(v_n\), \(\hat v_n\)	Measured scalar amplitude and predicted amplitude at sample \(n\)
\(L_{\rm wave}=192\)	Samples per independently generated trajectory
\(W_{\rm win}=32\)	Number of observed samples supplied to every predictor
\(o\)	First unobserved index; the input is \(v_{o-W_{\rm win}},\ldots,v_{o-1}\)
\(F=64\), \(k=0,\ldots,F-1\)	Rollout length and step index; horizon \(k+1\) predicts sample \(o+k\)
\(A_j\), \(f_j\), \(\varphi_j\), \(j=1,2\)	Amplitude, frequency in cycles/sample, and phase of each sinusoid
\(\varepsilon_n\), \(\sigma_{\rm noise}=0.03\)	Independent zero-mean Gaussian measurement noise and its standard deviation
\(\mu_{\rm train}\), \(\sigma_{\rm train}\)	Mean and standard deviation fitted using training trajectories only
\(C=8\), \(q\), \(\kappa\)	Convolutional feature channels, feature-position index, and filter offset; distinct from raw sample index \(n\)
MSE, \(\lambda_{\rm ridge}\)	Mean squared error; the linear model’s fixed regularization coefficient

1. Make trajectories before making windows

Each record is a sum of two sinusoids plus noise:

\[v_n=\sum_{j=1}^{2} A_j\sin(2\pi f_j n+\varphi_j)+\varepsilon_n.\]

Frequencies are fixed at 0.025 and 0.065 cycles/sample. Amplitudes and phases vary independently between trajectories. Thus test trajectories contain new phases, amplitudes, and noise, while the frequency family stays familiar. This is not a test of unseen frequencies or real audio.

We generate 32 training, 8 validation, and 16 test trajectories with independent random streams. Only then do we extract windows. Splitting overlapping windows from one trajectory at random could put nearly identical contexts and shared measurements on both sides of the boundary.

# @title Forecast setup: fixed data, model, training, and evaluation choices
from torch import nn
from IPython.display import Markdown
import copy
import hashlib

torch.set_num_threads(1)
wave_started = time.perf_counter()
WAVE_WINDOW, WAVE_HORIZON, WAVE_LENGTH = 32, 64, 192
WAVE_NOISE, WAVE_RIDGE = .03, .001
WAVE_EPOCHS, WAVE_BATCH, WAVE_LR = 30, 256, .003
WAVE_SEEDS = (6406, 6407, 6408)
WAVE_NEURAL = ("MLP", "RNN", "GRU", "1D CNN", "CNN → RNN")
WAVE_KINDS = ("Persistence", "Linear AR", *WAVE_NEURAL)
WAVE_ORIGINS = (32, 96)
WAVE_EXAMPLE_INDEX = 1  # test trajectory 0, origin 96; fixed before scoring

def make_wave_trajectories(count, seed, frequencies=(.025, .065)):
    rng = np.random.default_rng(seed)
    amplitudes = rng.uniform([.6, .2], [1., .5], size=(count, 2))
    phases = rng.uniform(0, 2*np.pi, size=(count, 2))
    frequencies = np.asarray(frequencies, dtype=float)
    assert frequencies.shape == (2,)
    index = np.arange(WAVE_LENGTH)
    angles = 2*np.pi*frequencies[None, :, None]*index + phases[:, :, None]
    clean = (amplitudes[:, :, None]*np.sin(angles)).sum(axis=1)
    return clean + rng.normal(0, WAVE_NOISE, size=clean.shape)

wave_splits = {"train": make_wave_trajectories(32, 6411),
               "validation": make_wave_trajectories(8, 6412),
               "test": make_wave_trajectories(16, 6413)}
wave_mean = wave_splits["train"].mean()
wave_scale = wave_splits["train"].std()
assert wave_scale > 0
wave_scaled = {name: (series-wave_mean)/wave_scale for name, series in wave_splits.items()}

# @title Verification: reproducible, disjoint whole trajectories
wave_hashes = {name: {hashlib.sha256(row.tobytes()).hexdigest() for row in series}
               for name, series in wave_splits.items()}
assert all(len(wave_hashes[name]) == len(wave_splits[name]) for name in wave_splits)
assert not (wave_hashes["train"] & wave_hashes["validation"])
assert not (wave_hashes["train"] & wave_hashes["test"])
assert not (wave_hashes["validation"] & wave_hashes["test"])
assert np.array_equal(wave_splits["train"], make_wave_trajectories(32, 6411))
assert all(origin+WAVE_HORIZON <= WAVE_LENGTH for origin in WAVE_ORIGINS)

plot_wave_window(origin=96)

Figure 25. Slide the observation window without moving information across the boundary. Blue measurements enter the model. The orange continuation is available as a training label or an evaluation target, and is excluded from that prediction’s input. Change origin in the plotting call to inspect another valid window; the scoring origins remain fixed.

Checkpoint. Would an MLP gain an information advantage by flattening this window?

Show explanation — the same evidence in different representations

No. Flattening preserves the order and values of the same 32 samples. An RNN reads the sequence, and an MLP assigns parameters to its window positions. Their computations and inductive biases differ, but neither sees extra observations here. A model reading the entire record would be answering a different comparison.

2. Pair each window with the next sample

The input ends just before the target. sliding_window_view gives consecutive spans of 33 values: the first 32 are input, and the last is the next-sample target. Overlap within a split is allowed; overlap across splits is prevented by the trajectory boundary.

Inputs have shape \((B,32,1)\) and targets \((B,)\). We standardize using training statistics for optimization, then report squared error in the original amplitude units. Each trajectory contributes the same number of windows, so their average contributions are equal. These overlapping windows are not independent experimental replicates.

def wave_windows(series):
    spans = np.lib.stride_tricks.sliding_window_view(series, WAVE_WINDOW+1, axis=1)
    inputs = np.ascontiguousarray(spans[..., :WAVE_WINDOW].reshape(-1, WAVE_WINDOW, 1))
    targets = np.ascontiguousarray(spans[..., WAVE_WINDOW].reshape(-1))
    return torch.tensor(inputs, dtype=torch.float32), torch.tensor(targets, dtype=torch.float32)

wave_data = {name: wave_windows(series) for name, series in wave_scaled.items()}
wave_train_x, wave_train_y = wave_data["train"]
wave_valid_x, wave_valid_y = wave_data["validation"]
assert wave_train_x.shape == (5120, 32, 1) and wave_train_y.shape == (5120,)
np.testing.assert_allclose(wave_train_x[0, :, 0], wave_scaled["train"][0, :32], rtol=1e-6, atol=1e-7)
np.testing.assert_allclose(wave_train_y[0], wave_scaled["train"][0, 32], rtol=1e-6, atol=1e-7)

# @title Verification: future measurements are excluded from input windows
altered_wave = wave_scaled["validation"].copy()
altered_wave[:, WAVE_WINDOW:] += 100  # change only measurements after each first context
altered_inputs, altered_targets = wave_windows(altered_wave)
first_window = np.arange(len(altered_wave))*(WAVE_LENGTH-WAVE_WINDOW)
torch.testing.assert_close(altered_inputs[first_window], wave_valid_x[first_window], rtol=0, atol=0)
assert not torch.equal(altered_targets[first_window], wave_valid_y[first_window])
print("Changing future measurements preserves each trajectory's first input window.")

Changing future measurements preserves each trajectory's first input window.

3. Start with two inexpensive predictors

Persistence repeats the last observation. It is useful when a signal changes slowly, but will draw a flat line in a recursive forecast.

The linear autoregressive (AR) predictor fits one coefficient per window position and an intercept. “Autoregressive” means it predicts a signal from earlier values of that same signal. We fit the next sample by minimizing mean squared error plus \(\lambda_{\rm ridge}\) times the squared coefficient norm; the intercept is unpenalized. The fixed coefficient is 0.001, and fitting uses training windows only.

For a fixed-frequency sinusoid, the next noise-free sample is a linear combination of the previous two. A sum of two fixed frequencies also has a finite linear recurrence. Consequently, a linear model is a serious competitor on this task. No architecture should be required to beat it for the tutorial to make sense.

wave_linear_started = time.perf_counter()
wave_design = np.column_stack((wave_train_x[:, :, 0].numpy(), np.ones(len(wave_train_x))))
penalty = np.eye(WAVE_WINDOW+1)*WAVE_RIDGE
penalty[-1, -1] = 0  # do not regularize the intercept
wave_coefficients = np.linalg.solve(
    wave_design.T @ wave_design/len(wave_design) + penalty,
    wave_design.T @ wave_train_y.numpy()/len(wave_design))

wave_linear_seconds = time.perf_counter()-wave_linear_started

def predict_persistence(windows):
    return windows[:, -1, 0]

def predict_linear_wave(windows):
    return windows[:, :, 0] @ wave_coefficients[:-1] + wave_coefficients[-1]

assert np.isfinite(wave_coefficients).all()

4. Change how the same window is represented

Every neural model receives the same \((B,32,1)\) window and outputs one scalar per example. The linear model has 33 fitted scalars; persistence has none.

Predictor	Representation and readout	Parameters
MLP	32 values → 16 ReLU features → scalar	545
RNN	20 tanh coordinates; read the last state	481
GRU	12 gated coordinates; read the last state	553
1D CNN	Two temporal convolutions → flatten all features → scalar	457
CNN → RNN	Temporal convolution → 16-coordinate RNN → last-state readout	481

Similar parameter counts do not equalize computation, representation widths, or optimization. The CNN and hybrid extend a previously inspected benchmark; they were specified before their own test evaluation.

All five neural models use the same 30 epochs, batch orders, Adam at 0.003, and clipping at 1. Train on next-sample MSE; select checkpoints by validation next-sample MSE. No rollout score selects an architecture, seed, or epoch. This fixed budget does not establish convergence.

The architecture map below traces the implemented routes. Each route predicts one next sample. The later diagnostic grids vary window length; the topology remains the same, while sequence lengths and the MLP/CNN flattened input sizes change.

class WaveWindowModel(nn.Module):
    def __init__(self, kind):
        super().__init__()
        self.kind = kind
        if kind == "MLP":
            self.network = nn.Sequential(nn.Linear(WAVE_WINDOW, 16), nn.ReLU(), nn.Linear(16, 1))
        else:
            width = 20 if kind == "RNN" else 12
            recurrent = nn.RNN if kind == "RNN" else nn.GRU
            self.recurrent = recurrent(1, width, batch_first=True)
            self.readout = nn.Linear(width, 1)

    def forward(self, windows):
        if self.kind == "MLP":
            return self.network(windows[:, :, 0]).squeeze(-1)
        states, _ = self.recurrent(windows)  # fresh zero state on every call
        return self.readout(states[:, -1]).squeeze(-1)

Share filters along time, then choose a readout

A 1D convolution slides the same learned weights over consecutive measurements. Our first layer has eight filters, each reading five samples. Input axes change from \((B,32,1)\) to \((B,1,32)\) because PyTorch Conv1d expects batch, channels, then time.

With stride one and no padding, each filter produces 28 positions. For the standardized input \(x\), its feature after ReLU is

\[u_{b,c,q}=\operatorname{ReLU}\!\left(b_c+\sum_{\kappa=0}^{4}W_{c,0,\kappa}x_{b,q+\kappa,0}\right).\]

The filter is used as stored, without reversing its offsets: PyTorch implements cross-correlation, following Week 05’s convention. \(q=0,\ldots,27\) indexes a feature’s leftmost input sample.

Route	Feature shapes after each operation
1D CNN	\((B,1,32)\) → Conv5/ReLU \((B,8,28)\) → Conv3/ReLU \((B,8,26)\) → flatten \((B,208)\) → scalar
CNN → RNN	\((B,1,32)\) → Conv5/ReLU \((B,8,28)\) → transpose \((B,28,8)\) → RNN \((B,28,16)\) → last state → scalar

Each feature after the CNN’s second convolution depends on seven adjacent raw samples (five plus two). Its flattened head combines all 26 feature positions, so the complete predictor can use the full 32-sample window. The hybrid processes 28 overlapping five-sample features in order; its state restarts at every call.

All feature positions lie inside the already observed window. We predict only the sample after that window, not an output at every raw time position; these valid convolutions therefore receive no unknown future measurement.

Checkpoint. Does the CNN’s seven-sample feature receptive field mean that its final prediction sees only seven samples? Trace the readout before answering.

Show explanation — local features and the final observation window

Each second-layer CNN feature uses seven adjacent measurements, but the readout combines all 26 feature positions. Their union covers the full observed window. The hybrid likewise processes features spanning the whole window; neither predictor gets additional history.

class WaveConvModel(nn.Module):
    def __init__(self, kind):
        super().__init__()
        if kind not in ("1D CNN", "CNN → RNN"):
            raise ValueError("Unknown convolutional predictor")
        self.kind = kind
        self.conv_in = nn.Conv1d(1, 8, kernel_size=5)  # valid: 32 -> 28
        if kind == "1D CNN":
            self.conv_out = nn.Conv1d(8, 8, kernel_size=3)  # 28 -> 26
            self.readout = nn.Linear(8*(WAVE_WINDOW-6), 1)
        else:
            self.recurrent = nn.RNN(8, 16, batch_first=True)
            self.readout = nn.Linear(16, 1)

    def forward(self, windows):
        features = F.relu(self.conv_in(windows.transpose(1, 2)))
        if self.kind == "1D CNN":
            features = F.relu(self.conv_out(features))
            return self.readout(features.flatten(1)).squeeze(-1)
        states, _ = self.recurrent(features.transpose(1, 2))
        return self.readout(states[:, -1]).squeeze(-1)

# @title Check the architecture map against the implemented tensor routes
with torch.random.fork_rng():
    torch.manual_seed(6400)
    sample = torch.zeros(2, 32, 1)
    for kind in WAVE_NEURAL:
        model = WaveConvModel(kind) if kind in ("1D CNN", "CNN → RNN") else WaveWindowModel(kind)
        assert model(sample).shape == (2,)
        if kind == "MLP":
            assert model.network[:2](sample[:, :, 0]).shape == (2, 16)
        elif kind in ("RNN", "GRU"):
            states, _ = model.recurrent(sample)
            assert states.shape == (2, 32, 20 if kind == "RNN" else 12)
        else:
            features = F.relu(model.conv_in(sample.transpose(1, 2)))
            assert features.shape == (2, 8, 28)
            if kind == "1D CNN":
                assert F.relu(model.conv_out(features)).flatten(1).shape == (2, 208)
            else:
                states, _ = model.recurrent(features.transpose(1, 2))
                assert states.shape == (2, 28, 16)
print("All five architecture routes and scalar outputs match the diagram.")

All five architecture routes and scalar outputs match the diagram.

plot_wave_architectures()

Architecture map — five ways to represent the same history. Read each row left to right. MLP combines positions with dense weights; CNN shares filters across positions before combining them; RNN and GRU repeatedly update a state. The hybrid first extracts local features, then processes those features in order. All five predict a scalar amplitude, not a token distribution.

The references are simpler: persistence copies the final input value; linear AR weights all 32 observations and adds an intercept. Neither has a hidden representation. The architecture map describes the original \(W_{\rm win}=32\) comparison; the shifted-frequency test later selects different windows using validation.

# @title Verify temporal filter arithmetic, axes, shapes, and gradients
for kind, count in (("1D CNN", 457), ("CNN → RNN", 481)):
    torch.manual_seed(6400)
    probe = WaveConvModel(kind)
    inputs = wave_valid_x[:3].clone().requires_grad_(True)
    first = probe.conv_in(inputs.transpose(1, 2))
    assert first.shape == (3, 8, 28)
    # Direct scalar check for filter 0 at feature position 2: exactly raw samples 2..6.
    expected = (inputs[0, 2:7, 0]*probe.conv_in.weight[0, 0]).sum()+probe.conv_in.bias[0]
    torch.testing.assert_close(first[0, 0, 2], expected)
    changed = inputs.detach().clone(); changed[:, 7:, 0] += 100
    torch.testing.assert_close(probe.conv_in(changed.transpose(1, 2))[0, 0, 2], expected.detach())
    if kind == "1D CNN":
        assert probe.conv_out(F.relu(first)).shape == (3, 8, 26)
    prediction = probe(inputs)
    assert prediction.shape == (3,)
    assert sum(p.numel() for p in probe.parameters()) == count
    prediction.square().mean().backward()
    assert inputs.grad is not None and torch.isfinite(inputs.grad).all()
    assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in probe.parameters())
print("Filter offsets, feature shapes, parameter counts, and finite backward paths passed.")

Filter offsets, feature shapes, parameter counts, and finite backward paths passed.

One epoch still follows predict, measure, update

For a batch, the loss is the mean of \((\hat v-v)^2\) in standardized units. Multiplying an evaluated loss by \(\sigma_{\rm train}^2\) converts it back to original squared-amplitude units. The same mean/scale conventions apply to every predictor.

def wave_training_epoch(model, optimizer, inputs, targets, order):
    model.train()
    for indices in order.split(WAVE_BATCH):
        optimizer.zero_grad(set_to_none=True)
        loss = F.mse_loss(model(inputs[indices]), targets[indices])
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1.)
        optimizer.step()


def fit_wave_model(kind, seed, training_pair, validation_pair):
    torch.manual_seed(seed)
    model = WaveConvModel(kind) if kind in ("1D CNN", "CNN → RNN") else WaveWindowModel(kind)
    optimizer = torch.optim.Adam(model.parameters(), lr=WAVE_LR)
    inputs, targets = training_pair
    best_mse, best_epoch, best_state = float("inf"), None, None
    start = time.perf_counter()
    for epoch in range(WAVE_EPOCHS):
        order = torch.randperm(len(inputs), generator=torch.Generator().manual_seed(6500+epoch))
        wave_training_epoch(model, optimizer, inputs, targets, order)
        model.eval()
        with torch.no_grad():
            validation_mse = F.mse_loss(model(validation_pair[0]), validation_pair[1]).item()
        if validation_mse < best_mse:
            best_mse, best_epoch = validation_mse, epoch+1
            best_state = copy.deepcopy(model.state_dict())
    model.load_state_dict(best_state)
    model.eval()
    return dict(model=model, epoch=best_epoch, validation_mse=best_mse*wave_scale**2,
                parameters=sum(p.numel() for p in model.parameters()), seconds=time.perf_counter()-start)

wave_runs = {}
for seed in WAVE_SEEDS:
    for kind in WAVE_NEURAL:
        wave_runs[(kind, seed)] = fit_wave_model(kind, seed, wave_data["train"], wave_data["validation"])
assert set(wave_runs) == {(kind, seed) for kind in WAVE_NEURAL for seed in WAVE_SEEDS}
assert all(wave_runs[(kind, WAVE_SEEDS[0])]["parameters"] == count
           for kind, count in (("MLP", 545), ("RNN", 481), ("GRU", 553), ("1D CNN", 457), ("CNN → RNN", 481)))
print("All fifteen validation-selected wave models ready; test performance not scored yet.")

All fifteen validation-selected wave models ready; test performance not scored yet.

# @title Verification: no hidden state carries between windows
for run in wave_runs.values():
    model = run["model"]
    with torch.no_grad():
        before = model(wave_valid_x[:3])
        model(wave_valid_x[20:24])
        after = model(wave_valid_x[:3])
    torch.testing.assert_close(before, after, rtol=0, atol=0)
print("All selected models reset context: intervening calls do not alter a prediction.")

All selected models reset context: intervening calls do not alter a prediction.

5. Stop observing and feed predictions back

For one-step evaluation, every test window contains actual preceding measurements. Each prediction asks only for the next sample.

For a recursive rollout, start with 32 actual measurements. Predict one value, remove the oldest value, append the prediction, and repeat. After the start, no actual future measurement enters the input. At step \(k\), the model’s own earlier guesses occupy part of the window; after 32 steps, the window consists entirely of guesses.

We fix two origins, 32 and 96, in each of the 16 test trajectories and evaluate 64 future samples from each origin. These origins are correlated within trajectories. The curves average errors over those 32 origins; bands describe variability across three neural initializations on these same data, not confidence intervals. The baselines have one deterministic fit each.

Checkpoint. Can a model have excellent one-step MSE and poor rollout MSE?

Show explanation — small errors can change the next input

Yes. One-step evaluation repeatedly restores the observed context. A rollout replaces observations with predictions, so amplitude or phase errors can alter later inputs and accumulate. The model was selected on next-sample validation error. Stable long rollouts are an additional requirement that this criterion does not guarantee.

plot_wave_feedback()

Feedback map — the rollout uses the same one-step model repeatedly. In recursive evaluation, predictions become inputs. In observed-context evaluation, new measured values replace them. Recurrent state resets inside every call: the feedback carries predicted amplitudes through the window, not a persistent hidden state.

def wave_model_predictor(model):
    def predict(windows):
        with torch.no_grad():
            return model(torch.tensor(windows, dtype=torch.float32)).numpy()
    return predict


def roll_wave(predict, initial_windows, horizon):
    window = initial_windows.copy()
    predictions = []
    for _ in range(horizon):
        next_value = predict(window)
        predictions.append(next_value)
        window = np.concatenate((window[:, 1:], next_value[:, None, None]), axis=1)
    return np.stack(predictions, axis=1)

wave_predictors = {("Persistence", None): predict_persistence, ("Linear AR", None): predict_linear_wave}
wave_predictors.update({key: wave_model_predictor(run["model"]) for key, run in wave_runs.items()})
assert len(wave_predictors) == 17  # both references plus all fifteen neural checkpoints

wave_test_origins = [(trajectory, origin) for trajectory in range(16) for origin in WAVE_ORIGINS]
wave_contexts = np.stack([wave_scaled["test"][i, o-WAVE_WINDOW:o]
                          for i, o in wave_test_origins]).astype(np.float32)[:, :, None]
wave_future = np.stack([wave_splits["test"][i, o:o+WAVE_HORIZON] for i, o in wave_test_origins])
assert wave_contexts.shape == (32, 32, 1) and wave_future.shape == (32, 64)
assert wave_test_origins[WAVE_EXAMPLE_INDEX] == (0, 96)
wave_test_x, wave_test_y = wave_data["test"]
wave_results = {}
for key, predict in wave_predictors.items():
    one_step_mse = np.mean((predict(wave_test_x.numpy())-wave_test_y.numpy())**2)*wave_scale**2
    scaled_predictions = roll_wave(predict, wave_contexts, WAVE_HORIZON)
    np.testing.assert_allclose(scaled_predictions[:, 0], predict(wave_contexts), rtol=1e-6, atol=1e-7)
    predictions = scaled_predictions*wave_scale + wave_mean
    squared_errors = (predictions-wave_future)**2
    assert np.isfinite(squared_errors).all()
    wave_results[key] = dict(one_step_mse=float(one_step_mse), rollout_mse=float(squared_errors.mean()),
                             horizon_mse=squared_errors.mean(axis=0), predictions=predictions)

6. Read the held-out evidence

The compact table compares observed-prefix next-sample MSE with mean error over a 64-step recursive rollout. Neural entries show mean ± sample SD over all three initializations; baselines have one fitted predictor. Open the ledger for every seed, selected epoch, parameter count, runtime, and errors at specific horizons. No seed is discarded.

The generator adds independent noise of standard deviation 0.03. Even a perfect prediction of the clean future has expected measured-target MSE \(0.03^2=0.0009\). This is a generator reference, not a strict lower bound on each finite test estimate. The models must also estimate the oscillation from noisy context, which can add error.

# @title Display the compact wave comparison
wave_lines = ["| Predictor | Next-sample MSE | 64-step rollout MSE |", "|---|---:|---:|"]
for kind in WAVE_KINDS:
    runs = [r for (name, seed), r in wave_results.items() if name == kind]
    columns = []
    for metric in ("one_step_mse", "rollout_mse"):
        values = np.array([r[metric] for r in runs])
        value = f"{values.mean():.4g}"
        if len(values) > 1:
            value += f" ± {values.std(ddof=1):.2g}"
        columns.append(value)
    wave_lines.append(f"| {kind} | " + " | ".join(columns) + " |")
display(Markdown("\n".join(wave_lines)))

Predictor	Next-sample MSE	64-step rollout MSE
Persistence	0.01963	0.7261
Linear AR	0.001049	0.0011
MLP	0.001181 ± 1.4e-05	0.001304 ± 8.2e-05
RNN	0.00358 ± 7.4e-05	0.3111 ± 0.14
GRU	0.003974 ± 0.00056	0.4066 ± 0.2
1D CNN	0.001156 ± 3.1e-05	0.001274 ± 9.4e-05
CNN → RNN	0.00319 ± 0.0007	0.3873 ± 0.43

# @title Display the full forecasting ledger
wave_ledger = ["| Predictor | Seed | Epoch | Params | Seconds | Next MSE | Rollout MSE | h=1 MSE | h=8 MSE | h=32 MSE | h=64 MSE |",
               "|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|"]
for (kind, seed), result in wave_results.items():
    run = wave_runs.get((kind, seed), {})
    count = run.get("parameters", 33 if kind == "Linear AR" else 0)
    errors = " | ".join(f"{result['horizon_mse'][h-1]:.5g}" for h in (1, 8, 32, 64))
    seconds = f"{run['seconds']:.2f}" if run else (f"{wave_linear_seconds:.2f}" if kind == "Linear AR" else "—")
    wave_ledger.append(f"| {kind} | {seed if seed is not None else '—'} | {run.get('epoch', '—')} | {count} | "
                       f"{seconds} | {result['one_step_mse']:.5g} | {result['rollout_mse']:.5g} | {errors} |")
display(Markdown("<details>\n<summary>Full forecasting ledger — every seed and selected horizons</summary>\n\n"
                 + "\n".join(wave_ledger) + "\n\n</details>"))

Full forecasting ledger — every seed and selected horizons

Predictor	Seed	Epoch	Params	Seconds	Next MSE	Rollout MSE	h=1 MSE	h=8 MSE	h=32 MSE	h=64 MSE
Persistence	—	—	0	—	0.019634	0.72613	0.022455	0.60993	0.4536	1.1039
Linear AR	—	—	33	0.00	0.0010491	0.0010995	0.0009415	0.00086083	0.0011084	0.0010579
MLP	6406	29	545	0.19	0.0011693	0.0012317	0.0013088	0.00091924	0.0013296	0.001406
RNN	6406	30	481	1.63	0.0036596	0.46555	0.0031421	0.14211	0.61146	0.85756
GRU	6406	29	553	3.15	0.004444	0.50472	0.0030055	0.21914	0.4942	0.71276
1D CNN	6406	25	457	2.63	0.0011805	0.0013763	0.00098257	0.0010389	0.0017094	0.0017855
CNN → RNN	6406	30	481	2.22	0.0037775	0.8683	0.0023898	0.19806	1.056	1.0859
MLP	6407	29	545	0.17	0.0011961	0.0013935	0.0011582	0.0010456	0.0017395	0.001819
RNN	6407	30	481	1.61	0.0035136	0.26413	0.0030582	0.107	0.31633	0.63092
GRU	6407	29	553	3.12	0.0041189	0.53474	0.0028134	0.22072	0.70441	0.61583
1D CNN	6407	22	457	2.67	0.0011667	0.0012514	0.00094675	0.0010229	0.0012735	0.0012651
CNN → RNN	6407	27	481	2.21	0.00241	0.032296	0.0021292	0.026943	0.049859	0.064889
MLP	6408	30	545	0.17	0.0011784	0.0012866	0.0011936	0.0011001	0.0011206	0.0013639
RNN	6408	30	481	1.60	0.0035662	0.20363	0.0029865	0.12214	0.25432	0.33521
GRU	6408	30	553	3.13	0.0033582	0.18027	0.002586	0.098908	0.18672	0.2984
1D CNN	6408	30	457	2.68	0.0011219	0.0011933	0.00098206	0.00098159	0.0014167	0.0012883
CNN → RNN	6408	29	481	2.23	0.0033831	0.26144	0.0022015	0.1194	0.23386	0.45622

plot_wave_rollout()

Figure 26. Compare continuations from the same observed window. Four panels share axes and repeat the same context and actual future so curves remain legible. The example is trajectory 0 at origin 96, fixed before scoring; neural predictions use seed 6406. Once the dotted boundary is crossed, each predictor uses its own earlier predictions. One example illustrates behavior but does not establish the overall ranking.

plot_wave_horizon_errors()

Figure 27. Ask how error changes with distance into the future. Linear AR and MLP appear in both panels to connect the comparisons. At each horizon, error is averaged over the same fixed test origins. Neural lines average three initializations, with sample-SD bands; baselines have no initialization band. The vertical axis is logarithmic and begins at one third of the noise reference. Shading below that display boundary is clipped; mean-minus-SD is not a lower bound on error. Small observed error at one horizon is not a guarantee about the next, and correlated origins are not independent replicates.

What this tiny comparison can teach

Conclusion. Linear AR retains the lowest mean next-sample and rollout errors. The 1D CNN and MLP both forecast well: mean rollout MSE is about 0.00127 for CNN and 0.00130 for MLP. This small difference, relative to the initialization variation, does not establish a reliable CNN advantage.

The CNN → RNN hybrid has mean rollout MSE about 0.39 with sample SD about 0.43. Its three runs vary widely; adding a convolutional front end does not deliver consistently stable recurrent forecasts here. RNN and GRU also drift, with mean rollout MSE about 0.31 and 0.41. Read each seed in the ledger alongside the mean.

The fixed frequencies admit a linear recurrence, making a linear predictor well suited to this task. The shared 32-sample window is sufficient for strong MLP and CNN forecasts here. RNNs provide a mechanism for updating a representation as observations arrive, but that mechanism does not guarantee better predictions. This comparison resets their state at every window, so it does not test carrying information through a longer stream.

The diagnostic lab below tests training budget, window length, and independent trajectory count, then evaluates one shifted frequency family. Persistent-state forecasting and gradual frequency drift remain open questions.

Interpretation checkpoint. Why does persistence become less wrong again around some later horizons, although it predicts a flat line?

Show explanation — an oscillation can return toward the last value

An oscillation revisits earlier values, so horizon error need not increase monotonically. The lower-frequency component has a period of 40 samples; near that distance it returns toward the phase of the last observation. The second component and measurement noise prevent an exact return in general. Persistence still draws a flat line rather than following the wave, and its mean error across the whole rollout is large. Read the full horizon curve, not only one favorable distance.

Design checkpoint. How would you test whether extra history helps, without giving one model the whole record?

Show explanation — vary a shared observation budget

Choose several window lengths and give every model the same observations at each length. Declare model sizes, validation selection, compute budgets, and test origins before evaluation; reserve fresh test trajectories for the new study. Report next-sample and rollout errors separately. If you instead carry an RNN state over the entire observed prefix, state that additional history explicitly and compare against predictors with access to it.

Evidence checkpoint. The CNN’s mean rollout MSE is slightly smaller than the MLP’s. Does that establish that CNN is better? Compare all three seeds and the initialization spread.

Show explanation — a small mean difference needs more evidence

Both fit this familiar signal family well. The CNN is not better in every same-seed run, and the mean difference is small relative to the displayed spread. These three initializations on fixed trajectories support similar performance under this protocol, not a general ranking. Establishing a broader claim needs a separately specified study across new data and training conditions.

print(f"Forecasting mini-lab executed in {time.perf_counter()-wave_started:.2f} seconds.")
print("Whole-trajectory splits; identical 32-sample contexts; no future observations in rollouts.")

Forecasting mini-lab executed in 32.12 seconds.
Whole-trajectory splits; identical 32-sample contexts; no future observations in rollouts.

Diagnostic lab — does this mean RNNs do not work?

Pause before running. The MLP and CNN follow this wave well, while the recurrent forecasts drift. Does this mean RNNs do not work? Write down two explanations that the table alone cannot distinguish. What measurement would separate them?

The experiment contains two fixed-frequency sinusoids, with new phases, amplitudes, and noise. It supplies every model a finite window and resets recurrent state at each call. This task needs a good next-sample rule; it does not require retaining a cue outside that window. Earlier, the robot’s state intervention showed that a trained RNN did use its memory at delay 8. Its failures at larger delays also required explanation.

Here we investigate the wave result rather than demand a particular winner. Our hypotheses are: the original training budget constrained the fit; errors accumulate when predictions replace observations; window length changes the fit; and more independent trajectories change generalization. These are hypotheses until the corresponding comparison is checked.

Start with a rule we can verify

A clean sinusoid of frequency \(f\) obeys \(v_n^{\rm clean}=2\cos(2\pi f)v_{n-1}^{\rm clean}-v_{n-2}^{\rm clean}\). Applying the two frequency factors gives a fourth-order recurrence for their sum. Let \(c_j=\cos(2\pi f_j)\):

\[\begin{aligned} v_n^{\rm clean}&=2(c_1+c_2)v_{n-1}^{\rm clean}-(2+4c_1c_2)v_{n-2}^{\rm clean}\\ &\quad+2(c_1+c_2)v_{n-3}^{\rm clean}-v_{n-4}^{\rm clean}. \end{aligned}\]

The four coefficients below multiply lags 1, 2, 3, and 4. This check uses a clean wave and the known generator frequencies. It explains why a linear predictor is appropriate; it is not an additional trained competitor. With independent measurement noise, directly applying these four coefficients adds variance \(\sigma_{\rm noise}^2(1+\sum_j a_j^2)\), where \(a_j\) are the lag coefficients. The printed reference is much larger than the original AR error: longer windows can help estimate the underlying signal. This is the error of that naive rule, not a lower bound for forecasting.

# A clean two-frequency wave satisfies the stated linear rule.
f = np.array([.025, .065])
c1, c2 = np.cos(2*np.pi*f)
lag_coefficients = np.array([2*(c1+c2), -(2+4*c1*c2), 2*(c1+c2), -1.])
clean_wave = np.sin(2*np.pi*f[0]*np.arange(192)+.7) + .4*np.sin(2*np.pi*f[1]*np.arange(192)+1.2)
recurrence_prediction = sum(lag_coefficients[j]*clean_wave[3-j:191-j] for j in range(4))
np.testing.assert_allclose(recurrence_prediction, clean_wave[4:], rtol=0, atol=2e-13)
print("Lag coefficients:", np.round(lag_coefficients, 5))
print("Maximum clean-wave discrepancy:", f"{np.max(np.abs(recurrence_prediction-clean_wave[4:])):.2e}")
# Independent noise at the target and four lags adds its squared coefficients.
naive_four_lag_noise_mse = .03**2*(1+np.square(lag_coefficients).sum())
print("Naive four-lag rule: expected noisy-target MSE", f"{naive_four_lag_noise_mse:.4f}")

Lag coefficients: [ 3.81089 -5.62582  3.81089 -1.     ]
Maximum clean-wave discrepancy: 1.72e-14
Naive four-lag rule: expected noisy-target MSE 0.0564

Separate fitting from error accumulation

A learning curve tracks prediction error as training proceeds. We extend the original training paths from 30 to 60 and 150 epochs, using the same data, seeds, optimizer, and batch order. At each budget, the selected checkpoint minimizes validation next-sample MSE. The first 30 epochs reproduce the original models. A separate selection rule chooses the lowest validation rollout error on the same training path.

For the matched horizon comparison, observed-context prediction receives actual preceding observations at every horizon. Recursive prediction receives the initial window and then its own generated samples. Observed-context evaluation assumes new measurements become available; it is a diagnostic rather than a forecast made entirely at the origin. Both are evaluated at the same origins; their first predictions agree. This comparison measures the consequence of replacing observations with predictions.

Diagnostic protocol — what is controlled?

\(W\) is window length, \(N\) is the number of independent training trajectories, and \(u\) is optimizer updates. \(F=64\) retains the mini-lab’s forecast-horizon convention; \(H\) continues to denote recurrent hidden width.

Comparison	Changed	Held fixed
Training budget	30/60/150 epochs; 600/1,200/3,000 updates	Original 32 trajectories, \(W=32\), five neural architectures, three seeds
Checkpoint selection	Validation next-sample or rollout MSE	Same training path and available budget
Window	\(W=8,16,32,64\)	MLP/RNN/GRU, 32 trajectories, 3,000 updates, identical target endpoints
Data	\(N=8,32,128\) nested independent trajectories	\(W=32\), MLP/RNN/GRU, 3,000 updates, batch size 256

The window and data grids share a newly generated training pool; they use endpoints 64–191, so shorter windows do not gain extra targets. Evaluation uses origins 64 and 96. MLP parameter count changes with \(W\); recurrent widths remain fixed. Each data subset fits its own training-only scaling. Equal updates give 768,000 processed examples but different effective epoch counts. Overlapping windows are not independent trajectories.

Original validation trajectories guide development. All configurations, trained checkpoints, and validation evidence are frozen before generating 32 fresh confirmation trajectories (seed 6799), evaluated at origins 64/96. These are new samples from the same frequency distribution. Three initialization seeds measure initialization variability on fixed data; they do not establish statistical significance. Each trajectory contributes the mean of its two correlated origins.

The original test comparison has already motivated this follow-up. Keep its scores as historical evidence, and use the new confirmation sample for the frozen contrasts. No choices change after confirmation scoring. All predefined runs are retained, including unfavorable results.

Run or inspect the diagnostics

The notebook retains the complete synthetic result record and regenerates the figures from it. Default Run all checks that record and executes the small recurrence calculation. Full diagnostic retraining is optional: change RETRAIN_WAVE_DIAGNOSTICS to True in the folded runner below. It uses the same embedded engine as the repository’s scripts/week06-wave-diagnostics.py, needs no downloads, and writes checkpoints/configuration into a temporary directory. Full diagnostic retraining took 13.7 minutes on the development laptop; default notebook execution remains about 4–6 minutes. Allow another 20–30 minutes for discussing the diagnostic questions and plots.

Verified 78 run records / 153 frozen checkpoint reports; first horizon agrees.
Measured complete diagnostic retraining: 13.7 minutes on this laptop.

# @title Optional: replay the full training grid

RETRAIN_WAVE_DIAGNOSTICS = False
# The full engine is visible here; the same source lives in scripts/week06-wave-diagnostics.py.
diagnostic_engine_source = r'''"""Week 06 fixed-frequency forecasting diagnostics; no downloads or new dependencies.
Run: python3 scripts/week06-wave-diagnostics.py --output /tmp/week06-diagnostics
Notebook embeds the same engine for optional standalone/Colab replay.
"""
import copy
import hashlib
import json
import time
from pathlib import Path
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F

SEEDS = (6406, 6407, 6408)
KINDS = ("MLP", "RNN", "GRU", "1D CNN", "CNN → RNN")
WAVE_WINDOW = 32
FREQUENCIES = np.array([.025, .065])

def generate(count, seed, noise=.03):
    rng = np.random.default_rng(seed)
    amplitudes = rng.uniform([.6, .2], [1., .5], size=(count, 2))
    phases = rng.uniform(0, 2*np.pi, size=(count, 2))
    angles = 2*np.pi*FREQUENCIES[None,:,None]*np.arange(192) + phases[:,:,None]
    clean = (amplitudes[:,:,None]*np.sin(angles)).sum(axis=1)
    return clean + rng.normal(0, noise, size=clean.shape)

def tensors(series, window, endpoints):
    # Every comparison explicitly supplies the same target endpoints.
    assert 0 < window <= min(endpoints) and max(endpoints) < series.shape[1]
    x = np.stack([series[:, n-window:n] for n in endpoints], axis=1)
    y = series[:, endpoints]
    return torch.tensor(x.reshape(-1,window,1), dtype=torch.float32), torch.tensor(y.reshape(-1), dtype=torch.float32)

class WaveWindowModel(nn.Module):
    def __init__(self, kind):
        super().__init__()
        self.kind = kind
        if kind == "MLP":
            self.network = nn.Sequential(nn.Linear(WAVE_WINDOW, 16), nn.ReLU(), nn.Linear(16, 1))
        else:
            width = 20 if kind == "RNN" else 12
            recurrent = nn.RNN if kind == "RNN" else nn.GRU
            self.recurrent = recurrent(1, width, batch_first=True)
            self.readout = nn.Linear(width, 1)

    def forward(self, windows):
        if self.kind == "MLP":
            return self.network(windows[:, :, 0]).squeeze(-1)
        states, _ = self.recurrent(windows)  # fresh zero state on every call
        return self.readout(states[:, -1]).squeeze(-1)
class WaveConvModel(nn.Module):
    def __init__(self, kind):
        super().__init__()
        if kind not in ("1D CNN", "CNN → RNN"):
            raise ValueError("Unknown convolutional predictor")
        self.kind = kind
        self.conv_in = nn.Conv1d(1, 8, kernel_size=5)  # valid: 32 -> 28
        if kind == "1D CNN":
            self.conv_out = nn.Conv1d(8, 8, kernel_size=3)  # 28 -> 26
            self.readout = nn.Linear(8*(WAVE_WINDOW-6), 1)
        else:
            self.recurrent = nn.RNN(8, 16, batch_first=True)
            self.readout = nn.Linear(16, 1)

    def forward(self, windows):
        features = F.relu(self.conv_in(windows.transpose(1, 2)))
        if self.kind == "1D CNN":
            features = F.relu(self.conv_out(features))
            return self.readout(features.flatten(1)).squeeze(-1)
        states, _ = self.recurrent(features.transpose(1, 2))
        return self.readout(states[:, -1]).squeeze(-1)


def digest(model):
    h = hashlib.sha256()
    for p in model.state_dict().values():
        h.update(p.cpu().numpy().tobytes())
    return h.hexdigest()

@torch.no_grad()
def paired(model, raw, mean, scale, window, origins):
    """Per-trajectory errors; horizons share origins and first observations."""
    scaled = (raw-mean)/scale
    contexts = np.stack([scaled[i,o-window:o] for i in range(len(raw)) for o in origins])
    actual = np.stack([raw[i,o:o+64] for i in range(len(raw)) for o in origins])
    observed = np.stack([scaled[i,o+h-window:o+h] for i in range(len(raw)) for o in origins for h in range(64)])
    model.eval()
    teacher = model(torch.tensor(observed[:,:,None],dtype=torch.float32)).numpy().reshape(-1,64)*scale+mean
    context = torch.tensor(contexts[:,:,None],dtype=torch.float32)
    predicted = []
    for _ in range(64):
        nxt = model(context)
        predicted.append(nxt.numpy()*scale+mean)
        context = torch.cat((context[:,1:],nxt[:,None,None]),dim=1)
    recursive = np.stack(predicted,axis=1)
    assert np.isfinite(recursive).all()
    np.testing.assert_allclose(teacher[:,0],recursive[:,0],rtol=1e-5,atol=2e-6)
    result = {}
    for name, predictions in (("observed",teacher),("recursive",recursive)):
        errors = ((predictions-actual)**2).reshape(len(raw),len(origins),64)
        result[name] = float(errors.mean())
        result[name+'_horizon'] = errors.mean(axis=(0,1)).tolist()
        result[name+'_trajectory'] = errors.mean(axis=(1,2)).tolist()
    return result


def train(kind, seed, raw_train, raw_valid, window, endpoints, updates, origins,
          budgets, selection_interval, stage, save_dir):
    global WAVE_WINDOW
    WAVE_WINDOW = window
    torch.manual_seed(seed)
    model = WaveConvModel(kind) if kind in ("1D CNN","CNN → RNN") else WaveWindowModel(kind)
    optimizer = torch.optim.Adam(model.parameters(),lr=.003)
    mean, scale = float(raw_train.mean()),float(raw_train.std())
    tx,ty = tensors((raw_train-mean)/scale,window,endpoints)
    vx,vy = tensors((raw_valid-mean)/scale,window,endpoints)
    best, rollout_best = float('inf'),float('inf')
    best_state,roll_state = None,None
    best_at,roll_at = None,None
    history, records = [],[]
    start = time.perf_counter()
    order, cursor, epoch = None,0,0
    for update in range(1,updates+1):
        if order is None or cursor >= len(order):
            order = torch.randperm(len(tx),generator=torch.Generator().manual_seed(6500+epoch))
            cursor=0; epoch+=1
        indices=order[cursor:cursor+256];cursor+=256
        assert len(indices)==256  # all predefined datasets are multiples of batch size
        optimizer.zero_grad(set_to_none=True)
        loss=F.mse_loss(model(tx[indices]),ty[indices])
        loss.backward()
        assert all(p.grad is None or torch.isfinite(p.grad).all() for p in model.parameters())
        nn.utils.clip_grad_norm_(model.parameters(),1.)
        optimizer.step()
        if update % selection_interval==0 or update in budgets:
            model.eval()
            with torch.no_grad():
                valid=float(F.mse_loss(model(vx),vy))*scale**2
                training=float(F.mse_loss(model(tx),ty))*scale**2
            history.append(dict(update=update,train=training,validation=valid))
            if valid<best:
                best,best_at,best_state=valid,update,copy.deepcopy(model.state_dict())
            if stage=='budget':
                val_pair=paired(model,raw_valid,mean,scale,window,origins)
                if val_pair['recursive']<rollout_best:
                    rollout_best,roll_at,roll_state=val_pair['recursive'],update,copy.deepcopy(model.state_dict())
        if update in budgets:
            current=copy.deepcopy(model.state_dict())
            for rule,state,at in [('next',best_state,best_at)]+([('rollout',roll_state,roll_at)] if stage=='budget' else []):
                model.load_state_dict(state)
                record=dict(stage=stage,kind=kind,seed=seed,window=window,n=len(raw_train),updates=update,
                            selection=rule,selected_update=at,parameters=sum(p.numel() for p in model.parameters()),
                            mean=mean,scale=scale,endpoints=[int(endpoints[0]),int(endpoints[-1])],origins=list(origins),
                            examples_processed=update*256,effective_epochs=update*256/len(tx),
                            digest=digest(model),validation=paired(model,raw_valid,mean,scale,window,origins))
                if stage=='budget' and update==600 and rule=='next':
                    with torch.no_grad(): record['legacy_next_mse']=float(F.mse_loss(model(vx),vy))*scale**2
                path=save_dir / f"{stage}-{kind.replace(' → ','-').replace(' ','_')}-{seed}-w{window}-n{len(raw_train)}-u{update}-{rule}.pt"
                torch.save(dict(state=state,record=record),path)
                record['checkpoint']=path.name
                records.append(record)
            model.load_state_dict(current)
            model.train()
        else:
            model.train()
    return dict(stage=stage,kind=kind,seed=seed,window=window,n=len(raw_train),seconds=time.perf_counter()-start,
                history=history,records=records)


def run_diagnostics(output):
    torch.set_num_threads(1)
    output=Path(output);output.mkdir(parents=True,exist_ok=True)
    start=time.perf_counter()
    config=dict(seeds=list(SEEDS),kinds=list(KINDS),frequencies=FREQUENCIES.tolist(),noise=.03,batch=256,lr=.003,
                clip=1.,budget_epochs=[30,60,150],windows=[8,16,32,64],trajectory_counts=[8,32,128],grid_updates=3000,
                legacy_data_seeds=[6411,6412],pool_seed=6701,confirmation_seed=6799,
                grid_endpoints=[64,191],grid_origins=[64,96],budget_origins=[32,96],selection_interval_grid=100)
    (output/'config.json').write_text(json.dumps(config,indent=2))
    clean=generate(3,6700,noise=0)
    c1,c2=np.cos(2*np.pi*FREQUENCIES)
    coeff=np.array([2*(c1+c2),-(2+4*c1*c2),2*(c1+c2),-1.])
    reconstructed=sum(coeff[j]*clean[:,3-j:191-j] for j in range(4))
    np.testing.assert_allclose(reconstructed,clean[:,4:],atol=2e-13,rtol=0)
    old_train,valid=generate(32,config['legacy_data_seeds'][0]),generate(8,config['legacy_data_seeds'][1])
    pool=generate(max(config['trajectory_counts']),config['pool_seed'])
    hashes=lambda rows:{hashlib.sha256(row.tobytes()).hexdigest() for row in rows}
    assert not hashes(pool)&hashes(valid)
    runs=[]
    for seed in config['seeds']:
        for kind in config['kinds']:
            run=train(kind,seed,old_train,valid,32,np.arange(32,192),20*max(config['budget_epochs']),tuple(config['budget_origins']),
                      tuple(20*e for e in config['budget_epochs']),20,'budget',output)
            runs.append(run)
            print('budget',kind,seed,round(run['seconds'],1),flush=True)
    # Predefined common-endpoint window grid and nested independent-trajectory grid.
    for stage,conditions in [('window',[(w,32) for w in config['windows']]),('data',[(32,n) for n in config['trajectory_counts']])]:
        for window,n in conditions:
            # Grid N=32 uses the same nested pool across both experiments; reuse exact duplicate.
            for seed in config['seeds']:
                for kind in config['kinds'][:3]:
                    if stage=='data' and n==32:
                        base=next(r for r in runs if r['stage']=='window' and r['window']==32 and r['kind']==kind and r['seed']==seed)
                        duplicate=copy.deepcopy(base);duplicate['stage']='data';duplicate['seconds']=0.
                        for rec in duplicate['records']:rec['stage']='data'
                        runs.append(duplicate);continue
                    run=train(kind,seed,pool[:n],valid,window,np.arange(config['grid_endpoints'][0],config['grid_endpoints'][1]+1),
                              config['grid_updates'],tuple(config['grid_origins']),(config['grid_updates'],),
                              config['selection_interval_grid'],stage,output)
                    runs.append(run)
                    print(stage,window,n,kind,seed,round(run['seconds'],1),flush=True)
    # Freeze trained states and validation evidence BEFORE generating confirmation trajectories.
    frozen=dict(config=config,recurrence_max_error=float(np.max(np.abs(reconstructed-clean[:,4:]))),runs=runs,
                development_seconds=time.perf_counter()-start)
    (output/'frozen-validation.json').write_text(json.dumps(frozen,indent=2))
    confirmation=generate(32,config['confirmation_seed'])
    assert not hashes(confirmation)&(hashes(pool)|hashes(valid)|hashes(old_train))
    for run in runs:
        for rec in run['records']:
            WAVE_WINDOW=rec['window']
            # Models use global window only during construction.
            globals()['WAVE_WINDOW']=rec['window']
            torch.manual_seed(rec['seed'])
            model=WaveConvModel(rec['kind']) if rec['kind'] in KINDS[3:] else WaveWindowModel(rec['kind'])
            checkpoint=torch.load(output/rec['checkpoint'],weights_only=True)
            model.load_state_dict(checkpoint['state'])
            assert digest(model)==rec['digest']
            rec['confirmation']=paired(model,confirmation,rec['mean'],rec['scale'],rec['window'],(64,96))
    result=dict(config=config,runs=runs,recurrence_max_error=frozen['recurrence_max_error'],seconds=time.perf_counter()-start,
                torch_version=str(torch.__version__))
    assert len(runs)==78 and sum(len(r['records']) for r in runs)==153
    (output/'results.json').write_text(json.dumps(result,indent=2))
    print('COMPLETE',round(result['seconds'],2),flush=True)
    return result

def verify_saved(output):
    """Replay frozen checkpoints without repeating the training grid."""
    global WAVE_WINDOW
    output=Path(output)
    result=json.loads((output/'results.json').read_text())
    frozen=json.loads((output/'frozen-validation.json').read_text())
    assert result['config']==frozen['config']
    confirmation=generate(32,result['config']['confirmation_seed'])
    torch.set_num_threads(1)
    count=0
    frozen_records=[r for run in frozen['runs'] for r in run['records']]
    for run in result['runs']:
        for rec in run['records']:
            assert {k:v for k,v in rec.items() if k != 'confirmation'}==frozen_records[count]
            WAVE_WINDOW=rec['window']
            model=WaveConvModel(rec['kind']) if rec['kind'] in KINDS[3:] else WaveWindowModel(rec['kind'])
            saved=torch.load(output/rec['checkpoint'],weights_only=True)
            model.load_state_dict(saved['state'])
            assert digest(model)==rec['digest']
            replay=paired(model,confirmation,rec['mean'],rec['scale'],rec['window'],(64,96))
            for key,value in replay.items():
                np.testing.assert_allclose(value,rec['confirmation'][key],rtol=0,atol=0)
            count+=1
    assert count==153
    print('Exact frozen-checkpoint confirmation replay:',count,flush=True)

'''
if RETRAIN_WAVE_DIAGNOSTICS:
    import tempfile
    diagnostic_namespace = {'__name__': 'week06_diagnostics'}
    exec(diagnostic_engine_source, diagnostic_namespace)
    diagnostic_output = Path(tempfile.mkdtemp(prefix='week06-diagnostics-'))
    wave_diagnostics_record = diagnostic_namespace['run_diagnostics'](diagnostic_output)
    print('Config, checkpoints, frozen validation, and metrics:', diagnostic_output)

# @title Auxiliary: diagnostic learning curves and controlled grids

def diagnostic_records(stage, **filters):
    return [r for run in wave_diagnostics_record['runs'] for r in run['records']
            if r['stage']==stage and all(r[k]==v for k,v in filters.items())]

def plot_wave_diagnostic_budget():
    fig, flat_axes = plt.subplots(4,1,figsize=(8,14))
    axes = np.array(flat_axes).reshape(2,2).T
    for col,kind in enumerate(('RNN','GRU')):
        runs=[r for r in wave_diagnostics_record['runs'] if r['stage']=='budget' and r['kind']==kind]
        ax=axes[0,col]
        for metric,color in [('train',BLUE),('validation','#c67c19')]:
            values=np.array([[h[metric] for h in r['history']] for r in runs])
            epochs=np.array([h['update']/20 for h in runs[0]['history']])
            ax.plot(epochs,values.mean(axis=0),color=color,label=metric)
            for row in values:ax.plot(epochs,row,color=color,alpha=.18,lw=.7)
        ax.axvline(30,color=GRAY,ls=':')
        ax.set(title=f'{kind}: continued training',xlabel='Epoch',ylabel='Observed-context MSE',yscale='log')
        ax.legend(fontsize=9)
        ax=axes[1,col]
        for updates,color in [(600,GRAY),(3000,BLUE)]:
            recs=diagnostic_records('budget',kind=kind,updates=updates,selection='next')
            for mode,style in [('observed','--'),('recursive','-')]:
                errors=np.array([r['confirmation'][mode+'_horizon'] for r in recs])
                ax.plot(np.arange(1,65),errors.mean(axis=0),color=color,ls=style,
                        label=f'{updates//20} epochs, {mode}')
        ax.axhline(.03**2,color='#c67c19',ls=':',label='noise reference')
        ax.set(title=f'{kind}: matched fresh origins',xlabel='Samples ahead',ylabel='Confirmation MSE',yscale='log')
    handles, labels = axes[1,0].get_legend_handles_labels()
    fig.legend(handles, labels, loc='lower center', ncol=2, fontsize=9)
    fig.tight_layout(rect=(0,.07,1,1));save_wave_comparison(fig,'28-wave-budget-diagnostics');plt.show()

def plot_wave_diagnostic_grids():
    fig,flat_axes=plt.subplots(4,1,figsize=(8,14))
    axes=np.array(flat_axes).reshape(2,2).T
    for col,(stage,key,choices,label) in enumerate([('window','window',[8,16,32,64],'Observed samples W'),
                                                   ('data','n',[8,32,128],'Independent training trajectories N')]):
        for row,mode in enumerate(('observed','recursive')):
            ax=axes[row,col]
            for kind in ('MLP','RNN','GRU'):
                errors=np.array([[r['confirmation'][mode] for r in diagnostic_records(stage,kind=kind,**{key:v})]
                                 for v in choices])
                ax.plot(range(len(choices)),errors.mean(axis=1),marker='o',color=WAVE_COLORS[kind],label=kind)
                for seed_index in range(3):
                    ax.plot(range(len(choices)),errors[:,seed_index],color=WAVE_COLORS[kind],alpha=.2,lw=.8)
            ax.axhline(.03**2,color=GRAY,ls=':',label='noise reference')
            ax.set(yscale='log',xlabel=label,ylabel='Confirmation MSE',title=f'{stage.capitalize()}: {mode}')
            ax.set_xticks(range(len(choices)));ax.set_xticklabels(choices)
            ax.legend(fontsize=8)
    fig.tight_layout();save_wave_comparison(fig,'29-wave-window-data-diagnostics');plt.show()

Read the frozen confirmation evidence

Every error below uses the same 32 fresh trajectories and origins 64/96. Observed and recursive MSE average all 64 horizons; entries are mean ± sample SD across three initializations. These evaluation locations differ from the historical table, so its numbers need not match even when the 30-epoch weights match exactly.

# @title Display diagnostic budget and control comparisons

def diagnostic_summary(records, mode, split='confirmation'):
    values=np.array([r[split][mode] for r in records])
    return f'{values.mean():.4g} ± {values.std(ddof=1):.2g}'

budget_lines=['| Predictor | Epoch budget | Observed MSE | Recursive MSE |', '|---|---:|---:|---:|']
for kind in ('MLP','RNN','GRU','1D CNN','CNN → RNN'):
    for updates in (600,1200,3000):
        recs=diagnostic_records('budget',kind=kind,updates=updates,selection='next')
        budget_lines.append(f'| {kind} | {updates//20} | {diagnostic_summary(recs,"observed")} | {diagnostic_summary(recs,"recursive")} |')
display(Markdown('\n'.join(budget_lines)))

selection_lines=['| Predictor | Next-sample selection | Rollout selection |','|---|---:|---:|']
for kind in ('MLP','RNN','GRU','1D CNN','CNN → RNN'):
    next_recs=diagnostic_records('budget',kind=kind,updates=3000,selection='next')
    roll_recs=diagnostic_records('budget',kind=kind,updates=3000,selection='rollout')
    selection_lines.append(f'| {kind} | {diagnostic_summary(next_recs,"recursive")} | {diagnostic_summary(roll_recs,"recursive")} |')
display(Markdown('<details>\n<summary>Same training path — compare checkpoint selection</summary>\n\n'+'\n'.join(selection_lines)+'\n\n</details>'))

ledger=['| Stage | Model | Seed | W | N | Updates | Selection | Selected update | Params | Effective epochs | Observed MSE | Recursive MSE |',
        '|---|---|---:|---:|---:|---:|---|---:|---:|---:|---:|---:|']
for run in wave_diagnostics_record['runs']:
    for r in run['records']:
        ledger.append(f'| {r["stage"]} | {r["kind"]} | {r["seed"]} | {r["window"]} | {r["n"]} | {r["updates"]} | {r["selection"]} | {r["selected_update"]} | {r["parameters"]} | {r["effective_epochs"]:.2f} | {r["confirmation"]["observed"]:.5g} | {r["confirmation"]["recursive"]:.5g} |')
display(Markdown('<details>\n<summary>Full diagnostic ledger — all 153 reports</summary>\n\n'+'\n'.join(ledger)+'\n\n</details>'))

Predictor	Epoch budget	Observed MSE	Recursive MSE
MLP	30	0.001135 ± 1.6e-05	0.001171 ± 7.5e-05
MLP	60	0.001086 ± 7.7e-06	0.001093 ± 5.6e-07
MLP	150	0.001063 ± 1.3e-05	0.001092 ± 2.5e-05
RNN	30	0.003566 ± 5.4e-05	0.3226 ± 0.14
RNN	60	0.003314 ± 4.2e-05	0.1568 ± 0.014
RNN	150	0.0017 ± 0.00017	0.008856 ± 0.0036
GRU	30	0.003914 ± 0.00053	0.4132 ± 0.22
GRU	60	0.002883 ± 0.00058	0.1641 ± 0.12
GRU	150	0.001954 ± 0.00039	0.01518 ± 0.013
1D CNN	30	0.001124 ± 2.6e-05	0.001209 ± 7.7e-05
1D CNN	60	0.00111 ± 3.5e-06	0.001127 ± 3.5e-05
1D CNN	150	0.001094 ± 8.1e-06	0.00111 ± 3.3e-05
CNN → RNN	30	0.003189 ± 0.00073	0.3886 ± 0.43
CNN → RNN	60	0.002457 ± 0.00037	0.02943 ± 0.0045
CNN → RNN	150	0.001763 ± 0.00014	0.009194 ± 0.0035

Same training path — compare checkpoint selection

Predictor	Next-sample selection	Rollout selection
MLP	0.001092 ± 2.5e-05	0.001069 ± 1.7e-05
RNN	0.008856 ± 0.0036	0.008621 ± 0.005
GRU	0.01518 ± 0.013	0.01282 ± 0.0089
1D CNN	0.00111 ± 3.3e-05	0.001123 ± 2.4e-05
CNN → RNN	0.009194 ± 0.0035	0.007677 ± 0.0017

Full diagnostic ledger — all 153 reports

Stage	Model	Seed	W	N	Updates	Selection	Selected update	Params	Effective epochs	Observed MSE	Recursive MSE
budget	MLP	6406	32	32	600	next	580	545	30.00	0.0011178	0.0011189
budget	MLP	6406	32	32	600	rollout	580	545	30.00	0.0011178	0.0011189
budget	MLP	6406	32	32	1200	next	1200	545	60.00	0.0010917	0.0010927
budget	MLP	6406	32	32	1200	rollout	1200	545	60.00	0.0010917	0.0010927
budget	MLP	6406	32	32	3000	next	2920	545	150.00	0.0010704	0.0011116
budget	MLP	6406	32	32	3000	rollout	1480	545	150.00	0.0010832	0.0010732
budget	RNN	6406	32	32	600	next	600	481	30.00	0.0036258	0.4785
budget	RNN	6406	32	32	600	rollout	500	481	30.00	0.0035976	0.19096
budget	RNN	6406	32	32	1200	next	1120	481	60.00	0.0033341	0.15216
budget	RNN	6406	32	32	1200	rollout	1140	481	60.00	0.0032507	0.14954
budget	RNN	6406	32	32	3000	next	3000	481	150.00	0.0016175	0.0054854
budget	RNN	6406	32	32	3000	rollout	3000	481	150.00	0.0016175	0.0054854
budget	GRU	6406	32	32	600	next	580	553	30.00	0.0043673	0.52499
budget	GRU	6406	32	32	600	rollout	140	553	30.00	0.0079682	0.27763
budget	GRU	6406	32	32	1200	next	1200	553	60.00	0.0030979	0.20092
budget	GRU	6406	32	32	1200	rollout	1180	553	60.00	0.0030896	0.14273
budget	GRU	6406	32	32	3000	next	2880	553	150.00	0.0020569	0.0089375
budget	GRU	6406	32	32	3000	rollout	2860	553	150.00	0.0020285	0.0096931
budget	1D CNN	6406	32	32	600	next	500	457	30.00	0.0011531	0.0012918
budget	1D CNN	6406	32	32	600	rollout	500	457	30.00	0.0011531	0.0012918
budget	1D CNN	6406	32	32	1200	next	940	457	60.00	0.0011132	0.0011646
budget	1D CNN	6406	32	32	1200	rollout	940	457	60.00	0.0011132	0.0011646
budget	1D CNN	6406	32	32	3000	next	2200	457	150.00	0.0011028	0.0011464
budget	1D CNN	6406	32	32	3000	rollout	2200	457	150.00	0.0011028	0.0011464
budget	CNN → RNN	6406	32	32	600	next	600	481	30.00	0.0038169	0.86193
budget	CNN → RNN	6406	32	32	600	rollout	60	481	30.00	0.022551	0.35111
budget	CNN → RNN	6406	32	32	1200	next	1180	481	60.00	0.0026421	0.031072
budget	CNN → RNN	6406	32	32	1200	rollout	1180	481	60.00	0.0026421	0.031072
budget	CNN → RNN	6406	32	32	3000	next	3000	481	150.00	0.0018299	0.013086
budget	CNN → RNN	6406	32	32	3000	rollout	2940	481	150.00	0.0018111	0.0072626
budget	MLP	6407	32	32	600	next	580	545	30.00	0.0011494	0.0012573
budget	MLP	6407	32	32	600	rollout	560	545	30.00	0.0011687	0.001287
budget	MLP	6407	32	32	1200	next	1120	545	60.00	0.0010776	0.0010931
budget	MLP	6407	32	32	1200	rollout	1120	545	60.00	0.0010776	0.0010931
budget	MLP	6407	32	32	3000	next	2920	545	150.00	0.001048	0.0010637
budget	MLP	6407	32	32	3000	rollout	1880	545	150.00	0.0010635	0.0010833
budget	RNN	6407	32	32	600	next	600	481	30.00	0.0035215	0.27275
budget	RNN	6407	32	32	600	rollout	500	481	30.00	0.0033645	0.14468
budget	RNN	6407	32	32	1200	next	1100	481	60.00	0.0032666	0.14563
budget	RNN	6407	32	32	1200	rollout	900	481	60.00	0.0031883	0.12161
budget	RNN	6407	32	32	3000	next	2880	481	150.00	0.0015871	0.0083861
budget	RNN	6407	32	32	3000	rollout	2960	481	150.00	0.0016307	0.0060003
budget	GRU	6407	32	32	600	next	580	553	30.00	0.0040461	0.5546
budget	GRU	6407	32	32	600	rollout	280	553	30.00	0.0048177	0.31808
budget	GRU	6407	32	32	1200	next	1200	553	60.00	0.0033225	0.25778
budget	GRU	6407	32	32	1200	rollout	1140	553	60.00	0.0032878	0.14391
budget	GRU	6407	32	32	3000	next	2940	553	150.00	0.0022862	0.030583
budget	GRU	6407	32	32	3000	rollout	2780	553	150.00	0.0023124	0.0229
budget	1D CNN	6407	32	32	600	next	440	457	30.00	0.0011147	0.0011392
budget	1D CNN	6407	32	32	600	rollout	600	457	30.00	0.0010965	0.00109
budget	1D CNN	6407	32	32	1200	next	680	457	60.00	0.0011064	0.0011208
budget	1D CNN	6407	32	32	1200	rollout	680	457	60.00	0.0011064	0.0011208
budget	1D CNN	6407	32	32	3000	next	2300	457	150.00	0.0010871	0.0010825
budget	1D CNN	6407	32	32	3000	rollout	2160	457	150.00	0.0011133	0.0010979
budget	CNN → RNN	6407	32	32	600	next	540	481	30.00	0.0023876	0.030345
budget	CNN → RNN	6407	32	32	600	rollout	540	481	30.00	0.0023876	0.030345
budget	CNN → RNN	6407	32	32	1200	next	1160	481	60.00	0.002028	0.024379
budget	CNN → RNN	6407	32	32	1200	rollout	640	481	60.00	0.0021487	0.014548
budget	CNN → RNN	6407	32	32	3000	next	2860	481	150.00	0.0016014	0.0082744
budget	CNN → RNN	6407	32	32	3000	rollout	1540	481	150.00	0.001786	0.0095472
budget	MLP	6408	32	32	600	next	600	545	30.00	0.001138	0.0011379
budget	MLP	6408	32	32	600	rollout	600	545	30.00	0.001138	0.0011379
budget	MLP	6408	32	32	1200	next	1200	545	60.00	0.0010901	0.0010938
budget	MLP	6408	32	32	1200	rollout	1000	545	60.00	0.0010882	0.0010729
budget	MLP	6408	32	32	3000	next	2780	545	150.00	0.0010702	0.0010993
budget	MLP	6408	32	32	3000	rollout	1500	545	150.00	0.0010737	0.0010509
budget	RNN	6408	32	32	600	next	600	481	30.00	0.0035509	0.21656
budget	RNN	6408	32	32	600	rollout	220	481	30.00	0.0041361	0.17045
budget	RNN	6408	32	32	1200	next	1120	481	60.00	0.0033425	0.17249
budget	RNN	6408	32	32	1200	rollout	880	481	60.00	0.0033005	0.13842
budget	RNN	6408	32	32	3000	next	2960	481	150.00	0.0018949	0.012697
budget	RNN	6408	32	32	3000	rollout	1820	481	150.00	0.0019921	0.014377
budget	GRU	6408	32	32	600	next	600	553	30.00	0.0033294	0.15997
budget	GRU	6408	32	32	600	rollout	600	553	30.00	0.0033294	0.15997
budget	GRU	6408	32	32	1200	next	1200	553	60.00	0.0022301	0.033706
budget	GRU	6408	32	32	1200	rollout	1180	553	60.00	0.0021979	0.03065
budget	GRU	6408	32	32	3000	next	2980	553	150.00	0.0015203	0.006031
budget	GRU	6408	32	32	3000	rollout	2880	553	150.00	0.0015401	0.0058796
budget	1D CNN	6408	32	32	600	next	600	457	30.00	0.0011041	0.0011949
budget	1D CNN	6408	32	32	600	rollout	380	457	30.00	0.001103	0.0011598
budget	1D CNN	6408	32	32	1200	next	1060	457	60.00	0.0011112	0.0010964
budget	1D CNN	6408	32	32	1200	rollout	1060	457	60.00	0.0011112	0.0010964
budget	1D CNN	6408	32	32	3000	next	1800	457	150.00	0.0010918	0.0011017
budget	1D CNN	6408	32	32	3000	rollout	2040	457	150.00	0.0011159	0.0011242
budget	CNN → RNN	6408	32	32	600	next	580	481	30.00	0.0033626	0.27359
budget	CNN → RNN	6408	32	32	600	rollout	580	481	30.00	0.0033626	0.27359
budget	CNN → RNN	6408	32	32	1200	next	1120	481	60.00	0.0026997	0.032852
budget	CNN → RNN	6408	32	32	1200	rollout	1120	481	60.00	0.0026997	0.032852
budget	CNN → RNN	6408	32	32	3000	next	2840	481	150.00	0.0018584	0.0062206
budget	CNN → RNN	6408	32	32	3000	rollout	2840	481	150.00	0.0018584	0.0062206
window	MLP	6406	8	32	3000	next	2600	161	187.50	0.0038605	0.42871
window	RNN	6406	8	32	3000	next	2900	481	187.50	0.0037419	0.22941
window	GRU	6406	8	32	3000	next	2600	553	187.50	0.0038243	0.22862
window	MLP	6407	8	32	3000	next	2600	161	187.50	0.003869	0.4285
window	RNN	6407	8	32	3000	next	2900	481	187.50	0.0037834	0.22223
window	GRU	6407	8	32	3000	next	2900	553	187.50	0.0038096	0.28385
window	MLP	6408	8	32	3000	next	2300	161	187.50	0.0038964	0.31792
window	RNN	6408	8	32	3000	next	2600	481	187.50	0.0037686	0.32165
window	GRU	6408	8	32	3000	next	2200	553	187.50	0.0038893	0.28939
window	MLP	6406	16	32	3000	next	2500	289	187.50	0.0016554	0.0028628
window	RNN	6406	16	32	3000	next	2300	481	187.50	0.0016944	0.0059089
window	GRU	6406	16	32	3000	next	2900	553	187.50	0.0017174	0.0067846
window	MLP	6407	16	32	3000	next	2900	289	187.50	0.0016801	0.0058991
window	RNN	6407	16	32	3000	next	2400	481	187.50	0.0016525	0.0056188
window	GRU	6407	16	32	3000	next	2900	553	187.50	0.0017262	0.0074877
window	MLP	6408	16	32	3000	next	3000	289	187.50	0.0016886	0.004695
window	RNN	6408	16	32	3000	next	2500	481	187.50	0.0017119	0.0069401
window	GRU	6408	16	32	3000	next	2900	553	187.50	0.0016769	0.0048707
window	MLP	6406	32	32	3000	next	2800	545	187.50	0.0010893	0.0011558
window	RNN	6406	32	32	3000	next	3000	481	187.50	0.0017282	0.0098296
window	GRU	6406	32	32	3000	next	3000	553	187.50	0.0020184	0.011643
window	MLP	6407	32	32	3000	next	2800	545	187.50	0.0010793	0.0011082
window	RNN	6407	32	32	3000	next	2700	481	187.50	0.0018189	0.0092404
window	GRU	6407	32	32	3000	next	2900	553	187.50	0.0022997	0.027111
window	MLP	6408	32	32	3000	next	2800	545	187.50	0.0010787	0.0010866
window	RNN	6408	32	32	3000	next	2600	481	187.50	0.0017962	0.01098
window	GRU	6408	32	32	3000	next	2600	553	187.50	0.001599	0.0095926
window	MLP	6406	64	32	3000	next	2100	1057	187.50	0.0010256	0.001032
window	RNN	6406	64	32	3000	next	2900	481	187.50	0.0016961	0.0089541
window	GRU	6406	64	32	3000	next	2900	553	187.50	0.0019265	0.0041544
window	MLP	6407	64	32	3000	next	1800	1057	187.50	0.0010139	0.0010148
window	RNN	6407	64	32	3000	next	2600	481	187.50	0.0014942	0.006832
window	GRU	6407	64	32	3000	next	3000	553	187.50	0.0015032	0.0049086
window	MLP	6408	64	32	3000	next	2100	1057	187.50	0.0010213	0.001021
window	RNN	6408	64	32	3000	next	2900	481	187.50	0.0018135	0.015112
window	GRU	6408	64	32	3000	next	2900	553	187.50	0.0014706	0.0079249
data	MLP	6406	32	8	3000	next	1900	545	750.00	0.0011706	0.0012378
data	RNN	6406	32	8	3000	next	2700	481	750.00	0.0020002	0.010272
data	GRU	6406	32	8	3000	next	2800	553	750.00	0.0022589	0.019505
data	MLP	6407	32	8	3000	next	2200	545	750.00	0.0011597	0.0012173
data	RNN	6407	32	8	3000	next	2900	481	750.00	0.0019522	0.015113
data	GRU	6407	32	8	3000	next	3000	553	750.00	0.0021796	0.019981
data	MLP	6408	32	8	3000	next	1600	545	750.00	0.0011797	0.0011682
data	RNN	6408	32	8	3000	next	2900	481	750.00	0.0021102	0.012295
data	GRU	6408	32	8	3000	next	2900	553	750.00	0.0017024	0.010283
data	MLP	6406	32	32	3000	next	2800	545	187.50	0.0010893	0.0011558
data	RNN	6406	32	32	3000	next	3000	481	187.50	0.0017282	0.0098296
data	GRU	6406	32	32	3000	next	3000	553	187.50	0.0020184	0.011643
data	MLP	6407	32	32	3000	next	2800	545	187.50	0.0010793	0.0011082
data	RNN	6407	32	32	3000	next	2700	481	187.50	0.0018189	0.0092404
data	GRU	6407	32	32	3000	next	2900	553	187.50	0.0022997	0.027111
data	MLP	6408	32	32	3000	next	2800	545	187.50	0.0010787	0.0010866
data	RNN	6408	32	32	3000	next	2600	481	187.50	0.0017962	0.01098
data	GRU	6408	32	32	3000	next	2600	553	187.50	0.001599	0.0095926
data	MLP	6406	32	128	3000	next	2600	545	46.88	0.0010738	0.0011516
data	RNN	6406	32	128	3000	next	2900	481	46.88	0.0019054	0.022343
data	GRU	6406	32	128	3000	next	3000	553	46.88	0.0020271	0.0088516
data	MLP	6407	32	128	3000	next	1800	545	46.88	0.0010792	0.0010865
data	RNN	6407	32	128	3000	next	3000	481	46.88	0.0016185	0.0089725
data	GRU	6407	32	128	3000	next	2900	553	46.88	0.0017523	0.0092091
data	MLP	6408	32	128	3000	next	2000	545	46.88	0.001074	0.0011168
data	RNN	6408	32	128	3000	next	3000	481	46.88	0.0017839	0.0083263
data	GRU	6408	32	128	3000	next	2900	553	46.88	0.0014946	0.0048939

plot_wave_diagnostic_budget()

Figure 28. Follow training, then separate two uses of the same predictor. Each model’s learning-curve panel uses the original training/validation splits; bold lines average three seeds and thin lines retain individual runs. The dotted line marks the old 30-epoch limit. Each following horizon panel uses fresh confirmation trajectories: dashed curves restore observations, solid curves feed predictions back. Both start from the same contexts. The noise line is a reference, not a finite-sample bound.

plot_wave_diagnostic_grids()

Figure 29. Change one observation or data budget at a time. Bold lines average all three seeds; thin lines show each seed. Each window condition uses identical target endpoints, and each data condition uses a nested subset of one training pool. All receive 3,000 updates. Errors use the same fresh confirmation trajectories and origins. These grids use a different training pool and endpoint protocol from Figure 28; compare conditions within each grid.

What changed, and what remains uncertain?

More training mattered. On fresh confirmation trajectories, extending 30 to 150 epochs reduces mean RNN recursive MSE from 0.323 to 0.00886, and GRU from 0.413 to 0.0152. Every paired seed improves. Training and validation learning curves also improve. The original budget constrained these configurations; the comparison did not establish convergence. MLP and CNN still have lower rollout error, around 0.0011, under the extended budget.

Recursive prediction remains harder. At 150 epochs, RNN observed-context MSE is about 0.00170, versus recursive MSE 0.00886; GRU is about 0.00195 versus 0.0152. Restoring actual observations keeps errors lower. Feeding predictions back magnifies the remaining errors. This evidence does not isolate phase drift, amplitude drift, or gradient problems as a unique cause. Selecting by validation rollout error gives mixed effects across seeds on confirmation data; it is not a guaranteed fix.

Window length has architecture-dependent effects. In the controlled grid, reducing 32 to 16 observations improves RNN and GRU rollout error in every seed: means fall from 0.0100 to 0.00616 and 0.0161 to 0.00638. Eight observations produce poor rollouts for all three models. Increasing to 64 helps the MLP and GRU relative to 32 in every seed, while RNN changes are mixed. Thus “the window is too wide” is an incomplete explanation. Window length changes both information and recurrent computation; MLP parameter count also changes.

Additional data help some configurations. With the same 3,000 updates, increasing 32 to 128 independent trajectories reduces GRU rollout MSE from 0.0161 to 0.00765, with all three seeds improving. RNN effects are mixed and its mean increases from 0.0100 to 0.0132. The MLP changes little between 32 and 128. More trajectories alone do not consistently remove the recurrent gap under this computation budget. This is one nested training pool and one confirmation sample, not a universal data-scaling law.

We can now explain the original result more carefully: a short training budget and sensitivity to generated inputs contributed to large recurrent errors; window length and data budget also affect the measured outcome. The fixed-frequency task supports a strong linear alternative. These findings assess the tested configurations rather than promise one architecture a win.

Make a claim that the evidence can support

Checkpoint. Return to your opening explanations. Which has evidence now? Does the poor 30-epoch wave result mean that RNNs do not work?

Show explanation — judge a configuration and protocol

The original table describes particular small models, a training budget, noisy fixed-frequency data, a window length, and a checkpoint rule. Continued training changes their errors. Matched observed/recursive predictions expose the effect of feeding predictions back. The window and data grids test additional specific explanations. Those findings narrow the claim; none warrants rejecting all RNNs. The robot experiment already demonstrates learned recall at delay 8 and a loss of performance under harder conditions.

An RNN offers a reusable state update as samples arrive. Whether training makes that state useful depends on the task and protocol. On this wave task, linear structure is a strong alternative. On a dependency extending outside a fixed input window, an architecture’s access to information becomes an additional experimental choice. Persistent-state forecasting has not been tested here.

TODO(student): Checked-effect investigation. Choose one budget or window contrast, predict its direction before running, and replay the diagnostic engine. Report all three seeds, observed-context error, and recursive error. State one conclusion and one limitation. If you change the protocol, use validation for development and reserve new confirmation trajectories rather than repeatedly choosing from test outcomes.

Reasoning checkpoint. Would a shorter window improving an RNN prove that long-term memory is harmful? Would more trajectories failing to improve this grid prove that data never help?

Show explanation — identify what each control changes

Window length changes both available observations and the number of recurrent updates within a prediction. It can affect denoising, optimization, and the learned representation. A shorter-window improvement alone does not diagnose vanishing gradients. The data grid fixes optimizer updates, so larger datasets receive fewer effective passes. It measures a particular data/computation tradeoff; it cannot establish every effect of additional training data.

TODO(student): Open hypothesis. Propose a test using a different frequency family, persistent recurrent state, or a multi-step training loss. Specify what changes, what remains comparable, the selection criterion, and a fresh evaluation split. Predict an outcome and a possible falsifying observation. The mini-test below checks one fixed frequency shift; these further extensions remain open, with no promised winning architecture.

Mini-test — change the test frequencies only

Predict before reading the result. Training and validation still use frequencies 0.025/0.065 cycles per sample. What happens if the test family uses 0.035/0.075? Will better next-sample predictions imply stable recursive forecasts?

This is a fixed frequency shift between training and test families, not gradual drift inside a trajectory. A change of 0.01 cycles/sample accumulates 0.64 cycles over a 64-sample forecast, so a small absolute frequency change can matter greatly over the rollout. Both test conditions have 32 new trajectories with seed 6806. Using the same seed preserves amplitudes, phases, and noise; only frequencies change. Contexts and future targets both belong to the shifted family. Models receive observed measurements but no frequency labels, updates, or test-fitted scaling.

We choose each architecture’s best previously trained configuration by mean validation rollout MSE across all three seeds, rescoring candidates at common validation origins 64/96. This harmonizes the earlier grids’ different origin sets. Selection is frozen before either new test condition is generated. Settings can differ between architectures: compare the complete recipes as well as the additional matched-window control.

Model	Selected window	Training trajectories	Updates	Checkpoint criterion
MLP	64	32, nested pool	3,000	Validation next-sample MSE
RNN / GRU	16	32, nested pool	3,000	Validation next-sample MSE
1D CNN / CNN → RNN	32	32, original pool	3,000	Validation rollout MSE

The original training-fit AR and persistence remain references at window 32. Error is averaged over 64 horizons at origins 64/96. Observed-context evaluation receives fresh real measurements at each horizon; recursive evaluation receives only the original context and subsequent predictions. The original comparison still uses the familiar frequencies.

# Shift ONLY the additional test family; all training/validation calls retain defaults.
shift_test_familiar = make_wave_trajectories(32, 6806)
shift_test_changed = make_wave_trajectories(32, 6806, frequencies=(.035, .075))
assert shift_test_familiar.shape == shift_test_changed.shape == (32, 192)
assert not np.array_equal(shift_test_familiar, shift_test_changed)
np.testing.assert_array_equal(make_wave_trajectories(32, 6411), wave_splits['train'])
print('Original training trajectories unchanged; paired test frequencies changed.')

Original training trajectories unchanged; paired test frequencies changed.

# @title Optional: repeat the shift test after diagnostic retraining

# Same replay engine as scripts/week06-wave-frequency-shift.py, supplied with the embedded diagnostics.
frequency_shift_engine_source = r'''"""Frozen-checkpoint frequency-shift mini-test; never selects models on test scores.
Run after week06-wave-diagnostics.py:
python3 scripts/week06-wave-frequency-shift.py --checkpoints /tmp/week06-diagnostics --output /tmp/week06-frequency-shift
"""
import argparse
import importlib.util
import json
from pathlib import Path
import numpy as np
import torch


def make_wave_trajectories(count, seed, frequencies=(.025,.065)):
    """Only the frequencies change: count/seed preserve amplitudes, phases, and noise."""
    rng=np.random.default_rng(seed)
    amplitudes=rng.uniform([.6,.2],[1.,.5],size=(count,2))
    phases=rng.uniform(0,2*np.pi,size=(count,2))
    frequencies=np.asarray(frequencies,dtype=float)
    assert frequencies.shape==(2,)
    angles=2*np.pi*frequencies[None,:,None]*np.arange(192)+phases[:,:,None]
    return (amplitudes[:,:,None]*np.sin(angles)).sum(axis=1)+rng.normal(0,.03,size=(count,192))

def load_model(record, checkpoint_dir):
    diag.WAVE_WINDOW=record['window']
    torch.manual_seed(record['seed'])
    model=diag.WaveConvModel(record['kind']) if record['kind'] in diag.KINDS[3:] else diag.WaveWindowModel(record['kind'])
    saved=torch.load(checkpoint_dir/record['checkpoint'],weights_only=True)
    model.load_state_dict(saved['state'])
    assert diag.digest(model)==record['digest']
    model.eval()
    return model

def run(checkpoint_dir, output):
    torch.set_num_threads(1)
    checkpoint_dir,output=Path(checkpoint_dir),Path(output)
    output.mkdir(parents=True,exist_ok=True)
    frozen=json.loads((checkpoint_dir/'frozen-validation.json').read_text())
    valid=make_wave_trajectories(8,6412)
    np.testing.assert_array_equal(valid,diag.generate(8,6412))
    candidates={}
    # Deduplicate reused grid records and score every validation candidate at common origins.
    seen=set()
    for run in frozen['runs']:
        for r in run['records']:
            if r['checkpoint'] in seen:continue
            seen.add(r['checkpoint'])
            model=load_model(r,checkpoint_dir)
            validation=diag.paired(model,valid,r['mean'],r['scale'],r['window'],(64,96))
            key=(r['kind'],r['stage'],r['window'],r['n'],r['updates'],r['selection'])
            candidates.setdefault(key,[]).append(dict(record=r,validation=validation))
    selected={}
    for kind in diag.KINDS:
        options=[(key,rows) for key,rows in candidates.items() if key[0]==kind]
        assert all(len(rows)==3 for _,rows in options)
        key,rows=min(options,key=lambda pair:np.mean([r['validation']['recursive'] for r in pair[1]]))
        selected[kind]=dict(configuration=list(key),validation_mean=float(np.mean([r['validation']['recursive'] for r in rows])),rows=rows)
    config=dict(original_frequencies=[.025,.065],shifted_frequencies=[.035,.075],test_seed=6806,
                trajectories=32,origins=[64,96],horizon=64,selection='mean validation recursive MSE, three seeds, common origins 64/96',selected=selected)
    # Durable freeze BEFORE either test condition is generated/scored.
    (output/'frozen-selection.json').write_text(json.dumps(config,indent=2))
    tests={name:make_wave_trajectories(32,6806,freq) for name,freq in [('familiar',(.025,.065)),('shifted',(.035,.075))]}
    rows=[]
    for kind,choice in selected.items():
        for run in choice['rows']:
            r=run['record'];model=load_model(r,checkpoint_dir)
            rows.append(dict(kind=kind,seed=r['seed'],window=r['window'],n=r['n'],updates=r['updates'],
                             selection=r['selection'],selected_update=r['selected_update'],stage=r['stage'],
                             digest=r['digest'],validation=run['validation'],
                             **{name:diag.paired(model,raw,r['mean'],r['scale'],r['window'],(64,96)) for name,raw in tests.items()}))
    # Original training-fit ridge AR and persistence remain references, not newly tuned models.
    train=make_wave_trajectories(32,6411);mean,scale=float(train.mean()),float(train.std())
    x,y=diag.tensors((train-mean)/scale,32,np.arange(32,192))
    design=np.column_stack((x[:,:,0].numpy(),np.ones(len(x))))
    penalty=np.eye(33)*.001;penalty[-1,-1]=0
    coeff=np.linalg.solve(design.T@design/len(design)+penalty,design.T@y.numpy()/len(design))
    class Baseline(torch.nn.Module):
        def __init__(self,kind):super().__init__();self.kind=kind
        def forward(self,window):
            if self.kind=='Persistence':return window[:,-1,0]
            values=window[:,:,0].numpy()@coeff[:-1]+coeff[-1]
            return torch.tensor(values,dtype=torch.float32)
    for kind in ['Persistence','Linear AR']:
        rows.append(dict(kind=kind,seed=None,window=32,n=32,
                         **{name:diag.paired(Baseline(kind),raw,mean,scale,32,(64,96)) for name,raw in tests.items()}))
    controls=[]
    for r in [r for run in frozen['runs'] if run['stage']=='window' and run['window']==16 and run['kind']=='MLP' for r in run['records']]:
        model=load_model(r,checkpoint_dir)
        controls.append(dict(kind='MLP',seed=r['seed'],window=16,parameters=r['parameters'],
                             **{name:diag.paired(model,raw,r['mean'],r['scale'],16,(64,96)) for name,raw in tests.items()}))
    result=dict(window_control=controls,window_control_status='Exploratory matched-window audit added after inspecting main results; existing validation-selected checkpoints, no tuning.',config={k:v for k,v in config.items() if k!='selected'},selected={k:{p:v for p,v in choice.items() if p!='rows'} for k,choice in selected.items()},rows=rows)
    (output/'results.json').write_text(json.dumps(result,indent=2))
    print('| Model | W | Familiar observed | Shifted observed | Familiar rollout | Shifted rollout |')
    print('|---|---:|---:|---:|---:|---:|')
    for kind in ['Persistence','Linear AR',*diag.KINDS]:
        group=[r for r in rows if r['kind']==kind]
        metrics=[]
        for mode in ['observed','recursive']:
            for condition in ['familiar','shifted']:
                a=np.array([r[condition][mode] for r in group]);value=f'{a.mean():.5g}'
                if len(a)>1:value+=f' ± {a.std(ddof=1):.2g}'
                metrics.append(value)
        print('| '+kind+' | '+str(group[0]['window'])+' | '+' | '.join(metrics)+' |')
    print('Frozen choices:',json.dumps(result['selected'],ensure_ascii=False))
    return result

'''
if RETRAIN_WAVE_DIAGNOSTICS:
    from types import ModuleType
    shift_diagnostics = ModuleType('embedded_wave_diagnostics')
    exec(diagnostic_engine_source, shift_diagnostics.__dict__)
    shift_namespace = {'__name__': 'wave_shift', 'diag': shift_diagnostics}
    exec(frequency_shift_engine_source, shift_namespace)
    frequency_shift_record = shift_namespace['run'](diagnostic_output, diagnostic_output/'frequency-shift')

# @title Display paired frequency-shift results

for mode, title in [('observed','Observed-context next-sample MSE'),('recursive','Recursive 64-step MSE')]:
    shift_lines = [f'**{title}**', '', '| Predictor | W | Familiar | Shifted |', '|---|---:|---:|---:|']
    for kind in ('Persistence','Linear AR','MLP','RNN','GRU','1D CNN','CNN → RNN'):
        rows = [r for r in frequency_shift_record['rows'] if r['kind']==kind]
        columns = []
        for condition in ('familiar','shifted'):
            values = np.array([r[condition][mode] for r in rows])
            value = f'{values.mean():.4g}'
            if len(values)>1:value += f' ± {values.std(ddof=1):.2g}'
            columns.append(value)
        shift_lines.append(f'| {kind} | {rows[0]["window"]} | '+' | '.join(columns)+' |')
    display(Markdown('\n'.join(shift_lines)))

Observed-context next-sample MSE

Predictor	W	Familiar	Shifted
Persistence	32	0.02097	0.03206
Linear AR	32	0.001003	0.4251
MLP	64	0.001003 ± 8.4e-06	0.7337 ± 0.0012
RNN	16	0.001684 ± 3.3e-05	0.02206 ± 0.00091
GRU	16	0.001682 ± 1.9e-05	0.02263 ± 0.0013
1D CNN	32	0.001072 ± 1e-06	0.492 ± 0.055
CNN → RNN	32	0.001841 ± 1.9e-05	0.02195 ± 0.0044

Recursive 64-step MSE

Predictor	W	Familiar	Shifted
Persistence	32	0.8228	0.689
Linear AR	32	0.001036	1.097
MLP	64	0.001023 ± 4.1e-06	0.9097 ± 0.016
RNN	16	0.00681 ± 0.00086	1.048 ± 0.044
GRU	16	0.006976 ± 0.0014	1.063 ± 0.016
1D CNN	32	0.001161 ± 1.9e-05	1.049 ± 0.028
CNN → RNN	32	0.008559 ± 0.0019	0.9105 ± 0.2

Conclusion. On the shifted family, observed-context MSE is about 0.0221 for RNN, 0.0226 for GRU, and 0.0220 for CNN → RNN, versus 0.734 for the selected MLP, 0.492 for CNN, and 0.425 for AR. These selected recurrent recipes tolerate this shift better locally; the matched-window control below qualifies the size of that advantage. Every neural model still has much larger shifted recursive error, around 0.91–1.06. Better local predictions therefore do not establish successful long-range forecasting.

A comparison with shared history. CNN and CNN → RNN both use 32 samples, the original training pool, 3,000 updates, and validation rollout selection. Their shifted observed-context errors are 0.492 and 0.0220: about a 22-fold difference. This is evidence of better local generalization for the trained hybrid under this shift. Its recursive error remains large; the comparison does not promise an architecture advantage for other frequency families.

Check the window explanation. The selected MLP uses 64 samples; RNN/GRU use 16. After inspecting the main results, we evaluated the already trained, validation-selected 16-sample MLP on the same paired test sample. Its shifted observed-context MSE is 0.0285 ± 0.00043, much closer to RNN/GRU, while shifted rollout MSE is 1.016 ± 0.060. The 16-sample MLP has 289 parameters, versus 481 for RNN and 553 for GRU, so equal history does not establish equal capacity. This secondary control is exploratory; no checkpoint was tuned on the shifted test. Much of the large initial gap depends on the selected window/recipe. At equal 16-sample history, RNN/GRU show a smaller local-prediction advantage on this sample, while all three rollouts fail to follow the shifted wave accurately.

Interpretation checkpoint. Why can a predictor excel at familiar frequencies and struggle when they change? Why do restored observations help more than a recursive forecast?

Show explanation — generalization and repeated feedback are separate tests

A predictor can fit relationships specific to the training frequencies. Changing frequency changes those relationships, including the clean linear recurrence coefficients. A longer-window predictor can fit a frequency-specific pattern very well and still extrapolate poorly. That explanation is consistent with these results; the comparison does not uniquely identify the learned internal mechanism.

Observed-context prediction repeatedly receives the actual shifted wave. Recursive prediction must sustain the motion using its own guesses, so local errors can change later inputs. The frequency-shift experiment supports a limited strength of the tested recurrent configurations for local prediction under this particular shift. It provides no general architecture ranking or evidence of within-trajectory adaptation.

TODO(student): Compare the matched 16-sample predictors, then propose a separately frozen test of gradual within-trajectory drift or a different shift magnitude. State which measurements would support robust local prediction and which would support sustained forecasting. Neither extension has been checked here.

Replay. Default execution validates the stored synthetic record. To retrain and repeat the complete selection/test protocol, enable RETRAIN_WAVE_DIAGNOSTICS above; the optional runner uses its newly saved checkpoints. From a repository checkout, the faster checkpoint-only test is python3 scripts/week06-wave-frequency-shift.py --checkpoints /tmp/week06-diagnostics --output /tmp/week06-frequency-shift after running the diagnostic script. The main test protocol was frozen before scoring; the 16-sample MLP control is explicitly exploratory.

Final synthesis

The robot experiment asks whether a predictor preserves a specific earlier cue. Recall accuracy isolates that dependency; all-token loss also includes unpredictable wandering. A successful model must use the available information, while representing the uncertainty that remains.

The wave experiment asks a different question: can a finite noisy history sustain a continuous forecast? Fixed frequencies favor a strong linear rule, MLP, and CNN. Longer training greatly improves the recurrent models, but predictions fed back into the input still accumulate error. On the shifted frequencies, recurrent recipes make better local predictions than the selected long-window MLP and CNN; the matched-window control narrows that advantage, and every learned 64-step rollout remains poor.

Takeaway. Judge a trained architecture together with its observation window, training budget, selection rule, and evaluation task. Separate next-sample accuracy from sustained forecasting. A useful follow-up changes one declared condition and reserves fresh data for the final check.

Reproducibility and scope

The notebook runs top-to-bottom without dataset downloads. Its diagrams and plots are generated from code, with PNG/SVG exports, provenance, and text alternatives. Optional diagnostic retraining reproduces the stored synthetic evidence. No assessment rules, private solutions, or external figure assets are included.

print(f"Complete tutorial executed in {time.perf_counter() - started:.2f} seconds")
print("27 validation-selected checkpoints evaluated on the reserved test episodes.")
print("Figures exported to:", FIGURES.relative_to(root) if root else FIGURES)
# Small public synthetic result record for reproducibility; no model weights or private data.
comparison_record = {"initialization_seeds": FINAL_SEEDS,
    "runs": [dict(model=k, delay=d, seed=s, epoch=final_runs[(k,d,s)]["best_epoch"],
                  parameters=final_runs[(k,d,s)]["parameters"], **test_results[(k,d,s)])
             for k in FINAL_KINDS for d in FINAL_DELAYS for s in FINAL_SEEDS]}
# The displayed tables above are the notebook's durable results.

Complete tutorial executed in 265.46 seconds
27 validation-selected checkpoints evaluated on the reserved test episodes.
Figures exported to: assets/figures/shared/week06