code is folded so the argument and saved outputs come first. expand any cell to inspect the implementation. saved outputs are historical evidence from the source notebook, not a new run or independent verification.
attention language model¶
Start from the same Tiny Shakespeare token stream as the bigram model, then run explicit multi-head causal self-attention on the embedded batch.
inspect code · cell 01
import random
import sys
from pathlib import Path
import torch
from torch import nn
from torch.nn import functional as F
repo_root = Path.cwd()
if not (repo_root / "data").exists():
repo_root = repo_root.parent
sys.path.insert(0, str(repo_root))
from experiments.notebook_theme import apply_notebook_theme
from src.dataset import get_batch, load_tiny_shakespeare_tokens, split_token_stream
apply_notebook_theme()
prepare token streams and batches¶
inspect code · cell 02
tokens, vocab, _ = load_tiny_shakespeare_tokens(repo_root / "data")
train_tokens, validation_tokens = split_token_stream(tokens)
vocab_size = len(vocab)
block_size = 8
batch_size = 32
n_embd = 32
num_heads = 2
head_size = n_embd // num_heads
random.seed(42)
x_batch, y_batch = get_batch("train", train_tokens, validation_tokens, block_size, batch_size)
x_batch = torch.tensor(x_batch, dtype=torch.long)
y_batch = torch.tensor(y_batch, dtype=torch.long)
embed the current batch¶
inspect code · cell 03
token_embedding_table = nn.Embedding(vocab_size, n_embd)
x = token_embedding_table(x_batch)
assert x.shape == (batch_size, block_size, n_embd)
tuple(x.shape)
multi-head causal self-attention¶
inspect code · cell 04
class MultiHeadAttention(nn.Module):
def __init__(self, n_head, d_model, d_k, d_v):
super().__init__()
self.n_head = n_head
self.d_k = d_k
self.d_v = d_v
self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)
self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)
self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False)
self.output_projection = nn.Linear(n_head * d_v, d_model, bias=False)
def forward(self, query, key, value):
d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
B, len_q, _ = query.shape
_, len_k, _ = key.shape
_, len_v, _ = value.shape
q = self.w_qs(query).view(B, len_q, n_head, d_k)
k = self.w_ks(key).view(B, len_k, n_head, d_k)
v = self.w_vs(value).view(B, len_v, n_head, d_v)
q = q.transpose(1, 2)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
# scaled dot-product attention
# q @ k^T: (B, h, T, d_k) @ (B, h, d_k, T)
# -> (B, h, T, T)
scores = q @ k.transpose(-2, -1)
scores = scores / (d_k**0.5)
# causal mask
causal_mask = torch.triu(
torch.ones(len_q, len_k, device=scores.device, dtype=torch.bool), diagonal=1
)
scores = scores.masked_fill(causal_mask, float("-inf"))
# normalize, then retrieve values
weights = torch.softmax(scores, dim=-1) # (B, h, T, T)
output = weights @ v # (B, h, T, d_v)
# put heads beside each other again
output = output.transpose(1, 2).contiguous() # (B, T, h, d_v)
output = output.view(B, len_q, n_head * d_v) # (B, T, h*d_v)
output = self.output_projection(output) # (B, T, C)
return output, weights
position-wise feed-forward network¶
inspect code · cell 05
class PositionWiseFeedForward(nn.Module):
def __init__(self, d_model, d_ff):
super().__init__()
self.w_1 = nn.Linear(d_model, d_ff)
self.w_2 = nn.Linear(d_ff, d_model)
def forward(self, x):
return self.w_2(F.relu(self.w_1(x)))
attention shape check¶
inspect code · cell 06
x = torch.randn(batch_size, block_size, n_embd)
mha = MultiHeadAttention(n_head=num_heads, d_model=n_embd, d_k=head_size, d_v=head_size)
out, weights = mha(x, x, x)
assert out.shape == x.shape
assert weights.shape == (batch_size, num_heads, block_size, block_size)
tuple(out.shape)
combine attention and feed-forward transformations with post-normalization: LayerNorm(x + Sublayer(x)), matching section 3.1.
the saved experiment results below were produced with the earlier pre-norm block. they have not been retrained after this change; use source revision 27127de to replay those checkpoints faithfully.
inspect code · cell 07
class TransformerBlock(nn.Module):
def __init__(self, d_model, n_head, d_k, d_v, d_ff):
super().__init__()
self.ln1 = nn.LayerNorm(d_model)
self.attn = MultiHeadAttention(n_head, d_model, d_k, d_v)
self.ln2 = nn.LayerNorm(d_model)
self.ff = PositionWiseFeedForward(d_model, d_ff)
def forward(self, x):
# Normalize after adding the attention update to the residual stream.
attn_out, weights = self.attn(x, x, x)
x = self.ln1(x + attn_out)
# Normalize after the feed-forward residual addition as well.
x = self.ln2(x + self.ff(x))
return x, weights
block shape check¶
inspect code · cell 08
block = TransformerBlock(
d_model=n_embd,
n_head=num_heads,
d_k=head_size,
d_v=head_size,
d_ff=4 * n_embd,
)
out, weights = block(x)
print(x.shape)
print(out.shape)
print(weights.shape)
train a one-block language model¶
the shared next-token loss trains all registered model parameters together. this notebook uses BPE tokens, not individual characters. define the loop here; the implemented one-block wrapper is in src/language_models.py. plain minibatch SGD exposes the update rule directly; the optimizer and learning rate are experimental choices.
inspect code · cell 09
def train_language_model(model, *, steps, learning_rate, eval_every=100, eval_batches=10):
"""Train a model mapping token ids [B, T] to raw vocabulary logits [B, T, V]."""
if min(steps, eval_every, eval_batches) < 1 or learning_rate <= 0:
raise ValueError("steps, evaluation counts, and learning rate must be positive")
device = next(model.parameters()).device
# plain minibatch SGD makes the update rule we discussed explicit.
# construct once: every registered trainable component shares this objective.
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
history = []
def batch_loss(split):
inputs, targets = get_batch(split, train_tokens, validation_tokens, block_size, batch_size)
inputs = torch.tensor(inputs, dtype=torch.long, device=device)
targets = torch.tensor(targets, dtype=torch.long, device=device)
logits = model(inputs)
# one prediction per position; flatten B and T without mixing vocabulary scores.
# cross_entropy takes raw logits, so do not apply softmax first.
return F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
model.train()
for step in range(1, steps + 1):
optimizer.zero_grad(set_to_none=True) # discard the previous batch's gradients
loss = batch_loss("train") # forward pass and shared next-token objective
loss.backward() # chain rule: calculate gradients throughout the model
optimizer.step() # SGD: change parameters using those gradients
if step == 1 or step % eval_every == 0 or step == steps:
model.eval() # evaluation behavior; this alone does not disable gradients
with torch.no_grad():
validation_loss = (
sum(batch_loss("validation").item() for _ in range(eval_batches)) / eval_batches
)
model.train()
# training loss is this update's batch loss, measured before the update.
history.append(
{"step": step, "train_batch_loss": loss.item(), "validation_loss": validation_loss}
)
print(f"step {step}: train batch {loss.item():.4f}, validation {validation_loss:.4f}")
return history
# the implemented wrapper lives in src/language_models.py:
# token ids -> token + position embeddings -> TransformerBlock (unpack its pair)
# -> final layer norm -> vocabulary projection -> raw logits.
# all components are registered so model.parameters() includes them.
# TODO: choose steps and learning_rate explicitly, then call train_language_model.
# rerunning the function continues the supplied model's current weights.
# TODO: inspect both losses; a lower training batch loss alone is not generalization.
compare bigram and one-block outputs¶
evaluate both trained models on identical sampled training and validation windows, then generate from the same prompt. delta means transformer minus bigram cross-entropy (nats per BPE token); negative favors the transformer. this is a between-model gap, not a before/after training change.
record each training budget and optimizer when interpreting results: the bigram notebook uses AdamW, while the loop above uses SGD. these models also have different parameter counts; this comparison alone does not isolate architecture.
inspect code · cell 10
def compare_language_models(
bigram_model, transformer_model, *, prompt_ids, new_tokens=100, loss_batches=10, seed=42
):
"""Compare trained models using shared BPE ids and identical loss-evaluation batches."""
from src.tokenizer import bpe_decode
if not prompt_ids or new_tokens < 0 or loss_batches < 1:
raise ValueError("provide a nonempty prompt and valid generation/evaluation counts")
models = {"bigram": bigram_model, "single transformer": transformer_model}
# use a local RNG so comparison does not change future training batch sampling.
rng = random.Random(seed)
batches = {}
for split, stream in [("train", train_tokens), ("validation", validation_tokens)]:
batches[split] = []
for _ in range(loss_batches):
starts = [rng.randrange(len(stream) - block_size) for _ in range(batch_size)]
inputs = [stream[i : i + block_size] for i in starts]
targets = [stream[i + 1 : i + block_size + 1] for i in starts]
batches[split].append((inputs, targets))
def get_logits(model, ids):
result = model(ids)
# the existing bigram returns (logits, loss); the new wrapper returns logits.
return result[0] if isinstance(result, tuple) else result
results = {}
for name, model in models.items():
device = next(model.parameters()).device
was_training = model.training
model.eval()
try:
with torch.no_grad():
losses = {}
for split, shared_batches in batches.items():
total = 0.0
for inputs, targets in shared_batches:
inputs = torch.tensor(inputs, dtype=torch.long, device=device)
targets = torch.tensor(targets, dtype=torch.long, device=device)
logits = get_logits(model, inputs)
total += F.cross_entropy(
logits.reshape(-1, logits.size(-1)), targets.reshape(-1)
).item()
losses[split] = total / loss_batches
ids = torch.tensor([prompt_ids], dtype=torch.long, device=device)
sampler = torch.Generator(device="cpu").manual_seed(seed)
for _ in range(new_tokens):
# keep the transformer inside its context window; bigram uses the last token.
logits = get_logits(model, ids[:, -block_size:])[:, -1, :]
probabilities = torch.softmax(logits, dim=-1).cpu()
next_id = torch.multinomial(probabilities, 1, generator=sampler).to(device)
ids = torch.cat((ids, next_id), dim=1)
results[name] = {
"train_loss": losses["train"],
"validation_loss": losses["validation"],
"text": bpe_decode(ids[0].tolist(), vocab, errors="replace"),
"parameters": sum(p.numel() for p in model.parameters()),
}
finally:
model.train(was_training)
print("same prompt, temperature 1, same sampling seed; one sample is qualitative evidence")
for name, result in results.items():
print(f"\n{name} ({result['parameters']:,} parameters):\n{result['text']}")
print(
f"train loss: {result['train_loss']:.4f}; validation loss: {result['validation_loss']:.4f}"
)
for metric in ("train_loss", "validation_loss"):
delta = results["single transformer"][metric] - results["bigram"][metric]
print(f"{metric} delta (transformer - bigram): {delta:+.4f} nats/token")
return results
# the next cell loads both trained models and calls this comparison.
paired control and treatment experiment¶
three paired seeds; identical batches and shared-component initialization; 10,000 uninterrupted AdamW updates each, rerun from the same initial seeds. the treatment adds learned positions, one transformer block, and final normalization. equal updates are not equal compute or parameter counts. validation loss is the primary outcome; these exploratory runs are not evidence of universal superiority. generated weights and results stay local under artifacts/.
only the training budget changed from the 1,000-update run. compare both budgets below: a widening, stable, or closing gap are all informative outcomes. full checkpoints include optimizer and random-generator state; the original short-run artifacts remain available.
inspect code · cell 11
import json
from IPython.display import Image, display
from src.language_models import BigramLanguageModel, SingleTransformerLanguageModel
# reproduce all three paired runs from the repository root with:
# uv run python -m experiments.compare_bigram_transformer --steps 10000 \
# --output artifacts/bigram-vs-transformer-10k
experiment_dir = repo_root / "artifacts" / "bigram-vs-transformer-10k"
report = json.loads((experiment_dir / "results.json").read_text())
print(report["config"])
# compare fixed evaluation windows at the two preselected training budgets.
for budget in (1000, 10000):
for split in ("train", "validation"):
means = {}
for name in ("bigram", "transformer"):
values = [
h[split]
for run in report["runs"]
if run["model"] == name
for h in run["history"]
if h["step"] == budget
]
means[name] = sum(values) / len(values)
print(
f"{budget:5d} updates, {split}: "
f"bigram={means['bigram']:.4f}, transformer={means['transformer']:.4f}, "
f"delta={means['transformer'] - means['bigram']:+.4f} nats/token"
)
for seed in report["config"]["seeds"]:
pair = {r["model"]: r for r in report["runs"] if r["seed"] == seed}
for split in ("train", "validation"):
control = pair["bigram"]["history"][-1][split]
treatment = pair["transformer"]["history"][-1][split]
print(
f"seed {seed} {split}: bigram={control:.4f}, transformer={treatment:.4f}, "
f"delta={treatment - control:+.4f} nats/token"
)
from experiments.compare_bigram_transformer import plot_loss_report
plot_loss_report(report, experiment_dir)
display(Image(filename=str(experiment_dir / "loss.png")))
# inspect the first preselected seed, rather than selecting the best-looking run.
seed = report["config"]["seeds"][0]
bigram_model = BigramLanguageModel(vocab_size, n_embd)
single_transformer_model = SingleTransformerLanguageModel(vocab_size, n_embd, block_size, num_heads)
bigram_model.load_state_dict(torch.load(experiment_dir / f"bigram-{seed}.pt", weights_only=True))
single_transformer_model.load_state_dict(
torch.load(experiment_dir / f"transformer-{seed}.pt", weights_only=True)
)
comparison = compare_language_models(
bigram_model, single_transformer_model, prompt_ids=train_tokens[:8]
)
text at 1k, 2.5k, 5k, and 10k updates¶
same prompt and sampling seed at each checkpoint, with 100 new BPE tokens. compare changes within each model as well as between models; the samples illustrate behavior, while the loss curves provide broader evidence.
inspect code · cell 12
from experiments.plot_checkpoint_text import plot_checkpoint_text
# reproduce the preselected seed's checkpoints, without selecting attractive samples:
# from experiments.compare_bigram_transformer import run
# run(steps=10000, seeds=(42,), output=str(repo_root / "artifacts/bigram-vs-transformer-milestones"))
milestone_dir = repo_root / "artifacts" / "bigram-vs-transformer-milestones"
figure_path = plot_checkpoint_text(milestone_dir, seed=42)
display(Image(filename=str(figure_path)))
compose transformer blocks¶
each block refines the preceding representation with separate parameters. depth and block dimensions are configurable. the stack returns final embeddings and a list of attention maps in layer order; vocabulary projection remains the language-model wrapper’s responsibility. the reusable version is in src/language_models.py.
inspect code · cell 13
class TransformerStack(nn.Module):
"""Compose independent blocks; return embeddings and attention maps in layer order."""
def __init__(self, n_layers, d_model, n_head, d_k, d_v, d_ff):
super().__init__()
if n_layers < 1:
raise ValueError("n_layers must be positive")
self.blocks = nn.ModuleList(
[TransformerBlock(d_model, n_head, d_k, d_v, d_ff) for _ in range(n_layers)]
)
def forward(self, x):
# x stays [batch, tokens, d_model] throughout the stack.
attention_weights = []
for block in self.blocks:
x, weights = block(x)
attention_weights.append(weights)
return x, attention_weights
inspect code · cell 14
# use the same block objects to verify sequential composition.
stack = TransformerStack(
n_layers=2,
d_model=n_embd,
n_head=num_heads,
d_k=head_size,
d_v=head_size,
d_ff=4 * n_embd,
)
# dummy embeddings let us check the wiring without loading text.
stack_input = torch.randn(2, block_size, n_embd)
stack_output, layer_weights = stack(stack_input)
first_output, _ = stack.blocks[0](stack_input)
explicit_output, _ = stack.blocks[1](first_output)
torch.testing.assert_close(stack_output, explicit_output)
assert stack_output.shape == stack_input.shape
assert len(layer_weights) == 2
# compare object identities, not parameter values: each layer learns independently.
first_block_parameter_ids = {id(parameter) for parameter in stack.blocks[0].parameters()}
second_block_parameter_ids = {id(parameter) for parameter in stack.blocks[1].parameters()}
shared_parameter_ids = first_block_parameter_ids.intersection(second_block_parameter_ids)
assert len(shared_parameter_ids) == 0, "blocks must have separate parameters"
print("stack output:", tuple(stack_output.shape))
attention_shapes = [tuple(weights.shape) for weights in layer_weights]
print("attention per layer:", attention_shapes)
print("composition and independent-parameter checks passed")
all treatments: matched 10k-update comparison¶
compare bigram, one transformer block, and two transformer blocks across seeds 42, 43, and 44. all models see the same batches and fixed evaluation windows. width 32, context 8, batch size 32; AdamW at 0.001 with weight decay 0.01. shared components start at identical values, with separate parameter objects. the second block adds parameters and compute.
question and pre-run hypothesis, controls, checkpoints, and reproduction command are in experiments/depth-comparison.md. validation loss is primary. the curves show the mean and full seed range, not confidence intervals. samples use the first preselected seed and are qualitative. full optimizer/RNG checkpoints and source hashes are saved locally.
inspect code · cell 15
import json
from pathlib import Path
from statistics import mean
from IPython.display import Image, display
comparison_root = Path.cwd()
if not (comparison_root / "artifacts").exists():
comparison_root = comparison_root.parent
all_treatments_directory = comparison_root / "artifacts" / "all-treatments-10k"
all_treatments_report = json.loads((all_treatments_directory / "results.json").read_text())
treatment_names = ["bigram", "transformer", "two_blocks"]
print("final cross-entropy: nats per BPE token; update time excludes evaluation")
for treatment_name in treatment_names:
treatment_runs = [
run for run in all_treatments_report["runs"] if run["model"] == treatment_name
]
training_losses = [run["history"][-1]["train"] for run in treatment_runs]
validation_losses = [run["history"][-1]["validation"] for run in treatment_runs]
update_times = [run["update_seconds"] for run in treatment_runs]
parameter_count = treatment_runs[0]["parameters"]
print(
f"{treatment_name}: parameters={parameter_count:,}, "
f"train={mean(training_losses):.4f}, validation={mean(validation_losses):.4f}, "
f"validation range=[{min(validation_losses):.4f}, {max(validation_losses):.4f}], "
f"mean update seconds={mean(update_times):.1f}"
)
print("\npaired validation losses and depth differences")
for seed in all_treatments_report["config"]["seeds"]:
losses_by_treatment = {}
for run in all_treatments_report["runs"]:
if run["seed"] == seed:
losses_by_treatment[run["model"]] = run["history"][-1]["validation"]
depth_difference = losses_by_treatment["two_blocks"] - losses_by_treatment["transformer"]
print(f"seed {seed}: {losses_by_treatment}; two minus one={depth_difference:+.4f}")
from experiments.compare_bigram_transformer import plot_loss_report
plot_loss_report(all_treatments_report, all_treatments_directory)
display(Image(filename=str(all_treatments_directory / "loss.png")))
print("\nseed 42 samples: same prompt, temperature 1, sampling seed 123")
for run in all_treatments_report["runs"]:
if run["seed"] == 42:
print(f"\n{run['model']}\n{run['sample']}")
extending depth to 4, 6, and 8 blocks¶
same three seeds, 10,000 updates, width 32, context 8, data, AdamW settings, and fixed evaluation windows. new models start from matching shared components and overlapping block prefixes. prior bigram/1/2-block runs are retained and loaded alongside these results; source and control compatibility are checked before combining.
the pre-run question is whether validation improvements flatten or reverse with additional depth. more blocks add parameters and compute. timings across the two runs are approximate. the plots show means and seed ranges, not confidence intervals. see experiments/depth-comparison.md for the experimental record and reproduction command.
inspect code · cell 16
import sys
from pathlib import Path
from IPython.display import Image, Markdown, display
comparison_root = Path.cwd()
if not (comparison_root / "artifacts").exists():
comparison_root = comparison_root.parent
sys.path.insert(0, str(comparison_root))
from experiments.plot_depth_comparison import (
load_depth_reports,
plot_depth_comparison,
summarize_depths,
)
depth_runs = load_depth_reports(comparison_root / "artifacts")
depth_summary = summarize_depths(depth_runs)
table_lines = [
"| treatment | parameters | train loss | validation loss | seed range | update seconds |",
"|---|---:|---:|---:|---|---:|",
]
for row in depth_summary:
table_lines.append(
f"| {row['treatment']} | {row['parameters']:,} | {row['train_loss']:.4f} | "
f"{row['validation_loss']:.4f} | "
f"{row['validation_min']:.4f}–{row['validation_max']:.4f} | "
f"{row['update_seconds']:.1f} |"
)
display(Markdown("\n".join(table_lines)))
display(
Markdown(
"losses are means over three seeds, in nats per BPE token. update time excludes evaluation and checkpointing."
)
)
depth_figure_path = comparison_root / "artifacts" / "depth-4-6-8-10k" / "combined-depths.png"
plot_depth_comparison(depth_runs, depth_figure_path)
display(Image(filename=str(depth_figure_path)))
10k treatment samples, side by side¶
same prompt and preselected training seed 42, sampling seed 123, temperature 1, and 100 generated BPE tokens. each panel shows its individual seed's validation loss; the table above reports three-seed means. spelling and line breaks are retained. these samples illustrate behavior, not overall quality.
inspect code · cell 17
from experiments.plot_depth_comparison import plot_treatment_text
treatment_text_path = plot_treatment_text(comparison_root / "artifacts", seed=42, step=10000)
display(Image(filename=str(treatment_text_path)))
context training and cross-evaluation¶
four blocks, width 32, contexts 8/32/128/256; three seeds and 10,000 updates. each update uses the same 256 target tokens, split at different sequence boundaries. evaluation scores the same 128 positions per split using the requested preceding history. loss units are nats per BPE token. these metrics use a different protocol from the earlier depth comparison.
at completion, each model is evaluated at every supported context length. weights stay fixed during cross-evaluation; cropped inputs use learned positions starting at zero. native-context curves show mean and seed range, not confidence intervals. tables retain each seed. generated samples use training seed 42, sampling seed 123, temperature 1, and a shared 256-token prompt truncated to the evaluation context.
experiment record: experiments/context-comparison.md. the cell below runs the experiment if the output directory has no results, or loads completed results. to rerun from scratch, choose a fresh output directory. results and interpretation will be walked through together; no conclusions are supplied here.
inspect code · cell 18
import sys
from pathlib import Path
from IPython.display import Image, Markdown, display
context_repo_root = Path.cwd()
if not (context_repo_root / "src").exists():
context_repo_root = context_repo_root.parent
sys.path.insert(0, str(context_repo_root))
from experiments.compare_context_lengths import run as run_context_comparison
from experiments.plot_context_comparison import (
cross_evaluation_table,
plot_measurements,
plot_text_panels,
read_report,
summary_table,
)
context_output_directory = context_repo_root / "artifacts" / "context-comparison-10k"
if not (context_output_directory / "results.json").exists():
# The runner loads data relative to the repository root.
import os
previous_directory = Path.cwd()
try:
os.chdir(context_repo_root)
run_context_comparison(steps=10000, seeds=(42, 43, 44), output=context_output_directory)
finally:
os.chdir(previous_directory)
context_report = read_report(context_output_directory)
display(Markdown("### native-context results: three-seed means"))
display(Markdown(summary_table(context_report)))
display(Markdown("### cross-evaluation: validation means"))
display(Markdown(cross_evaluation_table(context_report)))
for seed in context_report["config"]["seeds"]:
display(Markdown(f"### cross-evaluation: seed {seed}"))
display(Markdown(cross_evaluation_table(context_report, seed=seed)))
context_measurement_figure = plot_measurements(context_report, context_output_directory)
display(Image(filename=str(context_measurement_figure)))
inspect code · cell 19
native_context_text = plot_text_panels(context_report, context_output_directory)
display(Image(filename=str(native_context_text)))
frozen_model_text = plot_text_panels(
context_report, context_output_directory, frozen_long_model=True
)
display(Image(filename=str(frozen_model_text)))
reflection: context and loss¶
i expected the loss to go down as context increased, but it didn’t. my first thought was that there isn’t a tidy linear relationship between context and loss. more specifically, though, i realized i don’t yet understand what that relationship actually is.
the connection that makes sense to me is that you can have memories and still be wrong about how you apply them. having more information available doesn’t necessarily mean using it correctly. that gives me a way to think about the result, even if it doesn’t yet explain what happened in this experiment.
sinusoidal positional encoding¶
implementation companion (assistant): attention is all you need, section 3.5, page 6.
for position pos and feature pair i, even features use sin(pos / 10000**(2*i/d_model)) and odd features use cos(pos / 10000**(2*i/d_model)). the position vector has the same width as the token embedding, so the wrapper adds them before the transformer stack.
the module below computes fixed vectors on demand. its frequency scales are a registered buffer, not trainable parameters. both language-model wrappers accept position_encoding="sinusoidal"; "learned" remains the default for existing runs and checkpoints. select the same encoding when reconstructing a model for checkpoint loading.
this implements section 3.5 only: token-embedding scaling and the then-current pre-norm blocks were unchanged in that experiment. position ids still restart at zero when inputs are cropped. computing positions beyond the training length is possible; good predictions at those lengths are not guaranteed. the demonstration below is an untrained wiring check, not a new treatment result.
inspect code · cell 20
import importlib
import sys
from pathlib import Path
import torch
from IPython.display import Markdown, display
position_repo_root = Path.cwd()
if not (position_repo_root / "src").exists():
position_repo_root = position_repo_root.parent
sys.path.insert(0, str(position_repo_root))
import src.language_models as position_models
# Refresh the reusable implementation in an already-running notebook kernel.
position_models = importlib.reload(position_models)
inspect code · cell 21
position_encoder = position_models.SinusoidalPositionalEncoding(d_model=32)
position_ids = torch.arange(4)
position_vectors = position_encoder(position_ids)
# Show the first four sine/cosine pairs at positions 0 through 3.
position_table = "| position | " + " | ".join(f"feature {i}" for i in range(8)) + " |\n"
position_table += "| --- | " + " | ".join(["---"] * 8) + " |\n"
for position, vector in zip(position_ids.tolist(), position_vectors):
feature_values = " | ".join(f"{value:.4f}" for value in vector[:8].tolist())
position_table += f"| {position} | {feature_values} |\n"
display(Markdown(position_table))
# Synthetic token ids check the full path without loading data or training.
with torch.random.fork_rng(devices=[]):
torch.manual_seed(42)
sinusoidal_model = position_models.StackedTransformerLanguageModel(
vocab_size=16,
n_layers=4,
n_embd=32,
block_size=256,
position_encoding="sinusoidal",
)
with torch.no_grad():
sinusoidal_logits = sinusoidal_model(torch.arange(8).unsqueeze(0))
assert sinusoidal_logits.shape == (1, 8, 16)
assert torch.isfinite(sinusoidal_logits).all()
position_parameter_count = sum(
parameter.numel() for parameter in sinusoidal_model.position_embedding_table.parameters()
)
display(
Markdown(
f"fixed position parameters: **{position_parameter_count}**. "
f"untrained logits shape: **{tuple(sinusoidal_logits.shape)}** "
"(batch, tokens, vocabulary)."
)
)
bigram → six-block transformer → sinusoidal encoding¶
three paired arms: bigram; six transformer blocks without explicit positions; the same six blocks with sinusoidal positional encoding. context 32, width 32, two heads, seeds 42/43/44, 10,000 updates. each update uses the same 256 targets in eight sequences of 32 for every arm. token embeddings and vocabulary heads start identically; the transformers also share initial block and final-norm weights, with independent parameters.
the first contrast adds the transformer architecture, parameters, and access to history. the second isolates adding sinusoidal position vectors, with equal trainable parameter counts. causal masks remain in both transformers. the bigram uses only the final input token when scoring each target.
loss is cross-entropy in nats per BPE token on the same 1,024 sampled targets per split. tables show means, seed ranges (not confidence intervals), and paired validation differences. negative differences mean lower loss for the first named treatment. no figures or performance interpretation supplied. full methods and checkpoint schedule: experiment record.
run from the repository root:
uv run python -m experiments.compare_position_encodings --steps 10000 --output artifacts/position-encoding-six-blocks-10k
inspect code · cell 22
import json
import sys
from pathlib import Path
from IPython.display import Markdown, display
position_comparison_root = Path.cwd()
if not (position_comparison_root / "src").exists():
position_comparison_root = position_comparison_root.parent
sys.path.insert(0, str(position_comparison_root))
from experiments.compare_position_encodings import result_tables
position_result_path = (
position_comparison_root / "artifacts/position-encoding-six-blocks-10k/results.json"
)
position_comparison_report = json.loads(position_result_path.read_text())
position_summary, position_paired = result_tables(position_comparison_report)
display(Markdown("three-seed mean losses (nats per BPE token)\n\n" + position_summary))
display(Markdown("paired validation losses and planned contrasts\n\n" + position_paired))
matched text samples: bigram, six blocks, and positional encoding¶
same 32-token prompt, training seed 42, sampling seed 123, temperature 1, and 100 new BPE tokens from each 10,000-update checkpoint. labels show the individual seed's validation loss, not the three-seed mean. continuations are preserved verbatim with visual wrapping. raw samples and checkpoint references are saved beside the figure in artifacts/position-encoding-six-blocks-10k/text-samples.json.
inspect code · cell 23
import sys
from pathlib import Path
from IPython.display import Image, display
sample_repo_root = Path.cwd()
if not (sample_repo_root / "src").exists():
sample_repo_root = sample_repo_root.parent
sys.path.insert(0, str(sample_repo_root))
from experiments.plot_position_samples import plot_position_samples
position_sample_figure = plot_position_samples(
sample_repo_root / "artifacts/position-encoding-six-blocks-10k", seed=42
)
display(Image(filename=str(position_sample_figure)))