File size: 4,794 Bytes
495b681 5648744 495b681 5648744 495b681 5648744 495b681 5648744 495b681 5648744 495b681 5648744 495b681 5648744 495b681 5648744 495b681 5648744 495b681 5648744 495b681 5648744 495b681 5648744 495b681 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | ---
license: mit
language:
- code
library_name: sphere-attention
tags:
- poisson-attention
- hypersphere
- code-generation
- long-context
- sparse-attention
- research
datasets:
- codeparrot/codeparrot-clean
---
# POET-51M
51M-parameter POET (seq 512) — the reference model for the headline comparisons.
## Results
Matched parameters, data, optimizer, and token budget (983M tokens);
paired evaluations use identical batches, so the z-scores are large.
| model | non-emb params | best val ppl | paired vs this model |
|---|---|---|---|
| GPT-2-style baseline | 25.2M | 2.938 | POET **−2.9%** (z = −68) |
| nanoGPT (vendored) | 25.2M | 2.938 | POET **−3.0%** (z = −72) |
| Modern (RMSNorm/SwiGLU/RoPE) | 25.1M | 2.875 | POET **−0.7%** (z = −20) |
| HF `LlamaForCausalLM` | 25.1M | 2.849 | +0.1% (parity) |
| **POET-51M (this model)** | **25.2M** | **2.828** | — |
Replicated across 3 seeds (2.828 / 2.836 / ~2.85); seed spread is ~10x
smaller than the margin. FLOP-comparable to the baselines.
**Recipe:** seq 512, lr 6e-4, cosine schedule, 60k steps x 16,384
tokens/step, selective weight decay, EMA 0.999 (best-checkpoint numbers
above are non-EMA for like-for-like comparison; EMA best is 2.783).
## 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
```bash
pip install git+https://github.com/Grayblock-AI/spherical-attention
```
```python
import torch, tiktoken
from sphere_attention.hub import load_poet
model = load_poet("Grayblock-AI/POET-51M")
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
```bibtex
@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
```json
{
"model": "sphere-ff",
"dim": 512,
"depth": 8,
"heads": 8,
"seq_len": 512,
"batch_size": 32,
"steps": 60000,
"warmup": 1000,
"lr": 0.0006,
"solver_iters": 4,
"solver_tol": 0.001,
"eval_every": 1000,
"seed": 0,
"grad_accum": 2,
"no_bf16": false,
"r_init": 0.7,
"logit_scale": 10.0,
"tag": "1b",
"save_checkpoint": true,
"checkpoint_every": 0,
"s3_prefix": "",
"data_dir": "data1b",
"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
}
```
|