Inside Tiny: Residual Blocks, Attention and Sparse Autoencoders

In the previous post, we built a compact recursive model around one idea:

context + candidate + latent state
              ↓
          projection
              ↓
         reusable core
              ↓
       proposed update
              ↓
      z ← z + α · update
              ↓
            repeat

That architecture looked more sophisticated than MR.Q, EBT or SICQL because it introduced recurrence.

But the central idea of this series is that a model stops looking mysterious when we keep opening it.

So in this post we are going to open Tiny.

The recursive shell itself is not the whole model.

Inside it are smaller components:

Tiny
├── state fusion
├── residual block
│   ├── LayerNorm
│   ├── Linear
│   ├── GELU
│   ├── Linear
│   └── residual addition
├── optional attention block
│   ├── LayerNorm
│   ├── MultiheadAttention
│   └── residual addition
├── sparse autoencoder
│   ├── encoder
│   ├── sparse latent code
│   └── decoder
└── output heads

Each of those pieces is itself small enough to understand line by line.

That is the point of this post.

We are going one level deeper.

1. Why study the inside of Tiny separately?

Because model architecture is compositional.

If we only look at the top-level class, we see names such as:

• recursive latent state;
• attention core;
• sparse autoencoder;
• halting head;
• uncertainty head.

Those names can make the model sound complicated.

But they are not explanations.

A better question is:

What exact tensor operation happens next?

For Tiny, most of the model can be reduced to a handful of operations:

normalize
project
activate
add
repeat

Then optionally:

split into query/key/value
compare positions
mix information
add residual

Then:

compress
activate sparsely
reconstruct

Once we understand those pieces, the recursive model becomes ordinary composition.

2. Start with the smallest useful residual block

A simple Tiny-style block can be written as:

import torch
import torch.nn as 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, 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.mlp(self.norm(x))

This already contains a surprisingly important pattern:

x
│
├───────────────┐
│               │
▼               │
LayerNorm       │
▼               │
Linear          │
▼               │
GELU            │
▼               │
Linear          │
▼               │
update          │
│               │
└──── add ◄─────┘
      ↓
    output

The block does not replace the existing representation.

It proposes an update.

Then the residual connection adds that update back to the original state.

That distinction matters.

3. A residual block is a learned correction

Without the residual connection, we might write:

def forward(self, x):
    return self.mlp(self.norm(x))

The model must produce an entirely new representation at every layer.

With a residual connection:

return x + self.mlp(self.norm(x))

we can think of the block as learning:

new_state = old_state + correction

or mathematically:

x' = x + f(x)

This is an extremely useful way to reason about deep and recursive models.

The network does not need to reconstruct everything it already knew.

It can learn only the change.

4. Residual reasoning appears twice in Tiny

Tiny actually contains two related residual ideas.

Inside each block:

h' = h + MLP(LN(h))

And outside the block, during recursion:

z' = z + α · update

So we have residual updates at two scales:

local block residual
        ↓
core representation update
        ↓
recursive state residual

This matters because stability now depends on both.

A block may produce a reasonable update while the recursive step scale still makes the state unstable.

Or the recursive step scale may be small while the block itself creates pathological activations.

The two should be diagnosed separately.

5. Why normalize before the MLP?

The block uses:

self.norm = nn.LayerNorm(d_model)

and then:

self.mlp(self.norm(x))

This is a pre-normalization pattern.

The rough data flow is:

raw state
  ↓
normalize features
  ↓
learn update
  ↓
add original state

LayerNorm operates across the feature dimension for each example.

For a tensor:

[B, D]

normalization is performed over D.

For:

[B, T, D]

it is still performed over the last dimension D independently for each batch/time position.

This shape-preserving property makes it convenient inside residual blocks.

6. Do not confuse normalization with information

Normalization changes the scale and distribution of features.

It does not create new information.

If the state entering the block has collapsed, normalization cannot recover distinctions that are gone.

That means we should still inspect representation diversity.

For example:

def batch_feature_stats(x):
    return {
        "mean": x.mean().item(),
        "std": x.std().item(),
        "row_cosine_mean": torch.nn.functional.cosine_similarity(
            x[:-1], x[1:], dim=-1
        ).mean().item() if x.size(0) > 1 else float('nan'),
    }

If every example becomes nearly identical, the problem is representation collapse, not merely normalization.

7. Why expand to four times the dimension?

The MLP uses:

