# /// script # requires-python = ">=3.12" # dependencies = ["marimo==0.25.2", "matplotlib>=3.10,<4", "torch==2.14.1"] # /// """Open with `marimo edit notebook.py` from the extracted experiment bundle.""" import marimo __generated_with = "0.25.2" app = marimo.App(width="medium") @app.cell def _(): import hashlib import json from pathlib import Path import marimo as mo import matplotlib.pyplot as plt import torch import train_sms as sms return Path, hashlib, json, mo, plt, sms, torch @app.cell(hide_code=True) def _(mo): mo.md(r""" # What do the trained feature detectors read and write? Explore the article's **saved seed-42 model**, trained for 40 epochs on three word counts: **claim, prize, meeting**. No training or network access is needed to inspect these results. Keep the files from the download bundle beside this notebook. First examine the original responses. Then predict what will change when you disable a neuron, perform that intervention, and compare both the score and predicted class. The saved records are never overwritten. These are final-model activations, not snapshots from every training epoch. The real messages come from the [UCI SMS Spam Collection](https://archive.ics.uci.edu/dataset/228/sms%2Bspam%2Bcollection), Almeida & Hidalgo (2011), [CC BY 4.0](https://creativecommons.org/licenses/by/4.0/). """) return @app.cell def _(Path, hashlib, json, mo, sms, torch): notebook_dir = Path(mo.notebook_location()) saved = json.loads((notebook_dir / "example-results.json").read_text()) audit = json.loads((notebook_dir / "interpretability-results.json").read_text()) _archive_bytes = (notebook_dir / "interpretability-results-activations.json").read_bytes() assert hashlib.sha256(_archive_bytes).hexdigest() == audit["activation_archive"]["sha256"] archive = json.loads(_archive_bytes) assert archive["dataset_sha256"] == saved["dataset_sha256"] assert archive["parameters"] == saved["parameters"] records = archive["records"] assert len(records) == 5159 model = sms.ReadWriteClassifier() model.load_state_dict({key: torch.tensor(value) for key, value in saved["parameters"].items()}) model.eval() for _parameter in model.parameters(): _parameter.requires_grad_(False) _inputs = torch.tensor([row["counts"] for row in records]) torch.testing.assert_close(model.W(_inputs), torch.tensor([row["preactivations"] for row in records])) torch.testing.assert_close(model.W(_inputs).relu(), torch.tensor([row["activations"] for row in records])) torch.testing.assert_close(model(_inputs), torch.tensor([row["logit"] for row in records])) return archive, audit, model, notebook_dir, records, saved @app.cell(hide_code=True) def _(mo, saved): _m = saved["held_out"] mo.md(f""" ## The recorded result **{saved['train_count']:,} training / {_m['count']:,} held-out messages.** Accuracy: **{_m['accuracy']:.2%}**, compared with **{_m['always_nonspam_accuracy']:.2%}** for always predicting non-spam. Spam recall: **{_m['spam_recall']:.2%}** (25 of 128). **996 held-out messages have the same zero-count input.** They include 893 non-spam and all 103 missed spam messages. The model cannot distinguish messages with identical encodings. Its neuron labels describe its use of these three counts, not an understanding of message meaning. All saved activations and scores have been checked against the reconstructed model. """) return @app.cell def _(mo, plt, saved): _fig, _ax = plt.subplots(figsize=(7, 2.8)) _ax.plot(range(1, len(saved["train_loss"]) + 1), saved["train_loss"]) _ax.set(xlabel="Epoch", ylabel="Training cross-entropy", title="Recorded training loss") _fig.tight_layout() mo.vstack([ _fig, mo.ui.table([ {"Neuron": i, **dict(zip(saved["vocabulary"], row)), "Hidden bias": saved["parameters"]["W.bias"][i], "Outgoing weight": saved["parameters"]["V.weight"][0][i]} for i, row in enumerate(saved["parameters"]["W.weight"]) ], selection=None, label="Learned readers and writers"), mo.md(f"Output bias: **{saved['parameters']['V.bias'][0]:.4f}**.") ]) return @app.cell(hide_code=True) def _(mo): split = mo.ui.dropdown(["test", "train"], value="test", label="Data split") ordering = mo.ui.dropdown(["Neuron 0 activation", "Neuron 1 activation", "Original order"], value="Neuron 0 activation", label="Sort messages by") mo.vstack([mo.md("""## Find inputs that activate each neuron Choose a split and sort order, search the table, and select a row. The table contains every message in the selected split. Neuron 1's positive bias activates it even with no tracked words; meeting strengthens that response. """), mo.hstack([split, ordering])]) return ordering, split @app.cell def _(mo, ordering, records, split): split_records = [row for row in records if row["split"] == split.value] _ordered = sorted(split_records, key=lambda row: ( -row["activations"][0] if ordering.value == "Neuron 0 activation" else -row["activations"][1] if ordering.value == "Neuron 1 activation" else row["split_index"])) message_table = mo.ui.table([ {"Index": row["split_index"], "Message": row["message"], "Label": "spam" if row["label"] else "non-spam", "claim": row["counts"][0], "prize": row["counts"][1], "meeting": row["counts"][2], "h0": row["activations"][0], "h1": row["activations"][1], "Score": row["logit"], "Spam probability": row["probability"]} for row in _ordered ], selection="single", initial_selection=[0], page_size=10, label="Select a message to trace", show_column_summaries=False) message_table return message_table, split_records @app.cell def _(message_table, mo, model, plt, split_records): mo.stop(not message_table.value, mo.md("Select a message above.")) selected = next(row for row in split_records if row["split_index"] == message_table.value[0]["Index"]) _fig, _ax = plt.subplots(figsize=(7, 2.8)) _values = selected["contributions"] + [model.V.bias.item()] _ax.barh(["Neuron 0", "Neuron 1", "Output bias"], _values, color=["#d55e00", "#0072b2", "#666666"]) _ax.axvline(0, color="black", linewidth=.6) _ax.set(xlabel="Signed contribution to the spam score", title=f"Total score: {selected['logit']:.4f}") _fig.tight_layout() mo.vstack([ mo.plain_text(selected["message"]), mo.md(f"Counts: **{selected['counts']}**; true label: **{'spam' if selected['label'] else 'non-spam'}**."), mo.ui.table([ {"Neuron": i, "Before ReLU": selected["preactivations"][i], "After ReLU": selected["activations"][i], "Outgoing weight": model.V.weight[0, i].item(), "Contribution": selected["contributions"][i]} for i in range(2) ], selection=None), _fig ]) return (selected,) @app.cell(hide_code=True) def _(mo): expectation = mo.ui.text_area(label="Predict the effect on the score, probability, and class") intervention = mo.ui.dropdown(["None", "Neuron 0", "Neuron 1", "Both"], value="None", label="Disable") mo.vstack([mo.md("## Disable a neuron, with weights fixed"), expectation, intervention]) return expectation, intervention @app.cell def _(intervention, mo, model, selected, split_records, torch): removed = {"None": [], "Neuron 0": [0], "Neuron 1": [1], "Both": [0, 1]}[intervention.value] _h = torch.tensor([row["activations"] for row in split_records]) _changed_h = _h.clone() _changed_h[:, removed] = 0 intervention_logits = model.V(_changed_h).squeeze(-1) _original = torch.tensor([row["logit"] for row in split_records]) _labels = torch.tensor([row["label"] for row in split_records], dtype=torch.float32) _selected_index = next(i for i, row in enumerate(split_records) if row["split_index"] == selected["split_index"]) _after = intervention_logits[_selected_index].item() _removed_contribution = sum(selected["contributions"][i] for i in removed) assert abs(_after - (selected["logit"] - _removed_contribution)) < 1e-5 _comparison = [] for _name, _scores in [("Original", _original), ("Intervention", intervention_logits)]: _predictions = _scores > 0 _comparison.append({"Condition": _name, "Accuracy": (_predictions == _labels.bool()).float().mean().item(), "Spam recall": (_predictions & _labels.bool()).sum().item() / _labels.sum().item(), "Cross-entropy": torch.nn.functional.binary_cross_entropy_with_logits(_scores, _labels).item(), "Changed predictions": (_predictions != (_original > 0)).sum().item()}) mo.vstack([ mo.md(f"""Selected message: score **{selected['logit']:.4f} → {_after:.4f}**; spam probability **{selected['probability']:.2%} → {torch.sigmoid(torch.tensor(_after)).item():.2%}**. Predicted class: **{'spam' if selected['logit'] > 0 else 'non-spam'} → {'spam' if _after > 0 else 'non-spam'}**. The same intervention across the selected split:"""), mo.ui.table(_comparison, selection=None), mo.md("Disabling a neuron removes its signed contribution. A class can remain unchanged even when its probability changes substantially. No parameters are updated.") ]) return intervention_logits, removed @app.cell(hide_code=True) def _(mo): custom_message = mo.ui.text(value="Claim your prize at the meeting.", label="Your message", full_width=True) mo.vstack([mo.md("## Try a different combination of words"), custom_message]) return (custom_message,) @app.cell def _(custom_message, mo, model, records, sms): probe = sms.explain(model, [custom_message.value])[0] support = {name: sum(row["split"] == name and row["counts"] == probe["counts"] for row in records) for name in ("train", "test")} mo.vstack([ mo.md(f"""Counts: **{probe['counts']}**. Score: **{probe['logit']:.4f}**; spam probability: **{probe['spam_probability']:.2%}**. This count vector occurs **{support['train']}** times in training and **{support['test']}** times in testing. Only claim, prize, and meeting are recorded. Other words and word order cannot affect this model. A combination absent from the data is a synthetic probe; its score is not evidence of accuracy on real examples. """), mo.ui.table([{"Neuron": i, "Activation": probe["hidden_activations"][i], "Contribution": probe["neuron_contributions"][i]} for i in range(2)], selection=None) ]) return probe, support @app.cell(hide_code=True) def _(audit, mo): mo.vstack([ mo.md("""## Same predictions, different implementations These are saved measurements, not new training runs. The split and minibatch order were held fixed while initialization changed. A zero active count means inactive on the held-out inputs, not necessarily on every possible input. """), mo.ui.table([ {"Seed": row["initialization_seed"], "Accuracy": row["held_out"]["accuracy"], "Cross-entropy": row["held_out"]["cross_entropy"], "Active inputs, neuron 0": row["held_out_active_counts"][0], "Active inputs, neuron 1": row["held_out_active_counts"][1], "Outgoing weight 0": row["parameters"]["V.weight"][0][0], "Outgoing weight 1": row["parameters"]["V.weight"][0][1], "Output bias": row["parameters"]["V.bias"][0], "Different predictions": row["predictions_disagree_with_saved"]} for row in audit["fixed_split_initialization_sweep"] ], selection=None), mo.md("""Seed 42 has a positive-writing promotional neuron; seed 8 instead supplies spam evidence through a positive output bias. Both hidden neurons are off on its promotional inputs. Matching predictions do not establish matching mechanisms. This small experiment does not demonstrate superposition or show that nonlinear depth is needed: “claim or prize occurs” matches every held-out decision. It demonstrates how to inspect a trained read–activate–write computation. """) ]) return @app.cell(hide_code=True) def _(mo): mo.md(r""" ## Implementation and training details The cells above inspect saved results. The code below documents how they were obtained; opening this notebook does not rerun training. ### Two layers implement reading and writing ```python import torch from torch import nn from torch.nn import functional as F class ReadWriteClassifier(nn.Module): def __init__(self): super().__init__() self.W = nn.Linear(3, 2) # Two rows: one incoming weight vector per hidden neuron. self.V = nn.Linear(2, 1) # One row; one outgoing weight per hidden neuron. def forward(self, x): # x: [batch, 3]. Linear computes x @ W.weight.T + W.bias. h = F.relu(self.W(x)) # READ: [batch, 2] hidden activations. logits = self.V(h).squeeze(-1) # WRITE + SUM: [batch] classification scores. return logits # Positive favors spam; negative favors non-spam. ``` `model.W.weight` has shape `(2, 3)`: each row is a reader’s incoming weight vector. `model.V.weight` has shape `(1, 2)`: its two entries are the readers’ outgoing weights into the classifier. The biases have shapes `(2,)` and `(1,)`. The final `squeeze(-1)` removes only the singleton output axis, leaving one score per message even for a batch of size 1. Before training, the script checks the manual calculation with our chosen weights and zero biases, including $h=(4,0)$ and $s=4$. It also checks a non-spam message and examples on the decision boundary. Training then starts from fresh random parameters. ### Array dimensions describe how vectors are stored PyTorch stores a batch of examples as rows. For two messages and our three-word vocabulary, $$ X=\begin{bmatrix}2&1&0\\0&1&0\end{bmatrix}. $$ The shape is `(2, 3)`: two examples, each with three coordinates. The hidden activations have shape `(2, 2)`, and the output layer initially produces shape `(2, 1)`. Removing its last axis gives two scalar scores with shape `(2,)`. The matrix calculations are `X @ W.weight.T + W.bias` and `H @ V.weight.T + V.bias`. The transpose changes how the calculation is written for a batch, not what each neuron computes. ### Messages and labels become training batches The script downloads the [SMS Spam Collection](https://archive.ics.uci.edu/dataset/228/sms%2Bspam%2Bcollection), removes duplicate messages after normalizing case and whitespace, and reserves 20% of each class for testing. It counts the same three words as the manual example. Each message has one target: `0` for non-spam or `1` for spam. A `DataLoader` supplies batches of up to 64 examples. `word_counts` is a float tensor with shape `(batch, 3)`, and `labels` is a float tensor with shape `(batch,)`, as required by the binary loss. Labels come from the dataset, not from the presence of our selected words. Test messages are not used in optimizer updates. ### The annotated training loop Our tiny model runs on the CPU, so the tensors and parameters are already on the same device. ```python def train(model, train_loader, epochs, learning_rate): optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) loss_history = [] for epoch in range(epochs): model.train() loss_sum = 0.0 for word_counts, labels in train_loader: # 1. Reset: PyTorch accumulates gradients unless we clear them. optimizer.zero_grad() # 2. Forward: read with W, activate, write with V, add contributions. logits = model(word_counts) # 3. Loss: compare one raw score with each binary target (0 or 1). # This loss includes sigmoid internally; do not apply it first. loss = F.binary_cross_entropy_with_logits(logits, labels) # 4. Backward: calculate gradients for both matrices and biases. loss.backward() # 5. Update: Adam changes the parameters using those gradients. optimizer.step() loss_sum += loss.item() * len(labels) mean_loss = loss_sum / len(train_loader.dataset) loss_history.append(mean_loss) if epoch == 0 or (epoch + 1) % 10 == 0 or epoch + 1 == epochs: print(f"Epoch {epoch + 1:3d}: train cross-entropy = {mean_loss:.4f}") return loss_history ``` [`binary_cross_entropy_with_logits`](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.binary_cross_entropy_with_logits.html) compares one raw score with each binary target and includes the sigmoid calculation internally. Do not apply sigmoid before this loss. `loss.backward()` calculates derivatives, and `optimizer.step()` updates both weight matrices and both biases. `zero_grad()` clears accumulated gradients before the next batch. For predictions, the script uses `logits > 0`. This agrees with checking whether the sigmoid probability exceeds $0.5$. It assigns an exact zero to non-spam as a deterministic tie rule; the manual widgets instead label zero as the decision boundary. ### Inspect a trained neuron's contribution Training does not guarantee that the two neurons retain our hand-chosen meanings. We can still inspect what each one reads and writes. For the original message, calculate the two contributions explicitly: ```python model.eval() with torch.no_grad(): x = torch.tensor([[2., 1., 0.]]) h = F.relu(model.W(x)) contributions = h * model.V.weight[0] logits = contributions.sum(dim=1) + model.V.bias[0] torch.testing.assert_close(logits, model(x)) print("Activations:", h) print("Each neuron's contribution:", contributions) print("Output bias:", model.V.bias[0]) print("Classification score:", logits) print("Spam probability:", torch.sigmoid(logits)) ``` ### Run the example After downloading `train_sms.py`, run: ```bash python -m pip install torch python train_sms.py --epochs 40 --seed 42 --output results.json ``` The first run downloads the dataset. The script reports held-out accuracy, spam precision and recall, an always-non-spam baseline, and the number of messages whose three recorded counts are all zero. Its optional JSON report includes the learned parameters and decomposed predictions. ### How training adjusts incoming and outgoing weights The training loop updates both incoming and outgoing weights. Their gradients follow the [partial derivatives](https://indepth.dev/posts/2124/en/from-derivatives-to-backpropagation#multiple-inputs-and-partial-derivatives) and [chain rule](https://indepth.dev/posts/2124/en/from-derivatives-to-backpropagation#the-chain-rule-following-one-path) developed in the calculus article. To isolate the two weight updates, consider one reader with zero bias and an identity activation. Its contribution is $s=v(w^\top x)$, where $v$ is its scalar outgoing weight. Let $L$ be the loss and $g=\partial L/\partial s$. Holding $x$ fixed, the chain rule gives $$ \frac{\partial L}{\partial v}=(w^\top x)g, \qquad \nabla_w L=(vg)x. $$ The outgoing weight’s gradient is scaled by the current reader response. The incoming weights’ gradient depends on both the outgoing weight and the downstream loss derivative. Gradient descent subtracts a learning-rate-scaled gradient from each parameter, changing both what the neuron detects and what its response contributes. With a vector output, the same rules become $\nabla_v L=(w^\top x)g$ and $\nabla_w L=(v^\top g)x$, where $g=\nabla_z L$. The [Jacobian article](https://indepth.dev/posts/2131/en/jacobians-from-local-changes-to-neural-network-backpropagation#backpropagation-needs-jacobian-products) explains how these backward calculations combine derivatives through vector operations. In the ReLU model used by our training loop, the outgoing gradient is scaled by the activation, while the incoming gradient also includes the derivative of ReLU at the pre-activation. These gradients describe how the neuron receives a learning signal. Whether its learned response corresponds to a recognizable feature still needs evidence from the trained model. """) return if __name__ == "__main__": app.run()