tasmaikeni13/algebraic-transformer

The Algebraic Transformer: Transcendental-Free Language Modeling

The Algebraic Transformer is a causal language model whose entire forward, backward, and loss training path is strictly free of transcendental functions—avoiding exponentials (e^x), logarithms (ln(x)), trigonometric functions (sin(x), cos(x)), sigmoids (σ(x)), and error functions (erf(x)). Every operation in the network is arithmetic, comparison, or radical square root (ℚ[√·]).

Evaluated against the standard Transformer at 125M parameters on 2.5B tokens of FineWeb-Edu across 3 matched seeds, the Algebraic Transformer achieves lower held-out perplexity (22.78 vs 23.86), superior long-context retrieval, and high throughput using custom Pallas TPU kernels.

This repository provides the official model checkpoints, evaluation predictions, and comparison tables for all three evaluated architectures across seeds 42, 43, and 44.

---

Architecture Overview

Standard Transformers rely heavily on transcendental functions: softmax attention uses e^x, cross-entropy computes ln(x), position encodings use sinusoidal rotations (sin(x), cos(x)), and feed-forward networks use GELU/Swish activations (σ(x), erf(x)). In hardware, these operations require multi-cycle Special Function Units (SFUs) running CORDIC or polynomial approximations.

The Algebraic Transformer replaces these primitives with exact algebraic alternatives:

Subsystem	Standard Transformer	Pure Algebraic Transformer	Minimal-Algebra Hybrid
Attention Kernel	Softmax FlashAttention (exp)	Asinh-Power Attention ρ(s/8)¹⁶ with learned sink	Softmax FlashAttention (exp)
Attention Normalization	Standard Scaling (1/√d)	QK-Normalization (learned rsqrt)	QK-Normalization (learned rsqrt)
Position Encoding	Rotary Embeddings (RoPE: sin, cos)	Geometric Cayley Rotations	Rotary Embeddings (RoPE: sin, cos)
Activation Function	SwiGLU (x · σ(βx))	ALU-GLU (x · ALU₁/₂(y))	SwiGLU (x · σ(βx))
Normalization	RMSNorm (rsqrt) with learned gains	RMSNorm (rsqrt) with learned gains	RMSNorm (rsqrt) with learned gains
Vocabulary Head	Linear Logits + Softmax Cross-Entropy	Algebraic Maximum-Likelihood Head ρ(z/2)⁸	Algebraic Maximum-Likelihood Head ρ(z/2)⁸
Training Gradient	Transcendental (ln, exp)	Strictly Algebraic (ℚ[√·])	Strictly Algebraic (ℚ[√·])

Core Mathematical Components

1. Algebraic Maximum-Likelihood Head: Probabilities are defined entrywise by power scores:

$$p_j\propto \rho (z_j\mathrm{/}c{)}^n\text{where}\rho (x)=x+\sqrt{1+x^2}p\_j\; \backslash propto\; \backslash rho(z\_j\; /\; c)^n\; \backslash quad\; \backslash text\{where\}\; \backslash quad\; \backslash rho(x)\; =\; x\; +\; \backslash sqrt\{1\; +\; x^2\}$$

with $n=8$ and $c=2$. The log-likelihood loss gradient with respect to logit $z_j$ is rational-radical:

$$\frac{\mathrm{\partial }\mathcal{L}}{\mathrm{\partial }{z}_j}=\frac{n}{c}{(1+{\left(\frac{z_j}{c}\right)}^2)}^{-1\mathrm{/}2}(p_j-y_j)\backslash frac\{\backslash partial\; \backslash mathcal\{L\}\}\{\backslash partial\; z\_j\}\; =\; \backslash frac\{n\}\{c\}\; \backslash left(1\; +\; \backslash left(\backslash frac\{z\_j\}\{c\}\backslash right)^2\backslash right)^\{-1/2\}\; (p\_j\; -\; y\_j)$$

The model trains on exact likelihood without evaluating a logarithm. Furthermore, the entrywise asinh compression escapes the Softmax Rank Bottleneck (Yang et al., 2018), achieving infinite effective rank in log-probability space.