nn.Linear(d_model, 4 * d_model)

followed by:

nn.Linear(4 * d_model, d_model)

So if:

d_model = 256

then the hidden layer is:

1024

The block briefly moves into a wider feature space:

256 → 1024 → 256

This gives the nonlinear transformation more capacity while preserving the external interface.

The outside world still sees [B, 256].

The interior gets a larger working space.

8. Count the parameters instead of guessing

For the two linear layers:

D → 4D
4D → D

ignoring biases for a moment, the parameter count is approximately:

4D² + 4D² = 8D²

For D = 256:

8 × 256² = 524,288

That is already over half a million weights in what looks like a tiny MLP block.

The lesson is important:

A model can look small in source code while still contain substantial parameter capacity.

Always count.

def count_parameters(module):
    return sum(p.numel() for p in module.parameters())

9. GELU is just the nonlinearity in the middle

The block uses:

nn.GELU()

Without a nonlinear activation, stacking linear transformations gives another linear transformation.

Conceptually:

Linear(Linear(x))

can be collapsed into one linear mapping.

The activation prevents that collapse.

So the MLP is really:

project wider
    ↓
nonlinear transform
    ↓
project back

The exact activation is a design choice.

It should be tested if it matters materially.

10. Dropout is not reasoning

Dropout appears in the block:

nn.Dropout(dropout)

Its purpose is regularization.

It randomly zeros activations during training.

It should not be described as part of the model’s reasoning mechanism.

This distinction matters because architectural explanations often mix together:

• representation machinery;
• optimization machinery;
• regularization machinery.

Dropout belongs primarily to the third category.

11. Test the block before using it recursively

Before embedding this block in a recursive model, test it independently.

D = 64
block = TinyBlock(D)

x = torch.randn(8, D)
y = block(x)

assert y.shape == x.shape
assert torch.isfinite(y).all()

Then verify gradients:

loss = y.pow(2).mean()
loss.backward()

for name, p in block.named_parameters():
    print(name, p.grad is None)

If the component is broken by itself, recurrence will only make debugging harder.

Part II — Attention

12. What problem does attention solve here?

An MLP transforms features at one position.

Attention allows one position to condition its update on other positions.

That distinction only matters if there are multiple positions.

This gives us a critical rule:

Before adding attention, identify what is attending to what.

If you cannot answer that clearly, the attention layer may be decorative.

13. The common single-position trap

Suppose our latent representation is:

[B, D]

To feed it into nn.MultiheadAttention, we might convert it to:

x = x.unsqueeze(1)

which gives:

[B, 1, D]

There is only one sequence position.

Attention scores therefore have shape:

[B, H, 1, 1]

There is nothing to choose between.

The softmax over one element is always 1.

So while the projections inside the attention module still perform learned transformations, the mechanism is not performing meaningful routing between multiple positions.

This is a structural issue, not a training issue.

14. Attention becomes meaningful with actual slots

We can instead represent three distinct pieces of information as separate positions:

position 0 = context
position 1 = candidate
position 2 = current latent state

Then:

[B, 3, D]

Now the attention matrix is:

[B, H, 3, 3]

Each slot can assign different weight to the others.

That gives attention something real to do.

15. Build a slot attention core

A minimal block:

class SlotAttentionBlock(nn.Module):
    def __init__(self, d_model, n_heads=4, dropout=0.1):
        super().__init__()

        self.norm1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=n_heads,
            dropout=dropout,
            batch_first=True,
        )

        self.norm2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),
            nn.Linear(4 * d_model, d_model),
        )

    def forward(self, slots):
        h = self.norm1(slots)
        attn_out, attn_weights = self.attn(
            h, h, h,
            need_weights=True,
            average_attn_weights=False,
        )

        slots = slots + attn_out
        slots = slots + self.mlp(self.norm2(slots))

        return slots, attn_weights

Now create the slots:

slots = torch.stack([context, candidate, latent], dim=1)

Shape:

[B, 3, D]

Run attention:

slots, weights = block(slots)

Then extract the updated latent slot:

latent_update = slots[:, 2]

This gives us a much clearer architectural story.

16. What exactly are Q, K and V here?

Within multi-head self-attention, each slot is projected into:

query
key
value

The queries ask:

What information am I looking for?

Keys describe:

What kind of information do I contain?

Values provide:

What information should be transferred if I am attended to?

Mathematically, the core score is based on:

QKᵀ / √d

then softmax produces routing weights.

