Text And Rotating Dimensions In Token Space
or: how to make LMs bigger on the inside
Welcome to TARDITS, a library exploring (tiny) Complex-Valued LM architectures. It's my take on the awesome Let's build GPT tutorial. After finishing my GPT-net and updating it (RoPE, swiGLU, RMSNorm...), I got the idea to write a complex transformer. Mainly because I asked myself if they really need an activation function... turns out they don't 🤯
After a few iterations, the $\mathbb{C}$-net got surprisingly stable. Of course, I had to try quaternions ($\mathbb{H}$) next. They were built using the Cayley-Dickson construction — so I pass 2 complex values through the transformer and let PyTorch do all the backprop work. Then the concept escalated a little, and the Octonion ($\mathbb{O}$) and Sedenion ($\mathbb{S}$) nets also somehow happened. Both work, but they take a really long time to train. Especially on CPU :D
There may be better uses for CVNNs than predicting text... but why not Shakespeare?
🙃 Architecture
Click to expand 🔻
⛵ SAIL (Self-Activating Interference Layer)
The main idea. In standard $\mathbb{R}$-nets, non-linear activation functions (like GELU or SwiGLU) in the feed-forward layer are strictly necessary. However, multiplying two complex numbers is inherently non-linear in itself. SAIL is essentially SwiGLU without the "Swi": It splits the input via a Linear Layer into two parts and multiplies them together. That alone is sufficient to get a fully trainable complex transformer. Combined with residual connections z = z + SAIL(z), my guess is it could also be chaotic somehow (think Mandelbrot set).
🪢 GyRoPE - learnable RoPE, full circle
While implementing experimenting with RoPE in my $\mathbb{R}$-net, I found that a trainable RoPE initialized randomly from $-2\pi$ to $2\pi$ yielded results comparable to classical positive positional encodings. This proved to work nicely with the complex transformers. However the weights are few and they change really slowly during training.
- $\mathbb{C}$-RoPE rotates in the $2\mathcal{D}$ complex plane.
- $\mathbb{H}$-RoPE rotates around a unit 3-sphere ($S^3$).
- $\mathbb{O}$ and $\mathbb{S}$ rotate around their respective hyperspherical counterparts.
👀 Attention
How do you route attention in complex spaces? I experimented with various variants, but turns out: simplicity wins. Attention scores are computed via the Hermitian dot product ($Q K^\dagger$). By extracting only the real component before applying the causal mask and standard Softmax:
# In C-Space:
wei = (q @ k.adjoint()) * (2 * self.head_size) ** -0.5
wei = wei.real
wei = F.softmax(wei.masked_fill(self.tril[:T, :T] == 0, float("-inf")), dim=-1)
# Weighted value aggregation remains hypercomplex:
out = wei.to(v.dtype) @ v
The real part perfectly captures the phase alignment and magnitude between queries and keys after multiplication, acting as a natural scalar affinity score. The resulting real weights are then used to scale the full complex/hypercomplex value vectors ($V$). No other variant yielded a better validation loss so far.
🏋️ Weight Initialization
Attention and SAIL weights are initialized around the unit circle/sphere and scaled by:
$$\sqrt{\frac{S}{d_{in} + d_{out}}}$$
(where $S$ is $\mathbb{C}=1.0$, $\mathbb{H}=0.5$, $\mathbb{O}=0.25$, $\mathbb{S}=0.125$)
🛌 Embedding & LM-Head
To map tokens into hypercomplex spaces, I used separate embedding tables - one for each real component (e.g., 4 tables for $\mathbb{H}$, 16 for $\mathbb{S}$) and initialized them with a scaled normal distribution (Var = 1.0):
# Quaternion Embedding (4D)
z1 = torch.complex(emb_r1(x), emb_i1(x))
z2 = torch.complex(emb_r2(x), emb_i2(x))
---
nn.init.normal_(emb.weight, std=(1 / math.sqrt(4)))
To project back to vocabulary logits, the LMHead simply flattens all complex components back to $\mathbb{R}$:
features = torch.cat([z1.real, z1.imag, z2.real, z2.imag], dim=-1)
return self.lin(features)
🛡️ RMSNorm - Norming Amplitudes, Preserving Phases
🧨 remove for exploding transformer noises (NaN) 💥
Standard normalization layers can easily distort geometric phase information in complex spaces.
The custom RMSNorm calculates the Root Mean Square solely based on vector magnitudes ($\Vert{}q\Vert{}^2 = \vert{}z_1\vert{}^2 + \vert{}z_2\vert{}^2$):
# Calculate quaternion amplitude squared
ampsq = (z1 * z1.conj()).real + (z2 * z2.conj()).real
rms = torch.rsqrt((ampsq / 2.0).mean(dim=-1, keepdim=True) + self.eps)
return (z1 * rms) * self.weight, (z2 * rms) * self.weight
This scales feature amplitudes cleanly across layers while leaving phase angles completely intact.
🔥 Hot Curves (Benchmarks & Results)
Click to expand 🔻
🌊 SAIL Interference Patterns ($\mathbb{C}$)
🎬 2D SAIL Interference Evolution
🎬 3D SAIL Phase-Amplitude Topography
Figure: Evolution of selective phase interference across all 8 Transformer layers over training steps. Deep layers (6 & 7) demonstrate extreme destructive phase cancellation alongside razor-sharp resonance peaks (> 10,000 amplitude). That's why the Post-Norm design makes a lot of sense here
📉 Comparison Runs
shared parameters:
| batch_size | block_size | dropout | learning rate | vocab_size |
|---|---|---|---|---|
| 32 | 128 | 0.2 | 0.0005 | 65 |
$\mathbb{C}$ vs $\mathbb{R}$ / both ~180k params
| Model | Space | $n_{embd}$ | Layers | Heads | Params | Train Loss | Val Loss |
|---|---|---|---|---|---|---|---|
| ComplexGPT | $\mathbb{C}$ | 54 | 4 | 6 | ~179k | 1.472 | 1.704 |
| RealGPT | $\mathbb{R}$ | 60 | 4 | 6 | ~181k | 1.526 | 1.759 |
$\mathbb{R}$ vs the complex rest
Number of parameters scaled. Note: the higher dimensional runs take forever. A "Sedenion" needs 256 multiplications instead of one.
| Model | Space | $n_{embd}$ | Layers | Heads | Params | Train Loss | Val Loss |
|---|---|---|---|---|---|---|---|
| RealGPT | $\mathbb{R}$ | 96 | 5 | 6 | ~567k | 1.367 | 1.645 |
| ComplexGPT | $\mathbb{C}$ | 72 | 5 | 6 | ~384k | 1.359 | 1.633 |
| QuaternionGPT | $\mathbb{H}$ | 54 | 4 | 4 | ~356k | 1.367 | 1.633 |
| OctonionGPT | $\mathbb{O}$ | 32 | 4 | 4 | ~264k | 1.398 | 1.668 |
| SedenionGPT | $\mathbb{S}$ | 18 | 4 | 6 | ~184k | 1.460 | 1.702 |
v0.2.0 Bonus - Centumduodevigintunions 🔻
🚀 Quickstart & Usage
Install
pip install tardits
# or
uv add tardits
Optional - for MLFlow experiment tracking:
pip install tardits[mlflow]
Quickstart
Example configuration (save as .yaml):
device: "auto" # Choices: "auto", "cuda", "cpu", "mps"
seed: 1337
model:
level: 1 # Complex
block_size: 64
n_head: 6
n_embd: 42
n_layer: 4
dropout: 0.2
trainer:
batch_size: 32
max_iters: 5000
eval_interval: 500
learning_rate: 0.0005
eval_iters: 200
save_interval: 0
checkpoint: "saved_model.pt"
mlflow:
experiment: "" # add name to enable
log_interval: 10
deep_log_interval: 0
system_metrics: false
data:
input_path: "my_training_data.txt"
vocab_save_path: "saved/charset.txt"
You can also take a config from [examples/configs] and the shakespeare.txt form [examples/data]
Start Training
ttrain -c config.yaml
Generate Text
tgen -c config.yaml -p "To be or not to be" -l 200
tgen -h
options:
-h, --help show this help message and exit
--config, -c CONFIG Path to YAML configuration file
--checkpoint CHECKPOINT
Optional override for model checkpoint path
--vocab VOCAB Optional override for vocab path
--prompt, -p PROMPT Text prompt to start generation
--length, -l LENGTH Number of tokens to generate
--temp, -t TEMP Model sampling temperature
--top_k, -k TOP_K Top-k sampling threshold
--follow, -f Endless stream of tokens
--delay DELAY Delay between tokens in follow mode (seconds)
v0.2.0 (Current Release)
- Dynamic Dimension Scaling: Universal recursive Cayley-Dickson engine ($2^k$ dimensions, up to 128D Centumduodevigintunion)
- Package Refactoring: Clean library structure for distribution
- CLI Tools: Easy executions via
ttrainandtgen - PyPI Release: Official package published on PyPI
🧭 Outlook v0.3.0 - Tesseract
- Logit Distillation Engine:
DistillTrainersupporting dynamic teacher-student logit alignment (e.g., SmolLM2) - Synth Training Data: script to generate customized synthetic training data
- Interactive TUI:
Simpleamazing terminal-based chat interface (usingprompt-toolkit) - Basic Inference API + docker/compose
- Tesseract Model Series: Releasing small-footprint complex chat models (2D/4D) on HuggingFace
- Getting distracted by SAIL-experiments for possible future releases 🤫🧪⛵
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file tardits-0.2.1.tar.gz.
File metadata
- Download URL: tardits-0.2.1.tar.gz
- Upload date:
- Size: 586.8 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via:
uv/0.12.0 {"installer":{"name":"uv","version":"0.12.0","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Arch Linux","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
4d17d741ab2b6367c07b39257f50b7359de6e4fd9e202d52902d3e4fe8167f40
|
|
| MD5 |
1b673ac78037caf6f7853b0e670ff64e
|
|
| BLAKE2b-256 |
c924c8e94ae52f4dcfc19040148c572718feb5ad8e95bb2b4d6fad785d8786c8
|
File details
Details for the file tardits-0.2.1-py3-none-any.whl.
File metadata
- Download URL: tardits-0.2.1-py3-none-any.whl
- Upload date:
- Size: 25.6 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via:
uv/0.12.0 {"installer":{"name":"uv","version":"0.12.0","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Arch Linux","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
917595b1d453efe410d83ed30f23558db1db3f4ba4e41ea42bbe61c6500e5e8f
|
|
| MD5 |
5e60073d2ab4e8a3c37b27063ded1fbf
|
|
| BLAKE2b-256 |
36ea605ba20d02e44b94373c9ee1d11e41dc357d138483c8a1eec2d8ae7769b1
|