Editing hidden states

import torch, torch.nn.functional as F, numpy as np

@torch.no_grad()
def scale_sweep(model, prompt, target_id, factors=(0.25, 0.5, 0.8, 1.0, 1.25, 2.0, 4.0)):
    results = np.zeros((model.config.n_layer, len(factors)))

    for L in range(model.config.n_layer):
        for j, f in enumerate(factors):
            with intervene(model.transformer.h[L], scale(f)):
                logits = model(**prompt).logits[0, -1]
            # Negative log-probability of the target: lower is better.
            results[L, j] = -F.log_softmax(logits, -1)[target_id].item()
    return results, factors

surface, factors = scale_sweep(model, prompt, paris_id)

baseline_col = factors.index(1.0)
delta = surface - surface[:, [baseline_col]]     # change relative to no edit

for L, row in enumerate(delta):
    worst = np.abs(row).max()
    print(f"layer {L:2}  max |Δloss| {worst:5.2f}  "
          f"{'SENSITIVE' if worst > 0.5 else ''}")

# Note that h[L] returns a tuple, so `scale` must edit output[0] — the
# intervene() helper above already handles that.

import torch

def add_noise(sigma_ratio=0.5, seed=0):
    """Noise scaled to the activation's own std, so it is comparable across layers."""
    def fn(x):
        g = torch.Generator(device=x.device).manual_seed(seed)
        noise = torch.randn(x.shape, generator=g, device=x.device, dtype=x.dtype)
        return x + sigma_ratio * x.std() * noise
    return fn

def shuffle_positions(seed=0):
    """Same values, different positions: tests positional binding only."""
    def fn(x):
        g = torch.Generator(device=x.device).manual_seed(seed)
        perm = torch.randperm(x.size(1), generator=g, device=x.device)
        return x[:, perm]
    return fn

def shuffle_features(seed=0):
    """Same norm, scrambled direction: tests directional structure only."""
    def fn(x):
        g = torch.Generator(device=x.device).manual_seed(seed)
        perm = torch.randperm(x.size(-1), generator=g, device=x.device)
        return x[..., perm]
    return fn

# Robustness curve on BERT's masked prediction:
for sigma in [0.0, 0.1, 0.25, 0.5, 1.0, 2.0]:
    with intervene(bert.encoder.layer[6], add_noise(sigma)):
        acc = masked_accuracy(bert, tok, sentences)
    print(f"sigma {sigma:4.2f}  accuracy {acc:.3f}")
# The level at which accuracy collapses is the model's precision requirement
# at that depth — and it is usually much higher than people expect.

import torch, numpy as np

@torch.no_grad()
def bias_direction(model, tok, pairs, layer):
    """pairs: [(sentence_a, sentence_b), ...] differing only in the attribute."""
    diffs = []
    for a, b in pairs:
        ha = hidden_at(model, tok, a, layer)      # [dim], pooled or [CLS]
        hb = hidden_at(model, tok, b, layer)
        diffs.append(ha - hb)
    d = torch.stack(diffs).mean(0)
    return d / d.norm()

def project_out(direction):
    """Remove the component along `direction` from every position."""
    d = direction / direction.norm()
    def fn(x):
        # x: [B, T, dim]; subtract the projection onto d.
        return x - (x @ d).unsqueeze(-1) * d
    return fn

d = bias_direction(bert, tok, gender_pairs, layer=8)

before = bias_score(bert, tok, probe_sentences)
with intervene(bert.encoder.layer[8], project_out(d)):
    after = bias_score(bert, tok, probe_sentences)
    # The control that makes it a finding rather than damage:
    control_after = general_accuracy(bert, tok, neutral_sentences)

print(f"bias {before:.3f} -> {after:.3f}   "
      f"unrelated accuracy {control_after:.3f}")

# Repeat per layer: where in depth is the attribute linearly available, and
# where does removing it actually change behaviour? Those two depths need not
# be the same, and the gap between them is interesting.

import torch, numpy as np

@torch.no_grad()
def capture_all_layers(model, prompt):
    out = model(**prompt, output_hidden_states=True)
    return [h.clone() for h in out.hidden_states]

@torch.no_grad()
def patching_map(model, clean, corrupted, answer_id, wrong_id):
    """For every (layer, position): patch clean -> corrupted and measure recovery."""
    clean_acts = capture_all_layers(model, clean)

    def logit_diff(logits):
        # The correct metric: difference between the two candidate answers.
        # It isolates the decision and cancels prompt-level shifts.
        return (logits[0, -1, answer_id] - logits[0, -1, wrong_id]).item()

    base_clean = logit_diff(model(**clean).logits)
    base_corr  = logit_diff(model(**corrupted).logits)

    n_layers = model.config.n_layer
    n_pos = corrupted["input_ids"].size(1)
    recovery = np.zeros((n_layers, n_pos))

    for L in range(n_layers):
        for p in range(n_pos):
            donor = clean_acts[L + 1]         # +1: index 0 is the embeddings
            with intervene(model.transformer.h[L], patch_from(donor, position=p)):
                got = logit_diff(model(**corrupted).logits)
            # 0 = no recovery, 1 = fully restored to the clean behaviour.
            recovery[L, p] = (got - base_corr) / (base_clean - base_corr + 1e-9)

    return recovery

# Indirect object identification.
clean     = tok("When John and Mary went to the shop, John gave a drink to",
                return_tensors="pt")
corrupted = tok("When John and Mary went to the shop, Mary gave a drink to",
                return_tensors="pt")

heat = patching_map(model, clean, corrupted,
                    answer_id=tok(" Mary")["input_ids"][0],
                    wrong_id=tok(" John")["input_ids"][0])

# Read the map: high values early at the differing token, then moving to the
# final position in middle layers. That trace IS the information flow.

import torch, torch.nn.functional as F

def identity(x):
    """Return the input unchanged — the block adds nothing to the stream."""
    return x

@torch.no_grad()
def skip_profile(model, prompt, target_id):
    base = -F.log_softmax(model(**prompt).logits[0, -1], -1)[target_id].item()

    rows = []
    for L in range(model.config.n_layer):
        # Skipping means the block's contribution is zero, i.e. output = input.
        with intervene(model.transformer.h[L], lambda x: x * 0.0 + x):
            loss = -F.log_softmax(model(**prompt).logits[0, -1], -1)[target_id].item()
        rows.append((L, loss - base))
    return base, rows

base, rows = skip_profile(model, prompt, paris_id)
for L, delta in rows:
    bar = "#" * int(max(delta, 0) * 10)
    print(f"layer {L:2}  Δloss {delta:+6.3f}  {bar}")

# Cumulative skipping is the harsher test: remove the least-important layers
# one at a time and find where performance finally falls off a cliff. That
# curve is the empirical basis for depth pruning.