v0.5.1 PyTorch 2.0+ MIT Built with PCG

Attention, solved—not softmaxed.

xaker fuses Exclusive Self Attention with kernel ridge regression and solves the regularised system with Preconditioned Conjugate Gradient. One PyTorch module. Deterministic benches, typed API, paper-grade rubrics.

Get started → Read the paper View source
xaker / Fused · kernel ridge regression · PCG v0.5.1 
# A mathematically grounded attention block.
from xaker import Config, Model
import torch

cfg  = Config(dim=512, heads=8, kernel="exp", precond="fast")
m    = Model(cfg, num_layers=6, vocab_size=32000,
                max_seq_len=512, attention_type="fused")

x      = torch.randint(0, 32000, (2, 128))
logits = m(x)

print(f"parameters: {sum(p.numel() for p in m.parameters()):,}")
print(f"logits shape: {logits.shape}")

01 / Problem framing

Softmax made attention easy to write. It also constrains the math.

Scaled dot-product attention scales scores by softmax and produces a positive, doubly-stochastic interaction. The construction is clean, but the matrix it operates on is dense, the kernel it builds is not positive semidefinite, and the spectral structure can deteriorate on long sequences.

The empirical consequence

On a sequence of length 128 with dim=64, the condition number of the score matrix reaches ~7.5 × 104. A regularised formulation working in the same basis reaches the high single digits. The ratio is not a constant; it widens with length.

xaker takes the second path. Attention is rewritten as a kernel ridge regression: (K + λI) α = v, solved iteratively. The diagonal of K is removed (Exclusive Self Attention) so each token's self-aligned component no longer dominates its own output. A preconditioner chosen from a fixed strategy set accelerates the iterate.

02 / The system

One package, six registries, one typed driver.

xaker is built around polymorphism, not mode flags. Four attention variants, four kernels, four preconditioners, three XSA modes, two iterative solvers, and a typed benchmark driver — all on the same nn.Module surface.

attention / block.py

BLOCK

Polymorphic registry for attention variants: standard, xsa, fused, linear. Adding a variant is one class plus one entry.

attention / kernel.py

Kernels

exp, rbf, linear, cosine. Stateless ops plus a learnable stateful wrapper for Fused.

solver / precond.py

Preconditioners

Identity, Diagonal, Fast, Cccp. One factory, Make(config), dispatches by string.

attention / xsa.py

XSA modes

Projection, Zero, Mask. Strategy triple driven by XsaStrategy(config, scale).

solver / cg.py

PCG solver

pcg and richardson over the regularised operator (K + λI). Returns a Solve dataclass with full history.

bench / bench.py

Bench driver

Typed Spec → run → write producing schema- stable JSON, environment block, per-seed statistics, git_sha.

cli /

Four CLIs

xaker-train, xaker-eval, xaker-bench, xaker-validate. Save and load a checkpoint; the train→eval round-trip is on by default.

rubric /

Rubric gate

Six dimensions enforce paper-worthiness in CI: novelty, repro, correctness, efficiency, stability, usability.

03 / How xaker works

The pipeline, end to end.

One pass through Fused.attend per head. The tokens become a kernel; the kernel becomes a regularised system; the system is solved.

01
Project
Q, K, V = Qkv(x)
Bias-free Q/K/V projection per head.
02
Kernelise
K = k(q, k)
Configurable: exp, rbf, linear, cosine.
03
XSA diagonal removal
K ← K − diag(K)
Each token loses its self-aligned component.
04
Ridge regularise
A(α) = K α + λ α
λ = softplus(raw_λ) + ε; guaranteed positive.
05
Preconditioner
P = Make(config)
Identity · Diagonal · Fast · Cccp.
06
Solve (PCG)
α ← pcg(K, v, λ, P)
Returns a Solve dataclass; dense fallback on miss.
07
Output
Y = rms(xsa.apply(α, V))
Clamp + RMS-norm + XSA strategy projection.
x.Fused.attend(q, k, v)  =  XsaStrategy.apply( α , v )  // α = (K + λI)⁻¹ v, solved by PCG(P)

Source: xaker/attention/fused.py. Math derivation: Mathematical foundations.

04 / Evidence

One number, then a chart, then a derived claim.

The headline condition-number ratio is reproducible from a single command. The chart on the right shows how the gap widens with sequence length. Every table on this page is regenerated from paper_runs/ JSONs on every commit.

