Models From First Principles 04: HRM — Hierarchical Reasoning With Fast and Slow Recurrent State
HRM — Hierarchical Reasoning With Fast and Slow Recurrent State
The previous models in this series were mostly one-pass models.
MR.Q took two embeddings and produced one score.
EBT kept the same basic structure but added several heads.
SICQL made those heads explicit components.
The architecture grew, but the shape of the computation was still familiar:
input
↓
encoder
↓
representation
↓
heads
↓
outputs
HRM changes the question.
Instead of asking:
What can we predict from one representation?
we ask:
What if the representation itself is allowed to evolve through repeated computation?
That is the architectural jump in this post.
We are going to build a small hierarchical recurrent model with two coupled latent states:
- a low-level state that updates frequently;
- a high-level state that updates less frequently.
The low-level state performs several fine-grained updates before the high-level state is allowed to change.
Conceptually:
input
↓
input projector
↓
x̃
│
▼
┌────────────────────┐
│ low-level updates │
│ │
│ zL → zL → zL → zL │
└─────────┬──────────┘
│
▼
high-level update
│
zH → zH'
│
└──────────────┐
│
repeat cycle
This looks more sophisticated than the previous models.
But our rule for this series is unchanged:
If the architecture looks complicated, keep decomposing it.
By the time we reach the bottom, HRM will turn out to be a small number of ordinary PyTorch operations repeated according to a schedule.
1. What changed from SICQL?
SICQL gave us a useful composition:
context + candidate
↓
encoder
↓
zsa
┌─────┼─────┐
↓ ↓ ↓
Q V Policy
The obvious limitation is that zsa is computed once.
The model has one opportunity to transform its input before the heads make their predictions.
That may be enough.
And this is important:
recurrence is not automatically better.
If a feed-forward model solves the task, recurrence can simply add latency, optimization difficulty and more failure modes.
So HRM should not begin with the claim:
repeated computation is reasoning.
Instead we begin with a testable hypothesis:
some tasks may benefit from allowing the latent representation to be updated repeatedly before prediction.
That is much more precise.
2. The smallest possible recurrent model
Before building a hierarchy, build one recurrent state.
Suppose the input is:
x.shape == [B, D]
and we maintain a hidden state:
z.shape == [B, H]
A recurrent model repeatedly applies a function:
z_next = f(x, z_previous)
That is recurrence.
Nothing mystical has happened.
A minimal PyTorch version could be:
import torch
import torch.nn as nn
class SimpleRecurrentState(nn.Module):
def __init__(self, input_dim, hidden_dim):
super().__init__()
self.update = nn.Linear(input_dim + hidden_dim, hidden_dim)
def forward(self, x, z):
combined = torch.cat([x, z], dim=-1)
return torch.tanh(self.update(combined))
Then:
B = 8
D = 128
H = 64
x = torch.randn(B, D)
z = torch.zeros(B, H)
cell = SimpleRecurrentState(D, H)
for _ in range(4):
z = cell(x, z)
The same parameters are reused on every step.
That is a key distinction from simply stacking four different layers.
A deep feed-forward stack might be:
Layer1 → Layer2 → Layer3 → Layer4
A recurrent computation is:
Cell → Cell → Cell → Cell
with the same cell parameters reused.
3. Why use a GRUCell?
The simple tanh recurrence above can work, but recurrent optimization has well-known difficulties.
The hidden state needs to decide:
- what to preserve;
- what to overwrite;
- how much new information to incorporate.
A gated recurrent unit gives the model learned gates for this state update.
PyTorch exposes the basic recurrent operation directly:
nn.GRUCell(input_size, hidden_size)
A GRUCell consumes:
current input
+
previous hidden state
and returns:
next hidden state
So our recurrent block can become extremely small:
class RecurrentBlock(nn.Module):
def __init__(self, input_dim, hidden_dim):
super().__init__()
self.cell = nn.GRUCell(input_dim, hidden_dim)
def forward(self, z_prev, x):
return self.cell(x, z_prev)
Again:
advanced recurrent block
↓
GRUCell
We can continue descending inside the GRU if we want to understand reset/update gates, but for the architecture-level discussion this is already a useful boundary.
4. Add normalization
Repeated state updates can produce unstable activation scales.
One simple stabilization mechanism is RMS normalization.
RMSNorm does not subtract the mean like LayerNorm.
It rescales a vector according to its root-mean-square magnitude.
For a vector x:
rms(x) = sqrt(mean(x²) + ε)
and:
normalized = x / rms(x)
with a learned per-feature scale.
A compact implementation is:
class RMSNorm(nn.Module):
def __init__(self, dim, eps=1e-6):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def forward(self, x):
dtype = x.dtype
xf = x.float()
normed = xf * torch.rsqrt(
xf.pow(2).mean(dim=-1, keepdim=True) + self.eps
)
return normed.to(dtype) * self.weight
Now our recurrent block becomes:
class RecurrentBlock(nn.Module):
def __init__(self, input_dim, hidden_dim):
super().__init__()
self.cell = nn.GRUCell(input_dim, hidden_dim)
self.norm = RMSNorm(hidden_dim)
def forward(self, z_prev, x):
z_next = self.cell(x, z_prev)
return self.norm(z_next)
So one major HRM component is simply:
input + previous state
↓
GRUCell
↓
RMSNorm
↓
next state
5. The hierarchy
Now we add the architectural idea that makes HRM interesting.
Instead of one recurrent state, create two:
zL = low-level state
zH = high-level state
The low-level state updates many times.
The high-level state updates once per low-level cycle.
For example:
cycle 1
L step 1
L step 2
L step 3
L step 4
H step 1
cycle 2
L step 1
L step 2
L step 3
L step 4
H step 2
If there are:
N = number of high-level cycles
T = number of low-level steps per cycle
then the low-level module executes approximately:
N × T
updates while the high-level module executes:
N
updates.
That gives us two computational timescales.
6. Why call them fast and slow?
Not because one runs on a faster processor.
They operate at different update frequencies.
The low-level state receives repeated opportunities to process the current projected input in the context of the current high-level state.
The high-level state changes only after the low-level loop has completed.
So we can think of:
zL = frequently changing latent state
zH = slowly changing latent state
This is an architectural interpretation.
Do not automatically translate it into psychological claims such as:
zL = intuition
zH = conscious reasoning
unless experiments justify those interpretations.
The code only guarantees two differently scheduled hidden states.
7. Project the input first
Our input embedding might be large:
x.shape == [B, 2048]
We probably do not want every recurrent operation to work directly in 2048 dimensions.
So first project the input into a smaller hidden space:
class InputProjector(nn.Module):
def __init__(self, input_dim, hidden_dim, dropout=0.1):
super().__init__()
self.proj = nn.Linear(input_dim, hidden_dim)
self.norm = RMSNorm(hidden_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
return self.norm(self.dropout(self.proj(x)))
If:
input_dim = 2048
h_dim = 256
then:
[B, 2048]
↓ Linear
[B, 256]
↓ RMSNorm
[B, 256]
Call the projected input:
x̃
8. Building the low-level module
The low-level module should see:
- the projected input
x̃; - the current high-level state
zH.
If both have dimension h_dim, concatenate them:
l_input = torch.cat([x_tilde, zH], dim=-1)
So:
x̃ [B, H]
zH [B, H]
concat
↓
[B, 2H]
The low-level recurrent block therefore has:
low = RecurrentBlock(
input_dim=2 * h_dim,
hidden_dim=l_dim,
)
Its update is:
zL = low(zL, l_input)
Repeated T times.
9. Building the high-level module
After the low-level loop finishes, the high-level module receives:
final low-level state zL
+
previous high-level state zH
So:
h_input = torch.cat([zL, zH], dim=-1)
If:
zL: [B, L]
zH: [B, H]
then:
h_input: [B, L + H]
The high-level recurrent module is:
high = RecurrentBlock(
input_dim=l_dim + h_dim,
hidden_dim=h_dim,
)
and one high-level update is:
zH = high(zH, h_input)
That is the entire hierarchy.
10. The complete rollout
Now write the schedule explicitly.
def rollout(x_tilde, low, high, n_cycles, t_steps, l_dim, h_dim):
B = x_tilde.size(0)
device = x_tilde.device
zL = torch.zeros(B, l_dim, device=device)
zH = torch.zeros(B, h_dim, device=device)
for _ in range(n_cycles):
for _ in range(t_steps):
l_input = torch.cat([x_tilde, zH], dim=-1)
zL = low(zL, l_input)
h_input = torch.cat([zL, zH], dim=-1)
zH = high(zH, h_input)
return zL, zH
Read it slowly.
There is no hidden magic.
The hierarchy is mostly the loop structure.
This is one of the central lessons of the series:
architecture is often not a new primitive. It is a new composition of familiar primitives.
11. Put it inside a model
Now package the components.
class MiniHRM(nn.Module):
def __init__(
self,
input_dim=2048,
h_dim=256,
l_dim=128,
n_cycles=4,
t_steps=4,
dropout=0.1,
):
super().__init__()
self.input_dim = input_dim
self.h_dim = h_dim
self.l_dim = l_dim
self.n_cycles = n_cycles
self.t_steps = t_steps
self.input_projector = InputProjector(
input_dim,
h_dim,
dropout=dropout,
)
self.low = RecurrentBlock(
2 * h_dim,
l_dim,
)
self.high = RecurrentBlock(
l_dim + h_dim,
h_dim,
)
self.final_norm = RMSNorm(h_dim)
self.score_head = nn.Linear(h_dim, 1)
def rollout(self, x_tilde):
B = x_tilde.size(0)
device = x_tilde.device
zL = torch.zeros(B, self.l_dim, device=device)
zH = torch.zeros(B, self.h_dim, device=device)
for _ in range(self.n_cycles):
for _ in range(self.t_steps):
l_input = torch.cat([x_tilde, zH], dim=-1)
zL = self.low(zL, l_input)
h_input = torch.cat([zL, zH], dim=-1)
zH = self.high(zH, h_input)
return zL, self.final_norm(zH)
def forward(self, x):
x_tilde = self.input_projector(x)
zL, zH = self.rollout(x_tilde)
score_logit = self.score_head(zH).squeeze(-1)
return score_logit
We have now built the core architecture.
12. Parameter reuse changes the compute/parameter relationship
A recurrent architecture can perform more computation without increasing its parameters proportionally.
Suppose a recurrent block has P parameters.
Running it once costs roughly one application of those parameters.
Running it sixteen times still uses P learned parameters.
But compute becomes much larger.
So:
parameter count != amount of computation
This distinction matters enormously when comparing architectures.
A model may be relatively small in parameters yet expensive in latency because it repeatedly applies those parameters.
For HRM, useful measurements include:
parameters
forward latency
number of recurrent updates
memory use
training throughput
Do not use parameter count alone as a proxy for cost.
13. Count the recurrent steps
For our schedule:
n_cycles = 4
t_steps = 4
we execute:
low-level updates = 4 × 4 = 16
high-level updates = 4
Total recurrent cell applications:
20
That means one prediction contains twenty recurrent updates before the final head.
Compare that against a feed-forward baseline.
This is exactly the kind of architectural cost that should be measured rather than inferred from the number of model parameters.
14. A feed-forward baseline is mandatory
Before claiming the hierarchy helps, build a baseline with roughly comparable hidden dimensions.
class FeedForwardBaseline(nn.Module):
def __init__(self, input_dim=2048, hidden_dim=256):
super().__init__()
self.net = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.LayerNorm(hidden_dim),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, 1),
)
def forward(self, x):
return self.net(x).squeeze(-1)
Then compare:
validation performance
ranking performance
calibration
latency
memory
training stability
If the recurrent model does not beat the baseline on the metric that matters, the hierarchy may not be justified.
That is an architectural result too.
15. What does more recurrence buy us?
This is an experimental question.
Because our model can expose the number of cycles, we can test:
1 cycle
2 cycles
4 cycles
8 cycles
while keeping the trained weights fixed.
A useful interface is:
def rollout(self, x_tilde, max_cycles=None):
cycles = self.n_cycles if max_cycles is None else max_cycles
Then evaluate score quality as a function of recurrent compute.
This gives us a compute curve:
quality
↑
│ •
│ •
│ •
│ •
└──────────────────→ cycles
Possible outcomes:
- more cycles improve performance;
- performance saturates quickly;
- performance becomes worse;
- different examples benefit from different numbers of cycles.
Those are much more informative than simply saying the model is “hierarchical”.
16. Track the trajectory
A recurrent model produces intermediate states.
Do not throw them away while debugging.
Record them.
def rollout_with_trace(self, x_tilde):
B = x_tilde.size(0)
device = x_tilde.device
zL = torch.zeros(B, self.l_dim, device=device)
zH = torch.zeros(B, self.h_dim, device=device)
trace = []
for cycle in range(self.n_cycles):
for step in range(self.t_steps):
l_input = torch.cat([x_tilde, zH], dim=-1)
zL = self.low(zL, l_input)
trace.append({
"cycle": cycle,
"step": step,
"zL_norm": zL.norm(dim=-1).mean().item(),
"zH_norm": zH.norm(dim=-1).mean().item(),
})
h_input = torch.cat([zL, zH], dim=-1)
zH = self.high(zH, h_input)
return zL, zH, trace
Now you can inspect whether the latent state is actually changing.
17. State-change diagnostics
A useful measurement is:
||z_t - z_(t-1)||
If this rapidly approaches zero, later recurrent steps may be doing almost nothing.
def state_delta(a, b):
return (a - b).norm(dim=-1).mean().item()
Record:
cycle 0, step 0: 4.21
cycle 0, step 1: 1.37
cycle 0, step 2: 0.22
cycle 0, step 3: 0.03
That would suggest convergence within a cycle.
Perhaps four low-level steps are unnecessary.
Or perhaps convergence is exactly the behavior we want.
Either way, we now have evidence.
18. Cosine trajectory diagnostics
Norm change is not enough.
Two states can have similar magnitude while pointing in different directions.
Measure cosine similarity between successive states:
import torch.nn.functional as F
def cosine_change(prev, cur):
return F.cosine_similarity(prev, cur, dim=-1).mean().item()
Values near:
1.0
mean the state direction barely changes.
Lower values mean the representation is changing more substantially.
This gives another window into recurrent dynamics.
19. Track the maximum high-level state?
One design used by HRM-style architectures is to accumulate evidence over the high-level trajectory.
For example:
zH_max = torch.maximum(zH_max, zH)
on every high-level cycle.
Then the model has access to:
final high-level state
+
maximum activation over high-level trajectory
This is not the same as remembering the entire trajectory.
It is a simple summary statistic.
That distinction matters.
A max-reduced trajectory can preserve whether some feature became strongly active, but it discards:
- ordering;
- how long activation lasted;
- oscillation;
- intermediate sign changes.
So again:
evidence accumulation is a hypothesis about what summary of the trajectory is useful.
20. Add several diagnostic heads
Once we have a final high-level representation, we can reuse the lesson from EBT and SICQL.
One representation can support many small heads.
For example:
self.score_head = nn.Linear(h_dim, 1)
self.logvar_head = nn.Linear(h_dim, 1)
self.aux3_head = nn.Linear(h_dim, 3)
self.disagree_head = nn.Linear(h_dim, 1)
self.ood_head = nn.Linear(h_dim, 1)
self.temp_head = nn.Linear(h_dim, 1)
self.recon_head = nn.Linear(h_dim, h_dim)
Architecturally:
zH
│
┌───────┬───────┼────────┬────────┐
↓ ↓ ↓ ↓ ↓
score logvar aux3 disagree OOD
│
└──────── temperature calibration
zH
↓
reconstruction
This can look impressive on a diagram.
But remember the evidence boundary from the beginning of this series:
A head name does not establish what the head measures.
An ood_head is just a scalar neural network until training data and evaluation demonstrate out-of-distribution predictive validity.
21. A calibrated score head
Suppose the model predicts a raw quality logit:
score_logit = self.score_head(zH)
We can also predict a temperature:
tau_raw = self.temp_head(zH)
and force the temperature positive:
tau = 0.5 + 0.5 * F.softplus(tau_raw)
Then:
score = torch.sigmoid(score_logit / tau)
This gives each example a learned scaling of the logit.
But be precise.
A learned temperature head does not automatically mean the resulting probability is calibrated.
Calibration must be measured.
Useful metrics include:
Brier score
expected calibration error
reliability curves
negative log likelihood
22. Aleatoric uncertainty head
A common approach is to predict log variance:
log_var = self.logvar_head(zH)
The model can then use a heteroscedastic regression objective.
For target y, prediction mu, and log variance log_var:
def heteroscedastic_loss(mu, log_var, y):
precision = torch.exp(-log_var)
return 0.5 * (precision * (y - mu).pow(2) + log_var).mean()
This allows the model to represent examples as having different expected noise levels.
But again:
predicted variance
is not automatically equivalent to:
all forms of epistemic uncertainty
The training objective defines the operational meaning.
23. Three-way auxiliary classification
Another head can predict coarse classes:
bad
medium
good
with:
aux3_logits = self.aux3_head(zH)
and:
loss_aux3 = F.cross_entropy(aux3_logits, labels)
Why add an auxiliary classification task if we already have a scalar score?
Possible hypothesis:
coarse class supervision may encourage representations that separate qualitatively different response regions.
But the auxiliary head may also interfere with the main task.
Measure both possibilities.
24. Disagreement prediction
Suppose multiple judges evaluate the same candidate.
Then disagreement itself can become a target.
For example:
judge scores = [0.91, 0.89, 0.25, 0.22]
has high disagreement.
You can construct a target from score variance or another disagreement statistic and train:
disagree_logit = self.disagree_head(zH)
with an appropriate objective.
Now the head has an empirically defined meaning.
Without such a target, the name is just a name.
25. Reconstruction as a comprehension proxy
Another architecture idea is to ask the final state to reconstruct the projected input.
x_recon = self.recon_head(zH)
Then compare:
recon_sim = F.cosine_similarity(
x_recon,
x_tilde,
dim=-1,
)
The intuition is:
if the recurrent state has preserved useful information about its input, reconstruction should remain possible.
This is reasonable as an auxiliary constraint.
But do not overclaim it.
High reconstruction similarity does not prove semantic comprehension.
It proves the final representation contains information useful for reconstructing the chosen target representation.
That wording is less exciting and much more accurate.
26. Consistency under perturbation
We can also perturb the hidden representation:
mask = (
torch.rand_like(zH) < 0.1
).float()
zH_masked = zH * (1.0 - mask)
Then calculate similarity:
consistency_target = F.cosine_similarity(
zH,
zH_masked,
dim=-1,
)
A consistency head can be trained to predict this robustness signal.
This is useful because it turns an abstract word—“consistency”—into a concrete experimental definition.
The definition may or may not correspond to the kind of robustness your application cares about.
That is something you test.
27. Finite-difference sensitivity
A particularly useful diagnostic is to slightly perturb the input and see how much the score changes.
Create:
eps = 1e-3
x_eps = x + eps * F.normalize(
torch.randn_like(x),
dim=-1,
)
Run the model again:
score_eps = model(x_eps)
and estimate:
sensitivity = (
(score_eps - score).abs() / eps
)
This is not a full Jacobian.
It is a directional finite-difference approximation.
But it gives a practical signal for local sensitivity.
The decomposition is again simple:
input
↓ perturb
input'
↓ model
score'
↓ compare
local sensitivity
28. A complete multi-head HRM
Here is a compact but complete version:
class HRM(nn.Module):
def __init__(
self,
input_dim=2048,
h_dim=256,
l_dim=128,
n_cycles=4,
t_steps=4,
dropout=0.1,
):
super().__init__()
self.h_dim = h_dim
self.l_dim = l_dim
self.n_cycles = n_cycles
self.t_steps = t_steps
self.input_projector = InputProjector(
input_dim,
h_dim,
dropout,
)
self.low = RecurrentBlock(
2 * h_dim,
l_dim,
)
self.high = RecurrentBlock(
l_dim + h_dim,
h_dim,
)
self.final_norm = RMSNorm(h_dim)
self.head_drop = nn.Dropout(dropout)
self.score_head = nn.Linear(h_dim, 1)
self.logvar_head = nn.Linear(h_dim, 1)
self.aux3_head = nn.Linear(h_dim, 3)
self.disagree_head = nn.Linear(h_dim, 1)
self.ood_head = nn.Linear(h_dim, 1)
self.temp_head = nn.Linear(h_dim, 1)
self.recon_head = nn.Linear(h_dim, h_dim)
def _rollout(self, x_tilde, max_cycles=None):
B = x_tilde.size(0)
device = x_tilde.device
cycles = (
self.n_cycles
if max_cycles is None
else min(max_cycles, self.n_cycles)
)
zL = torch.zeros(B, self.l_dim, device=device)
zH = torch.zeros(B, self.h_dim, device=device)
zH_max = torch.zeros_like(zH)
for _ in range(cycles):
for _ in range(self.t_steps):
l_input = torch.cat([x_tilde, zH], dim=-1)
zL = self.low(zL, l_input)
h_input = torch.cat([zL, zH], dim=-1)
zH = self.high(zH, h_input)
zH_max = torch.maximum(zH_max, zH)
zH = self.final_norm(zH)
zH_max = self.final_norm(zH_max)
return zL, zH, zH_max
def forward(self, x, max_cycles=None):
x_tilde = self.input_projector(x)
zL, zH, zH_max = self._rollout(
x_tilde,
max_cycles=max_cycles,
)
h = self.head_drop(zH)
score_logit = self.score_head(h)
tau_raw = self.temp_head(h)
tau = 0.5 + 0.5 * F.softplus(tau_raw)
score = torch.sigmoid(score_logit / tau)
return {
"score": score.squeeze(-1),
"score_logit": score_logit.squeeze(-1),
"log_var": self.logvar_head(h).squeeze(-1),
"aux3_logits": self.aux3_head(h),
"disagree_logit": self.disagree_head(h).squeeze(-1),
"ood_logit": self.ood_head(h).squeeze(-1),
"temperature": tau.squeeze(-1),
"reconstruction": self.recon_head(h),
"zL": zL,
"zH": zH,
"zH_max": zH_max,
"x_tilde": x_tilde,
}
Now the entire architecture fits in one screen-sized mental model.
29. The tensor contracts
Always write the shape contracts down.
Suppose:
B = batch size
D = input dimension
H = high-level dimension
L = low-level dimension
Then:
x [B, D]
x_tilde [B, H]
zH [B, H]
zL [B, L]
l_input [B, 2H]
h_input [B, L + H]
score_logit [B, 1]
aux3_logits [B, 3]
These contracts are more useful than architecture names when debugging.
30. Add assertions
def assert_hrm_shapes(model, x, out):
B = x.size(0)
assert out["zH"].shape == (B, model.h_dim)
assert out["zL"].shape == (B, model.l_dim)
assert out["score"].shape == (B,)
assert out["aux3_logits"].shape == (B, 3)
This is cheap protection against silent shape drift.
31. Batch size one remains dangerous
Avoid:
.squeeze()
for scalar heads.
If:
[B, 1]
and:
B = 1
then .squeeze() can produce a zero-dimensional tensor.
Prefer:
.squeeze(-1)
so:
[1, 1] → [1]
not:
[1, 1] → []
32. Device-safe state initialization
A recurring bug in custom recurrent code is allocating the hidden state on CPU while the model input lives on CUDA.
Bad:
zH = torch.zeros(B, H)
Better:
zH = torch.zeros(
B,
H,
device=x.device,
dtype=x.dtype,
)
If mixed precision is involved, dtype alignment matters as well as device alignment.
33. Do not store a stale device string
It is tempting to do:
self.device = "cuda"
inside __init__.
Then later:
model.to("cpu")
and the stored string is wrong.
Prefer deriving device from the input or parameters:
device = x.device
or:
device = next(self.parameters()).device
This keeps state allocation aligned with the model’s real runtime location.
34. Gradient flow through repeated computation
Because recurrent parameters are reused, gradients from every recurrent application accumulate into the same parameters.
The computational graph looks conceptually like:
W
↙ ↓ ↘
step 1 step 2 step 3
\ | /
final loss
So deeper recurrence can produce different gradient behavior than a one-pass model.
Track gradient norms.
def gradient_report(model):
for name, p in model.named_parameters():
if p.grad is None:
print(name, "NO GRAD")
else:
print(
name,
float(p.grad.norm()),
)
Watch especially:
low.cell.*
high.cell.*
input_projector.*
35. Gradient clipping
Recurrent models can be sensitive to large gradients.
A practical safety mechanism is:
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0,
)
But clipping is not a substitute for understanding instability.
If gradients constantly exceed the clipping threshold, investigate:
- learning rate;
- recurrent depth;
- initialization;
- normalization;
- loss scaling;
- problematic data.
36. Prove parameters actually update
Before a long training run:
before = {
name: p.detach().clone()
for name, p in model.named_parameters()
}
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
for name, p in model.named_parameters():
delta = (
p.detach() - before[name]
).abs().max().item()
print(name, delta)
If recurrent cells have gradients but do not change, inspect optimizer membership and learning rate.
37. Overfit a tiny batch
The debugging rule from the PyTorch series still applies.
Take perhaps:
8 examples
and train repeatedly on them.
If the model cannot substantially overfit a tiny deterministic batch, do not start tuning large-scale hyperparameters.
Possible causes include:
- broken targets;
- detached graph;
- incorrect state initialization;
- optimizer omission;
- impossible objective;
- dropout noise;
- recurrent instability.
38. Train with a simple target first
Before training seven diagnostic heads simultaneously, verify the recurrent core with one target.
pred = model(x)["score_logit"]
loss = F.mse_loss(pred, target)
Only after this pipeline works should you introduce:
uncertainty loss
auxiliary classification
reconstruction
consistency
OOD
This isolates failures.
Architectural sophistication should be introduced incrementally.
39. Multi-task loss
Eventually we may use:
loss = (
lambda_score * loss_score
+ lambda_var * loss_var
+ lambda_aux3 * loss_aux3
+ lambda_disagree * loss_disagree
+ lambda_recon * loss_recon
+ lambda_ood * loss_ood
)
Every coefficient matters.
Print the unweighted losses first.
Otherwise one task can dominate simply because its numerical scale is larger.
40. Measure per-task gradients
Shared recurrence means auxiliary tasks can change the same recurrent parameters.
Inspect their gradient directions.
For losses A and B:
gA = torch.autograd.grad(
loss_A,
model.high.parameters(),
retain_graph=True,
allow_unused=True,
)
gB = torch.autograd.grad(
loss_B,
model.high.parameters(),
retain_graph=True,
allow_unused=True,
)
Flatten and compare cosine similarity.
If it is negative, the tasks are pushing the shared recurrent core in opposing directions on that batch.
This does not prove multi-task learning is bad.
It tells you there is an optimization trade-off to investigate.
41. Ablate the hierarchy
A useful hierarchy ablation is not merely:
HRM vs no model
Compare:
feed-forward baseline
single recurrent state
hierarchical recurrent state
That tells you where gains appear.
For example:
FFN accuracy 0.78
single RNN accuracy 0.79
HRM accuracy 0.82
would be more informative than only reporting HRM.
But also compare cost:
FFN 0.4 ms
single RNN 1.1 ms
HRM 4.9 ms
Architecture is always a trade-off.
42. Ablate the number of low-level steps
Test:
T = 1
T = 2
T = 4
T = 8
while keeping high-level cycles fixed.
Question:
does repeated low-level computation add useful information?
If not, the hierarchy may be unnecessarily expensive.
43. Ablate the number of high-level cycles
Similarly test:
N = 1
N = 2
N = 4
N = 8
Plot:
quality vs cycles
latency vs cycles
The most useful architecture may not be the largest setting.
44. Randomize the hierarchy schedule
A stronger experiment asks whether the exact hierarchical schedule matters.
For example, compare:
4 low updates → 1 high update
against:
1 low → 1 high → repeat
or another compute-matched schedule.
If all schedules perform similarly, perhaps hierarchy is not the source of the improvement.
Maybe repeated compute alone is enough.
That is exactly the kind of alternative explanation we should test.
45. Compute-matched baselines
This is especially important.
Suppose HRM performs twenty recurrent operations.
A fair comparison should include a non-hierarchical model with a similar compute budget.
Otherwise:
HRM beats baseline
may simply mean:
more compute beats less compute
not:
hierarchical scheduling is beneficial
A good experiment separates these possibilities.
46. Parameter-matched baselines
The opposite comparison also matters.
Because recurrence reuses weights, HRM can execute many steps with relatively few parameters.
Compare against:
same parameter count
and:
same compute budget
These answer different questions.
Both are useful.
47. Does the high-level state actually differ from the low-level state?
A two-state architecture does not guarantee functional specialization.
We should test whether the states encode different information.
Possible probes:
train a linear probe on zL
train a linear probe on zH
for several targets.
If both states support exactly the same predictions equally well, specialization may be weak.
If different information is linearly accessible from each, that is evidence of representational differentiation.
Not proof of “fast thinking” and “slow thinking”—but actual measurable differentiation.
48. Probe each cycle
Do not only inspect final zH.
Save:
zH_1
zH_2
zH_3
zH_4
and train simple probes or evaluate the main head at each cycle.
Then we can ask:
When does useful predictive information emerge?
Perhaps:
cycle 1: 0.71
cycle 2: 0.79
cycle 3: 0.81
cycle 4: 0.81
That suggests the fourth cycle may be unnecessary.
49. Score convergence
Another simple diagnostic:
run the score head after every high-level cycle.
cycle 1: 0.62
cycle 2: 0.74
cycle 3: 0.76
cycle 4: 0.761
The score appears to converge.
That opens the door to adaptive computation.
Instead of always running four cycles, perhaps some examples can halt earlier.
But we should first demonstrate that score convergence is reliable.
50. Halting is another model
Suppose we want adaptive computation.
We might add:
self.halt_head = nn.Linear(h_dim, 1)
and predict:
halt_prob = torch.sigmoid(
self.halt_head(zH)
)
Now we have another model inside the model.
It needs:
- a target;
- an objective;
- evaluation;
- a decision threshold.
Simply producing a halting probability does not mean the model knows when it has “finished reasoning.”
That semantic claim must be earned.
51. A simple convergence-based halting rule
Before training a halting head, test a deterministic rule.
For example:
if (score_cur - score_prev).abs().max() < 1e-3:
break
or latent convergence:
if (zH_cur - zH_prev).norm(dim=-1).mean() < threshold:
break
This is a useful baseline for any learned halting mechanism.
Never compare a learned controller only against “always run everything.”
Compare against simple heuristics too.
52. Evaluation should include compute-aware metrics
For recurrent models, report more than predictive quality.
A useful table might include:
model quality params recurrent steps latency
FFN 0.78 1.2M 0 0.4 ms
single RNN 0.80 0.8M 16 2.9 ms
HRM 0.82 0.9M 20 4.6 ms
Now architecture choices become visible.
53. Checkpoint the schedule
A recurrent model checkpoint is not fully described by weights alone.
Store configuration:
checkpoint = {
"model_state": model.state_dict(),
"input_dim": model.input_dim,
"h_dim": model.h_dim,
"l_dim": model.l_dim,
"n_cycles": model.n_cycles,
"t_steps": model.t_steps,
}
Why?
Because changing:
n_cycles
changes the computation even though tensor shapes remain compatible.
Configuration is part of the model contract.
54. Reproducibility is harder with recurrence
Small implementation changes can alter repeated computation significantly.
Record:
- seed;
- architecture config;
- dropout;
- optimizer;
- learning rate;
- gradient clipping;
- number of cycles;
- low-level steps;
- embedding version;
- loss weights.
A recurrent architecture creates more places where hidden differences can accumulate.
55. The debugger view
When HRM produces a strange output, do not stare at the final score.
Trace the pipeline:
input finite?
↓
projected input finite?
↓
zL after each step finite?
↓
zH after each cycle finite?
↓
state norms stable?
↓
state deltas meaningful?
↓
head inputs finite?
↓
head outputs finite?
This is exactly why the previous PyTorch debugging series matters.
Once an architecture becomes iterative, runtime evidence becomes even more valuable.
56. A reusable trajectory debugger
@torch.no_grad()
def debug_hrm(model, x):
x_tilde = model.input_projector(x)
B = x.size(0)
zL = torch.zeros(
B,
model.l_dim,
device=x.device,
dtype=x.dtype,
)
zH = torch.zeros(
B,
model.h_dim,
device=x.device,
dtype=x.dtype,
)
print("x", x.shape, torch.isfinite(x).all().item())
print(
"x_tilde",
x_tilde.shape,
float(x_tilde.norm(dim=-1).mean()),
)
for cycle in range(model.n_cycles):
for step in range(model.t_steps):
prev = zL.clone()
l_input = torch.cat([x_tilde, zH], dim=-1)
zL = model.low(zL, l_input)
print(
"L",
cycle,
step,
"norm",
float(zL.norm(dim=-1).mean()),
"delta",
float((zL - prev).norm(dim=-1).mean()),
)
prev_h = zH.clone()
h_input = torch.cat([zL, zH], dim=-1)
zH = model.high(zH, h_input)
print(
"H",
cycle,
"norm",
float(zH.norm(dim=-1).mean()),
"delta",
float((zH - prev_h).norm(dim=-1).mean()),
)
That tells us far more than:
model output looks wrong
57. The model inside HRM
Now step back.
What is HRM actually made of?
HRM
│
├── InputProjector
│ ├── Linear
│ ├── Dropout
│ └── RMSNorm
│
├── Low-level RecurrentBlock
│ ├── GRUCell
│ └── RMSNorm
│
├── High-level RecurrentBlock
│ ├── GRUCell
│ └── RMSNorm
│
├── rollout schedule
│ ├── low-level inner loop
│ └── high-level outer loop
│
└── prediction heads
├── score
├── variance
├── auxiliary class
├── disagreement
├── OOD
├── temperature
└── reconstruction
That is the architecture.
The intimidating name has disappeared.
58. What is genuinely new here?
Compared with SICQL, the genuine architectural novelty is not the heads.
We already understand heads.
The new idea is:
reused recurrent computation
+
two latent states
+
different update frequencies
Everything else is composition around that core.
This is exactly how we should read unfamiliar models:
- identify the new mechanism;
- separate it from familiar supporting machinery;
- test whether the new mechanism matters.
59. What HRM does not prove
The architecture does not by itself prove:
reasoning
hierarchical cognition
uncertainty awareness
out-of-distribution awareness
comprehension
adaptive intelligence
It provides mechanisms that can be trained and tested for operational definitions of those concepts.
That distinction keeps our claims attached to evidence.
60. When might hierarchical recurrence be useful?
Possible cases include tasks where:
- repeated refinement helps;
- a compact parameter budget is desirable;
- variable compute may be useful;
- intermediate latent trajectories contain signal;
- multiple auxiliary objectives can shape one recurrent representation.
But the architecture should earn its complexity against simpler baselines.
61. When might it be the wrong choice?
If:
- latency is critical;
- a shallow scorer already performs well;
- the dataset is small enough that recurrence overfits;
- intermediate computation does not improve predictions;
- recurrent training is unstable;
then HRM may be an unnecessary complication.
“More sophisticated architecture” is not itself a goal.
62. The crucial experiment
If I had to choose one experiment for this architecture, it would be:
Does additional recurrent computation improve held-out performance after controlling for parameter count and compute?
That question cuts through most of the mythology.
Test:
feed-forward
single recurrent
hierarchical recurrent
under:
parameter-matched
compute-matched
latency-reported
conditions.
Then inspect performance as cycles increase.
That would tell us much more than the architecture diagram alone.
63. Where the architecture wants to go next
HRM gives us powerful machinery.
But it is not small.
It contains:
- two recurrent systems;
- multiple update scales;
- many diagnostic heads;
- repeated computation;
- several auxiliary objectives.
That naturally creates the next question:
Can we preserve the useful idea of iterative latent refinement in a smaller, simpler architecture?
Instead of maintaining two recurrent states:
zL
zH
what if we maintain one latent state:
z
and repeatedly refine it using:
input + candidate + previous z
That takes us toward Tiny.
Conceptually:
HRM
x̃ + zH → zL → zL → zL
↓
zL + zH → zH
↓
repeat
becomes something closer to:
Tiny
x + y + z
↓
block
↓
z'
↓
residual update
↓
repeat
The hierarchy becomes recursion.
And the model becomes much easier to make parameter-efficient.
That is the next stage of Models From First Principles.
Final mental model
If you remember only one diagram from this post, remember this:
input embedding
↓
projection
↓
x̃
│
├─────────────────────────────┐
│ │
▼ │
[x̃ + zH] → low recurrent block │
↓ │
zL │
↓ repeat T times │
zL │
│ │
▼ │
[zL + zH] → high recurrent block │
↓ │
zH ────────────────────────────┘
│ repeat N cycles
▼
final representation
│
┌────┼───────────────┐
↓ ↓ ↓
score uncertainty diagnostics
HRM looks like a large model.
But once decomposed, it is mostly:
Linear
GRUCell
RMSNorm
concatenate
repeat
small heads
The architecture is in how those pieces are scheduled and connected.
That is the real lesson.
