Investigating token embeddings

import numpy as np, matplotlib.pyplot as plt

def cosine_matrix(X, center=True):
    """Pairwise cosine similarity. center=True removes the shared mean
    direction, without which contextual embeddings all look ~0.95 similar."""
    X = np.asarray(X, dtype=np.float64)
    if center:
        X = X - X.mean(axis=0, keepdims=True)
    Xn = X / np.linalg.norm(X, axis=1, keepdims=True)
    return Xn @ Xn.T

def null_distribution(shape, n=200, center=True):
    """What similarity looks like with no structure at all."""
    out = []
    for _ in range(n):
        R = np.random.randn(*shape)
        S = cosine_matrix(R, center)
        out.append(S[np.triu_indices_from(S, k=1)])
    return np.concatenate(out)

S = cosine_matrix(embeddings)                 # [n_tokens, n_tokens]
null = null_distribution(embeddings.shape)

observed = S[np.triu_indices_from(S, k=1)]
print(f"observed mean {observed.mean():+.3f}   null mean {null.mean():+.3f}")
print(f"95th pct of null: {np.percentile(np.abs(null), 95):.3f}")
# Only pairs exceeding that percentile are worth discussing.

fig, ax = plt.subplots(figsize=(7, 6))
im = ax.imshow(S, cmap="RdBu_r", vmin=-1, vmax=1)   # diverging map, symmetric
ax.set_xticks(range(len(labels))); ax.set_xticklabels(labels, rotation=90)
ax.set_yticks(range(len(labels))); ax.set_yticklabels(labels)
plt.colorbar(im)

import numpy as np, matplotlib.pyplot as plt
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA
from sklearn.cluster import DBSCAN
from sklearn.neighbors import NearestNeighbors

# PCA first: denoises and makes t-SNE much faster on 768 dimensions.
X = PCA(n_components=50, random_state=0).fit_transform(embeddings)

# Cluster in the ORIGINAL space, not in the 2-D projection. Clustering the
# projection means clustering an artefact of the projection.
nn = NearestNeighbors(n_neighbors=5).fit(X)
d, _ = nn.kneighbors(X)
plt.plot(np.sort(d[:, -1]))          # elbow of this curve is a good eps

labels_ = DBSCAN(eps=chosen_eps, min_samples=5).fit_predict(X)
print(f"{len(set(labels_)) - (1 in labels_)} clusters, "
      f"{(labels_ == -1).sum()} noise points")

# Then project only for viewing, and check stability across perplexity.
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
for ax, perp in zip(axes, [5, 30, 50]):
    Y = TSNE(n_components=2, perplexity=perp, init="pca",
             random_state=0).fit_transform(X)
    ax.scatter(Y[:, 0], Y[:, 1], c=labels_, cmap="tab10", s=12)
    ax.set_title(f"perplexity {perp}")
    ax.set_xticks([]); ax.set_yticks([])
# If your clusters only appear at one perplexity, they are not a finding.

# UMAP is the usual alternative: faster, retains a little more global
# structure, and has exactly the same interpretive caveats.

import numpy as np
from scipy.stats import spearmanr

def rdm(X):
    """Representational dissimilarity matrix: 1 - centred cosine similarity."""
    Xc = X - X.mean(axis=0, keepdims=True)
    Xn = Xc / np.linalg.norm(Xc, axis=1, keepdims=True)
    return 1 - Xn @ Xn.T

def rsa(X, Y):
    """X and Y: embeddings of the SAME stimuli from two models.
    Dimensions need not match — that is the whole point."""
    a, b = rdm(X), rdm(Y)
    iu = np.triu_indices_from(a, k=1)          # upper triangle only
    return spearmanr(a[iu], b[iu]).statistic

print("word2vec vs GPT-2:", rsa(w2v_vectors, gpt2_vectors))

# Laminar profile: how does each layer's structure relate to word2vec's?
profile = [rsa(w2v_vectors, layer_vectors[L]) for L in range(n_layers)]
# Typically high early (token identity dominates), falling with depth as
# contextual and task-specific structure takes over.

# Category selectivity by layer: is the RDM block-diagonal by category?
def category_selectivity(X, categories):
    D = rdm(X)
    same = np.equal.outer(categories, categories)
    np.fill_diagonal(same, False)
    iu = np.triu_indices_from(D, k=1)
    within  = D[iu][same[iu]].mean()
    between = D[iu][~same[iu]].mean()
    return between - within        # larger = more categorical organisation

import networkx as nx, numpy as np

