Investigating layers

import torch, numpy as np

@torch.no_grad()
def get_qkv(model, tok, text, layer):
    """GPT-2 packs Q, K, V into one c_attn projection: split it in thirds."""
    captured = {}
    h = model.transformer.h[layer].attn.c_attn.register_forward_hook(
        lambda m, i, o: captured.__setitem__("qkv", o.detach()))
    ids = tok(text, return_tensors="pt")
    model(**ids)
    h.remove()

    dim = model.config.n_embd
    q, k, v = captured["qkv"][0].split(dim, dim=-1)      # each [T, dim]
    return q.numpy(), k.numpy(), v.numpy(), tok.convert_ids_to_tokens(ids["input_ids"][0])

def structure_profile(model, tok, text, n_layers):
    rows = []
    for L in range(n_layers):
        q, k, v, _ = get_qkv(model, tok, text, L)
        rows.append({
            "layer": L,
            "within_q": mean_offdiag(cosine_matrix(q)),
            "within_k": mean_offdiag(cosine_matrix(k)),
            "within_v": mean_offdiag(cosine_matrix(v)),
            # Do Q and K organise tokens the SAME way? RSA answers it.
            "rsa_qk": rsa(q, k),
            "rsa_qv": rsa(q, v),
        })
    return rows

def mean_offdiag(S):
    iu = np.triu_indices_from(S, k=1)
    return float(S[iu].mean())

# Reshape by head to ask the same questions per head rather than per layer:
#   q.reshape(T, n_heads, head_dim).transpose(1, 0, 2)
# Head-level structure is usually much more specific than layer-level.

import numpy as np
from sklearn.decomposition import PCA

def participation_ratio(X):
    """Effective number of dimensions carrying variance."""
    X = X - X.mean(0, keepdims=True)
    ev = np.linalg.svd(X, compute_uv=False) ** 2
    p = ev / ev.sum()
    return float(1.0 / (p ** 2).sum())

def variance_dim(X, threshold=0.95):
    """Components needed to reach a variance threshold — the other convention."""
    pca = PCA().fit(X - X.mean(0, keepdims=True))
    return int(np.searchsorted(np.cumsum(pca.explained_variance_ratio_), threshold) + 1)

for L, acts in enumerate(hidden_by_layer):      # each [n_tokens, dim]
    print(f"layer {L:2}  nominal {acts.shape[1]:4}  "
          f"PR {participation_ratio(acts):6.1f}  "
          f"95% var {variance_dim(acts):4}")

# Control: you cannot have more effective dimensions than samples. With 200
# tokens the ceiling is 200 regardless of width, so hold the token count fixed
# across layers or the profile is measuring your sample size.

import numpy as np
from sklearn.feature_selection import mutual_info_regression
from sklearn.metrics import mutual_info_score

def mi_binned(x, y, bins=16):
    """Fast, but sensitive to bin count — always report the count used."""
    c = np.histogram2d(x, y, bins=bins)[0]
    return mutual_info_score(None, None, contingency=c)

def mi_knn(x, y, k=3):
    """Kraskov-style estimator; better on continuous data."""
    return float(mutual_info_regression(x.reshape(-1, 1), y,
                                        n_neighbors=k, random_state=0)[0])

def mi_with_baseline(x, y, n=100):
    rng = np.random.default_rng(0)
    observed = mi_knn(x, y)
    null = [mi_knn(x, rng.permutation(y)) for _ in range(n)]
    return observed, float(np.mean(null)), float(np.percentile(null, 95))

# MI vs covariance: a non-monotonic relationship shows the difference clearly.
x = np.random.randn(2000)
y = x ** 2 + 0.1 * np.random.randn(2000)
print("corr", np.corrcoef(x, y)[0, 1])      # ~0 — blind to it
print("MI  ", mi_knn(x, y))                 # clearly positive

# Pairwise MI through the network: does information about a token persist,
# and how does it decay with distance?
def mi_by_distance(hidden, max_dist=20):
    out = {}
    for d in range(1, max_dist + 1):
        vals = [mi_knn(hidden[i], hidden[i + d])
                for i in range(len(hidden) - d)]
        out[d] = float(np.mean(vals))
    return out
# A slow decay means long-range dependence is being maintained; a sharp drop
# means the layer is largely local.

import torch

@torch.no_grad()
def logit_lens(model, tok, prompt, k=5):
    ids = tok(prompt, return_tensors="pt")
    out = model(**ids, output_hidden_states=True)

    ln_f = model.transformer.ln_f      # MUST apply this first
    head = model.lm_head

    for layer, h in enumerate(out.hidden_states):
        logits = head(ln_f(h[0, -1]))            # last position only
        probs = torch.softmax(logits, dim=-1)
        top = probs.topk(k)
        preds = " ".join(f"{tok.decode([i])!r}({p:.2f})"
                         for p, i in zip(top.values, top.indices))
        print(f"layer {layer:2}  {preds}")

logit_lens(model, tok, "The Eiffel Tower is located in the city of")
# Early layers: generic frequent tokens.
# Middle layers: 'Paris' appears and climbs.
# Late layers: sharpens onto it.

# Track one target token's probability through depth — the cleanest summary:
@torch.no_grad()
def trace_token(model, tok, prompt, target):
    tid = tok(target)["input_ids"][0]
    out = model(**tok(prompt, return_tensors="pt"), output_hidden_states=True)
    return [torch.softmax(model.lm_head(model.transformer.ln_f(h[0, -1])), -1)[tid].item()
            for h in out.hidden_states]