The weighted values are then mixed.

The important point is that this operation acts over positions, not magically over abstract concepts.

We decide what the positions mean by how we construct the input.

17. Heads divide the feature space

If:

D = 256
H = 4

then each head typically works with:

64 dimensions

Conceptually:

256-dimensional slot
        ↓
4 parallel 64-dimensional attention views
        ↓
concatenate
        ↓
256-dimensional output

Different heads may learn different routing patterns.

But again:

Having multiple heads does not prove that they specialize meaningfully.

Inspect the behavior.

18. Attention weights are evidence, but weak evidence

We can inspect:

weights.shape

which for batch size B, heads H, and three slots is:

[B, H, 3, 3]

We can average across batches:

mean_weights = weights.mean(dim=0)

and examine whether the latent slot attends differently to context and candidate.

But attention weights should not automatically be treated as explanations.

They show routing inside one mechanism.

They do not prove causal importance.

A stronger test is intervention.

19. Attention ablation is stronger than attention visualization

Suppose the latent slot attends strongly to the candidate slot.

A useful experiment is to remove the candidate information:

candidate_zero = torch.zeros_like(candidate)

Then compare model behavior.

Or shuffle candidates across the batch:

perm = torch.randperm(candidate.size(0))
shuffled = candidate[perm]

If predictions barely change, the attention pattern was not doing much useful work.

This is the recurring theme of the series:

visualization < intervention

20. MLP vs attention must be an experiment

We can create two models:

Tiny-MLP
Tiny-Attention

Keep approximately matched:

• parameter count;
• recursion depth;
• training budget;
• embedding inputs;
• evaluation protocol.

Then measure:

• ranking accuracy;
• calibration;
• robustness;
• throughput;
• memory;
• state convergence.

Only then can we say whether attention helped.

Part III — Sparse Autoencoders

21. Why add a bottleneck at all?

After several recursive updates, Tiny has a latent state:

z_final ∈ R^D

We may want to know whether that state can be represented through a smaller set of active features.

A sparse autoencoder introduces:

z_final
   ↓
encoder
   ↓
code c
   ↓
decoder
   ↓
reconstruction

The hope is that the code becomes:

• compressed;
• sparse;
• useful for analysis.

But those are separate claims.

22. Build the smallest autoencoder first

class Autoencoder(nn.Module):
    def __init__(self, d_model, d_code):
        super().__init__()

        self.encoder = nn.Linear(d_model, d_code)
        self.decoder = nn.Linear(d_code, d_model)

    def forward(self, x):
        code = self.encoder(x)
        recon = self.decoder(code)
        return code, recon

Train using reconstruction loss:

code, recon = model(x)
loss = torch.nn.functional.mse_loss(recon, x)

This is compression.

It is not yet sparse.

23. Add a nonlinearity

A common simple version uses ReLU:

code = torch.relu(self.encoder(x))

Now activations are non-negative and some may become exactly zero.

That can create sparse activation patterns.

A compact module:

class SparseAutoencoder(nn.Module):
    def __init__(self, d_model, d_code):
        super().__init__()

        self.encoder = nn.Sequential(
            nn.Linear(d_model, d_code),
            nn.ReLU(),
        )

        self.decoder = nn.Linear(d_code, d_model)

    def forward(self, x):
        code = self.encoder(x)
        recon = self.decoder(code)
        return code, recon

24. ReLU alone does not guarantee useful sparsity

We should measure sparsity explicitly.

def sparsity_fraction(code, eps=1e-8):
    return (code.abs() <= eps).float().mean().item()

Possible output:

0.74

meaning 74% of activations are approximately zero.

That is useful descriptive evidence.

It still does not tell us what the active features mean.

25. Add an explicit sparsity penalty

We can penalize activation magnitude:

sparsity_loss = code.abs().mean()

Then:

loss = recon_loss + alpha * sparsity_loss

This creates an explicit trade-off:

reconstruct accurately
       vs
use fewer / smaller active features

As alpha increases, sparsity may increase while reconstruction worsens.

That trade-off should be measured.

26. The sparse autoencoder is itself a model-selection problem

Important hyperparameters include:

• code dimension;
• sparsity coefficient;
• activation function;
• decoder normalization;
• reconstruction objective;
• training data;
• whether the main model is frozen.

So adding an SAE does not end the model-design problem.

It creates another one.

27. Dense bottleneck vs sparse bottleneck

Always compare against a dense control.

