trojan0x/ultron-small-baseline

Ultron-Small Baseline

Recurrent-Depth Transformer pretrained on FineWeb-Edu

Architecture

Ultron is a Recurrent-Depth Transformer (RDT) that loops a small set of transformer layers multiple times per forward pass, enabling "thinking in latent space" — implicit chain-of-thought reasoning without generating intermediate tokens.

Input → [Embedding + RoPE] → [Prelude (2 layers)] → [LN(e)] → [RecurrentBlock ×8 loops (4 layers/loop)] → [C·h] → [Coda (2 layers)] → [LM Head]

Property	Value
Total Parameters	88.9M
Non-embedding Parameters	50.4M
Effective Depth	36 layers (8 loops × 4 layers + 2 prelude + 2 coda)
Attention	GQA (12 heads, 4 KV heads)
Hidden Dimension	768
Sequence Length	1024
MoE	No (dense FFN)
Spectral Radius ρ(A)	< 1 by construction (LTI stability)

Training

• Dataset: FineWeb-Edu (sample-10BT)
• Optimizer: AdamW (β1=0.9, β2=0.95, ε=1e-8, weight decay=0.1)
• LR Schedule: Linear warmup (1000 steps) + cosine decay (3e-4 → 3e-5)
• Batch Size: 8 × 8 gradient accumulation = 64 effective
• Precision: bf16
• Depth Sampling: Parcae geometric per-sequence (μ=8)
• Gradient Checkpointing: Enabled

Key Research Papers

• Parcae (2026) — LTI-stable injection
• Loop, Think, & Generalize (2025) — Depth extrapolation
• Saunshi et al. (2025) — Looped = implicit CoT

Usage

from ultron.model import Ultron, UltronConfig
import torch

# Load from checkpoint
ckpt = torch.load("ultron_final.pt", weights_only=False)
cfg = UltronConfig(**ckpt["config"])
model = Ultron(cfg)
model.load_state_dict(ckpt["model_state_dict"])
model.eval()

# Inference with depth extrapolation
logits = model(input_ids, n_loops=8)   # standard
logits = model(input_ids, n_loops=16)  # deeper reasoning

Repository

trojan0x/ultron — Full source code, training scripts, and notebook