File size: 4,631 Bytes
209ab40 1078a5f 209ab40 1078a5f 209ab40 1078a5f 209ab40 1078a5f 209ab40 1078a5f 209ab40 1078a5f 209ab40 1078a5f 209ab40 1078a5f 209ab40 1078a5f 209ab40 1078a5f 209ab40 1078a5f 209ab40 1078a5f 209ab40 | 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 | ---
license: mit
language:
- code
library_name: sphere-attention
tags:
- poisson-attention
- hypersphere
- code-generation
- long-context
- sparse-attention
- research
datasets:
- codeparrot/codeparrot-clean
---
# POET-124M
124M-parameter POET (GPT-2 small size) — the scaling-trend checkpoint: the advantage grows, not shrinks, with size.
## Results
Second rung of the scaling study, parameter-matched to a GPT-2-style
baseline (123.62M vs 124.01M total; 85.02M non-embedding each), same 983M
tokens and schedule.
| scale | baseline best ppl | POET | paired margin |
|---|---|---|---|
| 51M | 2.938 | 2.828 | −2.9% (z = −68) |
| **124M (this model)** | **2.72** | **2.75 final / 2.60 EMA** | **−3.3% (z = −74)** |
The margin **held and slightly grew** across a 2.4x size jump — the opposite
of the usual pattern for efficient-attention variants, which tend to fade
with scale. A 354M rung is in progress.
**Recipe:** dim 768 / depth 12 / heads 12, seq 512, lr 6e-4, 60k steps,
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
```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-124M")
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": 768,
"depth": 12,
"heads": 12,
"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": "125m",
"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,
"cluster_reg": 0.0,
"abort_divergence": 1.5
}
```
|