For example:

Model A:
z → Linear → ReLU → Linear → recon

Model B:
z → Linear → GELU → Linear → recon

or:

Model C:
z → Linear → Linear → recon

If the sparse variant performs similarly to the dense variant but provides much cleaner activation structure, that is interesting.

If the sparse variant performs worse and the features are unstable, the extra complexity may not be justified.

28. Reconstruction is necessary, not sufficient

A good SAE should reconstruct the latent representation reasonably well.

Measure:

recon_error = ((recon - x) ** 2).mean()

or cosine similarity:

cos = torch.nn.functional.cosine_similarity(
    recon, x, dim=-1
).mean()

But good reconstruction alone does not imply interpretability.

A dense random rotation can preserve information while being difficult to interpret.

29. What would make a feature interesting?

Suppose feature c[:, 17] activates strongly.

We can search for examples where that activation is high:

values = code[:, 17]
indices = values.argsort(descending=True)[:20]

Then inspect the corresponding inputs.

Questions:

• Do high-activation examples share a semantic property?
• Does the property recur on held-out data?
• Is the feature stable across training seeds?
• Is it redundant with other features?

These are much stronger questions than simply plotting the activation.

30. Intervention is stronger than correlation

If feature 17 appears related to a property, intervene on it.

For example:

code_modified = code.clone()
code_modified[:, 17] = 0

Then decode:

recon_modified = decoder(code_modified)

Feed the modified representation into downstream heads and measure what changes.

Or amplify it:

code_modified[:, 17] *= 2.0

A feature whose manipulation systematically changes relevant downstream behavior is much more interesting than one that merely correlates with examples.

31. Dead features matter

A sparse model can become too sparse.

Some features may never activate.

Measure per-feature activation frequency:

def activation_frequency(code, eps=1e-6):
    return (code.abs() > eps).float().mean(dim=0)

Then:

freq = activation_frequency(code)
dead = (freq == 0).sum().item()

If half the code dimensions are permanently dead, the effective capacity is lower than the declared dimension.

That may be fine.

But we should know it.

32. Feature redundancy matters too

Two features may activate together almost all the time.

We can inspect correlations:

centered = code - code.mean(dim=0, keepdim=True)
cov = centered.T @ centered / max(code.size(0) - 1, 1)

Highly redundant features may indicate that the SAE is not using its dictionary efficiently.

Again, interpretation requires measurement.

Part IV — Combining the components

33. A Tiny-style core from independent pieces

Now we can assemble a recursive unit explicitly.

class RecursiveCore(nn.Module):
    def __init__(
        self,
        d_model,
        step_scale=0.1,
        use_attention=False,
        n_heads=4,
    ):
        super().__init__()

        self.d_model = d_model
        self.step_scale = step_scale
        self.use_attention = use_attention

        self.fuse = nn.Linear(3 * d_model, d_model)

        if use_attention:
            self.core = SlotAttentionBlock(d_model, n_heads=n_heads)
        else:
            self.core = TinyBlock(d_model)

    def forward(self, context, candidate, latent):
        if self.use_attention:
            slots = torch.stack(
                [context, candidate, latent],
                dim=1,
            )

            slots, weights = self.core(slots)
            proposal = slots[:, 2]

        else:
            fused = torch.cat(
                [context, candidate, latent],
                dim=-1,
            )

            proposal = torch.tanh(self.fuse(fused))
            proposal = self.core(proposal)
            weights = None

        latent = latent + self.step_scale * proposal

        return latent, weights

Notice that this module is not mysterious anymore.

It contains components we have already built independently.

34. Add the SAE after recursion

class TinyWithSAE(nn.Module):
    def __init__(
        self,
        d_model=256,
        d_code=128,
        n_steps=6,
        step_scale=0.1,
    ):
        super().__init__()

        self.n_steps = n_steps

        self.core = RecursiveCore(
            d_model=d_model,
            step_scale=step_scale,
        )

        self.sae = SparseAutoencoder(
            d_model=d_model,
            d_code=d_code,
        )

        self.score_head = nn.Linear(d_model, 1)

    def forward(self, context, candidate):
        latent = torch.zeros_like(context)

        trajectory = []

        for _ in range(self.n_steps):
            latent, _ = self.core(
                context,
                candidate,
                latent,
            )

            trajectory.append(latent)

        code, recon = self.sae(latent)

        head_state = latent + recon
        score = self.score_head(head_state).squeeze(-1)

        return {
            "score": score,
            "latent": latent,
            "code": code,
            "recon": recon,
            "trajectory": trajectory,
        }

