跳至主要内容

Scoring

這些是 Model 上的「只做 forward，不做 sampling、也不進入 decode loop」的方法。它們適合用在評估、可解釋性分析，以及自訂採樣管線。

Model.score(...)​

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

輸入：

  • tokens：numpy.ndarray[int32] 或相容的 array-like

回傳形狀：

[seq_len, vocab_size]

適用於：

  • eval loops
  • parity tests
  • 完整位置分析

Model.score_last(...)​

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

回傳形狀：

[vocab_size]

適用於：

  • 下一個 token 排序
  • decode loops
  • 服務端自訂採樣

Model.score_logit_lens(...)​

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

這是一個逐迭代的「logit lens」：它會在每個 loop 深度 l ∈ [1, loops]，把最後一個位置的 hidden state 經過模型的最終 LayerNorm 與 head 投影成 logits。回傳陣列中的第 l - 1 列，代表如果模型在第 l 次 loop iteration 結束時就停止，它會輸出什麼 logits。

適合用於：

  • 可解釋性分析（模型對下一個 token 的預測如何隨深度變化？）
  • L-variance uncertainty scoring：例如 KL(softmax(lens[-1]) ‖ softmax(lens[L_lo])) 可以成為內建信心訊號
  • self-consistency gating：若 lens[L_lo] 與 lens[L_hi] 的 argmax 不同，就升到更深的 L

目前實作會用 n_loops 次序列 forward（總計算量約 O(L²)）。未來路線圖已追蹤一個更便宜的版本：在一次 forward 中直接擷取每輪 hidden state，再批次做 head projection。

這個 API 之所以少見，是因為它依賴權重共享 loop 的結構；傳統深層 Transformer 無法以這麼自然的形式提供。

Model.score_trajectory(...)​

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

這個 API 會在一次 forward pass 中，回傳每個 forward 階段之後、最後一個位置的 hidden state（embedding、每個 pre-block、每次 loop iteration）。其中 loop-block 這一段所形成的序列，就是同一組共享權重作用在離散動力系統上的 trajectory，而這恰好是只有權重共享 looped Transformer 才具備的結構。

簽名​

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"]
}

參數​

  • tokens（np.ndarray[int32] 或相容 array-like）：輸入 prompt，一個長度為 seq_len 的一維 token id 陣列。此呼叫會擷取 tokens[-1] 在每個階段後的 hidden state，因此前面位置的內容仍會透過注意力影響最終被觀察到的 hidden。tokens 本身不會被修改。
  • n_loops（int，預設 8）：要跑的 loop 深度。需滿足 1 <= n_loops <= max_seq_len - seq_len，讓 decode cache 有足夠空間。loop block 會被套用 n_loops 次，因此你會拿到每次 iteration 一列 hidden state。

回傳​

一個 Python dict，包含五個鍵：

  • labels（List[str]）：有序的階段標籤，長度為 1 + n_pre + n_loops。
    • labels[0] 永遠是 "embed"，代表 embedding + embed projection 之後、尚未進入任何 block 前的 hidden state。
    • labels[1 : 1 + n_pre] 依序為 "pre.0"、"pre.1"、…，表示每個 pre-block 之後的 hidden state。n_pre = model.config()["n_pre_blocks"]，TinyLoop 參考架構預設為 2。
    • labels[1 + n_pre :] 則是 "loop.0" 到 "loop.{n_loops-1}"，也就是每次 loop-block iteration 之後的 hidden state，這一段就是 loop trajectory。
  • hidden（np.ndarray[float32]，形狀為 [len(labels), dim]）：每個階段最後一個位置的 hidden state，型別為 FP32。之所以是 FP32 而不是 FP16，是因為擷取點來自每個 block 結束後的 FP32 residual stream。
  • dim（int）：hidden dimension，方便你直接取用，等同 model.config()["dim"]。
  • n_loops（int）：原樣回傳輸入參數。
  • n_pre（int）：模型的 pre-block 數，方便切片。若要只取 loop trajectory，直接用 hidden[1 + n_pre : 1 + n_pre + n_loops]。

可能拋出的錯誤​

  • RuntimeError：模型已關閉（例如已呼叫 model.close() 或 with 區塊已結束）
  • RuntimeError：內部 trace_hidden_forward 失敗（OOM、非法 token id、n_loops 超出 cache 範圍）

