Models From First Principles 05: Tiny — Recursive Reasoning With a Small Neural Network
Tiny — Recursive Reasoning With a Small Neural Network
The previous post introduced a much more ambitious architecture.
Instead of taking one representation and predicting from it once, the Hierarchical Reasoning Model repeatedly updated two latent states:
input
↓
low-level state
↓ ↓ ↓
high-level state
↓
repeat
That gave us something genuinely new:
computation could continue without adding a new set of parameters for every step.
The same recurrent modules were reused.
But it also gave us more machinery:
- two latent states;
- two recurrent modules;
- two update frequencies;
- a scheduling rule;
- recurrent state initialization;
- hierarchical coupling;
- several diagnostic heads.
So the next question is obvious.
How much of the useful idea can we keep if we make the architecture much smaller?
That is the motivation for the model in this post.
We will call it Tiny.
Not because the idea is trivial.
Because the architecture deliberately compresses the iterative computation into one latent state and a small set of reused blocks.
The core idea is:
x = goal / context embedding
y = candidate / response embedding
z = current latent state
[x, y, z]
↓
projection
↓
core block
↓
proposed update
↓
z = z + α · update
↓
repeat
That is the heart of the model.
Everything else in this post grows from it.
1. What changed from HRM?
The HRM-style architecture had two hidden states:
zL = low-level state
zH = high-level state
Tiny collapses that hierarchy to one state:
z
Instead of asking:
How should the fast state update?
How should the slow state update?
How often should each update?
we ask:
Given the goal,
the candidate,
and what the model currently knows,
what should the latent state become next?
In mathematical shorthand:
u_t = f(x, y, z_t)
z_{t+1} = z_t + α u_t
where:
x = context embedding
y = response embedding
z_t = current latent state
u_t = proposed state update
α = step scale
The same function f is used at every recursion.
That reuse is crucial.
A model that performs six recursive steps is not necessarily six times larger than a one-step model.
It is usually using the same parameters six times.
This is the same parameter-count-versus-compute distinction we introduced in HRM.
2. Start with the smallest possible recursive model
Let us strip the architecture down to almost nothing.
import torch
from torch import nn
class MinimalRecursiveModel(nn.Module):
def __init__(self, d_model: int = 256, step_scale: float = 0.1):
super().__init__()
self.d_model = d_model
self.step_scale = step_scale
self.update = nn.Linear(d_model * 3, d_model)
def forward(self, x, y, z, steps: int = 6):
for _ in range(steps):
fused = torch.cat([x, y, z], dim=-1)
dz = torch.tanh(self.update(fused))
z = z + self.step_scale * dz
return z
That is already a recursive neural model.
There is no magic hidden behind the name.
Each step is:
x [B,D]
y [B,D]
z [B,D]
│
└──── concatenate ────> [B,3D]
│
▼
Linear(3D,D)
│
▼
tanh
│
▼
dz
│
α · dz
│
▼
z ← z + αdz
Then we do it again.
3. Why include the previous latent state?
Suppose we removed z from the update function.
We would have:
u = f(x, y)
If x and y never change, then every recursion receives exactly the same input.
Unless the block contains its own recurrent state, there is no reason for step 5 to know what happened at step 4.
Including z_t changes that:
u_t = f(x, y, z_t)
Now each computation depends on the accumulated result of previous computation.
That makes the recursion stateful.
Step 2 is not simply a replay of step 1.
It sees a changed latent state.
4. The state is not automatically “reasoning”
This distinction matters throughout this series.
We may describe z informally as a reasoning state because it is repeatedly refined.
But the architecture only guarantees this:
zis a latent vector updated recursively from the inputs and its previous value.
It does not guarantee that:
- each step corresponds to a logical inference;
- later steps are more correct than earlier steps;
- the state contains interpretable propositions;
- more recursions always improve performance;
- the model has learned an algorithm.
Those are empirical questions.
Later in the post we will build the experiments needed to test them.
5. Why use a residual state update?
We could replace the state entirely:
z = dz
Tiny instead uses:
z = z + step_scale * dz
This makes the recursion look like incremental refinement.
If:
α = 0.1
then every update modifies the existing state rather than replacing it wholesale.
That can make recursive dynamics easier to control.
Conceptually:
z0
│
+ small correction
│
z1
│
+ small correction
│
z2
│
+ small correction
│
z3
instead of:
z0 → completely new z1 → completely new z2 → ...
This resembles residual connections elsewhere in deep learning.
The important point is not the analogy.
The important point is that step size becomes an architectural hyperparameter.
6. Step scale is really a dynamics control
Consider:
z = z + alpha * dz
If alpha is tiny:
0.001
then the state changes very slowly.
If alpha is large:
1.0
then each recursion can radically alter the state.
We should measure this.
@torch.no_grad()
def state_delta(old_z, new_z):
return (new_z - old_z).norm(dim=-1).mean().item()
A useful recursion trace is:
step 0: ||Δz|| = 0.091
step 1: ||Δz|| = 0.074
step 2: ||Δz|| = 0.051
step 3: ||Δz|| = 0.028
step 4: ||Δz|| = 0.014
step 5: ||Δz|| = 0.008
That would suggest convergence.
But if we see:
0.2
0.4
0.8
1.6
3.2
then we have a very different dynamical system.
7. Put a real processing block inside the recursion
The linear projection is intentionally minimal.
A richer version uses a residual MLP block.
class TinyBlock(nn.Module):
def __init__(self, d_model: int, dropout: float = 0.1):
super().__init__()
self.norm = nn.LayerNorm(d_model)
self.mlp = nn.Sequential(
nn.Linear(d_model, d_model * 4),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(d_model * 4, d_model),
nn.Dropout(dropout),
)
def forward(self, x):
return x + self.mlp(self.norm(x))
Again, recursively decompose it:
TinyBlock
↓
LayerNorm
↓
Linear D → 4D
↓
GELU
↓
Dropout
↓
Linear 4D → D
↓
Dropout
↓
residual addition
There is nothing irreducible here.
It is ordinary PyTorch machinery.
8. Why expand to 4D inside the MLP?
This is a common feed-forward pattern:
D → 4D → D
For example:
256 → 1024 → 256
The expansion gives the nonlinear transformation more intermediate capacity.
The second projection returns to the model dimension so the residual addition remains valid:
input: [B, D]
output: [B, D]
That stable interface is extremely useful.
It means we can replace the block without changing the recursive machinery.
9. Stable interfaces let us replace the core
Suppose every core block satisfies:
[B,D] → [B,D]
Then the recursive loop does not care whether the implementation is:
- one linear layer;
- a two-layer MLP;
- three residual MLP blocks;
- self-attention;
- a convolution;
- a mixture-of-experts module.
The loop remains:
z_next = self.core(z_next)
z = z + alpha * z_next
That is a recurring theme in this series.
Good component boundaries make architectural experimentation cheap.
10. Building the recursive core
Now we can write a stronger version.
class RecursiveCore(nn.Module):
def __init__(
self,
d_model: int = 256,
n_layers: int = 2,
dropout: float = 0.1,
):
super().__init__()
self.blocks = nn.Sequential(
*[
TinyBlock(d_model, dropout=dropout)
for _ in range(n_layers)
]
)
def forward(self, x):
return self.blocks(x)
Then the full recursive update becomes:
class RecursiveLatentModel(nn.Module):
def __init__(
self,
d_model: int = 256,
n_layers: int = 2,
n_recursions: int = 6,
step_scale: float = 0.1,
):
super().__init__()
self.d_model = d_model
self.n_recursions = n_recursions
self.step_scale = step_scale
self.z_proj = nn.Linear(d_model * 3, d_model)
self.core = RecursiveCore(d_model, n_layers)
self.final_norm = nn.LayerNorm(d_model)
def forward(self, x, y, z=None):
if z is None:
z = torch.zeros_like(x)
for _ in range(self.n_recursions):
fused = torch.cat([x, y, z], dim=-1)
z_next = torch.tanh(self.z_proj(fused))
z_next = self.core(z_next)
z = z + self.step_scale * z_next
return self.final_norm(z)
We now have the fundamental Tiny architecture.
11. Trace the shapes
Assume:
B = 32
D = 256
Then:
x [32, 256]
y [32, 256]
z [32, 256]
cat([x,y,z]) [32, 768]
z_proj [32, 256]
core [32, 256]
state update [32, 256]
final_norm [32, 256]
Put assertions in the model while developing it.
assert x.ndim == 2
assert y.ndim == 2
assert x.shape == y.shape
assert x.shape[-1] == self.d_model
Inside the loop:
assert fused.shape[-1] == self.d_model * 3
assert z_next.shape == z.shape
Shape contracts turn many bugs into immediate failures.
12. Recursion increases compute, not parameter count
Suppose the recursive model contains one million parameters.
Running:
1 recursion
and:
8 recursions
uses essentially the same learned parameter set.
But the second case performs the core computation eight times.
So we have two independent axes:
capacity ≈ parameter count
computation ≈ recursions × per-step compute
This gives us an interesting experimental lever.
We can increase test-time computation without increasing parameter count.
But we must not automatically assume that more computation means better output.
13. Does more recursion actually help?
This is one of the most important experiments in the entire post.
Evaluate the same trained model with different recursion counts:
@torch.no_grad()
def evaluate_by_steps(model, loader, step_values):
results = {}
for steps in step_values:
scores = []
for x, y, target in loader:
pred = model(x, y, steps=steps)
scores.append(metric(pred, target))
results[steps] = sum(scores) / len(scores)
return results
Possible result:
steps accuracy
1 0.712
2 0.741
3 0.754
4 0.756
6 0.755
8 0.749
That would be much more informative than saying:
recursion improves reasoning.
It tells us that, for this task and model, additional computation helps until roughly four steps and then saturates.
14. Allow the caller to choose the recursion depth
This is useful for experiments and possibly adaptive inference.
def _resolve_steps(self, steps):
if steps is None:
return self.n_recursions
return max(1, min(int(steps), self.n_recursions))
Then:
for _ in range(self._resolve_steps(steps)):
...
Now one trained model supports:
model(x, y, steps=1)
model(x, y, steps=3)
model(x, y, steps=6)
This is much better than training a separate model for every compute budget.
15. Add a halting signal
If the model refines its state repeatedly, we may want some estimate of whether further computation is useful.
A simple learned head is:
self.halt_head = nn.Linear(d_model, 1)
At every step:
halt_logit = self.halt_head(self.final_norm(z_next))
We can track the maximum:
halt_logits = torch.maximum(halt_logits, halt_logit)
Then:
halt_prob = torch.sigmoid(halt_logits)
But be careful with the language.
A head named halt_head does not automatically know when the model should stop.
It needs a meaningful training target or a validated relationship with marginal value from additional computation.
16. A better halting experiment
We can directly measure the benefit of one more recursion.
For each sample:
score_t = prediction after t steps
score_t+1 = prediction after t+1 steps
Define improvement relative to the target:
gain_t = error_t - error_t+1
If:
gain_t <= threshold
then extra computation did not help much.
A learned halting head could be evaluated against that signal.
That is stronger than simply plotting sigmoid(halt_logit) and calling it confidence.
17. Deterministic halting baselines first
Before adding learned halting, try simple rules.
State-delta rule
Stop when:
||z_t - z_{t-1}|| < ε
Score-delta rule
Stop when:
|score_t - score_{t-1}| < ε
Fixed-budget rule
Always run exactly N steps.
Patience rule
Stop if the score changes less than ε for k consecutive steps.
If a learned halting head cannot outperform those baselines, it may not be earning its complexity.
18. Add the prediction surface
So far our model returns a latent vector.
We need task outputs.
Start with one score:
self.score_head = nn.Linear(d_model, 1)
Then:
score_logit = self.score_head(z)
score = torch.sigmoid(score_logit)
Now the architecture is:
x,y
↓
recursive latent refinement
↓
z_final
↓
score head
↓
score
This looks conceptually similar to MR.Q again.
The difference is that the representation is no longer computed in one pass.
It is iteratively refined.
19. Add multiple diagnostic heads
Once we have a shared final state, additional heads are cheap.
For example:
self.score_head = nn.Linear(d_model, 1)
self.logvar_head = nn.Linear(d_model, 1)
self.aux3_head = nn.Linear(d_model, 3)
self.disagree_head = nn.Linear(d_model, 1)
self.ood_head = nn.Linear(d_model, 1)
self.temp_head = nn.Linear(d_model, 1)
self.recon_head = nn.Linear(d_model, d_model)
That produces a familiar pattern:
z_final
│
┌────────────┼────────────┐
↓ ↓ ↓
score uncertainty OOD
↓ ↓ ↓
calibration disagree reconstruction
Again:
the existence of a head is not proof of the semantic claim in its name.
Each head requires an objective and evaluation.
20. Heteroscedastic uncertainty
Suppose the model predicts:
μ = predicted score
log σ² = predicted log variance
A common Gaussian-style heteroscedastic loss is:
L = 0.5 * exp(-logvar) * (target - μ)^2
+ 0.5 * logvar
In code:
def heteroscedastic_loss(mean, log_var, target):
precision = torch.exp(-log_var)
return (
0.5 * precision * (target - mean).pow(2)
+ 0.5 * log_var
).mean()
Usually we clamp log_var for numerical safety:
log_var = log_var.clamp(-5.0, 5.0)
But the key evaluation question remains:
Are larger predicted variances actually associated with larger errors?
That should be measured.
21. Temperature calibration
Tiny can also predict a sample-dependent temperature.
tau_raw = self.temp_head(z)
tau = 0.5 + 0.5 * F.softplus(tau_raw)
Then:
score = torch.sigmoid(score_logit / tau)
A larger tau softens the score.
A smaller tau makes it more decisive.
This is an architectural mechanism.
Whether it improves calibration is an empirical question.
Measure:
- expected calibration error;
- Brier score;
- negative log likelihood;
- reliability curves;
- ranking metrics separately.
Do not collapse calibration and ranking into one number.
22. Add consistency under perturbation
Suppose we want the latent representation to be stable when a small portion of it is masked.
mask = (torch.rand_like(z) < 0.1).float()
z_masked = z * (1.0 - mask)
We can compute cosine similarity:
cos = F.cosine_similarity(z, z_masked, dim=-1)
consistency_target = (cos + 1.0) * 0.5
Then a consistency head can attempt to predict this quantity.
But notice what we have really defined:
stability under this specific perturbation process.
We have not defined universal robustness.
That distinction is important.
23. Finite-difference sensitivity
We can perturb the response embedding slightly:
eps = 1e-3
noise = F.normalize(torch.randn_like(y), dim=-1)
y_eps = y + eps * noise
Then compare scores:
score = model(x, y)
score_eps = model(x, y_eps)
sensitivity = (score_eps - score).abs() / eps
This gives a local directional sensitivity estimate.
It is not the full Jacobian.
It is not a proof of causal understanding.
It is one measurable response to a defined perturbation.
24. Now introduce the sparse autoencoder
The next component is especially interesting because it is another model inside the model.
After recursive refinement we have:
z_final [B,D]
We introduce an encoder:
D → D/2
and a decoder:
D/2 → D
For example:
self.sae_encoder = nn.Sequential(
nn.Linear(d_model, d_model // 2),
nn.ReLU(),
nn.LayerNorm(d_model // 2),
)
self.sae_decoder = nn.Linear(d_model // 2, d_model)
Then:
concepts = self.sae_encoder(z_final)
reconstruction = self.sae_decoder(concepts)
We can feed the reconstructed information back into the head representation:
z_head = z_final + reconstruction
Now the model contains another learned subsystem.
25. What is a sparse autoencoder trying to do?
An autoencoder tries to reconstruct an input through a constrained intermediate representation.
z
↓
encoder
↓
c
↓
decoder
↓
z_hat
If the bottleneck is narrower:
256 → 128 → 256
then the model cannot simply copy every dimension directly.
If we also encourage sparse activations in c, the hope is that individual latent features become more selective.
A simple objective might be:
L_sae = reconstruction_loss + λ * sparsity_penalty
For example:
def sae_loss(z, z_hat, concepts, alpha=0.05):
recon = F.mse_loss(z_hat, z)
sparsity = concepts.abs().mean()
return recon + alpha * sparsity
26. ReLU is not proof of interpretability
This is another semantic boundary worth stating explicitly.
A sparse bottleneck does not automatically produce human-interpretable concepts.
To justify an interpretability claim, we would need evidence such as:
- feature activation examples;
- feature selectivity;
- intervention tests;
- causal feature steering;
- stability across checkpoints;
- correspondence with known dataset factors;
- comparison against random or dense bottlenecks.
The architecture creates an opportunity for interpretability.
It does not prove interpretability.
27. Measure sparsity directly
Useful metrics include:
Mean absolute activation
concepts.abs().mean()
Fraction of exact zeros
With ReLU:
(concepts == 0).float().mean()
Active features per sample
(concepts > threshold).float().sum(dim=-1).float().mean()
Dead features
A feature that never activates:
active_anywhere = (concepts > threshold).any(dim=0)
dead_fraction = (~active_anywhere).float().mean()
These metrics do not tell us what a feature means.
They do tell us whether the representation is actually sparse.
28. Intervention is stronger than inspection
Suppose concept dimension 37 appears to activate on poor factual answers.
That observation is interesting.
A stronger test is:
concepts[:, 37] = 0
Decode again and run the prediction heads.
Does the score change systematically?
Or amplify it:
concepts[:, 37] *= 2.0
Do downstream predictions move in the expected direction?
Interventions help distinguish:
correlation
from:
functional influence inside the model
29. A complete Tiny-style model
Now we can combine the pieces.
import torch
import torch.nn.functional as F
from torch import nn
class TinyBlock(nn.Module):
def __init__(self, d_model: int, dropout: float = 0.1):
super().__init__()
self.norm = nn.LayerNorm(d_model)
self.mlp = nn.Sequential(
nn.Linear(d_model, d_model * 4),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(d_model * 4, d_model),
nn.Dropout(dropout),
)
def forward(self, x):
return x + self.mlp(self.norm(x))
class TinyModel(nn.Module):
def __init__(
self,
d_model: int = 256,
n_layers: int = 2,
n_recursions: int = 6,
step_scale: float = 0.1,
dropout: float = 0.1,
):
super().__init__()
self.d_model = d_model
self.n_recursions = n_recursions
self.step_scale = step_scale
self.z_proj = nn.Linear(d_model * 3, d_model)
self.core = nn.Sequential(
*[
TinyBlock(d_model, dropout)
for _ in range(n_layers)
]
)
self.final_norm = nn.LayerNorm(d_model)
self.halt_head = nn.Linear(d_model, 1)
self.score_head = nn.Linear(d_model, 1)
self.logvar_head = nn.Linear(d_model, 1)
self.aux3_head = nn.Linear(d_model, 3)
self.disagree_head = nn.Linear(d_model, 1)
self.ood_head = nn.Linear(d_model, 1)
self.temp_head = nn.Linear(d_model, 1)
self.recon_head = nn.Linear(d_model, d_model)
self.sae_encoder = nn.Sequential(
nn.Linear(d_model, d_model // 2),
nn.ReLU(),
nn.LayerNorm(d_model // 2),
)
self.sae_decoder = nn.Linear(d_model // 2, d_model)
self.head_drop = nn.Dropout(dropout)
def _resolve_steps(self, steps):
if steps is None:
return self.n_recursions
return max(1, min(int(steps), self.n_recursions))
def recur(self, x, y, z=None, steps=None, return_trace=False):
assert x.ndim == 2
assert y.ndim == 2
assert x.shape == y.shape
assert x.shape[-1] == self.d_model
if z is None:
z = torch.zeros_like(x)
halt_logits = torch.full(
(x.shape[0], 1),
-1e9,
device=x.device,
dtype=x.dtype,
)
trace = []
for step in range(self._resolve_steps(steps)):
z_before = z
fused = torch.cat([x, y, z], dim=-1)
dz = torch.tanh(self.z_proj(fused))
dz = self.core(dz)
step_halt = self.halt_head(self.final_norm(dz))
halt_logits = torch.maximum(halt_logits, step_halt)
z = z + self.step_scale * dz
if return_trace:
trace.append({
"step": step,
"z": z,
"dz": dz,
"state_delta": (z - z_before).norm(dim=-1),
"halt_logit": step_halt,
})
z_final = self.final_norm(z)
concepts = self.sae_encoder(z_final)
sae_recon = self.sae_decoder(concepts)
z_head = self.head_drop(z_final + sae_recon)
return z_final, z_head, concepts, halt_logits, trace
def forward(self, x, y, z=None, steps=None, return_trace=False):
z_final, z_head, concepts, halt_logits, trace = self.recur(
x,
y,
z=z,
steps=steps,
return_trace=return_trace,
)
score_logit = self.score_head(z_head)
tau_raw = self.temp_head(z_head)
tau = 0.5 + 0.5 * F.softplus(tau_raw)
tau = tau.clamp_min(1e-2)
score = torch.sigmoid(score_logit / tau)
log_var = self.logvar_head(z_head).clamp(-5.0, 5.0)
aux3_logits = self.aux3_head(z_head)
disagree_logit = self.disagree_head(z_head)
ood_logit = self.ood_head(z_head)
recon = self.recon_head(z_head)
return {
"score": score,
"score_logit": score_logit,
"log_var": log_var,
"aux3_logits": aux3_logits,
"disagree_logit": disagree_logit,
"ood_logit": ood_logit,
"temperature": tau,
"reconstruction": recon,
"halt_logits": halt_logits,
"halt_prob": torch.sigmoid(halt_logits),
"z_final": z_final,
"concepts": concepts,
"trace": trace,
}
That is a substantial model.
But it is still made out of pieces we already understand.
30. The model inside the model, again
Decompose it:
TinyModel
│
├── state fusion
│ └── Linear(3D,D)
│
├── recursive core
│ ├── TinyBlock
│ │ ├── LayerNorm
│ │ ├── Linear
│ │ ├── GELU
│ │ ├── Linear
│ │ └── residual
│ └── TinyBlock
│
├── sparse autoencoder
│ ├── encoder
│ └── decoder
│
├── score head
├── uncertainty head
├── OOD head
├── disagreement head
├── calibration head
├── reconstruction head
└── halting head
A complicated architecture has become a tree of small modules.
That is the entire thesis of this series in one diagram.
31. Attention as an optional core
So far every recursion processes one vector per sample:
[B,D]
We can still wrap it as a sequence of length one and use self-attention:
class TinyAttentionBlock(nn.Module):
def __init__(
self,
d_model: int,
n_heads: int = 4,
dropout: float = 0.1,
):
super().__init__()
self.norm = nn.LayerNorm(d_model)
self.attn = nn.MultiheadAttention(
embed_dim=d_model,
num_heads=n_heads,
dropout=dropout,
batch_first=True,
)
self.drop = nn.Dropout(dropout)
self.ff = TinyBlock(d_model, dropout)
def forward(self, x):
squeeze = False
if x.ndim == 2:
x = x.unsqueeze(1)
squeeze = True
q = k = v = self.norm(x)
h, _ = self.attn(q, k, v, need_weights=False)
x = x + self.drop(h)
x = self.ff(x)
if squeeze:
x = x.squeeze(1)
return x
But this raises an important question.
32. Attention over a sequence of length one
If the input is:
[B,1,D]
then the attention matrix is effectively:
[1 × 1]
There is only one position to attend to.
So self-attention cannot express meaningful token-to-token or slot-to-slot routing in the usual sense.
The projections still transform the representation, but the attention mechanism has little relational structure to exploit.
This means we should distinguish:
attention module present
from:
attention is doing useful multi-position routing
That is exactly the sort of architectural claim this series should interrogate.
33. Give attention something to attend over
A more meaningful variant could preserve the three inputs as slots:
slot 0 = x
slot 1 = y
slot 2 = z
Stack them:
tokens = torch.stack([x, y, z], dim=1)
Shape:
[B,3,D]
Now self-attention can learn relationships among:
- context;
- candidate;
- current latent state.
After attention we can pool or select the latent slot.
For example:
updated = attention(tokens)
z_next = updated[:, 2]
That is a materially different architecture from pretending a single vector is a meaningful sequence.
34. Slot-based recursive attention
Here is a compact version:
class SlotRecursiveBlock(nn.Module):
def __init__(self, d_model=256, n_heads=4):
super().__init__()
self.norm = nn.LayerNorm(d_model)
self.attn = nn.MultiheadAttention(
d_model,
n_heads,
batch_first=True,
)
self.ff = TinyBlock(d_model)
def forward(self, x, y, z):
slots = torch.stack([x, y, z], dim=1)
h = self.norm(slots)
h, _ = self.attn(h, h, h, need_weights=False)
slots = self.ff(slots + h)
return slots[:, 2]
Now attention has a clear role.
We can ablate:
MLP fusion
vs
slot attention fusion
That is a meaningful experiment.
35. MLP vs attention should be measured, not assumed
Compare:
Model A: MLP recursive core
Model B: slot-attention recursive core
Control:
- parameter count;
- recursion depth;
- training budget;
- data;
- embedding model;
- optimizer;
- evaluation split.
Measure:
- ranking accuracy;
- calibration;
- throughput;
- latency;
- memory;
- sensitivity to recursion depth;
- robustness to input perturbation.
Attention is not automatically better because it is attention.
36. Train the score head first
When a model has many heads, debugging everything simultaneously becomes difficult.
Start with the simplest objective.
out = model(x, y)
score = out["score"].squeeze(-1)
loss = F.mse_loss(score, target)
Or binary quality:
loss = F.binary_cross_entropy_with_logits(
out["score_logit"].squeeze(-1),
target.float(),
)
Prove that the recursive core learns anything before adding seven auxiliary losses.
37. Then add auxiliary objectives one at a time
For example:
L_total =
L_score
+ λ_var L_var
+ λ_aux L_aux3
+ λ_ood L_ood
+ λ_recon L_recon
+ λ_sae L_sae
Do not add everything at once and then wonder why the model stopped learning.
For every new loss:
- establish a baseline;
- add one objective;
- measure the primary metric;
- measure the auxiliary metric;
- inspect gradient interaction;
- keep it only if it earns its complexity.
38. Gradient conflicts become more likely
All these heads share z_head.
Their gradients eventually reach the recursive core.
One task may push the representation one way while another pushes it another way.
We can measure cosine similarity between task gradients.
def flat_grad(loss, params):
grads = torch.autograd.grad(
loss,
params,
retain_graph=True,
allow_unused=True,
)
pieces = []
for p, g in zip(params, grads):
if g is None:
pieces.append(torch.zeros_like(p).reshape(-1))
else:
pieces.append(g.reshape(-1))
return torch.cat(pieces)
Then:
g_score = flat_grad(score_loss, shared_params)
g_recon = flat_grad(recon_loss, shared_params)
cos = F.cosine_similarity(
g_score.unsqueeze(0),
g_recon.unsqueeze(0),
).item()
Interpretation:
+1 same direction
0 orthogonal-ish
-1 directly conflicting
This is much better than saying multi-task learning is beneficial by default.
39. Trace the recursive trajectory
The most useful debugging surface in Tiny is often not the final output.
It is the sequence:
z0, z1, z2, ... zN
Track:
- state norm;
- state delta;
- cosine similarity to previous state;
- score at each step;
- uncertainty at each step;
- halt probability;
- concept sparsity.
For example:
@torch.no_grad()
def summarize_trace(trace):
rows = []
prev = None
for item in trace:
z = item["z"]
row = {
"step": item["step"],
"z_norm": z.norm(dim=-1).mean().item(),
"delta": item["state_delta"].mean().item(),
"halt": torch.sigmoid(item["halt_logit"]).mean().item(),
}
if prev is not None:
row["cos_prev"] = F.cosine_similarity(
prev,
z,
dim=-1,
).mean().item()
rows.append(row)
prev = z
return rows
40. Detect a model that is not really using recursion
A recursive architecture can silently collapse into effectively one-step behavior.
Symptoms:
z1 ≈ z2 ≈ z3 ≈ z4
or:
score_1 ≈ score_6
or:
performance(1 step) ≈ performance(6 steps)
Then the recursion may not be adding useful computation.
That is not necessarily a bug.
But it means the claim that recursive refinement matters has not been supported.
41. Detect unstable recursion
The opposite failure is exploding state dynamics.
Track:
z.norm(dim=-1).mean()
across steps.
If we see:
1.1
1.8
3.5
7.9
18.2
something is wrong.
Possible causes:
- step scale too high;
- poorly conditioned projection;
- too many recursive steps;
- unstable auxiliary losses;
- insufficient normalization.
Try:
- smaller
step_scale; - gradient clipping;
- normalized updates;
- fewer recursions;
- stronger normalization;
- explicit update magnitude penalties.
42. Update magnitude regularization
One optional control is:
L_update = mean(||dz||²)
In code:
update_loss = dz.pow(2).mean()
This encourages smaller proposed changes.
But again, it changes the objective.
Do not add it because it sounds sensible.
Test whether it improves:
- stability;
- primary metric;
- convergence across steps;
- robustness.
43. Tiny-batch overfit test
Before any large training run, make the model memorize a tiny dataset.
For example:
8 or 16 examples
Train repeatedly on those exact examples.
If the model cannot drive the training loss down substantially, do not tune the data loader or add more recursions.
Something is structurally wrong.
Possible causes:
- detached graph;
- wrong target;
- frozen parameters;
- incorrect loss;
- optimizer missing parameters;
- score saturation;
- dropout too strong;
- output/target shape broadcasting.
This diagnostic remains one of the highest-value tests in the entire series.
44. Prove the recursive core receives gradients
loss.backward()
for name, p in model.named_parameters():
if "z_proj" in name or "core" in name:
print(
name,
None if p.grad is None else p.grad.norm().item(),
)
If the heads receive gradients but the recursive core does not, then the architecture may be disconnected from the objective.
45. Prove the parameters actually change
Before the step:
before = {
name: p.detach().clone()
for name, p in model.named_parameters()
}
Then:
optimizer.step()
Compare:
for name, p in model.named_parameters():
delta = (p.detach() - before[name]).abs().max().item()
print(name, delta)
The optimizer existing is not evidence that the intended parameters changed.
Runtime evidence is better.
46. Recursion changes gradient depth
Even though parameters are reused, autograd still sees a deeper computation graph.
For six recursions:
parameter use
↓
parameter use
↓
parameter use
↓
parameter use
↓
parameter use
↓
parameter use
Gradients accumulate contributions through each use.
This can increase:
- gradient variance;
- memory use;
- exploding/vanishing tendencies;
- sensitivity to recursion count.
Monitor gradient norms by step-count configuration.
47. Gradient clipping
A practical guard is:
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0,
)
But clipping is not a substitute for understanding unstable dynamics.
If every batch is being aggressively clipped, inspect:
- recursion depth;
- step scale;
- loss weights;
- initialization;
- target scaling.
48. Compare against a feed-forward baseline
A fair experiment needs a non-recursive model.
class FeedForwardBaseline(nn.Module):
def __init__(self, d_model=256):
super().__init__()
self.net = nn.Sequential(
nn.Linear(d_model * 2, d_model),
nn.ReLU(),
nn.LayerNorm(d_model),
nn.Linear(d_model, d_model),
nn.ReLU(),
)
self.head = nn.Linear(d_model, 1)
def forward(self, x, y):
z = self.net(torch.cat([x, y], dim=-1))
return self.head(z).squeeze(-1)
If Tiny does not beat this baseline under appropriate controls, recursion may not be helping.
49. Parameter-matched baseline
Tiny may have fewer parameters than a large feed-forward network but more compute because parameters are reused.
So compare against a feed-forward model with approximately the same parameter count.
This asks:
Is the recursive structure better than using the same number of parameters once?
50. Compute-matched baseline
Also compare against a larger/deeper model with similar total FLOPs or latency.
This asks:
Is repeated parameter reuse better than spending the same compute on a deeper one-pass network?
These are different questions.
Both matter.
51. Recursion-depth ablation
Train or evaluate:
1
2
3
4
6
8
steps.
Plot:
quality metric vs recursion depth
latency vs recursion depth
memory vs recursion depth
The useful result is not necessarily monotonic.
You may find a clear elbow point.
That can define the default inference budget.
52. Step-scale ablation
Try:
0.01
0.05
0.10
0.25
0.50
1.00
Measure:
- training stability;
- final metric;
- state delta;
- state norm;
- gradient norm;
- sensitivity to recursion depth.
Step scale is not a cosmetic parameter.
It controls the dynamics of the latent trajectory.
53. Does the state accumulate information?
One test is to train probes on z_t at each recursion.
Suppose the task has known labels:
correctness
relevance
style
risk
Freeze the model.
For each step:
z1
z2
z3
...
train the same lightweight probe.
If some information becomes progressively easier to decode, we have evidence that the recursive state changes functionally across computation.
That is stronger than merely seeing that the vector moved.
54. Does later computation correct earlier mistakes?
A very strong analysis is per-example transition tracking.
For each sample:
step 1 prediction
step 2 prediction
...
step N prediction
Count:
wrong → right
right → wrong
wrong → wrong
right → right
If later recursion genuinely refines decisions, we should see useful wrong → right transitions.
If right → wrong is equally common, more recursion may just add noise.
55. Candidate ordering across steps
For ranking tasks, do not only inspect absolute scores.
Suppose one context has candidates A, B and C.
Track:
step 1: B > A > C
step 2: A > B > C
step 3: A > B > C
If A is the preferred candidate, the recursion corrected the ranking at step 2.
This is often more informative than a tiny change in MSE.
56. The sparse-autoencoder ablation
Compare:
Tiny without SAE
Tiny with dense bottleneck
Tiny with sparse bottleneck
Measure:
- primary task quality;
- calibration;
- reconstruction;
- representation sparsity;
- latency;
- parameter count;
- feature stability.
If the sparse autoencoder adds no measurable benefit, keep that result.
Architecture should earn its place.
57. Residual SAE vs bottleneck-only heads
There are at least two designs.
Residual
z_head = z_final + sae_decoder(concepts)
This preserves a direct path from z_final.
Bottleneck-only
z_head = sae_decoder(concepts)
Now every head must depend on information that passed through the bottleneck.
The second gives the bottleneck much more pressure to preserve task-relevant information.
But it may reduce performance.
That is a useful ablation.
58. Reconstruction target matters
What should the SAE reconstruct?
Possible targets include:
z_final
response embedding y
context embedding x
pair representation
Each objective encourages different information.
If we reconstruct z_final, the SAE compresses the model’s own latent state.
If we reconstruct y, it is encouraged to preserve response information.
Do not treat these objectives as interchangeable.
59. Length effects
If response length is available, we can inspect whether predictions depend systematically on it.
For example:
length_feature = torch.tanh(
sequence_length.float() / 512.0
)
But be careful.
Length can be:
- a legitimate signal;
- a shortcut;
- a dataset artifact.
A useful test is to stratify evaluation by response length.
If quality collapses on short or long responses, the aggregate metric may hide an important failure mode.
60. Detect shortcut learning
Suppose the model predicts quality using superficial properties of embeddings rather than the relationship between context and candidate.
A strong test is to shuffle the context-candidate pairing.
Keep the response embeddings and labels, but randomize contexts.
If performance barely changes, the model may not be using context meaningfully.
Similarly:
shuffle candidate
zero context
zero candidate
replace context with mean embedding
These simple ablations can reveal surprising shortcuts.
61. Input-ablation matrix
Evaluate:
| Input condition | What it tests |
|---|---|
full x,y,z |
normal model |
x=0 |
reliance on context |
y=0 |
reliance on candidate |
z=0 every step |
reliance on recursive state |
shuffled x |
context/candidate alignment |
shuffled y |
candidate specificity |
| one recursion | recursion contribution |
The exact numbers matter more than the architectural story we hope is true.
62. Zeroing z every step is a powerful test
Normal recursion:
z_{t+1} depends on z_t
Ablated recursion:
z input is always zero
The core is still run repeatedly, but it cannot accumulate state.
If performance is unchanged, then the recurrent state may not be doing meaningful work.
That is a particularly strong architectural ablation.
63. Randomize recursion order? There is only one state
HRM gave us a scheduling question between low- and high-level states.
Tiny removes that degree of freedom.
That is part of what makes it simpler.
The remaining schedule is mostly:
how many times do we apply the same update rule?
This reduction in architectural choices is one of Tiny’s strengths.
64. Simpler architecture means easier falsification
With HRM, a poor result could come from:
- low-state width;
- high-state width;
- L/H coupling;
- number of low steps;
- number of high cycles;
- recurrent cell choice.
With Tiny, the core questions are narrower:
Does recursive state help?
How many recursions help?
What step scale works?
Does the core need attention?
Does the SAE help?
Do the auxiliary heads help?
A simpler model can sometimes be scientifically easier to understand.
65. Parameter counting
Always count parameters explicitly.
def count_parameters(model):
return sum(
p.numel()
for p in model.parameters()
if p.requires_grad
)
Break them down by component:
def parameter_breakdown(model):
groups = {}
for name, p in model.named_parameters():
prefix = name.split(".")[0]
groups[prefix] = groups.get(prefix, 0) + p.numel()
return groups
You may discover that most parameters live in the MLP core rather than the many small diagnostic heads.
66. Compute accounting
A rough conceptual model is:
compute ≈ recursions × core_cost + head_cost
The heads are usually evaluated once.
The recursive core is evaluated repeatedly.
So a small increase in core width may be more expensive than adding another head.
This is why parameter count alone is not enough.
67. Throughput benchmark
import time
def benchmark(model, x, y, steps, runs=100):
model.eval()
if x.is_cuda:
torch.cuda.synchronize()
start = time.perf_counter()
with torch.inference_mode():
for _ in range(runs):
model(x, y, steps=steps)
if x.is_cuda:
torch.cuda.synchronize()
elapsed = time.perf_counter() - start
return elapsed / runs
Measure:
steps=1
steps=2
steps=4
steps=6
Performance gains need to be considered against latency cost.
68. Adaptive compute becomes possible
If some samples benefit from more recursion and others do not, we can imagine:
easy example → 2 steps
hard example → 6 steps
But adaptive compute should be justified by measured marginal benefit.
A simple offline policy is:
run 2 steps
if state delta > threshold:
continue
Then compare against fixed budgets.
69. Halting must save real compute
A halting mechanism that computes all six steps and merely reports that it would have stopped after step three saves nothing.
To matter operationally, inference must actually break:
for step in range(max_steps):
...
if should_stop:
break
For batched inference this becomes harder because different samples may want to stop at different times.
That introduces masking or dynamic batching complexity.
Again, the architectural idea is simpler than the production system around it.
70. Batch-wise adaptive halting
One simple implementation stops the whole batch when every sample meets a criterion.
if (delta < threshold).all():
break
This is easy but inefficient when one difficult sample keeps the rest running.
More advanced implementations maintain an active mask.
That is a useful later optimization, but it should not complicate the first model.
71. Active-mask recursion
Conceptually:
active = torch.ones(B, dtype=torch.bool, device=x.device)
for step in range(max_steps):
if not active.any():
break
# compute updates
...
z = torch.where(
active.unsqueeze(-1),
z_new,
z,
)
active = active & (delta >= threshold)
Whether this actually improves wall-clock performance depends on hardware and implementation.
Benchmark it.
72. Checkpoint the architecture contract
A checkpoint should include more than the state dict.
checkpoint = {
"model_state": model.state_dict(),
"d_model": model.d_model,
"n_recursions": model.n_recursions,
"step_scale": model.step_scale,
"embedding_model": embedding_model_name,
"embedding_dim": embedding_dim,
"training_objective": objective_name,
"sae_enabled": True,
"version": 1,
}
The weights are only meaningful relative to the architecture and embedding contract that produced them.
73. Embedding drift still matters
Tiny consumes embeddings.
If you change the upstream embedding model while keeping dimensions identical, the code may still run.
That does not mean the semantics are compatible.
A 256-dimensional vector from embedding system A is not automatically interchangeable with a 256-dimensional vector from embedding system B.
Version the embedding source.
Evaluate after any change.
74. Cache embeddings when they are frozen
If upstream embeddings are fixed, cache them.
Training then becomes:
load x embedding
load y embedding
run Tiny
backprop through Tiny only
This can dramatically reduce training cost when text encoding is expensive.
It also gives cleaner experiments because the representation input is held fixed across architecture comparisons.
75. Separate representation research from recursive-model research
A good experimental strategy is:
Phase 1:
freeze one embedding system
compare recursive architectures
Phase 2:
freeze the best recursive architecture
compare embedding systems
If both change simultaneously, attribution becomes difficult.
This evidence-first separation is especially important when many components are generated or modified by LLMs.
76. LLMs make architectural instrumentation more important
An LLM can generate this model in seconds.
It can also generate:
- another SAE;
- another halting head;
- another loss term;
- an attention variant;
- another optimizer;
- five more diagnostic heads.
That makes architectural complexity extremely cheap to create.
The scarce resource becomes something else:
evidence that each component is doing useful work.
For Tiny, that means traces, ablations and controlled comparisons matter more than the ability to generate another module.
77. A practical experiment ladder
I would evaluate Tiny in this order.
Experiment 1 — one-pass baseline
context + candidate → MLP → score
Experiment 2 — recursive state
same parameter scale
1/2/4/6 recursions
Experiment 3 — step scale
0.05 / 0.1 / 0.25
Experiment 4 — MLP vs attention core
Prefer a meaningful slot-attention version if testing attention.
Experiment 5 — SAE
none / dense bottleneck / sparse bottleneck
Experiment 6 — auxiliary heads
Add one at a time.
Experiment 7 — halting
Compare learned halting with deterministic delta thresholds.
That experiment sequence makes causal attribution much easier.
78. The complete training loop
Here is a deliberately simple score-only training loop.
def train_epoch(model, loader, optimizer, device):
model.train()
total_loss = 0.0
for x, y, target in loader:
x = x.to(device)
y = y.to(device)
target = target.to(device).float()
optimizer.zero_grad(set_to_none=True)
out = model(x, y)
logits = out["score_logit"].squeeze(-1)
loss = F.binary_cross_entropy_with_logits(
logits,
target,
)
loss.backward()
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0,
)
optimizer.step()
total_loss += loss.item() * x.shape[0]
return total_loss / len(loader.dataset)
Start here.
Then add complexity only after it works.
79. Validation by recursion depth
@torch.no_grad()
def validate_steps(model, loader, device, step_values=(1, 2, 4, 6)):
model.eval()
results = {}
for steps in step_values:
correct = 0
total = 0
for x, y, target in loader:
x = x.to(device)
y = y.to(device)
target = target.to(device)
out = model(x, y, steps=steps)
pred = (out["score"].squeeze(-1) >= 0.5)
correct += (pred == target.bool()).sum().item()
total += target.numel()
results[steps] = correct / max(total, 1)
return results
This should be a standard report for any recursive model.
80. Trajectory report for one sample
@torch.no_grad()
def inspect_sample(model, x, y):
model.eval()
out = model(
x.unsqueeze(0),
y.unsqueeze(0),
return_trace=True,
)
for item in out["trace"]:
print(
f"step={item['step']} "
f"delta={item['state_delta'].item():.6f} "
f"halt={torch.sigmoid(item['halt_logit']).item():.4f}"
)
print("final score:", out["score"].item())
Do this on:
- obvious positives;
- obvious negatives;
- borderline examples;
- examples the model gets wrong;
- distribution-shifted examples.
The trajectory can reveal behavior hidden by aggregate metrics.
81. What would count as evidence for iterative refinement?
Strong evidence would look something like:
- multi-step evaluation beats one-step evaluation;
- the gain survives parameter-matched baselines;
- the gain survives compute-matched baselines;
- removing state accumulation removes the gain;
- later steps correct meaningful earlier mistakes;
- state probes show useful information emerging across steps;
- gains reproduce across seeds;
- gains hold on a true held-out distribution.
That would justify stronger language about recursive refinement.
Without those tests, we should say only what the architecture actually does.
82. What would falsify the recursion story?
Suppose we find:
1 step = 0.782 accuracy
6 steps = 0.783 accuracy
and:
zero recurrent state = 0.782
Then the simplest interpretation is:
the recursive mechanism is not contributing much on this task.
That is a useful result.
It may mean:
- the task does not require iterative computation;
- the training objective does not reward it;
- the embedding already solves most of the task;
- the model needs a different recurrent structure.
Evidence can tell us not to keep complexity too.
83. Tiny vs HRM
We can now compare the architectural ideas directly.
HRM
two states
zL + zH
multiple update frequencies
GRU-based recurrent blocks
explicit hierarchy
Tiny
one state
z
one repeated update rule
residual refinement
optional MLP or attention core
optional sparse bottleneck
HRM asks whether hierarchy helps.
Tiny asks whether simple recursive refinement is enough.
Neither is universally superior.
That is an experimental question.
84. Tiny is a useful scientific baseline for HRM
This is easy to overlook.
If HRM beats a feed-forward model, we still do not know whether the benefit came from:
recurrence
or:
hierarchical recurrence
Tiny gives us the missing control.
Compare:
feed-forward
single-state recursion
hierarchical recursion
Now we can begin separating the effect of recurrence from the effect of hierarchy.
That makes Tiny valuable even if HRM ultimately performs better.
85. The architecture ladder so far
We started with MR.Q:
pair → representation → score
Then EBT:
pair → representation → Q / V / policy
Then SICQL:
pair → explicit modular Q / V / policy components
Then HRM:
pair → hierarchical recurrent state → diagnostic surface
Now Tiny:
pair + latent state
↓
recursive refinement
↓
sparse bottleneck
↓
diagnostic surface
Each stage adds one architectural pressure rather than inventing an unrelated model.
86. Why this model is called Tiny
The important interpretation of “Tiny” is not merely file size or parameter count.
It is architectural economy.
The model tries to get leverage from:
parameter reuse
instead of endlessly adding new blocks.
The same small core can be applied repeatedly.
That makes computation a tunable resource.
It also makes the model easier to instrument because the same transformation is revisited over a visible trajectory.
87. But repeated compute is not free
A 5-million-parameter model run ten times can be slower than a 20-million-parameter model run once.
So “small model” must not be confused with:
cheap inference
Always measure:
- parameters;
- FLOPs;
- wall-clock latency;
- memory;
- throughput;
- quality.
Recursive models make this distinction especially important.
88. What Tiny adds conceptually
The important additions are:
1. One explicit latent state
z_t
2. Recursive state refinement
z_{t+1} = z_t + α f(x,y,z_t)
3. Parameter reuse across computation
The same core runs every step.
4. Variable computation depth
One trained model can be evaluated with different recursion budgets.
5. Optional sparse latent bottleneck
A second model compresses the final state before prediction.
6. Rich diagnostic surface
Multiple heads can observe the same recursively refined representation.
Those are the ideas worth carrying forward.
89. What Tiny does not prove
Tiny does not prove that:
- recursion equals reasoning;
- SAE features are interpretable;
- halting probability identifies completed thought;
- OOD heads detect all distribution shift;
- uncertainty heads are calibrated;
- attention improves the recursive core;
- more inference steps improve decisions.
Those are all hypotheses.
The architecture gives us places to test them.
90. The broader lesson
A model that initially sounds complicated:
parameter-efficient recursive multi-task quality model with halting, uncertainty, sparse concept bottleneck and optional attention
reduces to:
concatenate three vectors
↓
linear projection
↓
small residual network
↓
add update to state
↓
repeat
↓
compress/reconstruct state
↓
a few linear heads
That is the point of Models From First Principles.
We are not making the model less sophisticated.
We are making the sophistication inspectable.
The full compact implementation
For reference, here is a smaller production-shaped version of the core model without every diagnostic convenience from earlier sections.
import torch
import torch.nn.functional as F
from torch import nn
class TinyBlock(nn.Module):
def __init__(self, d_model, dropout=0.1):
super().__init__()
self.norm = nn.LayerNorm(d_model)
self.net = nn.Sequential(
nn.Linear(d_model, 4 * d_model),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(4 * d_model, d_model),
nn.Dropout(dropout),
)
def forward(self, x):
return x + self.net(self.norm(x))
class Tiny(nn.Module):
def __init__(
self,
d_model=256,
n_layers=2,
n_recursions=6,
step_scale=0.1,
dropout=0.1,
):
super().__init__()
self.d_model = d_model
self.n_recursions = n_recursions
self.step_scale = step_scale
self.fuse = nn.Linear(3 * d_model, d_model)
self.core = nn.Sequential(
*[
TinyBlock(d_model, dropout)
for _ in range(n_layers)
]
)
self.final_norm = nn.LayerNorm(d_model)
self.sae_enc = nn.Sequential(
nn.Linear(d_model, d_model // 2),
nn.ReLU(),
nn.LayerNorm(d_model // 2),
)
self.sae_dec = nn.Linear(d_model // 2, d_model)
self.score = nn.Linear(d_model, 1)
self.logvar = nn.Linear(d_model, 1)
self.halt = nn.Linear(d_model, 1)
def forward(self, x, y, steps=None):
if x.shape != y.shape:
raise ValueError(
f"x/y shape mismatch: {x.shape} vs {y.shape}"
)
if x.shape[-1] != self.d_model:
raise ValueError(
f"expected d_model={self.d_model}, got {x.shape[-1]}"
)
steps = self.n_recursions if steps is None else int(steps)
steps = max(1, min(steps, self.n_recursions))
z = torch.zeros_like(x)
halt_logit = torch.full(
(x.shape[0], 1),
-1e9,
device=x.device,
dtype=x.dtype,
)
for _ in range(steps):
fused = torch.cat([x, y, z], dim=-1)
dz = torch.tanh(self.fuse(fused))
dz = self.core(dz)
step_halt = self.halt(self.final_norm(dz))
halt_logit = torch.maximum(halt_logit, step_halt)
z = z + self.step_scale * dz
z = self.final_norm(z)
concepts = self.sae_enc(z)
z_head = z + self.sae_dec(concepts)
score_logit = self.score(z_head)
log_var = self.logvar(z_head).clamp(-5.0, 5.0)
return {
"score_logit": score_logit,
"score": torch.sigmoid(score_logit),
"log_var": log_var,
"halt_logit": halt_logit,
"halt_prob": torch.sigmoid(halt_logit),
"concepts": concepts,
"state": z,
}
There is enough here to train, measure and falsify.
That is what we want.
Where do we go next?
Tiny introduced the sparse autoencoder as one component inside a larger recursive model.
It deserves a deeper look of its own.
Because once you have a latent state like:
z = [0.17, -0.82, 0.04, ...]
one of the biggest questions becomes:
Can we make the internal representation more inspectable without destroying the task performance that made it useful?
That leads directly to the next post:
Models From First Principles 06
Inside Tiny — Residual Blocks, Attention and Sparse Autoencoders
We will pull the model apart one level further and ask which of those internal components actually earn their place.