We can now inspect every stage.

35. Why add the reconstruction residually?

One pattern is:

head_state = latent + recon

This means the SAE decoder is not replacing the latent state.

It contributes an additional learned reconstruction term.

That gives the downstream head access to:

original latent
      +
SAE-decoded representation

This is a design choice.

Another option is:

head_state = recon

which forces the prediction heads to rely entirely on the bottleneck.

These two designs answer different questions.

36. Soft bottleneck vs hard bottleneck

Soft bottleneck

head_state = latent + recon

The SAE influences the prediction but cannot fully block information outside the code.

Hard bottleneck

head_state = recon

All downstream information must pass through the code.

If our goal is interpretability, the hard bottleneck is a stronger test.

If our goal is preserving task performance while adding an auxiliary sparse representation, the soft version may be safer.

Do not confuse the two.

37. Train reconstruction separately from the main task

One useful experiment is:

1. train the recursive model normally;
2. freeze it;
3. collect latent states;
4. train the SAE separately;
5. analyze the sparse code.

This isolates the SAE from the main task.

Another experiment is joint training:

main loss
  +
reconstruction loss
  +
sparsity penalty

Joint training may change the representation itself.

That can be good or bad.

The distinction should be explicit.

38. A simple joint objective

Suppose our main task is binary quality prediction.

main_loss = torch.nn.functional.binary_cross_entropy_with_logits(
    output["score"],
    labels.float(),
)

Reconstruction:

recon_loss = torch.nn.functional.mse_loss(
    output["recon"],
    output["latent"].detach(),
)

Sparsity:

sparse_loss = output["code"].abs().mean()

Total:

loss = (
    main_loss
    + 0.1 * recon_loss
    + 0.01 * sparse_loss
)

The coefficients are not truths.

They are hyperparameters.

Measure sensitivity to them.

39. Be careful with .detach()

In the reconstruction target above:

output["latent"].detach()

we prevent the reconstruction loss from pushing directly on the recursive representation.

That means:

SAE learns to represent latent

rather than:

latent learns to become easy for SAE to reconstruct

If we remove .detach(), the training dynamics change.

Neither choice is universally correct.

But they are scientifically different.

Part V — What to measure

40. Component-level parameter counts

Do not only count the whole model.

def named_parameter_counts(model):
    for name, module in model.named_children():
        n = sum(p.numel() for p in module.parameters())
        print(f"{name:20s} {n:,}")

This tells us where capacity actually lives.

A tiny output head may be irrelevant to parameter cost compared with the 4× MLP expansion.

41. Runtime cost by component

Parameter count is not runtime cost.

Attention cost depends strongly on sequence length.

Recursion multiplies repeated compute.

SAE cost is paid every forward pass if it is in the inference path.

Benchmark pieces independently.

import time


def benchmark(fn, n=100):
    start = time.perf_counter()

    for _ in range(n):
        fn()

    return (time.perf_counter() - start) / n

On CUDA, synchronize around timing.

42. State trajectory diagnostics

Record latent states:

trajectory = torch.stack(output["trajectory"], dim=1)

Shape:

[B, steps, D]

Then compute step deltas:

deltas = (
    trajectory[:, 1:] - trajectory[:, :-1]
).norm(dim=-1)

If deltas stay huge forever, the state may not be converging.

If they collapse immediately, later recursion may be wasted compute.

43. Prediction trajectory matters more than latent movement alone

A changing latent state does not necessarily imply changing decisions.

Run the score head at every step:

step_scores = []

for z in trajectory.unbind(dim=1):
    step_scores.append(
        torch.sigmoid(model.score_head(z)).squeeze(-1)
    )

step_scores = torch.stack(step_scores, dim=1)

Now measure:

score step 1
score step 2
score step 3
...

Interesting cases include:

wrong → right
right → wrong
right → right
wrong → wrong

Recursion should justify its compute by improving something measurable.

44. Measure sparsity by step too

If the SAE is applied to intermediate states, we can ask whether sparse structure emerges gradually.

For each recursive step:

code, _ = model.sae(z)

Measure:

• active feature count;
• reconstruction error;
• feature stability;
• feature overlap between steps.

This can tell us whether recursion changes only the dense representation or also the sparse concept structure.

45. Measure feature stability across seeds

