メインコンテンツまでスキップ

Scoring

Methods on Model that run one or more forward passes and return logits or hidden state — no sampling, no decode loop. Use these for evaluation, interpretability, and custom sampling pipelines.

Model.score(...)

logits = model.score(tokens, loops=8)

Expected input:

  • tokens: numpy.ndarray[int32] or compatible array-like

Return shape:

[seq_len, vocab_size]

Use this for:

  • eval loops
  • parity tests
  • full-position analysis

Model.score_last(...)

last_logits = model.score_last(tokens, loops=8)

Return shape:

[vocab_size]

Use this for:

  • next-token ranking
  • decode loops
  • service-side sampling code

Model.score_logit_lens(...)

lens = model.score_logit_lens(tokens, loops=8) # shape [loops, vocab_size]

Per-iteration "logit lens": projects the last-position hidden state through the model's final LayerNorm and head at every loop depth l ∈ [1, loops]. Row l - 1 of the returned array is the logit vector the model would emit if generation stopped after l loop iterations.

Useful for:

  • interpretability (how does the predicted next token evolve with depth?)
  • L-variance uncertainty scoring: KL(softmax(lens[-1]) ‖ softmax(lens[L_lo])) gives a built-in confidence signal
  • self-consistency gating — if lens[L_lo] and lens[L_hi] disagree on argmax, escalate to more loops

The current implementation uses n_loops sequential forward passes (O(L²) total compute). A single-forward variant that captures per-iteration hidden state and batches the head projection is tracked as an optimization in the production roadmap.

No other runtime provides this endpoint as a first-class API — it relies on the weight-shared loop structure, so traditional deep transformers can't expose it cheaply.

Model.score_trajectory(...)

report = model.score_trajectory(tokens, n_loops=8)

Run one forward pass and return the last-position hidden state after each forward-pass stage (embed, each pre-block, each loop iteration). The loop-block portion of the result is the trajectory of a discrete dynamical system under a shared weight, which is a structure only weight-shared looped transformers have.

Signature

report: Dict[str, Any] = Model.score_trajectory(
tokens: np.ndarray[int32], # shape [seq_len]
n_loops: int = 8,
) -> {
"labels": List[str], # length 1 + n_pre + n_loops
"hidden": np.ndarray[float32], # shape [1 + n_pre + n_loops, dim]
"dim": int, # hidden dimension
"n_loops": int, # echoes the argument
"n_pre": int, # model.config()["n_pre_blocks"]
}

Parameters

  • tokens (np.ndarray[int32] or compatible array-like) — The prompt, a 1-D array of token ids of length seq_len. The call captures the hidden state at the final position (tokens[-1]) after each stage, so the content of preceding positions affects which hidden state is observed via the attention interaction. tokens is not modified.
  • n_loops (int, default 8) — Loop depth to run during capture. Must satisfy 1 <= n_loops <= max_seq_len - seq_len for the decode cache to have room. The loop block is applied n_loops times; you get one hidden-state row per iteration.

Returns

