Test a Recurrent Middle Against Hidden States, Rollouts, and Cost

Flow

Fail-Closed Replacement Pipeline

A mismatch or failed acceptance gate returns the experiment to analysis rather than promoting the candidate.

1. 1Declare boundary⌄

   Choose one exact half-open layer interval and tensor contract.

2. 2Version trace schema⌄

   Bind every sample to architecture and preprocessing identity.

3. 3Capture reference⌄

   Store start state, end state, masks, positions, and next-token labels.

4. 4Validate corpus⌄

   Reject identity drift, shape mismatch, duplicates, and split leakage.

5. 5Train candidate⌄

   Optimize the repeated transition with reproducible settings.

6. 6Evaluate locally⌄

   Measure hidden-state loss and geometry.

7. 7Evaluate downstream⌄

   Compare logits, top tokens, calibration, and rollouts.

8. 8Benchmark resources⌄

   Measure latency, throughput, and memory at matched workloads.

9. 9Gate promotion⌄

   Advance only when every declared tolerance has observed evidence.

A transformer layer stack can be split into three regions:

Illustrative anonymized example
tokens
  -> embedding / early layers
  -> middle layer interval
  -> late layers
  -> logits
  -> sampling

The experiment focuses on a middle interval:

Illustrative anonymized example
layers [layer_start, layer_end)

The reference transition is:

# Illustrative anonymized example
def transformer_middle(hidden_start, layers):
    hidden = hidden_start

    for layer in layers:
        hidden = layer(hidden)

    return hidden

The candidate transition is a recurrent block:

# Illustrative anonymized example
class RecurrentMiddleBlock:
    def __init__(self, cell, steps: int) -> None:
        self.cell = cell
        self.steps = steps

    def __call__(self, hidden_start):
        hidden = hidden_start

        for _ in range(self.steps):
            hidden = self.cell(hidden)

        return hidden

The question is operational:

Can this candidate transition replace this specific transformer interval for this specific model under measured accuracy, distribution, rollout, latency, and memory tolerances?

That requires trace data from the reference model.

# Illustrative anonymized example
from dataclasses import dataclass
from typing import Sequence

@dataclass(frozen=True)
class HiddenStateTrace:
    trace_id: str
    sample_id: str

    model_revision: str
    tokenizer_revision: str
    architecture_revision: str
    layout_revision: str
    precision: str

    layer_start: int
    layer_end: int

    input_token_ids: Sequence[int]
    attention_mask_hash: str
    position_ids_hash: str
    position_encoding_revision: str | None

    hidden_start_ref: str
    hidden_end_ref: str
    next_token_id: int

    dataset_id: str
    split: str
    trace_format_revision: str

The trace says:

For this input and this exact model boundary, the original transformer transformed hidden_start into hidden_end.

It does not say the candidate is equivalent.

Equivalence has to be tested.

The layer boundary should be treated like an interface contract.

# Illustrative anonymized example
@dataclass(frozen=True)
class LayerBoundary:
    model_revision: str
    tokenizer_revision: str
    architecture_revision: str
    layout_revision: str
    precision: str

    hidden_size: int
    sequence_axis: int
    feature_axis: int

    layer_start: int
    layer_end: int
    position_encoding_revision: str | None

Validation should fail closed:

# Illustrative anonymized example
def validate_boundary(trace: HiddenStateTrace, boundary: LayerBoundary) -> None:
    checks = {
        "model_revision": trace.model_revision == boundary.model_revision,
        "tokenizer_revision": trace.tokenizer_revision == boundary.tokenizer_revision,
        "architecture_revision": trace.architecture_revision
        == boundary.architecture_revision,
        "layout_revision": trace.layout_revision == boundary.layout_revision,
        "precision": trace.precision == boundary.precision,
        "layer_start": trace.layer_start == boundary.layer_start,
        "layer_end": trace.layer_end == boundary.layer_end,
        "position_encoding_revision": (
            trace.position_encoding_revision
            == boundary.position_encoding_revision
        ),
    }

    failed = [name for name, passed in checks.items() if not passed]

    if failed:
        raise ValueError(f"layer boundary mismatch: {failed}")

This looks similar to cache identity because it is the same correctness pattern:

Reuse is safe only when identity proves the object still means the same thing.

Once the boundary is valid, evaluation can compare the reference and candidate.

Basic hidden-state metrics:

# Illustrative anonymized example
def hidden_state_mse(reference, candidate):
    return ((reference - candidate) ** 2).mean()

def hidden_state_cosine_similarity(reference, candidate):
    reference_flat = reference.reshape(reference.shape[0], -1)
    candidate_flat = candidate.reshape(candidate.shape[0], -1)

    return torch.nn.functional.cosine_similarity(
        reference_flat,
        candidate_flat,
        dim=-1,
    ).mean()

Next-token agreement is useful:

# Illustrative anonymized example
def top1_match_rate(reference_logits, candidate_logits):
    reference_top1 = reference_logits.argmax(dim=-1)
    candidate_top1 = candidate_logits.argmax(dim=-1)

    return (reference_top1 == candidate_top1).float().mean()

But it is not enough.

Two distributions can have the same top token while disagreeing everywhere else.

So measure distribution drift:

# Illustrative anonymized example
def kl_reference_to_candidate(reference_logits, candidate_logits):
    reference_log_probs = torch.log_softmax(reference_logits, dim=-1)
    candidate_log_probs = torch.log_softmax(candidate_logits, dim=-1)

    return torch.nn.functional.kl_div(
        input=candidate_log_probs,
        target=reference_log_probs,
        log_target=True,
        reduction="batchmean",
    )

The candidate should also be tested through the late layers and language head.

# Illustrative anonymized example
def final_position_logits(model, hidden_end, attention_mask):
    final_positions = attention_mask.sum(dim=1) - 1
    batch_index = torch.arange(hidden_end.shape[0], device=hidden_end.device)

    return model.lm_head(hidden_end[batch_index, final_positions])

The evaluation should separate:

• teacher-forced one-step checks
• free rollout checks
• latency checks
• memory checks

A minimal acceptance schema:

Illustrative anonymized example
{
  "acceptance_gates": {
    "hidden_state_error": null,
    "hidden_state_cosine": null,
    "next_token_match_rate": null,
    "distribution_kl": null,
    "rollout_match_rate": null,
    "rollout_quality_review": null,
    "latency_delta": null,
    "memory_delta": null
  }
}

The null values matter. They keep the article from pretending that measurements already exist.

The builder-level claim is:

A recurrent block is a candidate transition. It becomes interesting only after it survives boundary validation, hidden-state comparison, distribution comparison, rollout evaluation, and runtime measurement.

Illustrative anonymized example
Illustrative anonymized example:

# Illustrative anonymized example
with no_grad():
    reference = model.forward_prefix(batch)
    transitioned = adapter.forward_middle(reference.hidden_state)
    candidate = model.forward_suffix(transitioned)
    agreement = compare_logits(reference.logits, candidate.logits)
require(agreement >= policy.minimum_agreement)