def similarity_graph(S, labels, k=5):
    """k-nearest-neighbour graph — more stable than a global threshold,
    because it adapts to locally varying density."""
    G = nx.Graph()
    G.add_nodes_from(labels)
    np.fill_diagonal(S, -np.inf)
    for i, label in enumerate(labels):
        for j in np.argsort(S[i])[-k:]:
            G.add_edge(label, labels[j], weight=float(S[i, j]))
    return G

G = similarity_graph(cosine_matrix(embeddings), labels, k=5)

print(nx.number_connected_components(G))
print(sorted(nx.degree_centrality(G).items(), key=lambda kv: -kv[1])[:5])
print([sorted(c)[:6] for c in nx.community.greedy_modularity_communities(G)])

# High-degree nodes are usually generic or high-frequency tokens; they connect
# to everything and are worth excluding before interpreting communities.

import numpy as np

def semantic_axis(model, positive_pairs):
    """positive_pairs: [('hot','cold'), ('warm','cool'), ...]
    Average many differences so the axis reflects the attribute, not the words."""
    diffs = [model[a] - model[b] for a, b in positive_pairs]
    axis = np.mean(diffs, axis=0)
    return axis / np.linalg.norm(axis)

temperature = semantic_axis(glove, [("hot", "cold"), ("warm", "cool"),
                                     ("boiling", "freezing")])

for word in ["fire", "ice", "soup", "snow", "desk"]:
    v = glove[word]
    print(f"{word:6} {float(v @ temperature / np.linalg.norm(v)):+.3f}")
# 'desk' should score near zero. If it does not, the axis is contaminated —
# usually by frequency or by a topic the pairs share.

def analogy_accuracy(model, quadruples, k=1):
    """a:b :: c:d — score over MANY quadruples, not one."""
    hits = 0
    for a, b, c, d in quadruples:
        if any(w not in model for w in (a, b, c, d)):
            continue
        got = model.most_similar(positive=[b, c], negative=[a], topn=k)
        hits += int(d in [w for w, _ in got])
    return hits / len(quadruples)

import torch, numpy as np

@torch.no_grad()
def occurrence_vectors(model, tok, sentences, target, layer=-1):
    """One vector per occurrence of `target`, in context."""
    out = []
    for s in sentences:
        ids = tok(s, return_tensors="pt")
        hs = model(**ids, output_hidden_states=True).hidden_states[layer]
        pieces = tok.convert_ids_to_tokens(ids["input_ids"][0])
        for i, p in enumerate(pieces):
            if target in p:
                out.append((s, hs[0, i].numpy()))
    return out

def knn(query, bank, k=5):
    V = np.stack([v for _, v in bank])
    V = V / np.linalg.norm(V, axis=1, keepdims=True)
    q = query / np.linalg.norm(query)
    return [bank[i][0] for i in np.argsort(-(V @ q))[:k]]

def neighbourhood_overlap(bank_a, bank_b, k=10):
    """Jaccard overlap of the k-NN sets two models give for each item.
    A per-layer profile of this is a clean cross-architecture comparison."""
    scores = []
    for i in range(len(bank_a)):
        a = set(knn(bank_a[i][1], bank_a, k))
        b = set(knn(bank_b[i][1], bank_b, k))
        scores.append(len(a & b) / len(a | b))
    return float(np.mean(scores))

import numpy as np, matplotlib.pyplot as plt

def spectrum(X, center=True):
    X = np.asarray(X, dtype=np.float64)
    if center:
        X = X - X.mean(axis=0, keepdims=True)
    return np.linalg.svd(X, compute_uv=False)

def effective_dim(s):
    """Participation ratio: how many directions genuinely matter.
    Equals n for a flat spectrum, ~1 when one direction dominates."""
    p = s ** 2 / (s ** 2).sum()
    return float(1.0 / (p ** 2).sum())

for name, subset in categories.items():
    s = spectrum(subset)
    print(f"{name:12} eff.dim {effective_dim(s):5.1f}  "
          f"top-10 variance {(s[:10]**2).sum() / (s**2).sum():.2%}")

plt.semilogy(spectrum(embeddings), label="real")
plt.semilogy(spectrum(np.random.randn(*embeddings.shape)), label="random")
plt.legend(); plt.xlabel("component"); plt.ylabel("singular value")
# The gap between the two curves IS the structure. A random matrix's spectrum
# (the Marchenko-Pastur law) is the null you are comparing against.

# Project two related sets onto shared leading directions to see how they differ.
both = np.vstack([set_a, set_b])
U, s, Vt = np.linalg.svd(both - both.mean(0), full_matrices=False)
A, B = set_a @ Vt[:2].T, set_b @ Vt[:2].T