A Python dict with five keys:

  • labels (List[str]) — Ordered stage labels, length 1 + n_pre + n_loops.
    • labels[0] is always "embed" — the hidden state after embedding + embed projection but before any block.
    • labels[1 : 1 + n_pre] are "pre.0", "pre.1", … — the hidden state after each pre-block. There are n_pre = model.config()["n_pre_blocks"] of these (default 2 for TinyLoop's reference architecture).
    • labels[1 + n_pre :] are "loop.0" through "loop.{n_loops-1}" — the hidden state after each loop-block iteration. These rows form the loop trajectory.
  • hidden (np.ndarray[float32], shape [len(labels), dim]) — The captured hidden state for each stage, at the final input position (i.e. tokens[-1]), in FP32. The dtype is FP32 (not FP16) because the capture reads from the model's FP32 residual stream after the final LayerNorm of each block.
  • dim (int) — Hidden dimension. Provided for convenience; equals model.config()["dim"].
  • n_loops (int) — Echoes the input argument.
  • n_pre (int) — Number of pre-blocks in the model, for easy slicing. Use hidden[1 + n_pre : 1 + n_pre + n_loops] to isolate the loop trajectory.

Raises

  • RuntimeError — If the model was closed (via model.close() or the with block exited) before the call.
  • RuntimeError — If trace_hidden_forward fails internally (out-of-memory, invalid token ids, n_loops out of range for the cache buffers).

Numerical details

  • The returned FP32 rows are the residual stream state after each block's MLP output+residual add (for pre blocks and loop blocks) and after the embedding+embed-proj combination (for the "embed" row). They are pre-LN of the next block's attention, not post-LN.
  • The last-position capture uses a single-token write from the [seq_len, dim] block-output buffer; no per-position materialisation of the full sequence is performed, so memory use is (1 + n_pre + n_loops) × dim × 4 bytes regardless of seq_len.
  • Trajectory capture runs the same CUDA kernels as score(), so numerical output is bit-exact against a score() call at the same n_loops for row-0 / last-position alignment.

Cost

One forward pass per call (≈ score() cost). The per-stage hidden-state write is an additional dim × 4 bytes copy per stage — negligible versus the forward itself.

Example

import numpy as np
import tinyloop_py as tl

with tl.Model("model.tinyloop", max_seq_len=512) as model:
tokens = np.array([12, 34, 567, 890], dtype=np.int32)
report = model.score_trajectory(tokens, n_loops=16)
h = report["hidden"] # [19, dim] for n_pre=2, n_loops=16

# Slice out the loop trajectory
n_pre = report["n_pre"]
loop = h[1 + n_pre : 1 + n_pre + report["n_loops"]] # [16, dim]

# Per-iteration relative change rho_l = ||h^l - h^{l-1}|| / ||h^l||
rho = np.linalg.norm(np.diff(loop, axis=0), axis=1) / np.linalg.norm(loop[1:], axis=1)

# Consecutive-iteration cosine similarity
num = np.sum(loop[1:] * loop[:-1], axis=1)
denom = np.linalg.norm(loop[1:], axis=1) * np.linalg.norm(loop[:-1], axis=1)
cos_prev = num / np.clip(denom, 1e-12, None)

# Drift coherence: are consecutive deltas pointing in the same direction?
delta = np.diff(loop, axis=0)
num2 = np.sum(delta[1:] * delta[:-1], axis=1)
denom2 = np.linalg.norm(delta[1:], axis=1) * np.linalg.norm(delta[:-1], axis=1)
drift_cos = num2 / np.clip(denom2, 1e-12, None)

print(f"rho at L={len(rho)}: {rho[-1]:.4f} "
f"cos_prev: {cos_prev[-1]:.4f} "
f"drift_cos mean: {drift_cos.mean():.4f}")

Typical values (measured on 1B-effective, 5-prompt calibration)

For n_loops = 32, the last-iteration metrics on the 1B-effective checkpoint are approximately rho ≈ 0.033, cos(h^L, h^{L-1}) ≈ 0.9998, drift_cos ≈ 0.95. These describe the loop block's contractive dynamics; see the KV Cache Modes page for how this enables the h_mode compression.

Batch version

There is no batched variant — score_trajectory takes a single prompt at a time, and calling it N times gives N independent reports. If you need per-prompt statistics over a calibration set, use the driver below, which subprocess-isolates each call and aggregates means and standard deviations.

Driver script

tools/measure_trajectory.py is a ready-made wrapper that runs score_trajectory on a list of calibration prompts across multiple n_loops values and dumps a JSON with per-prompt stats plus per-iteration aggregates (mean and standard deviation of rho, cos_prev, drift_cos, and the anchor distance). The output JSON is what fed the fig_trajectory.{png,pdf} figure in the paper and the numbers cited in the KV Cache Modes page.

Usage:

python3 tools/measure_trajectory.py \
--model /workspace/model_1b_effective.tinyloop \
--loops 4 8 16 32 \
--output trajectory.json

Why this API exists only for weight-shared looped transformers

A standard deep transformer's per-layer hidden states are not a trajectory. Each layer uses different weights, so the {h^l} sequence follows a different update rule at each step. A weight-shared looped transformer applies the same block n_loops times, so the loop portion of the hidden array is a trajectory of a discrete dynamical system — h^{l+1} = g(h^l; x_{<t}) with g fixed across iterations. That structure is what makes rho_l → 0 and cos(h^l, h^{l-1}) → 1 meaningful as statements about the block's contractive behaviour, and what makes K and V reconstructible from h via the same W_k / W_v at every iteration (see KV Cache Modes).

Model.score_with_uncertainty(...)

logits, kl = model.score_with_uncertainty(tokens, L_lo=4, L_hi=8)
# logits : np.ndarray [vocab_size] (last-token logits at L_hi)
# kl : float (KL(p_L_hi || p_L_lo) on last-token dist)

Runs the input at two loop depths and returns (a) the logits at L_hi and (b) the KL divergence between the two distributions. High KL means the model's prediction at L_hi would disagree with what it would emit at L_lo — a built-in confidence signal with zero auxiliary heads, zero auxiliary training.

Model.score_with_consistency(...)

logits, escalated, L_used = model.score_with_consistency(
tokens, L_lo=8, L_hi=16, escalate_L=32,
)
# logits : np.ndarray [vocab_size]
# escalated : bool (did we have to fall through to escalate_L?)
# L_used : int (final loop depth used)

Self-consistency quality gating:

  1. Forward at L_lo and L_hi; compare top-1 argmax.
  2. If they agree, return L_hi's logits with escalated=False — the model self-certified at the cheaper depth.
  3. If they disagree, forward again at escalate_L and return those deeper logits with escalated=True.

The escalated flag is a zero-cost monitoring signal: downstream code can track what fraction of tokens in a stream needed deep reasoning vs. trivially agreed at shallow L. Only weight-shared looped transformers can do this cheaply — a standard deep stack has no "same weights, different depth" primitive.