Spaces:
Running on Zero
Running on Zero
| """Reading a direction back as words: the inverse of the transport. | |
| `Lens` maps words to a direction. `Vocabulary` maps a direction back to the words whose own | |
| transported directions align with it, which is how an axis gets a name — and how you find out | |
| that it has none. | |
| voc = ulens.Lens("qwen3.5-4b", block=23).vocabulary() | |
| voc.top(axis, k=12) # the words this direction reads as | |
| print(voc.report(cone)) # every axis, both ends | |
| Cheaply. Ranking by cosine against every transported row looks like it needs the full | |
| transported vocabulary `D = unit((W_U * gamma) @ J)`, which is gigabytes, but it factorises: | |
| numerator_v = (W_U[v] * gamma) . (J @ u) one matvec, streamed | |
| denominator_v = || (W_U[v] * gamma) @ J || independent of u | |
| so only the denominators have to be precomputed, and they are one float per token: | |
| `sqrt(r_v @ (J J^T) @ r_v)` with `r_v = W_U[v] * gamma`, computed in chunks and cached at | |
| about a megabyte instead of gigabytes. After that each decode is one streamed pass over the | |
| unembedding. `cosine=False` skips the denominators altogether. | |
| Two caveats that the numbers make unavoidable: | |
| * **Transported rows share a large common component**, so raw cosines return much the same | |
| high-frequency words for unrelated directions. `purifier()` estimates the shared direction | |
| and its leading principal components, and `top(..., purify=...)` removes them. | |
| * **An axis decoding to nothing is a real outcome.** Compare any top-k against | |
| `sample_directions()` — what a random direction of the same vocabulary decodes to — before | |
| reading meaning into it. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import re | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| import numpy as np | |
| from ._math import unit | |
| from .models import ModelSpec | |
| from .models import get as _model | |
| #: words a decode will report: alphabetic, at least this long | |
| ALPHA = re.compile(r"^[A-Za-z]+$") | |
| MIN_LEN = 3 | |
| #: rows per streamed chunk. The denominator pass is O(chunk * hidden^2) and the decode pass | |
| #: O(chunk * hidden), so the former runs in smaller chunks to bound peak memory. | |
| NORM_CHUNK = 8192 | |
| DECODE_CHUNK = 65536 | |
| class Hit: | |
| """One decoded token.""" | |
| word: str | |
| score: float | |
| token_id: int | |
| def __repr__(self) -> str: | |
| return f"{self.word}({self.score:+.3f})" | |
| class AxisDecode: | |
| """Both ends of one axis, as words.""" | |
| name: str | |
| positive: list[Hit] | |
| negative: list[Hit] | |
| def words(self, sign: str = "+") -> list[str]: | |
| return [h.word for h in (self.positive if sign == "+" else self.negative)] | |
| def __repr__(self) -> str: | |
| return f"AxisDecode({self.name!r}, +{self.words('+')[:4]}, -{self.words('-')[:4]})" | |
| def common_component(D: np.ndarray, k: int = 3): | |
| """`(mu, V[k, hidden])`: the shared direction of unit rows `D`, and the leading principal | |
| components around it. Computed through the `hidden x hidden` covariance, so the number of | |
| rows only enters linearly.""" | |
| X = np.asarray(D, np.float32) | |
| mu = X.mean(0) | |
| if k <= 0: | |
| return unit(mu), np.zeros((0, X.shape[1]), np.float32) | |
| Xc = X - mu | |
| C = (Xc.T @ Xc) / max(len(Xc) - 1, 1) | |
| _, V = np.linalg.eigh(C) | |
| return unit(mu), np.ascontiguousarray(V[:, ::-1][:, :k].T, np.float32) | |
| def purify(v: np.ndarray, mu: np.ndarray, V: np.ndarray) -> np.ndarray: | |
| """Remove the shared direction and the leading principal components from `v`. | |
| Ranking a purified direction against the raw vocabulary is the same ordering as ranking | |
| the raw direction against a purified vocabulary, so this is applied to the direction. | |
| """ | |
| v = np.asarray(v, np.float32) | |
| mu = np.asarray(mu, np.float32) | |
| v = v - (v @ mu) * mu | |
| V = np.asarray(V, np.float32) | |
| return v - V.T @ (V @ v) if V.size else v | |
| class Vocabulary: | |
| """Every token of a model, as a transported direction, without materialising them.""" | |
| def __init__(self, model: str | ModelSpec, block: int | None = None, | |
| cache_dir: str | Path | None = None): | |
| self.spec = _model(model) | |
| self.block = self.spec.readout_block if block is None else int(block) | |
| self._cache_dir = Path(cache_dir) if cache_dir else None | |
| self._J: np.ndarray | None = None | |
| self._gamma: np.ndarray | None = None | |
| self._norms: np.ndarray | None = None | |
| self._n: int | None = None | |
| self._tok = None | |
| self._word: dict[int, str] = {} | |
| def __repr__(self) -> str: | |
| return f"Vocabulary({self.spec.key!r}, block={self.block})" | |
| # ------------------------------------------------------------------ resources | |
| def J(self) -> np.ndarray: | |
| if self._J is None: | |
| from . import lensfile as _lensfile | |
| self._J = _lensfile.matrix(self.spec.lens, self.block) | |
| return self._J | |
| def n_tokens(self) -> int: | |
| if self._n is None: | |
| from . import unembed | |
| self._n = unembed.unembedding_shape(self.spec.hf, self.spec.revision)[0] | |
| return self._n | |
| def _open(self): | |
| """`(handle, weight key, gamma)` for the unembedding, read lazily, never copied.""" | |
| from . import unembed | |
| from ._deps import require | |
| safe_open = require("safetensors").safe_open | |
| names, shard = unembed._keys_and_shard(self.spec.hf, self.spec.revision) | |
| w_key = (unembed._pick(names, "lm_head.weight") | |
| or unembed._pick(names, "embed_tokens.weight")) | |
| if self._gamma is None: | |
| g_key = unembed._pick(names, "model.norm.weight", ".norm.weight") | |
| with safe_open(shard(g_key), framework="pt") as f: | |
| self._gamma = f.get_tensor(g_key).float().numpy().astype(np.float32) | |
| return safe_open(shard(w_key), framework="pt"), w_key, self._gamma | |
| def _rows(self, chunk: int): | |
| """Yield `(start, (W_U[start:stop] * gamma))` as float32 `[c, hidden]`.""" | |
| handle, w_key, gamma = self._open() | |
| with handle as f: | |
| sl = f.get_slice(w_key) | |
| n = int(sl.get_shape()[0]) | |
| for i in range(0, n, chunk): | |
| R = sl[i:min(i + chunk, n), :].float().numpy().astype(np.float32) | |
| yield i, R * gamma[None, :] | |
| def cache_path(self) -> Path: | |
| from .basis import cache_dir | |
| root = self._cache_dir or cache_dir() | |
| return Path(root) / "vocab_norms" / f"{self.spec.key}-b{self.block}.npy" | |
| def norms(self) -> np.ndarray: | |
| """`[V]` row norms of the transported vocabulary. Computed once, then cached. | |
| `||r_v @ J||^2 = r_v @ (J J^T) @ r_v`, so the `V x hidden` product is never formed. | |
| """ | |
| if self._norms is not None: | |
| return self._norms | |
| p = self.cache_path | |
| if p.exists(): | |
| self._norms = np.load(p) | |
| return self._norms | |
| M = (self.J @ self.J.T).astype(np.float32) | |
| parts = [np.sqrt(np.maximum(((R @ M) * R).sum(1), 0.0)) | |
| for _, R in self._rows(NORM_CHUNK)] | |
| self._norms = np.concatenate(parts).astype(np.float32) | |
| p.parent.mkdir(parents=True, exist_ok=True) | |
| np.save(p, self._norms) | |
| return self._norms | |
| def word(self, token_id: int) -> str: | |
| """The decoded, stripped string of one token. Decoded lazily and cached, because a | |
| decode only ever inspects the few hundred candidates it scans.""" | |
| t = int(token_id) | |
| if t not in self._word: | |
| if self._tok is None: | |
| from ._deps import require | |
| AutoTokenizer = require("transformers").AutoTokenizer | |
| kw = {"revision": self.spec.revision} if self.spec.revision else {} | |
| self._tok = AutoTokenizer.from_pretrained(self.spec.hf, **kw) | |
| self._word[t] = self._tok.decode([t]).strip() | |
| return self._word[t] | |
| # ------------------------------------------------------------------ decoding | |
| def scores(self, direction, cosine: bool = True, purify_with=None) -> np.ndarray: | |
| """Alignment of every token's transported direction with `direction`, `[V]`. | |
| `purify_with` is a `(mu, V)` pair from `purifier()`; it is applied to the direction. | |
| """ | |
| u = np.asarray(direction, np.float32).reshape(-1) | |
| if u.shape[0] != self.spec.hidden_size: | |
| raise ValueError(f"direction is {u.shape[0]}-wide, " | |
| f"{self.spec.key} is {self.spec.hidden_size}") | |
| if purify_with is not None: | |
| u = purify(u, *purify_with) | |
| v = (self.J @ u).astype(np.float32) | |
| s = np.concatenate([R @ v for _, R in self._rows(DECODE_CHUNK)]).astype(np.float32) | |
| if cosine: | |
| s = s / (self.norms + 1e-9) / (np.linalg.norm(u) + 1e-9) | |
| return s | |
| def top(self, direction, k: int = 12, sign: str = "+", min_len: int = MIN_LEN, | |
| alpha_only: bool = True, scan: int = 40, cosine: bool = True, | |
| purify_with=None) -> list[Hit]: | |
| """The `k` tokens best aligned with `direction` (`sign="-"` for the other end). | |
| Candidates are scanned `k * scan` deep and filtered to alphabetic words of at least | |
| `min_len` characters, deduplicated case-insensitively, which is the convention the | |
| reported decodes use: a raw top-k is mostly fragments and punctuation. | |
| """ | |
| s = self.scores(direction, cosine=cosine, purify_with=purify_with) | |
| order = np.argsort(-s if sign == "+" else s) | |
| hits: list[Hit] = [] | |
| seen: set[str] = set() | |
| for t in order[: max(k * scan, k)]: | |
| w = self.word(t) | |
| if (alpha_only and not ALPHA.match(w)) or len(w) < min_len or w.lower() in seen: | |
| continue | |
| seen.add(w.lower()) | |
| hits.append(Hit(word=w, score=float(s[t]), token_id=int(t))) | |
| if len(hits) >= k: | |
| break | |
| return hits | |
| def decode(self, target, k: int = 12, **kw) -> list[AxisDecode]: | |
| """Both ends of an Axis, of every axis of a Cone, or of a bare direction.""" | |
| from .axes import Axis, Cone | |
| if isinstance(target, Axis): | |
| pairs = [(target.name, target.vector)] | |
| elif isinstance(target, Cone): | |
| pairs = [(n, target.U[:, i]) for i, n in enumerate(target.names)] | |
| else: | |
| pairs = [("direction", np.asarray(target, np.float32))] | |
| return [AxisDecode(name=n, positive=self.top(v, k=k, sign="+", **kw), | |
| negative=self.top(v, k=k, sign="-", **kw)) for n, v in pairs] | |
| def report(self, target, k: int = 12, **kw) -> str: | |
| """A printable decode of every axis, both ends.""" | |
| out = [f"{self.spec.key} block {self.block}, {self.n_tokens} tokens"] | |
| for d in self.decode(target, k=k, **kw): | |
| out.append(f" [{d.name:22s}] + {', '.join(d.words('+'))}") | |
| out.append(f" {'':24s} - {', '.join(d.words('-'))}") | |
| return "\n".join(out) | |
| # ------------------------------------------------------------------ nulls | |
| def sample_directions(self, n: int = 20000, seed: int = 0) -> np.ndarray: | |
| """`[n, hidden]` unit transported directions of a random token sample. | |
| Feeds `purifier`, and gives the null an axis's decode has to be read against: what a | |
| random direction of this same vocabulary returns. | |
| """ | |
| rng = np.random.default_rng(seed) | |
| idx = np.sort(rng.choice(self.n_tokens, size=min(n, self.n_tokens), replace=False)) | |
| pos = {int(t): j for j, t in enumerate(idx)} | |
| out = np.empty((len(idx), self.spec.hidden_size), np.float32) | |
| for start, R in self._rows(DECODE_CHUNK): | |
| want = idx[(idx >= start) & (idx < start + len(R))] | |
| if not len(want): | |
| continue | |
| out[[pos[int(t)] for t in want]] = R[want - start] @ self.J | |
| return unit(out) | |
| def purifier(self, k: int = 3, n_rows: int = 20000, seed: int = 0): | |
| """`(mu, V)` for `purify`, estimated from a token sample of this vocabulary.""" | |
| return common_component(self.sample_directions(n_rows, seed), k=k) | |
| def save_decode(decodes: list[AxisDecode], path: str | Path) -> Path: | |
| """Write decoded axes to json.""" | |
| path = Path(path) | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| blob = [{"name": d.name, | |
| "positive": [{"word": h.word, "score": h.score, "id": h.token_id} | |
| for h in d.positive], | |
| "negative": [{"word": h.word, "score": h.score, "id": h.token_id} | |
| for h in d.negative]} for d in decodes] | |
| path.write_text(json.dumps(blob, indent=1) + "\n") | |
| return path | |