Embedding trajectories through the model

import torch, numpy as np

@torch.no_grad()
def trajectory(model, tok, text, position):
    out = model(**tok(text, return_tensors="pt"), output_hidden_states=True)
    return np.stack([h[0, position].numpy() for h in out.hidden_states])

def angle_between(u, v):
    c = u @ v / (np.linalg.norm(u) * np.linalg.norm(v) + 1e-12)
    return float(np.degrees(np.arccos(np.clip(c, -1, 1))))

def describe(traj):
    rows = []
    for i in range(len(traj) - 1):
        step = traj[i + 1] - traj[i]
        rows.append({
            "layer": i,
            "norm": float(np.linalg.norm(traj[i])),
            # Relative step: scale-free, so comparable across depth.
            "rel_step": float(np.linalg.norm(step) / np.linalg.norm(traj[i])),
            # How much did this layer rotate the representation?
            "rotation": angle_between(traj[i], traj[i + 1]),
            # Did this step continue the previous direction, or turn?
            "turn": angle_between(traj[i] - traj[i - 1], step) if i else None,
        })
    return rows

def path_length(traj, normalise=True):
    steps = np.diff(traj, axis=0)
    lengths = np.linalg.norm(steps, axis=1)
    if normalise:
        lengths = lengths / np.linalg.norm(traj[:-1], axis=1)
    return float(lengths.sum())

# Total displacement vs. path length: their ratio measures directness.
# Close to 1 means steady refinement; much less means the representation
# wandered and was revised.
def directness(traj):
    return float(np.linalg.norm(traj[-1] - traj[0]) /
                 np.linalg.norm(np.diff(traj, axis=0), axis=1).sum())

import torch, numpy as np

class Decomposer:
    """Captures each layer's attention and MLP contributions separately."""
    def __init__(self, model):
        self.model, self.parts, self.handles = model, {}, []

    def __enter__(self):
        for i, block in enumerate(self.model.transformer.h):
            self.handles.append(block.attn.register_forward_hook(self._save(f"attn{i}")))
            self.handles.append(block.mlp.register_forward_hook(self._save(f"mlp{i}")))
        return self

    def _save(self, name):
        def hook(module, inputs, output):
            t = output[0] if isinstance(output, tuple) else output
            self.parts[name] = t.detach()
        return hook

    def __exit__(self, *exc):
        for h in self.handles:
            h.remove()

@torch.no_grad()
def logit_attribution(model, tok, prompt, target, position=-1):
    """How much did each component push toward the target token?"""
    tid = tok(target)["input_ids"][0]
    # The unembedding row for the target IS the direction that raises its logit.
    direction = model.lm_head.weight[tid]

    with Decomposer(model) as dec:
        model(**tok(prompt, return_tensors="pt"))

    rows = []
    for name, contribution in dec.parts.items():
        rows.append((name, float(contribution[0, position] @ direction)))
    rows.sort(key=lambda kv: -abs(kv[1]))
    return rows

for name, score in logit_attribution(model, tok,
        "The Eiffel Tower is in", " Paris")[:8]:
    print(f"{name:8} {score:+.3f}")

# Sanity check: because the stream is additive, the contributions plus the
# embedding term should reconstruct the final logit up to the layernorm's
# non-linearity. If they do not, you are capturing the wrong tensor.

import spacy, numpy as np, pandas as pd
import statsmodels.formula.api as smf

nlp = spacy.load("en_core_web_sm")

def collect(model, tok, sentences):
    rows = []
    for sentence in sentences:
        doc = nlp(sentence)
        ids = tok(sentence, return_tensors="pt")
        pieces = tok.convert_ids_to_tokens(ids["input_ids"][0])

        for token in doc:
            span = find_span(pieces, token.text, tok)
            if span is None:
                continue
            traj = trajectory(model, tok, sentence, span[1] - 1)   # last piece
            rows.append({
                "pos": token.pos_,
                "path": path_length(traj),
                "rotation": angle_between(traj[0], traj[-1]),
                "directness": directness(traj),
                # Controls — every one of these confounds the comparison.
                "n_tokens": span[1] - span[0],
                "log_freq": np.log1p(corpus_frequency.get(token.text.lower(), 1)),
                "position": span[0] / len(pieces),
            })
    return pd.DataFrame(rows)

df = collect(model, tok, sentences)
print(df.groupby("pos")[["path", "rotation", "directness"]].agg(["mean", "count"]))

# The honest analysis: does part of speech still predict path length once the
# confounds are in the model?
fit = smf.ols("path ~ C(pos) + n_tokens + log_freq + position", data=df).fit()
print(fit.summary().tables[1])
# If the C(pos) terms lose significance here, the raw group difference was a
# frequency or length effect. That is the result, and it is worth reporting.