POET-51M-2048

POET-51M-2048

51M-parameter POET trained at 2048 context on repository-grouped code — best perplexity of the family, and the long-context stability result.

Results

Paired eval vs modern-2k: −0.9% ppl (z = −9.6). Best in the family on both perplexity and generation correctness (+25% relative exact-match over GPT-2 on next-line completion, 200 held-out samples, greedy decode).

* Long-context stability finding. On this data the GPT-2 recipe diverged twice at the standard learning rate (auto-aborted at steps ~22k and 13k) and only completed at 3x lower LR — so its column carries that handicap. POET trained cleanly at the standard LR; its attention scores are bounded by construction (cosines through a bounded kernel), which appears to prevent the blow-up. The modern RoPE+RMSNorm recipe was also stable.

Certified sparse inference (prefix 1792 keys, budgeted cap attention):

Discarding ~93% of prefix attention costs ~2% perplexity.

Recipe: seq 2048, lr 4e-4, warmup 2000, 60k steps, repo-grouped data, selective weight decay, EMA 0.999.

What POET is

POET (POisson attEntion Transformer) replaces softmax attention with the closed-form spherical Poisson kernel. Queries and keys are L2-retracted onto the unit hypersphere and scored by

P_r(<q,k>) = (1 - r^2) / (1 - 2 r <q,k> + r^2)^(d/2)

with a learnable per-head radius r (resolution, replacing temperature), rotary positions (rotations are isometries of the sphere), and a QUEST-style query-norm sharpness term. Equivalently, each attention weight is the harmonic measure of where a random walk started at the query's interior point first exits the sphere at that key.

Because keys live on a sphere, attention admits a geodesic cap decomposition that supports budgeted sparse attention with closed-form per-query error certificates.

Architecture is otherwise a standard pre-LN transformer (GELU MLP, tied embeddings, GPT-2 BPE). Custom code — not transformers-compatible.

Usage

pip install git+https://github.com/Grayblock-AI/spherical-attention

import torch, tiktoken
from sphere_attention.hub import load_poet

model = load_poet("Grayblock-AI/POET-51M-2048")
enc = tiktoken.get_encoding("gpt2")
ids = torch.tensor([enc.encode_ordinary("def quicksort(arr):")])
out = model.generate(ids, max_new_tokens=64, top_k=20)
print(enc.decode(out[0].tolist()))

Training data

codeparrot/codeparrot-clean (permissively licensed Python), GPT-2 BPE. The long-context variants use repository-grouped packing: files from the same repo are packed contiguously so long sequences contain genuine cross-file structure (imports, call sites, definitions).

Limitations

• Research checkpoint, small scale. Not instruction-tuned; a raw next-token model. Python only.
• Trained on ~983M tokens, which is under compute-optimal for the larger sizes — absolute perplexities are compressed, though all comparisons in the results are matched arm-for-arm.
• Exact long-range identifier retrieval is near-zero at these model sizes (for POET and all baselines) — a capacity limit, not architecture.
• Sparse-attention certificates are sound (coverage 1.0) but conservative in the tail; see POET-51M-certified for the regularized variant and the gating mechanism that bounds the tail operationally.
• Poisson attention has no fused kernel yet, so wall-clock is ~2x a flash-attention baseline at equal FLOPs.

Links

• Code, full experiment log, paper draft: https://github.com/Grayblock-AI/spherical-attention
• Experiment log: https://github.com/Grayblock-AI/spherical-attention/blob/main/EXPERIMENTS.md

Citation

@misc{poet2026,
  title  = {POET: Poisson Attention Transformers with Certified Sparsity on the Hypersphere},
  author = {Jerge, Michael},
  year   = {2026},
  url    = {https://github.com/Grayblock-AI/spherical-attention}
}

Config

{
  "model": "sphere-ff",
  "dim": 512,
  "depth": 8,
  "heads": 8,
  "seq_len": 2048,
  "batch_size": 8,
  "steps": 60000,
  "warmup": 2000,
  "lr": 0.0004,
  "solver_iters": 4,
  "solver_tol": 0.001,
  "eval_every": 1000,
  "seed": 0,
  "grad_accum": 4,
  "no_bf16": false,
  "r_init": 0.7,
  "logit_scale": 10.0,
  "tag": "2k",
  "save_checkpoint": true,
  "checkpoint_every": 0,
  "s3_prefix": "",
  "data_dir": "data2k",
  "wandb": false,
  "wandb_project": "spherical-attention",
  "wandb_entity": null,
  "cloudwatch": false,
  "cloudwatch_region": "us-east-2",
  "resume": "",
  "ngpt": true,
  "selective_wd": true,
  "grad_steps": 1,
  "ema": 0.999,
  "abort_divergence": 1.5
}