Blame for README.md

---
license: mit
tags:
  - looped-transformer
  - linear-systems
  - dynamical-systems
---

# linsys-gf3

A 0.4M-parameter looped transformer that solves square, full-rank linear systems `A x = b`
over GF(3) with 8 unknowns. One transformer block (d_model 256, 8 heads, MLP ratio 1) is
applied repeatedly; the token embeddings are re-injected every iteration, the loop state is
renormalised each step, and it starts from random noise on the 8 answer slots:

    x_{k+1} = LN(x_k + r(x_k, e)),   r(x, e) = block(x + e) - (x + e)

Because the start state is random, the model defines a map from initial latent to outcome. Its
answer settle time (the last iteration at which the decoded answer changes) is fractal on
instances that need Gaussian elimination, even where every initial condition reaches the
correct answer.

## Files

* `model.pt` — PyTorch state dict (fp32), training step 298000 of 300000. This is what
  `loopscape.linsys` downloads.
* `config.json` — architecture, task and training settings.
* `checkpoints/pk8_b1024/` — the full training run: `model_<step>.pt` every 2000 steps from 0 to
  298000, `model_best.pt` (best held-out accuracy, saved near the end of training),
  `metadata.json` (exact run arguments), and `train_state_156000.pt` (optimizer and loop-schedule
  state, needed to resume from step 156000).

There is no dataset: training systems are sampled on the fly (`loopscape.linsys.task`).
The checkpoints in `checkpoints/` use the original training code's parameter names
(`transformer.block.block.*`); `LoopedTransformer.from_checkpoint` maps them.

## Use

```python
import numpy as np
from loopscape.linsys import basin, task

model = basin.load(device="cuda")           # downloads this checkpoint
A, b, x = task.sample_instance(np.random.default_rng(10))
answers, settle = basin.settle_field(model, A, b, res=192, loops=256)
```

Code: [GilpinLab/loopscape](https://github.com/GilpinLab/loopscape), module `loopscape.linsys`.
`python -m loopscape.linsys.train --out <dir>` reproduces the training run (~4 h on one H100).

## Training

300k steps, batch 1024, AdamW (lr 2e-4, betas 0.9/0.99, weight decay 1e-3 on matrices),
2000-step warmup, constant LR, 1-sqrt cooldown over the final 5%, bf16. Loop count per step is
drawn from a Poisson with a log-normal rate (mean 16, sigma 0.3); gradients flow through the
last 8 iterations. Training systems draw a per-instance row-weight cap `kmax ~ U{1..8}` and then
each row's weight `~ U{1..kmax}`; 5% of instances are held out by a hash of `(A, b)`.