| --- |
| 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)`. |
|
|