Train several models with different random seeds.

Then compare SAE feature behavior.

If apparently meaningful features appear only in one seed and disappear completely in another, interpretability claims should be cautious.

A useful representation should ideally show some reproducible structure.

Feature identity may not align one-to-one across runs, so comparisons may require matching features by activation similarity or decoder direction.

46. Measure decoder direction similarity

Each SAE feature has a decoder direction.

If the decoder weight matrix is:

[d_model, d_code]

then each code dimension corresponds to one direction back into model space.

We can compare directions with cosine similarity.

This gives another way to detect redundant or unstable features.

47. Measure causal usefulness of features

For each feature:

1. identify examples where it activates strongly;
2. set it to zero;
3. recompute downstream predictions;
4. record the effect;
5. repeat on held-out examples.

This gives a feature-effect profile.

A feature that activates on a concept but whose intervention changes nothing downstream may be descriptive rather than causally important.

Part VI — Failure modes

48. Failure mode: attention over one token

Symptoms:

attention layer present
attention weights shape [..., 1, 1]

The layer may still learn projections, but attention is not routing between positions.

Fix:

Create meaningful slots or remove the attention layer.

49. Failure mode: residual explosion

If repeated updates grow too large:

||z_1|| < ||z_2|| < ||z_3|| << ||z_6||

inspect:

• step scale;
• block output norm;
• normalization;
• gradient norms.

A simple diagnostic:

for i, z in enumerate(trajectory):
    print(i, z.norm(dim=-1).mean().item())

50. Failure mode: residual inactivity

The opposite problem:

proposal norm ≈ 0

Then recursion does almost nothing.

Measure:

ratio = proposal.norm(dim=-1) / (
    latent.norm(dim=-1) + 1e-8
)

If the ratio is tiny at every step, later recursion may be decorative compute.

51. Failure mode: SAE learns identity-like dense codes

If almost every code feature activates for almost every example, the model is not meaningfully sparse.

Measure:

sparsity_fraction(code)

Do not assume ReLU automatically solved the problem.

52. Failure mode: SAE collapses to dead features

If most features never activate, sparsity pressure may be too strong.

Inspect:

activation_frequency(code)

Tune:

• sparsity coefficient;
• code width;
• learning rate;
• initialization.

53. Failure mode: reconstruction dominates the real task

If the total loss is:

main + λ_recon * reconstruction + λ_sparse * sparsity

and reconstruction gradients dominate, the model may optimize representation preservation rather than the task we care about.

Inspect gradient norms per objective.

54. Failure mode: task dominates and SAE becomes irrelevant

The opposite can happen too.

If the SAE receives tiny gradients, its code may remain noisy and unstructured.

Again, inspect gradient magnitudes rather than guessing.

55. Failure mode: interpreting attention as explanation

Attention weights are internal routing weights.

They are not automatically causal explanations.

Use:

• masking;
• slot ablation;
• permutation;
• counterfactual replacement.

56. Failure mode: interpreting sparse features by anecdotes

Finding three examples that seem related is not enough.

Use:

• held-out examples;
• top-k activation sets;
• negative examples;
• seed replication;
• intervention.

Interpretation should survive attempts to falsify it.

Part VII — Controlled experiments

57. Experiment 1: no residual vs residual

Compare:

MLP(x)

against:

x + MLP(LN(x))

Measure:

• optimization speed;
• gradient norms;
• final task performance;
• representation stability.

58. Experiment 2: different residual scales

Test:

α ∈ {0.01, 0.05, 0.1, 0.25, 0.5, 1.0}

Measure:

• convergence;
• instability;
• steps needed;
• score trajectory.

The best scale is empirical.

59. Experiment 3: MLP vs one-position attention

This is a useful negative control.

Compare:

TinyBlock

against:

MultiheadAttention over [B,1,D]

If they perform similarly, that is unsurprising.

The attention mechanism had no multi-position routing opportunity.

60. Experiment 4: MLP vs slot attention

Now compare MLP with meaningful three-slot attention.

[context, candidate, latent]

Measure:

• task quality;
• latency;
• attention patterns;
• robustness to candidate shuffling.

This is the more meaningful test.

61. Experiment 5: no bottleneck vs dense bottleneck vs SAE

Compare:

A: no bottleneck
B: dense autoencoder
C: sparse autoencoder

Measure:

• task performance;
• reconstruction;
• sparsity;
• feature stability;
• intervention effects.

