---
type: reference
title: "RingTransformer — Custom MLX Architecture for Nonce Prediction"
description: "Custom transformer on MLX: 17-ring input, 4D Continuous RoPE, cyclic loss, reinforcement training, 272k parameters."
tags: [mlx, transformer, rope, neural-network, hashfita, compute-optimization, ml-architecture]
timestamp: 2026-07-25
---

# 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

```python
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

```python
# 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

```python
# 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

```python
# 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

```python
# 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

```python
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

```python
# 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

```python
# 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