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

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
-
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support