2. Asinh-Power Attention: Attention scores are computed as:

$$A_{ij}\propto \rho {\left(\frac{q_i{k}_j^{\mathrm{\top }}}{8\sqrt{d}}\right)}^{16}A\_\{ij\}\; \backslash propto\; \backslash rho\backslash left(\backslash frac\{q\_i\; k\_j^\backslash top\}\{8\; \backslash sqrt\{d\}\}\backslash right)^\{16\}$$

coupled with a learned per-head attention sink. An exact log-free max shift (ρ(a) · ρ(−b) = ρ(a ⊖ b)) prevents numerical overflow in Pallas TPU tiles.

3. Minimal-Algebra Hybrid Transformer: Adds only the two algebraic components that carry the empirical gain—QK-normalization and the algebraic likelihood head—to a standard Transformer backbone. It captures ~86% of the full algebraic gain while retaining the standard Softmax FlashAttention kernel.

---

125M / 2.5B-Token Empirical Results

All three models were trained on 2,500,000,000 tokens of FineWeb-Edu (sequence length 2,048, global batch size 512, AdamW cosine schedule) across 3 random seeds (42, 43, 44) on a Google Cloud TPU v4-32 pod.

1. Language Modeling & Throughput (Mean ± SD across 3 seeds)

Metric	Standard Transformer	Minimal-Algebra Hybrid	Pure Algebraic Transformer	Advantage vs Standard
Held-Out NLL (nats/token)	3.1722 ± 0.0024	3.1323 ± 0.0011	3.1258 ± 0.0026	−0.0464 nats (t > 25, p < 10⁻⁶)
Held-Out Perplexity	23.86 ± 0.06	22.93 ± 0.03	22.78 ± 0.06	−1.08 PPL drop
Training Time per Run	25.9 minutes	29.4 minutes	30.2 minutes	—
Throughput (16 TPU v4 chips)	1.62M tokens/s	1.43M tokens/s	1.39M tokens/s	86% speed parity

Every algebraic run beats every standard run by a wide, non-overlapping margin. In scaling-law terms, a 0.0464 nat drop at 125M is equivalent to a ~35–45% increase in compute-equivalent capacity.

---

2. Public Zero-Shot Benchmarks (Mean ± SD across 3 seeds)

Mean continuation-token log-probability (LAMBADA evaluated with greedy word prefix match):

Benchmark	Standard Transformer	Minimal-Algebra Hybrid	Pure Algebraic Transformer	Architectural Lead
ARC-Easy	43.71% ± 0.60%	44.74% ± 0.95%	43.69% ± 0.43%	Hybrid (+1.03%)
ARC-Challenge	24.77% ± 0.26%	24.83% ± 0.39%	25.11% ± 0.50%	Pure Algebraic (+0.34%)
HellaSwag	28.60% ± 0.19%	29.12% ± 0.16%	29.32% ± 0.18%	Pure Algebraic (+0.72%)
PIQA	61.68% ± 0.31%	62.08% ± 0.86%	62.44% ± 0.69%	Pure Algebraic (+0.76%)
OpenBookQA	27.93% ± 1.51%	27.87% ± 0.76%	28.00% ± 0.35%	Pure Algebraic
BoolQ	58.86% ± 3.70%	58.03% ± 3.81%	58.80% ± 2.32%	Standard
WinoGrande	52.59% ± 0.20%	51.62% ± 0.65%	51.64% ± 0.53%	Standard
LAMBADA	17.81% ± 0.24%	17.24% ± 0.60%	17.41% ± 0.63%	Standard

---

3. Long-Context Recall Matrix (Mean ± SD across 3 seeds)

Evaluated under the frozen bounded-context recall protocol (256, 512, 1,024, and 2,048 tokens across 5 needle depths):

