Skip to main content

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 🔻

Figure: Complete dataflow inside the Transformer block, showing the transition from $\mathbb{R}$ to $\mathbb{Z}$ space, GyRoPE-infused attention, SAIL feed-forward, and residual routing.

⛵ 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

4 Layer SAIL 8 Layer SAIL

🎬 3D SAIL Phase-Amplitude Topography

3D 4 Layer SAIL 3D 8 Layer SAIL

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

Image
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.

Image *Figure: Training and validation loss trajectories across spaces ($\mathbb{R}, \mathbb{C}, \mathbb{H}, \mathbb{O}, \mathbb{S}$). Note how Complex-383k (green) consistently outperforms the larger Real-baseline (blue).*
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 🔻 image It's ridiculous, pointless and I can't even spell it... but it runs (very slowly). The most interesting usecases may be found on level 1-3 ;)



🚀 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 ttrain and tgen
  • PyPI Release: Official package published on PyPI

🧭 Outlook v0.3.0 - Tesseract

  • Logit Distillation Engine: DistillTrainer supporting dynamic teacher-student logit alignment (e.g., SmolLM2)
  • Synth Training Data: script to generate customized synthetic training data
  • Interactive TUI: Simple amazing terminal-based chat interface (using prompt-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

tardits-0.2.1.tar.gz (586.8 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

tardits-0.2.1-py3-none-any.whl (25.6 kB view details)

Uploaded Python 3

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

Hashes for tardits-0.2.1.tar.gz
Algorithm Hash digest
SHA256 4d17d741ab2b6367c07b39257f50b7359de6e4fd9e202d52902d3e4fe8167f40
MD5 1b673ac78037caf6f7853b0e670ff64e
BLAKE2b-256 c924c8e94ae52f4dcfc19040148c572718feb5ad8e95bb2b4d6fad785d8786c8

See more details on using hashes here.

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

Hashes for tardits-0.2.1-py3-none-any.whl
Algorithm Hash digest
SHA256 917595b1d453efe410d83ed30f23558db1db3f4ba4e41ea42bbe61c6500e5e8f
MD5 5e60073d2ab4e8a3c37b27063ded1fbf
BLAKE2b-256 36ea605ba20d02e44b94373c9ee1d11e41dc357d138483c8a1eec2d8ae7769b1

See more details on using hashes here.

Release history Release notifications | RSS feed

This release

0.2.1 This release

2 files

0.2.0

2 files

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page