數值細節​

  • 回傳的 FP32 row 表示的是每個 block MLP output + residual add 之後的 residual stream 狀態（對 pre-block 與 loop-block 都一樣）；"embed" 這一列則是 embedding + embed projection 之後的狀態。它們是下一個 block attention 看到的 pre-LN 狀態，而不是 post-LN。
  • 對最後一個位置的擷取只需要從 [seq_len, dim] 的 block-output buffer 裡複製一個 token，不需要把整段序列全量 materialize，因此額外記憶體成本永遠只有 (1 + n_pre + n_loops) × dim × 4 bytes，與 seq_len 無關。
  • Trajectory capture 用的就是 score() 會走的同一批 CUDA kernels，因此在相同 n_loops 下，其最後位置結果與 score() 在該位置上的對齊是位元級一致的。

成本​

每次呼叫只做一次完整 forward（成本約等於 score()）。每個階段多做一次 dim × 4 bytes 的 hidden-state 寫出，相較整個 forward 幾乎可忽略。

範例​

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

    # 取出 loop trajectory
    n_pre = report["n_pre"]
    loop = h[1 + n_pre : 1 + n_pre + report["n_loops"]]  # [16, dim]

    # 每次迭代的相對變化 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)

    # 連續兩次迭代的 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：相鄰 delta 是否朝相同方向
    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}")

典型值（1B-effective、5 個 prompt 校準）​

在 n_loops = 32 的情況下，1B-effective checkpoint 在最後一次迭代的典型量測約為：rho ≈ 0.033、cos(h^L, h^{L-1}) ≈ 0.9998、drift_cos ≈ 0.95。這些數字描述的是 loop block 的收縮動力學；也正是這個性質，讓 KV Cache Modes 中的 h_mode 壓縮成為可能。

批次版本​

目前沒有 batched 版本；score_trajectory 一次只處理一個 prompt，若要做 N 個 prompt 的統計，就必須呼叫 N 次。如果你需要對一整組 calibration set 收集統計，建議使用下方 driver script。

Driver script​

tools/measure_trajectory.py 是現成工具，可對一組 calibration prompts 在多個 n_loops 值下執行 score_trajectory，並輸出包含 per-prompt stats 與 per-iteration 聚合統計（rho、cos_prev、drift_cos、anchor distance 的平均與標準差）的 JSON。該 JSON 正是論文中的 fig_trajectory.{png,pdf} 圖，以及 KV Cache Modes 頁面數字的來源。

用法：

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

為什麼這個 API 只有權重共享 looped Transformer 才合理​

對標準深層 Transformer 而言，逐層 hidden state 並不構成真正的 trajectory，因為每一層都是不同權重，因此 {h^l} 序列在每一步都遵循不同的更新規則。權重共享 looped Transformer 則把同一個 block 套用 n_loops 次，因此 loop 這一段就是離散動力系統 h^{l+1} = g(h^l; x_{<t}) 的實例，其中 g 在每次迭代中都固定不變。這也是為什麼像 rho_l → 0 和 cos(h^l, h^{l-1}) → 1 這種描述，才真能被詮釋為「block 的收縮行為」；也正因此，K 與 V 才能透過相同的 W_k / W_v 在每次迭代下都從 h 重建（詳見 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)

此 API 會在兩個 loop 深度下跑相同輸入，並回傳：

  1. L_hi 下最後一個 token 的 logits
  2. 兩個分布之間的 KL divergence

KL 越高，代表模型在 L_hi 下的預測越可能與 L_lo 時不一致，也就是一個內建的不確定性訊號，而且不需要額外 head 或額外訓練。

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. 先分別在 L_lo 與 L_hi 下 forward，並比較 top-1 argmax
  2. 若兩者一致，直接回傳 L_hi 的 logits，並設 escalated=False，代表模型在較便宜深度下已完成自我驗證
  3. 若兩者不一致，則再往 escalate_L 做一次更深的 forward，並回傳那次 logits，同時設 escalated=True

escalated 本身就是一個零額外成本的監控訊號：你可以統計流量中有多少 token 需要深層推理，多少 token 在淺層就已足夠。這類能力之所以便宜，正是因為權重共享 looped Transformer 可以在「相同權重、不同深度」之間自由比較；標準深層 Transformer 沒有這個原語。