Benchmark Task	Metric	Standard Transformer	Minimal Hybrid	Pure Algebraic	Key Finding
Multi-Needle	4-Choice Recognition	93.12% ± 4.71%	97.29% ± 2.98%	99.84% ± 0.16%	Algebraic is near-flawless (+6.72%)
Associative Recall	4-Choice Recognition	81.04% ± 3.55%	82.71% ± 0.70%	84.01% ± 1.66%	Algebraic leads (+2.97%)
	Exact Retrieval (All)	1.82% ± 1.73%	0.47% ± 0.16%	3.07% ± 1.35%	Algebraic leads in generative retrieval
	Exact @ 2,048 Tokens	0.42% ± 0.36%	0.62% ± 0.00%	1.88% ± 1.88%	Algebraic maintains long context
Latest-Value Recall	Exact Retrieval (All)	2.03% ± 1.02%	2.19% ± 1.08%	3.49% ± 2.42%	Algebraic leads in recency tracking
	Exact @ 2,048 Tokens	2.71% ± 1.57%	1.46% ± 1.30%	4.58% ± 3.44%	Algebraic leads (+1.87%)
Needle Passkey	4-Choice Recognition	96.15% ± 0.09%	94.64% ± 2.80%	96.93% ± 0.55%	Algebraic leads (+0.78%)
	Exact @ 2,048 Tokens	10.42% ± 6.29%	10.94% ± 1.33%	6.04% ± 10.46%	Hybrid / Standard
	Absent Evidence (False Pos)	0.00% ± 0.00%	0.00% ± 0.00%	0.00% ± 0.00%	Zero hallucinations across all arms

---

Checkpoint Organization

Checkpoints are organized by architecture arm and seed:

checkpoints/
├── standard_s42/     # Standard baseline (Seed 42)
├── standard_s43/     # Standard baseline (Seed 43)
├── standard_s44/     # Standard baseline (Seed 44)
├── hybrid_s42/       # Minimal-Algebra Hybrid (Seed 42)
├── hybrid_s43/       # Minimal-Algebra Hybrid (Seed 43)
├── hybrid_s44/       # Minimal-Algebra Hybrid (Seed 44)
├── algebraic_s42/    # Pure Algebraic Transformer (Seed 42)
├── algebraic_s43/    # Pure Algebraic Transformer (Seed 43)
└── algebraic_s44/    # Pure Algebraic Transformer (Seed 44)

Each checkpoint directory contains:

• params.npz: Complete trained weight arrays in JAX/NumPy format.
• meta.json: Full model configuration and training hyperparameter arguments.
• result.json: Exact evaluation metrics, NLL, throughput, and training seconds.

Full post-training zero-shot and recall outputs are stored in evaluations/.

---

Loading & Inference

Weights can be loaded directly using NumPy and JAX:

import numpy as np
from huggingface_hub import hf_hub_download

# Download weights for the Pure Algebraic Transformer (Seed 42)
path = hf_hub_download(
    repo_id="tasmaikeni13/algebraic-transformer",
    filename="checkpoints/algebraic_s42/params.npz"
)

# Load weight dictionary
weights = dict(np.load(path))
print("Loaded parameter keys:", list(weights.keys())[:5])

---

Hardware & Pallas Kernels

The models were trained on Google Cloud TPU v4 hardware using custom Pallas TPU kernels:

• Fused Backward Attention: Computes $dQ, dK, dV$ simultaneously inside TPU Vector Memory (VMEM) tiles, eliminating high-bandwidth memory (HBM) materialization of attention matrices.
• Algebraic Speed Parity: Pure Algebraic attention achieves 6.32 ms forward+backward latency on TPU v4 (vs 4.28 ms for Standard Softmax and 22.14 ms for JAX default FlashAttention), reaching 86% end-to-end throughput parity with standard Transformers.

---

Acknowledgements

We thank Google's TPU Research Cloud (TRC) for generously providing the Cloud TPU v4-32 pod (16 chips) used to train and evaluate all models in this study.

---

Citation & License

This project is licensed under the Apache 2.0 License. Codebase: https://github.com/tasmaikeni13/algebraic-transformer