README.md

---
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
}
```