This isolates the benefit of sparsity from the benefit of simply adding another network.

62. Experiment 6: soft vs hard bottleneck

Compare:

soft: latent + recon
hard: recon only

If the hard bottleneck preserves task performance, the SAE code is carrying substantial task-relevant information.

That is stronger evidence than the soft case.

63. Experiment 7: freeze model, train SAE later

Train Tiny first.

Freeze it.

Train the SAE on recorded latent states.

Compare this with joint training.

This tells us whether interpretability pressure changes the underlying task representation.

64. Experiment 8: feature intervention

For the most active SAE features:

zero feature
amplify feature
replace feature with batch mean

Measure downstream score changes.

This is where the interpretability claim becomes experimentally interesting.

Part VIII — Debugging

65. Shape contracts first

Use explicit assertions:

B, D = context.shape

assert candidate.shape == (B, D)
assert latent.shape == (B, D)

For slot attention:

slots = torch.stack([context, candidate, latent], dim=1)
assert slots.shape == (B, 3, D)

For SAE:

code, recon = sae(latent)
assert recon.shape == latent.shape

Shape bugs are cheaper to catch immediately.

66. Inspect gradients by component

def grad_norm(module):
    total = 0.0

    for p in module.parameters():
        if p.grad is not None:
            total += p.grad.detach().pow(2).sum().item()

    return total ** 0.5

Then print:

print("core", grad_norm(model.core))
print("sae", grad_norm(model.sae))
print("score", grad_norm(model.score_head))

If one component receives no gradient, understand why before tuning hyperparameters.

67. Verify parameters actually update

before = {
    name: p.detach().clone()
    for name, p in model.named_parameters()
}

loss.backward()
optimizer.step()

for name, p in model.named_parameters():
    delta = (p.detach() - before[name]).abs().max().item()
    print(name, delta)

This remains one of the most useful debugging tests in the whole series.

68. Overfit a tiny dataset

Before running a large experiment, train on a tiny batch repeatedly.

If the model cannot memorize a handful of examples, suspect:

• broken objective;
• disconnected gradients;
• wrong labels;
• overly aggressive regularization;
• incorrect shape handling;
• optimizer mistakes.

Do not blame model capacity first.

69. Keep trajectory logs

Store per-step:

latent norm
proposal norm
score
halt probability
SAE sparsity
reconstruction error

This turns recursion from an opaque loop into observable computation.

Part IX — What this teaches us about model design

70. Tiny is not one model

Tiny is a composition of:

state transition model
      +
residual MLP model
      +
optional routing model
      +
sparse representation model
      +
prediction heads

Once separated, every component can be:

• tested;
• replaced;
• frozen;
• ablated;
• benchmarked.

That is the broader architectural lesson.

71. A reusable interface matters more than the component

Suppose the recursive core obeys:

(context, candidate, latent)
        →
new latent

Then we can swap:

MLP
GRU
attention
convolution
state-space block

without rewriting the surrounding model.

This is exactly why explicit component boundaries mattered in SICQL too.

Good interfaces turn architecture into experimentation.

72. Model sophistication often comes from scheduling

The individual Tiny block is simple.

The interesting behavior comes from using it repeatedly:

same parameters
new latent state
same inputs
repeat

This is a recurring lesson from HRM and Tiny:

Compute schedule can be as important as parameter count.

A small model run six times is not computationally equivalent to the same model run once.

73. Interpretability is another model layer, not a free property

Adding an SAE does not suddenly make the system interpretable.

It creates a new learned representation that can potentially be investigated.

Then we still need:

feature discovery
      ↓
replication
      ↓
negative controls
      ↓
intervention
      ↓
causal validation

Interpretability is an empirical program.

74. Attention should earn its place

Attention is useful when it can route information between meaningful positions.

If the sequence length is one, it may be unnecessary complexity.

If we create meaningful slots, it may become valuable.

But the model must demonstrate that value empirically.

75. Sparse autoencoders should earn their place too

An SAE adds:

• parameters;
• compute;
• an auxiliary objective;
• hyperparameters;
• another failure surface.

Its inclusion should produce something measurable:

• useful compression;
• stable sparse features;
• interpretable interventions;
• better diagnostics;
• improved robustness.

Otherwise it is simply more architecture.

Part X — A compact complete example

76. Full componentized implementation

import torch
import torch.nn as nn
import torch.nn.functional as F