κ ratio, L = 16
300×
kernel / softmax score
κ ratio, L = 32
591×
kernel / softmax score
κ ratio, L = 64
964×
kernel / softmax score
κ ratio, L = 128
1802×
kernel / softmax score

Condition number vs sequence length

log-scale. Lower is better. dim=64, lam=10.0, CPU.

Fused (kernel) Standard (score)
10⁵ 10⁴ 10³ 10² 16 32 64 128 256 kernel score

Reproduce with: python -m xaker.bench.condition --lam 10.0 --lengths 16 32 64 128 --out paper_runs/condition.json.

What changes when you switch

Side-by-side, in the same backbone, on the same hardware. Numbers below come from paper_runs/.

Conventional attention
xaker · Fused
/
kernel
softmax(QKT/√d)
learnable kernel k(q, k)
/
solver
single dense softmax
Preconditioned Conjugate Gradient
/
self-alignment
untreated (contributes to its own output)
diagonal removed (Projection · Zero · Mask)
/
κ at L = 128, dim = 64
~7.5 × 10⁴
~42
/
conditioning reg.
none
learnable λ; ridge via regularised operator
/
preconditioner
n/a
Identity · Diagonal · Fast · Cccp

05 / What ships

The package surface, as code.

Every public symbol is a single word. No _private names, no shim modules, no aliases. The polymorphic registry is the design.

Attention variants
  • StandardVaswani-style scaled dot-product attention.
  • XsaExclusive Self Attention with strategy dispatch.
  • FusedFlagship: XSA + kernel ridge regression via PCG.
  • LinearLinear-complexity baseline (Katharopoulos et al., 2020).
Kernel functions
  • expcosine-sim exponential, default; stateful + learnable.
  • rbfclassical Gaussian kernel σ = 1.
  • linearraw inner product; ridge compensates.
  • cosineL2-normalised inner product, range [−1, 1].
Preconditioners
  • IdentityP(r) = r; baseline.
  • Diagonallearned softplus-positive Jacobi.
  • Fastlearned low-rank + diagonal; cache-aware.
  • CccpTyler M-estimator; eigh-based inverse half-power.
Solvers and dispatch
  • pcgPreconditioned Conjugate Gradient; returns Solve.
  • richardsonfixed-iteration preconditioned Richardson.
  • BLOCKattention dispatch: standard · xsa · fused · linear.
  • Makepreconditioner dispatch by config.precond.
typed Config dataclass typed Spec / Result / Metrics four CLI entry points paper-worthiness rubric git_sha + schema-stable JSON single-word naming enforced in CI 12 paper_runs JSON artifacts

06 / Documentation

Pick a track. The doc system meets you where you are.

Four entry points. Each track tells you which docs to read first and which to skip.

Track 01

First-time user

I want to install xaker and run a forward pass.

Getting started Installation FAQ Troubleshooting Contributing
Track 02

Researcher

I want to understand the math, the pipeline, the numbers.

Tutorial Mathematical foundations Architecture Limitations Troubleshooting Contributing
Track 03

Contributor

I want to add a kernel, a preconditioner, a variant.

Recipes API reference Design decisions Troubleshooting Contributing
Track 04

Benchmark / validation

I want to evaluate xaker against my own baseline.

Paper rubric paper_runs/ RESULTS.md Troubleshooting Contributing

07 / Quickstart

Two commands to a benchmark. Three to a Transformer.

The working install is the source install (PyPI mirroring is planned). The headline benchmark is one reproducible command.

install · bash 
git clone https://github.com/sachncs/xaker.git
cd xaker
python3 -m venv .venv
source .venv/bin/activate
pip install -e '.[dev]'

The .[dev] is intentional; it means "this package, plus dev extras."

validate · bash 
xaker-validate
# novelty        2/3  novel=2/2 (fused.py,linear.py)
# repro          3/3  seeds=True bench=True paper_runs=12 cudnn=True
# correctness    3/3  fused=True cg=True dispatch=True
# efficiency     3/3  bench=True specs=True runs=12
# stability      3/3  seeds=True runs=True dtype=True
# usability      3/3  cli=True readme=True rubric=True
# Total          17/18 — PASS
forward pass · python 
from xaker import Config, Model
import torch

cfg = Config(dim=512, heads=8,
            kernel="exp", precond="fast")
m   = Model(cfg, num_layers=6, vocab_size=32000,
             max_seq_len=512, attention_type="fused")

logits = m(torch.randint(0, 32000, (2, 128)))
print(logits.shape)  # torch.Size([2, 128, 32000])

