OUROBOROS

OUROBOROS is a research prototype for a decoder-only language model that replaces additive residual updates with a learned geometric delta update:

x_next = x + beta * k * (v - <k, x>)

For a fixed unit direction k, the linear component has one eigenvalue 1 - beta along k and eigenvalue 1 on its orthogonal complement. This identity motivates the design; it is not evidence of better training, accuracy, forgetting, or gradient stability.

Important

This repository does not ship checkpoints, datasets, training recipes, or reproduced comparisons against standard residual networks. Empirical benefits remain unvalidated until controlled experiments are published.

What is implemented

• Hugging Face PreTrainedModel and PretrainedConfig integration
• causal scaled dot-product attention, RoPE, QK RMSNorm, and SwiGLU
• bounded per-token beta gates and scalar write targets
• optional causal one-token key/value shift reference
• tied token embedding / language-model head
• import, forward/backward, configuration, and block-size regression tests

Install and verify

Python 3.10+ is required.

python -m venv .venv
source .venv/bin/activate
pip install -e ".[test]"
pytest

Minimal example

import torch
from model import OuroborosConfig, OuroborosModel

config = OuroborosConfig(
    vocab_size=256,
    hidden_size=64,
    num_hidden_layers=2,
    num_attention_heads=4,
    head_dim=16,
    block_size=128,
)
model = OuroborosModel(config)
tokens = torch.randint(0, config.vocab_size, (2, 32))
logits, loss = model(tokens, targets=tokens)
loss.backward()
print(logits.shape, float(loss))

hidden_size must equal num_attention_heads * head_dim. Inputs longer than block_size fail explicitly. crop_block_size(n) only decreases that limit.

Opt-in residual diagnostics

Set ouroboros_collect_diagnostics=True in the config to collect per-layer metrics during a forward pass:

config = OuroborosConfig(
    vocab_size=256,
    hidden_size=64,
    num_hidden_layers=2,
    num_attention_heads=4,
    head_dim=16,
    ouroboros_collect_diagnostics=True,
)
model = OuroborosModel(config)
model(torch.randint(0, 256, (1, 16)))
print(model.get_residual_diagnostics())

Each attention and MLP residual reports beta range/mean, absolute gain along the updated axis |1-beta|, projection error before/after the write, and update norm. Collection is disabled by default because scalar extraction synchronizes accelerators. Metrics describe the realized update; they do not establish a quality improvement.

Allocation-aware update

The default geometric update uses torch.addcmul(x, k, delta) rather than materializing delta * k as a Python-level intermediate before addition. An explicit reference implementation remains in model.geometric, with output and backward-gradient parity covered by tests.

Measure the two paths on the target hardware:

python benchmarks/geometric_update.py --device cpu --iterations 200

On CUDA the script also reports peak allocated memory. The measurable acceptance target is equal outputs/gradients and no increase in peak allocation; this repository does not claim a hardware-independent latency improvement.

Research status and benchmark policy

The spectral statement above is an exact fixed-k, fixed-beta property. The full model is data-dependent and nonlinear. Any empirical claim should be backed by an additive-residual control with matched parameters/FLOPs, identical data order, multiple seeds, learning curves, downstream evaluation, and activation/gradient diagnostics.

The optional ShiftLinear is a zero-padded, one-token causal reference. Its streaming-cache interface is intentionally unsupported until cache correctness is tested.

License

Apache License 2.0. See LICENSE.