RingTransformer — Custom MLX Architecture

A custom transformer neural network built entirely on mlx.nn.Module for predicting Bitcoin nonce locations from SHA-256 ring features. Part of hashfita's Chronobreaker system.

---

Architecture

Input: 17 rings × 4 features (amplitude, phase, direction, frequency)
  ↓
Feature Projection: nn.Linear(4, 128, bias=False)
  ↓
4× AttentionBlock4D (pre-norm RMSNorm, 4 heads, head_dim=32)
  ↓
Output Head: 128 → 64 (LeakyReLU 0.1) → 1 (Sigmoid)
  ↓
Output: nonce_ratio ∈ [0, 1] → nonce = ratio × MAX_NONCE

Parameters: ~272k (referenced in benchmark_grokking.py)

Source: src/chronobreaker/model/transformer.py, lines 1-86

---

Component Details

Feature Projection

nn.Linear(4, d_model=128, bias=False)

Projects 4 ring features (amplitude, phase, direction, frequency) to 128-dimensional embeddings. No bias — the model relies entirely on the learned weight matrix for feature interaction.

AttentionBlock4D

# attention.py, lines 1-78
class AttentionBlock4D(nn.Module):
    # 4 heads, head_dim=32, d_model=128
    # Q/K/V/Out: nn.Linear(128, 128, bias=False)
    # Pre-norm: nn.RMSNorm(128)
    # FFN: 128 → 256 → 128 with nn.GELU
    # 4D RoPE applied to Q and K

Standard multi-head self-attention with custom 4D Rotary Position Embeddings. The QK-Clip mechanism optionally captures mx.max(mx.abs(scores), axis=(0, 2, 3)) for gradient stabilization (used by muonclip-optimizer).

4D Continuous Rotary Position Embedding

# rope.py, lines 1-106
# 4 axes: ring index, frequency, merkle pointer, iteration count
# Decomposes embedding dimension into 4 sub-axes
# Computes linear rotation angles per axis
# Applies: x1*cos - x2*sin, x1*sin + x2*cos

This is a custom extension of standard RoPE to support 4 continuous (non-integer) position axes. Standard RoPE operates on a single discrete sequence position; this version handles 4 continuous coordinates simultaneously.

Source: src/chronobreaker/model/rope.py, _compute_cos_sin (line 33), _apply (line 90)

Output Head

# transformer.py, lines 77-79
h1 = mx.maximum(h, h * 0.1)   # LeakyReLU(0.1) — prevents dead gradients
nonce_logits = mx.sigmoid(out)  # Constrained to [0, 1]

Two-stage projection with LeakyReLU between stages. The final sigmoid constrains output to a ratio of the nonce space. LeakyReLU (slope 0.1) was chosen over ReLU to prevent dead gradients — a known issue with small models on sparse features.

---

Cyclic Loss Function

# trainer.py, lines 32-86
theta_predicted = nonce_ratio * 2 * pi
theta_actual = actual_nonce / MAX_NONCE * 2 * pi
loss = 1.0 - mx.cos(theta_predicted - theta_actual)

The nonce space is circular (nonce wraps around at 2^32). Mapping to angles and using cosine distance is mathematically appropriate — it handles the wraparound correctly. A simple MSE would fail at the boundary (nonce 0 is close to nonce 2^32-1).

Solver Penalty

penalty = 2.0 - (solves_count * 0.5)
total_loss = loss * penalty

Biases gradient updates toward blocks where the model failed to find a valid nonce. As solves_count increases, the penalty decreases — the model focuses its learning budget on failures.

---

3-Pass Training Step

# trainer.py, lines 557-672
# Pass 1: Forward inference (MLX)
nonce_ratio = model(input_features)

# Pass 2: GPU validation (OpenCL) — crosses framework boundary
solved = batch_validate_nonces_multi_gpu(predicted_nonces, headers)

# Pass 3: Autograd (MLX)
loss, grads = nn.value_and_grad(model, loss_fn)(model, input, target, solved)
model_optimizer.update(model, grads)  # MuonClip update
qk_clip_pass()                        # Gradient stabilization

Critical architectural note: Pass 2 crosses the mlx-opencl-bridge boundary — block headers must traverse the 3-copy chain. This is the main performance bottleneck in training.

---

Training Pipeline

File: src/chronobreaker/train/trainer.py (813 lines), ReinforcementTrainer class

Feature	Implementation
LR Schedule	Linear warmup (1000 steps) + cosine decay (50000 steps)
Serialization	MLX native .safetensors + JSON sidecar
Continuity	accumulated_epochs counter persisted across sessions
Eval split	Default 10%
Smart sampling	--edges-sampling for difficulty transitions

Training Script

File: src/scripts/train_model.py (418 lines)

• Rich dashboard with real-time loss plots via plotext
• Auto-continuation: finds highest-epoch .safetensors file
• Weight cleanup: keeps only top-k checkpoints

---

Grokking Benchmark

# benchmark_grokking.py, lines 1-94
for batch in [64, 128, 256, ..., 32768]:
    # Forward + backward pass timing
    # OOM detection for Apple Silicon memory limits

Benchmarks the RingTransformer across batch sizes to find performance characteristics and memory boundaries.

---

Parameter Evolution

Version	Normalization	Activation	Notable
Early	LayerNorm	ReLU	Standard transformer
Current	RMSNorm	LeakyReLU(0.1)	Upgrade for performance + gradient stability

The LayerNorm → RMSNorm upgrade is consistent with modern LLM practice (Llama, Mistral). RMSNorm is ~10% faster and has slightly better training dynamics.

---

Cross-References

• compute-optimization — Hub
• muonclip-optimizer — The optimizer that trains this model
• mlx-opencl-bridge — The bottleneck in the training step
• floating-point-determinism — Precision considerations for cyclic loss
• gpu-async-pipelines — How training fits the async architecture