The same Config drives Standard, Xsa, Fused, and Linear through the BLOCK registry — no if mode == "fused" chains.

reproduce headline · bash 
python -m xaker.bench.condition \
    --lam 10.0 --lengths 16 32 64 128 \
    --out paper_runs/condition.json

# wrote paper_runs/condition.json
# ✓ headline condition numbers match the four metric cards above

Full tutorial: build a four-block Transformer step-by-step

08 / Design principles

Why the code looks the way it does.

A short manifesto on the choices behind xaker. These are defended in design_decisions.md; here are the four anchors.

i.

Polymorphism, not flags.

The four attention variants, four kernels, four preconditioners, and three XSA modes are all selected through a single dispatch. Adding a variant is one class plus one registry entry. Adding a fifth is the same amount of work.

ii.

The math is the spec.

docs/math.md and xaker/solver/precond.py agree byte for byte. If a derivation disagrees with the code, the code wins until a paper revision moves the goalposts.

iii.

Solver-backed, not softmax-backed.

Attention as a regularised system gives you a preconditioner knob. Identity, Diagonal, Fast, and Cccp cover the spectrum from debug-only to best-converges-on-ill-conditioned-kernels. Preconditioner choice is a config field, not a fork.

iv.

Numerical stability is a contract.

BOUND clamps. BOUND does not vary with the kernel. λ is guaranteed positive via softplus + ε. The dtype frontier is documented, not hidden.

v.

One public surface. One word per symbol.

_private prefixes are gone. Multi-word snake-case is gone. Aliases are gone. CI enforces this. The single-word rule keeps the API coherent across versions.

vi.

What we did not build, deliberately.

No FlashAttention. No sparse / Nyström. No AMP. No transformers-hub integration. Each of these is documented as absent and survives a paper revision. Roadmap, not TODO.

09 / Limitations

Where the math behaves, and where it doesn't.

Read this section before using xaker in production. The dtype frontier is a contract; the solver fallbacks are explicit; the unsupported environments are listed.

Where
Guard
dtype frontier
Notes
xaker/attention/linear.py
elu(x) + 1
fp16: x < −14 · bf16: x < −30
Per-dtype feature_clamp keeps the feature map strictly positive.
xaker/attention/func.py
exp(clamp(x, −100, 100))
fp16: x > ~50 · bf16: x > ~80
Same guard as PyTorch's softmax in low precision.
xaker/solver/precond.py
BOUND = 1e6
fp16-safe ceiling ~6.5e4
Lower BOUND to 1e4 for fp16-grade accuracy.
xaker/solver/cg.py
PCG fallback
matrix-dependent
Fused falls back to torch.linalg.solve on not converged ∧ finite.
xaker/solver/precond.py
Cccp cost
O(n³) per build
Beyond ~512 tokens prefer fast or diagonal.
PyTorch + MPS
Batched linalg
Apple Silicon
linalg.solve / eigh have shape bugs on 4-D inputs.
xaker/attention/linear.py
Position structure
task-dependent
Linear can't represent positions; fails on copy at length=32 (14%). Use Standard / Xsa / Fused for positional tasks.

Full limitations page

10 / Quality bar

This library is evaluated, not just presented.

The paper-worthiness rubric runs on every push to master and is enforced by CI. Each dimension below maps to a real grader that inspects the repo and emits an evidence string. The current run is presented here; it moves with the codebase.

noveltyof 3
2/3
novel=2/2 (fused.py, linear.py)
reproof 3
3/3
seeds=True bench=True paper_runs=12 cudnn=True
correctnessof 3
3/3
fused=True cg=True dispatch=True
efficiencyof 3
3/3
bench=True specs=True runs=12
stabilityof 3
3/3
seeds=True runs=True dtype=True
usabilityof 3
3/3
cli=True readme=True rubric=True
17 / 18 ● PASS Run on master; refreshed every CI build.

Rubric documentation

11 / Cite

Cite xaker.

The paper is in preparation. Until the arXiv identifier is assigned, cite the GitHub release matching the version you used. A canonical CITATION.cff ships at the repository root and is what GitHub's "Cite this repository" button reads.

cite · bibtex 
@software{xaker,
  title  = {xaker: A Mathematically Grounded Attention Framework},
  author = {sachin},
  year   = {2026},
  url    = {https://github.com/sachncs/xaker},
  note   = {Paper in preparation; arXiv ID will replace this entry on submission.}
}