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 whatloopscape.linsysdownloads.config.jsonโ architecture, task and training settings.checkpoints/pk8_b1024/โ the full training run:model_<step>.ptevery 2000 steps from 0 to 298000,model_best.pt(best held-out accuracy, saved near the end of training),metadata.json(exact run arguments), andtrain_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
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, 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).
- Downloads last month
- -