class TinyBlock(nn.Module):
    def __init__(self, d_model, dropout=0.1):
        super().__init__()

        self.norm = nn.LayerNorm(d_model)
        self.mlp = 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.mlp(self.norm(x))


class SparseAutoencoder(nn.Module):
    def __init__(self, d_model, d_code):
        super().__init__()

        self.encoder = nn.Linear(d_model, d_code)
        self.decoder = nn.Linear(d_code, d_model)

    def forward(self, x):
        code = F.relu(self.encoder(x))
        recon = self.decoder(code)
        return code, recon


class TinyFromPieces(nn.Module):
    def __init__(
        self,
        d_model=128,
        d_code=64,
        n_steps=6,
        step_scale=0.1,
    ):
        super().__init__()

        self.d_model = d_model
        self.n_steps = n_steps
        self.step_scale = step_scale

        self.fuse = nn.Linear(3 * d_model, d_model)
        self.block = TinyBlock(d_model)
        self.sae = SparseAutoencoder(d_model, d_code)
        self.score_head = nn.Linear(d_model, 1)

    def forward(self, context, candidate):
        assert context.ndim == 2
        assert candidate.shape == context.shape

        latent = torch.zeros_like(context)
        trajectory = []

        for _ in range(self.n_steps):
            fused = torch.cat(
                [context, candidate, latent],
                dim=-1,
            )

            proposal = torch.tanh(self.fuse(fused))
            proposal = self.block(proposal)

            latent = latent + self.step_scale * proposal
            trajectory.append(latent)

        code, recon = self.sae(latent)

        head_state = latent + recon
        score_logit = self.score_head(head_state).squeeze(-1)

        return {
            "score_logit": score_logit,
            "latent": latent,
            "code": code,
            "recon": recon,
            "trajectory": torch.stack(trajectory, dim=1),
        }

The complete model is still short.

That is not because the ideas are trivial.

It is because the complexity lives in composition and training behavior rather than code volume.

77. Tiny training example

torch.manual_seed(0)

B = 64
D = 128

context = torch.randn(B, D)
candidate = torch.randn(B, D)

# Synthetic target for demonstration only.
labels = (
    (context * candidate).sum(dim=-1) > 0
).float()

model = TinyFromPieces(
    d_model=D,
    d_code=64,
    n_steps=4,
)

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=3e-4,
)

for step in range(500):
    output = model(context, candidate)

    main_loss = F.binary_cross_entropy_with_logits(
        output["score_logit"],
        labels,
    )

    recon_loss = F.mse_loss(
        output["recon"],
        output["latent"].detach(),
    )

    sparse_loss = output["code"].abs().mean()

    loss = (
        main_loss
        + 0.1 * recon_loss
        + 0.01 * sparse_loss
    )

    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()

    if step % 50 == 0:
        with torch.no_grad():
            probs = torch.sigmoid(output["score_logit"])
            pred = (probs >= 0.5).float()
            acc = (pred == labels).float().mean().item()

            sparsity = (
                output["code"].abs() < 1e-8
            ).float().mean().item()

        print(
            step,
            "loss", round(loss.item(), 4),
            "acc", round(acc, 3),
            "sparsity", round(sparsity, 3),
        )

This is not a benchmark.

It is a microscope.

Its purpose is to let us inspect every component.

78. What we have learned

We started this series with a scalar pair scorer.

Then we added multiple heads.

Then we made the heads explicit modules.

Then we added hierarchical recurrence.

Then we collapsed that hierarchy into a compact recursive latent-state model.

Now we have opened that model again.

Inside Tiny we found:

residual MLP
attention router
sparse autoencoder
prediction heads

And inside those we found:

LayerNorm
Linear
GELU
softmax
matrix multiplication
residual addition

The sophisticated model eventually reduced back to ordinary tensor operations.

That is the central thesis of Models From First Principles:

Keep opening the model until every operation is understandable.

Then build back upward deliberately.

79. Where we go next

We have spent several posts changing what the model computes.

The next stage changes something different:

how the parameters move

We will leave model architecture for a moment and go underneath training itself.

The next post is:

Models From First Principles 07: PACS — Building an Optimizer From Gradient Statistics

We will start with a raw gradient and build:

gradient
   ↓
moving average
   ↓
running squared-gradient estimate
   ↓
diagonal preconditioner
   ↓
scaled parameter update

And once again, something that sounds sophisticated will turn out to be a small number of understandable operations composed carefully.