From 951bf724638445d8a0b84f526fbb1e1d66a1d4fb Mon Sep 17 00:00:00 2001
From: Kauna <16511995+klei22@users.noreply.github.com>
Date: Wed, 22 Jul 2026 18:55:58 -0700
Subject: [PATCH 1/2] Add QK attention projection functions
---
report/threejs/hf-vocab-sphere/app/main.py | 14 +-
.../hf-vocab-sphere/app/model_store.py | 309 +++++++++++++++++-
report/threejs/hf-vocab-sphere/app/schemas.py | 6 +
.../threejs/hf-vocab-sphere/app/static/app.js | 47 ++-
.../hf-vocab-sphere/app/templates/index.html | 12 +-
.../hf-vocab-sphere/app/vector_math.py | 103 ++++--
.../tests/test_frontend_contract.py | 9 +-
.../hf-vocab-sphere/tests/test_vector_math.py | 56 +++-
8 files changed, 510 insertions(+), 46 deletions(-)
diff --git a/report/threejs/hf-vocab-sphere/app/main.py b/report/threejs/hf-vocab-sphere/app/main.py
index 26079613a5..38a784b491 100644
--- a/report/threejs/hf-vocab-sphere/app/main.py
+++ b/report/threejs/hf-vocab-sphere/app/main.py
@@ -6,7 +6,6 @@
from pathlib import Path
import numpy as np
-
from fastapi import FastAPI, HTTPException, Query
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
@@ -19,12 +18,13 @@
list_local_models,
load_model,
model_status,
+ model_vector_function,
nearest_neighbors,
search_tokens,
- tokenize_text,
selected_token_rows,
selected_vectors,
token_record,
+ tokenize_text,
unload_model,
)
from .projections import nearest_neighbor_edges, project_vectors, projection_catalog
@@ -47,7 +47,7 @@
TokenSearchResponse,
TokenWindowResponse,
)
-from .vector_math import VectorExpressionResult, alias_for_index, evaluate_vector_expressions
+from .vector_math import VectorExpressionResult, alias_for_index, evaluate_vector_expressions, vector_dimension_metrics
BASE_DIR = Path(__file__).resolve().parent
INDEX_PATH = BASE_DIR / "templates" / "index.html"
@@ -237,7 +237,11 @@ def _prepare_projection_selection(
requested_ids.insert(0, anchor_id)
ids, base_vectors = selected_vectors(assets, requested_ids)
anchor_index = ids.index(anchor_id) if anchor_id is not None else None
- resultants = evaluate_vector_expressions(base_vectors, arithmetic_expressions or [])
+ resultants = evaluate_vector_expressions(
+ base_vectors,
+ arithmetic_expressions or [],
+ model_function=lambda name, args: model_vector_function(assets, name, args),
+ )
if resultants:
vectors = np.vstack([base_vectors, *(item.vector[None, :] for item in resultants)])
else:
@@ -262,6 +266,7 @@ def _projection_rows(
"kind": "token",
"alias": alias_for_index(index),
"label": row["display"],
+ **vector_dimension_metrics(base_vectors[index]),
"expression": None,
"referenced_aliases": [],
}
@@ -283,6 +288,7 @@ def _projection_rows(
"special": False,
"present_in_tokenizer": False,
"magnitude": item.magnitude,
+ **vector_dimension_metrics(item.vector),
"rank": None,
"cosine_similarity": None,
"angle_deg": None,
diff --git a/report/threejs/hf-vocab-sphere/app/model_store.py b/report/threejs/hf-vocab-sphere/app/model_store.py
index 1633a1f4a9..86410efcc0 100644
--- a/report/threejs/hf-vocab-sphere/app/model_store.py
+++ b/report/threejs/hf-vocab-sphere/app/model_store.py
@@ -6,9 +6,10 @@
import os
import re
import threading
+from collections.abc import Iterable
from dataclasses import dataclass
from pathlib import Path
-from typing import Any, Iterable
+from typing import Any
import numpy as np
import torch
@@ -102,7 +103,10 @@ def dtype(self) -> str:
@property
def memory_bytes(self) -> int:
- return int(self.weight.nelement() * self.weight.element_size() + self.magnitudes.nelement() * self.magnitudes.element_size())
+ return int(
+ self.weight.nelement() * self.weight.element_size()
+ + self.magnitudes.nelement() * self.magnitudes.element_size()
+ )
def token(self, token_id: int) -> TokenInfo:
if token_id < 0 or token_id >= self.vocab_size:
@@ -346,7 +350,9 @@ def _load_full_model(
try:
model = model_class.from_pretrained(model_name, **kwargs)
output = model.get_output_embeddings() if callable(getattr(model, "get_output_embeddings", None)) else None
- input_embedding = model.get_input_embeddings() if callable(getattr(model, "get_input_embeddings", None)) else None
+ input_embedding = (
+ model.get_input_embeddings() if callable(getattr(model, "get_input_embeddings", None)) else None
+ )
if requested_source == "output":
chosen, source = output, "output"
elif requested_source == "input":
@@ -354,7 +360,9 @@ def _load_full_model(
else:
chosen, source = (output, "output") if output is not None else (input_embedding, "input")
if chosen is None or getattr(chosen, "weight", None) is None:
- raise ValueError(f"{model_class.__name__} does not expose a usable {requested_source} vocabulary matrix.")
+ raise ValueError(
+ f"{model_class.__name__} does not expose a usable {requested_source} vocabulary matrix."
+ )
tensor = chosen.weight.detach().cpu().contiguous()
name = "get_output_embeddings().weight" if source == "output" else "get_input_embeddings().weight"
return tensor, name, source
@@ -975,3 +983,296 @@ def list_local_models() -> list[LocalModelInfo]:
modified = None
records[name] = LocalModelInfo(name, str(path), None, modified)
return sorted(records.values(), key=lambda item: item.model_name.casefold())
+
+
+_AUX_TENSOR_CACHE: dict[tuple[str, str, bool], dict[str, torch.Tensor]] = {}
+
+
+def _all_safetensor_tensors(model_name: str, *, revision: str, allow_download: bool = True) -> dict[str, torch.Tensor]:
+ cache_key = (model_name, revision, allow_download)
+ cached = _AUX_TENSOR_CACHE.get(cache_key)
+ if cached is not None:
+ return cached
+ tensors: dict[str, torch.Tensor] = {}
+ index_path = _try_repo_file(
+ model_name, "model.safetensors.index.json", revision=revision, allow_download=allow_download
+ )
+ if index_path is not None:
+ with index_path.open("r", encoding="utf-8") as handle:
+ data = json.load(handle)
+ weight_map = data.get("weight_map") or {}
+ files = sorted(set(str(value) for value in weight_map.values() if isinstance(value, str)))
+ else:
+ files = _list_safetensors_files(model_name, revision=revision, allow_download=allow_download)
+ for filename in files:
+ path = _try_repo_file(model_name, filename, revision=revision, allow_download=allow_download)
+ if path is None:
+ continue
+ with safe_open(str(path), framework="pt", device="cpu") as handle:
+ for name in handle.keys():
+ lname = str(name).casefold()
+ if (
+ lname.endswith(("norm.weight", "norm.bias", "layernorm.weight", "layernorm.bias"))
+ or ".q_proj.weight" in lname
+ or ".query.weight" in lname
+ or ".k_proj.weight" in lname
+ or ".key.weight" in lname
+ or ".v_proj.weight" in lname
+ or ".value.weight" in lname
+ or ".o_proj.weight" in lname
+ or ".out_proj.weight" in lname
+ ):
+ tensors[str(name)] = (
+ handle.get_tensor(str(name)).detach().cpu().to(dtype=torch.float32).contiguous()
+ )
+ _AUX_TENSOR_CACHE[cache_key] = tensors
+ return tensors
+
+
+def _layer_prefix(tensor_name: str, layer: int) -> str:
+ markers = (f"layers.{layer}.", f"h.{layer}.", f"blocks.{layer}.", f"layer.{layer}.")
+ for marker in markers:
+ if marker in tensor_name:
+ return tensor_name.split(marker, 1)[0] + marker
+ return ""
+
+
+def _find_tensor(
+ tensors: dict[str, torch.Tensor], needles: Iterable[str], *, layer: int | None = None, ndim: int | None = None
+) -> tuple[str, torch.Tensor]:
+ names = list(tensors)
+ if layer is not None:
+ names = [name for name in names if _layer_prefix(name, layer)]
+ needle_list = [needle.casefold() for needle in needles]
+ matches = [name for name in names if any(needle in name.casefold() for needle in needle_list)]
+ if ndim is not None:
+ matches = [name for name in matches if tensors[name].ndim == ndim]
+ if not matches:
+ scope = f" layer {layer}" if layer is not None else ""
+ raise ValueError(f"Could not find model tensor matching {', '.join(needles)}{scope}.")
+ return sorted(matches, key=lambda name: (len(name), name))[0], tensors[
+ sorted(matches, key=lambda name: (len(name), name))[0]
+ ]
+
+
+def _norm_weight_for_args(
+ tensors: dict[str, torch.Tensor], args: list[float | str | np.ndarray]
+) -> tuple[np.ndarray, torch.Tensor]:
+ if len(args) == 2 and str(args[1]).casefold() == "final":
+ vector = np.asarray(args[0], dtype=np.float64)
+ _name, weight = _find_tensor(
+ tensors, ("model.norm.weight", "final_layernorm.weight", "ln_f.weight", "norm.weight"), ndim=1
+ )
+ return vector, weight
+ if len(args) != 4:
+ raise ValueError("norm/invnorm expect (vector, layer, 'attn'|'ffn', 'input'|'output') or (vector, 'final').")
+
+ vector = np.asarray(args[0], dtype=np.float64)
+ layer = int(args[1])
+ block = str(args[2]).casefold()
+ position = str(args[3]).casefold()
+ if block in {"attn", "attention"} and position in {"input", "before", "in"}:
+ needles = ("input_layernorm.weight", "ln_1.weight", "attention_norm.weight")
+ elif block in {"attn", "attention"} and position in {"output", "after", "out"}:
+ needles = ("post_attention_layernorm.weight", "post_attention_norm.weight", "ln_2.weight")
+ elif block in {"ffn", "mlp"} and position in {"input", "before", "in"}:
+ needles = ("pre_feedforward_layernorm.weight", "post_attention_layernorm.weight", "ln_2.weight")
+ elif block in {"ffn", "mlp"} and position in {"output", "after", "out"}:
+ needles = ("post_feedforward_layernorm.weight", "post_feedforward_norm.weight")
+ else:
+ raise ValueError("norm/invnorm expect (vector, layer, 'attn'|'ffn', 'input'|'output') or (vector, 'final').")
+ _name, weight = _find_tensor(tensors, needles, layer=layer, ndim=1)
+ return vector, weight
+
+
+def _apply_rms_norm(vector: np.ndarray, weight: torch.Tensor) -> np.ndarray:
+ w = weight.numpy().astype(np.float64)
+ if w.shape[0] != vector.shape[0]:
+ raise ValueError(f"Norm width {w.shape[0]} does not match vector width {vector.shape[0]}.")
+ return vector * w / math.sqrt(float(np.mean(vector * vector)) + 1e-6)
+
+
+def _apply_inverse_rms_norm(vector: np.ndarray, weight: torch.Tensor) -> np.ndarray:
+ w = weight.numpy().astype(np.float64)
+ if w.shape[0] != vector.shape[0]:
+ raise ValueError(f"Norm width {w.shape[0]} does not match vector width {vector.shape[0]}.")
+ if np.any(np.abs(w) <= 1e-15):
+ raise ValueError("Cannot invert a norm with zero or near-zero weights.")
+ unweighted = vector / w
+ normalized_square_mean = float(np.mean(unweighted * unweighted))
+ if normalized_square_mean >= 1.0:
+ raise ValueError("Cannot invert this norm result because its implied pre-norm magnitude is not finite.")
+ scale = math.sqrt(1e-6 / max(1.0 - normalized_square_mean, 1e-15))
+ return unweighted * scale
+
+
+def _attention_head_counts(assets: ModelAssets) -> tuple[int, int]:
+ config_path = _try_repo_file(assets.model_name, "config.json", revision=assets.revision, allow_download=True)
+ if config_path is None:
+ return 1, 1
+ with config_path.open("r", encoding="utf-8") as handle:
+ config = json.load(handle)
+ attention_heads = int(config.get("num_attention_heads") or config.get("n_head") or 0)
+ key_value_heads = int(config.get("num_key_value_heads") or attention_heads or 0)
+ return max(attention_heads, 1), max(key_value_heads, 1)
+
+
+def _wov_head_slices(
+ *,
+ v_output_width: int,
+ o_input_width: int,
+ attention_heads: int,
+ key_value_heads: int,
+ head: int,
+) -> tuple[slice, slice, int]:
+ if attention_heads <= 0:
+ attention_heads = 1
+ if key_value_heads <= 0:
+ key_value_heads = attention_heads
+ if o_input_width % attention_heads != 0:
+ # Fall back to treating Wo as one concatenated head space when the config
+ # does not divide the tensor cleanly.
+ attention_heads = 1
+ head_dim = o_input_width // attention_heads
+ if head_dim <= 0:
+ raise ValueError("Wo input width is not usable for head slicing.")
+ if head < 0 or head >= attention_heads:
+ raise ValueError(f"head must be in [0, {attention_heads - 1}], got {head}.")
+
+ if v_output_width == o_input_width:
+ key_value_heads = attention_heads
+ elif v_output_width % key_value_heads != 0:
+ if v_output_width == head_dim:
+ key_value_heads = 1
+ elif v_output_width % head_dim == 0:
+ key_value_heads = v_output_width // head_dim
+ else:
+ raise ValueError("Wv output width cannot be partitioned into Wo-compatible heads.")
+
+ v_head_dim = v_output_width // key_value_heads
+ if v_head_dim != head_dim:
+ raise ValueError("Wv head width and Wo head width do not match.")
+ kv_head = min((head * key_value_heads) // attention_heads, key_value_heads - 1)
+ v_start = kv_head * head_dim
+ o_start = head * head_dim
+ return slice(v_start, v_start + head_dim), slice(o_start, o_start + head_dim), attention_heads
+
+
+def _qk_head_slices(
+ *,
+ q_output_width: int,
+ k_output_width: int,
+ attention_heads: int,
+ key_value_heads: int,
+ query_head: int,
+ key_head: int,
+) -> tuple[slice, slice]:
+ if attention_heads <= 0:
+ attention_heads = 1
+ if key_value_heads <= 0:
+ key_value_heads = attention_heads
+ if q_output_width % attention_heads != 0:
+ attention_heads = 1
+ q_head_dim = q_output_width // attention_heads
+ if q_head_dim <= 0:
+ raise ValueError("Wq output width is not usable for head slicing.")
+ if query_head < 0 or query_head >= attention_heads:
+ raise ValueError(f"query head must be in [0, {attention_heads - 1}], got {query_head}.")
+
+ if k_output_width == q_output_width:
+ key_value_heads = attention_heads
+ elif k_output_width % key_value_heads != 0:
+ if k_output_width == q_head_dim:
+ key_value_heads = 1
+ elif k_output_width % q_head_dim == 0:
+ key_value_heads = k_output_width // q_head_dim
+ else:
+ raise ValueError("Wk output width cannot be partitioned into Wq-compatible heads.")
+
+ k_head_dim = k_output_width // key_value_heads
+ if k_head_dim != q_head_dim:
+ raise ValueError("Wq head width and Wk head width do not match.")
+ if key_head < 0 or key_head >= key_value_heads:
+ raise ValueError(f"key head must be in [0, {key_value_heads - 1}], got {key_head}.")
+ q_start = query_head * q_head_dim
+ k_start = key_head * q_head_dim
+ return slice(q_start, q_start + q_head_dim), slice(k_start, k_start + q_head_dim)
+
+
+def model_vector_function(assets: ModelAssets, name: str, args: list[float | str | np.ndarray]) -> np.ndarray:
+ tensors = _all_safetensor_tensors(assets.model_name, revision=assets.revision, allow_download=True)
+ if name in {"norm", "invnorm", "inverse_norm"}:
+ vector, weight = _norm_weight_for_args(tensors, args)
+ if name == "norm":
+ return _apply_rms_norm(vector, weight)
+ return _apply_inverse_rms_norm(vector, weight)
+ if name == "wov":
+ if len(args) != 3:
+ raise ValueError("wov expects wov(vector, layer, head).")
+ vector = np.asarray(args[0], dtype=np.float64)
+ layer = int(args[1])
+ head = int(args[2])
+ _vn, v_weight = _find_tensor(
+ tensors,
+ ("self_attn.v_proj.weight", "attention.v_proj.weight", "attn.v_proj.weight", "value.weight"),
+ layer=layer,
+ ndim=2,
+ )
+ _on, o_weight = _find_tensor(
+ tensors,
+ ("self_attn.o_proj.weight", "attention.o_proj.weight", "attn.o_proj.weight", "out_proj.weight"),
+ layer=layer,
+ ndim=2,
+ )
+ v = v_weight.numpy().astype(np.float64)
+ o = o_weight.numpy().astype(np.float64)
+ if v.shape[1] != vector.shape[0] or o.shape[0] != vector.shape[0]:
+ raise ValueError("Attention projection widths do not match the selected vector space.")
+ attention_heads, key_value_heads = _attention_head_counts(assets)
+ v_slice, o_slice, _head_count = _wov_head_slices(
+ v_output_width=v.shape[0],
+ o_input_width=o.shape[1],
+ attention_heads=attention_heads,
+ key_value_heads=key_value_heads,
+ head=head,
+ )
+ # Multiplying by the selected Wo slice is equivalent to padding the
+ # selected Wv-head output with zeros before/after its head slot and then
+ # applying the full Wo matrix, but avoids materializing the padded vector.
+ return (vector @ v[v_slice, :].T) @ o[:, o_slice].T
+ if name in {"wqk", "wkq"}:
+ if len(args) != 4:
+ raise ValueError("wqk/wkq expect (vector, layer, query_head, key_head).")
+ vector = np.asarray(args[0], dtype=np.float64)
+ layer = int(args[1])
+ query_head = int(args[2])
+ key_head = int(args[3])
+ _qn, q_weight = _find_tensor(
+ tensors,
+ ("self_attn.q_proj.weight", "attention.q_proj.weight", "attn.q_proj.weight", "query.weight"),
+ layer=layer,
+ ndim=2,
+ )
+ _kn, k_weight = _find_tensor(
+ tensors,
+ ("self_attn.k_proj.weight", "attention.k_proj.weight", "attn.k_proj.weight", "key.weight"),
+ layer=layer,
+ ndim=2,
+ )
+ q = q_weight.numpy().astype(np.float64)
+ k = k_weight.numpy().astype(np.float64)
+ if q.shape[1] != vector.shape[0] or k.shape[1] != vector.shape[0]:
+ raise ValueError("Attention Q/K projection widths do not match the selected vector space.")
+ attention_heads, key_value_heads = _attention_head_counts(assets)
+ q_slice, k_slice = _qk_head_slices(
+ q_output_width=q.shape[0],
+ k_output_width=k.shape[0],
+ attention_heads=attention_heads,
+ key_value_heads=key_value_heads,
+ query_head=query_head,
+ key_head=key_head,
+ )
+ if name == "wqk":
+ return (vector @ q[q_slice, :].T) @ k[k_slice, :]
+ return (vector @ k[k_slice, :].T) @ q[q_slice, :]
+ raise ValueError(f"Unknown model function {name!r}.")
diff --git a/report/threejs/hf-vocab-sphere/app/schemas.py b/report/threejs/hf-vocab-sphere/app/schemas.py
index 7affb61d04..10e975830c 100644
--- a/report/threejs/hf-vocab-sphere/app/schemas.py
+++ b/report/threejs/hf-vocab-sphere/app/schemas.py
@@ -46,6 +46,9 @@ class TokenRecord(BaseModel):
special: bool = False
present_in_tokenizer: bool = True
magnitude: float | None = None
+ effective_dimension: float | None = None
+ dimensions_for_angle: int | None = Field(default=None, ge=0)
+ continuous_dimensions_for_angle: float | None = None
rank: int | None = Field(default=None, ge=1)
cosine_similarity: float | None = None
angle_deg: float | None = None
@@ -138,6 +141,9 @@ class ProjectionPoint(BaseModel):
special: bool = False
present_in_tokenizer: bool = True
magnitude: float | None = None
+ effective_dimension: float | None = None
+ dimensions_for_angle: int | None = Field(default=None, ge=0)
+ continuous_dimensions_for_angle: float | None = None
rank: int | None = Field(default=None, ge=1)
cosine_similarity: float | None = None
angle_deg: float | None = None
diff --git a/report/threejs/hf-vocab-sphere/app/static/app.js b/report/threejs/hf-vocab-sphere/app/static/app.js
index 8ce73b5d9e..75c6e93f0a 100644
--- a/report/threejs/hf-vocab-sphere/app/static/app.js
+++ b/report/threejs/hf-vocab-sphere/app/static/app.js
@@ -38,6 +38,8 @@ const state = {
showLabels: false,
showLabelAliases: true,
showLabelIds: true,
+ showLabelMagnitudes: false,
+ showLabelParticipation: false,
showEdgeLabels: false,
edgeColorMode: 'angle',
edgeWidth: 1.8,
@@ -363,14 +365,32 @@ function clearEdgeLabels() {
edgeLabelEntries = [];
}
+function labelMagnitudeText(point) {
+ return state.showLabelMagnitudes && Number.isFinite(Number(point.magnitude))
+ ? `‖v‖ ${formatNumber(point.magnitude, 3)}`
+ : '';
+}
+
+function labelParticipationText(point) {
+ return state.showLabelParticipation && Number.isFinite(Number(point.effective_dimension))
+ ? `PR ${formatNumber(point.effective_dimension, 2)}`
+ : '';
+}
+
function labelTextForPoint(point) {
- if (point.kind === 'resultant') return `${point.alias || 'R'}: ${point.label || point.display}`;
+ const magnitudeText = labelMagnitudeText(point);
+ const participationText = labelParticipationText(point);
+ if (point.kind === 'resultant') {
+ return [`${point.alias || 'R'}: ${point.label || point.display}`, magnitudeText, participationText]
+ .filter((value) => value.length)
+ .join(' · ');
+ }
const prefix = state.showLabelAliases && point.alias ? `${point.alias}: ` : '';
const tokenText = point.display !== null && point.display !== undefined ? String(point.display) : '';
const tokenId = state.showLabelIds && point.token_id !== null && point.token_id !== undefined
? String(point.token_id)
: '';
- const body = [tokenText, tokenId].filter((value) => value.length).join(' · ');
+ const body = [tokenText, tokenId, magnitudeText, participationText].filter((value) => value.length).join(' · ');
return `${prefix}${body}`.trim();
}
@@ -1354,7 +1374,8 @@ function renderArithmeticPanel() {
async function addArithmeticResult() {
const rows = selectedRows();
const slerpMode = $('arithmeticModeSelect').value === 'slerp';
- const fraction = Math.max(0, Math.min(1, Number($('slerpFractionInput').value) || 0));
+ const rawFraction = Number($('slerpFractionInput').value);
+ const fraction = Number.isFinite(rawFraction) ? rawFraction : 0;
const fromAlias = $('slerpFromSelect').value;
const toAlias = $('slerpToSelect').value;
const expression = slerpMode
@@ -1737,6 +1758,8 @@ function buildSettingsSnapshot() {
show_node_labels: state.showLabels,
show_alias_letters: state.showLabelAliases,
show_label_ids: state.showLabelIds,
+ show_label_magnitudes: state.showLabelMagnitudes,
+ show_label_participation: state.showLabelParticipation,
show_edge_labels: state.showEdgeLabels,
edge_color_mode: state.edgeColorMode,
edge_width: state.edgeWidth,
@@ -1848,7 +1871,7 @@ function normalizedWorkspace(snapshot) {
label: stringValue(editor.label, '', 120),
slerpFrom: stringValue(editor.slerp_from, 'A', 20),
slerpTo: stringValue(editor.slerp_to, 'B', 20),
- slerpFraction: finiteNumber(editor.slerp_fraction, 0.5, 0, 1),
+ slerpFraction: finiteNumber(editor.slerp_fraction, 0.5, -1e9, 1e9),
},
};
}
@@ -1944,6 +1967,8 @@ function applySettingsControls(snapshot) {
state.showLabels = booleanValue(appearance.show_node_labels, false);
state.showLabelAliases = booleanValue(appearance.show_alias_letters, true);
state.showLabelIds = booleanValue(appearance.show_label_ids, true);
+ state.showLabelMagnitudes = booleanValue(appearance.show_label_magnitudes, false);
+ state.showLabelParticipation = booleanValue(appearance.show_label_participation, false);
state.showEdgeLabels = booleanValue(appearance.show_edge_labels, false) && state.showEdges;
state.edgeColorMode = appearance.edge_color_mode === 'uniform' ? 'uniform' : 'angle';
state.edgeWidth = finiteNumber(appearance.edge_width, 1.8, 0.5, 8);
@@ -2015,6 +2040,8 @@ function syncAppearanceControls() {
$('showLabelsInput').checked = state.showLabels;
$('showLabelAliasesInput').checked = state.showLabelAliases;
$('showLabelIdsInput').checked = state.showLabelIds;
+ $('showLabelMagnitudesInput').checked = state.showLabelMagnitudes;
+ $('showLabelParticipationInput').checked = state.showLabelParticipation;
$('showEdgeLabelsInput').checked = state.showEdgeLabels;
$('pointSizeInput').value = String(state.pointSize);
$('pointSizeLabel').textContent = state.pointSize.toFixed(3);
@@ -2066,6 +2093,14 @@ function setShowLabelIds(value) {
state.showLabelIds = Boolean(value);
syncAppearanceControls();
}
+function setShowLabelMagnitudes(value) {
+ state.showLabelMagnitudes = Boolean(value);
+ syncAppearanceControls();
+}
+function setShowLabelParticipation(value) {
+ state.showLabelParticipation = Boolean(value);
+ syncAppearanceControls();
+}
function setShowEdgeLabels(value) {
state.showEdgeLabels = Boolean(value);
if (state.showEdgeLabels) state.showEdges = true;
@@ -2247,6 +2282,8 @@ $('showEdgeLabelsInput').addEventListener('change', (event) => setShowEdgeLabels
$('showLabelsInput').addEventListener('change', (event) => setShowLabels(event.target.checked));
$('showLabelAliasesInput').addEventListener('change', (event) => setShowLabelAliases(event.target.checked));
$('showLabelIdsInput').addEventListener('change', (event) => setShowLabelIds(event.target.checked));
+$('showLabelMagnitudesInput').addEventListener('change', (event) => setShowLabelMagnitudes(event.target.checked));
+$('showLabelParticipationInput').addEventListener('change', (event) => setShowLabelParticipation(event.target.checked));
$('resetCameraBtn').addEventListener('click', resetCamera);
$('toggleRotateBtn').addEventListener('click', () => setAutoRotate(!state.autoRotate));
$('toggleEdgesBtn').addEventListener('click', () => setShowEdges(!state.showEdges));
@@ -2269,6 +2306,8 @@ document.addEventListener('keydown', (event) => {
else if (event.key.toLowerCase() === 'l') setShowLabels(!state.showLabels);
else if (event.key.toLowerCase() === 't') setShowLabelAliases(!state.showLabelAliases);
else if (event.key.toLowerCase() === 'i') setShowLabelIds(!state.showLabelIds);
+ else if (event.key.toLowerCase() === 'm') setShowLabelMagnitudes(!state.showLabelMagnitudes);
+ else if (event.key.toLowerCase() === 'p') setShowLabelParticipation(!state.showLabelParticipation);
else if (event.key === 'Escape') {
state.pinnedIndex = null;
renderInspector();
diff --git a/report/threejs/hf-vocab-sphere/app/templates/index.html b/report/threejs/hf-vocab-sphere/app/templates/index.html
index 1b42f47543..12e2df6d14 100644
--- a/report/threejs/hf-vocab-sphere/app/templates/index.html
+++ b/report/threejs/hf-vocab-sphere/app/templates/index.html
@@ -261,15 +261,15 @@
Vector arithmetic
-
Interpolation fractiont = 0.50
-
+
Interpolation / extrapolation valuet = 0.50
+
-
SLERP follows the shortest geodesic between normalized full-dimensional vectors while linearly interpolating their original L2 magnitudes.
+
SLERP follows the shortest geodesic between normalized full-dimensional vectors while linearly interpolating their original L2 magnitudes; finite negative and positive t values are accepted for extrapolation.
-
Expressions support aliases, scalars, parentheses, + - * /, mean(...), and slerp(A,B,t). Example average: (A+B+C)/3 + D. Resultants are projected from the original model space, never by adding screen coordinates.
+
Expressions support aliases, scalars, parentheses, + - * /, mean(...), slerp(A,B,t), norm(A,0,'attn','input')/norm(A,'final'), invnorm(A,...), wov(A,0,0), wqk(A,0,0,0), and wkq(A,0,0,0). Functions can be nested. Example average: (A+B+C)/3 + D. Resultants are projected from the original model space, never by adding screen coordinates.
No resultants defined.
@@ -452,6 +452,8 @@
Appearance
+
+
@@ -468,6 +470,8 @@
Presenter hotkeys
LToggle node labelsTShow / hide A/B/C alias tagsIShow / hide token IDs in labels
+ MShow / hide vector magnitudes in labels
+ PShow / hide participation ratio in labelsEscClear pinned token
diff --git a/report/threejs/hf-vocab-sphere/app/vector_math.py b/report/threejs/hf-vocab-sphere/app/vector_math.py
index 99bd18f832..2088724be8 100644
--- a/report/threejs/hf-vocab-sphere/app/vector_math.py
+++ b/report/threejs/hf-vocab-sphere/app/vector_math.py
@@ -2,8 +2,8 @@
import ast
import math
+from collections.abc import Callable, Iterable
from dataclasses import dataclass
-from typing import Iterable
import numpy as np
@@ -43,7 +43,7 @@ def alias_map(vectors: np.ndarray) -> dict[str, np.ndarray]:
@dataclass(slots=True)
class _Value:
kind: str
- value: float | np.ndarray
+ value: float | str | np.ndarray
def spherical_linear_interpolation(a: np.ndarray, b: np.ndarray, t: float) -> np.ndarray:
@@ -56,17 +56,12 @@ def spherical_linear_interpolation(a: np.ndarray, b: np.ndarray, t: float) -> np
first = np.asarray(a, dtype=np.float64)
second = np.asarray(b, dtype=np.float64)
fraction = float(t)
- if not math.isfinite(fraction) or fraction < 0.0 or fraction > 1.0:
- raise ValueError("SLERP t must be a finite scalar in the interval [0, 1].")
+ if not math.isfinite(fraction):
+ raise ValueError("SLERP t must be a finite scalar.")
norm_a = float(np.linalg.norm(first))
norm_b = float(np.linalg.norm(second))
if norm_a <= 1e-15 or norm_b <= 1e-15:
raise ValueError("SLERP requires two non-zero vectors.")
- if fraction <= 0.0:
- return first.copy()
- if fraction >= 1.0:
- return second.copy()
-
unit_a = first / norm_a
unit_b = second / norm_b
dot = float(np.clip(np.dot(unit_a, unit_b), -1.0, 1.0))
@@ -85,19 +80,59 @@ def spherical_linear_interpolation(a: np.ndarray, b: np.ndarray, t: float) -> np
else:
angle = math.acos(dot)
sine = math.sin(angle)
- direction = (
- math.sin((1.0 - fraction) * angle) / sine * unit_a
- + math.sin(fraction * angle) / sine * unit_b
- )
+ direction = math.sin((1.0 - fraction) * angle) / sine * unit_a + math.sin(fraction * angle) / sine * unit_b
direction /= max(float(np.linalg.norm(direction)), 1e-15)
magnitude = (1.0 - fraction) * norm_a + fraction * norm_b
return direction * magnitude
+def vector_dimension_metrics(
+ vector: np.ndarray, max_angle_degrees: float = 5.0, eps: float = 1e-12
+) -> dict[str, float | int]:
+ values = np.asarray(vector, dtype=np.float64).reshape(-1)
+ energy = values * values
+ total = float(np.sum(energy))
+ if total <= eps:
+ return {
+ "effective_dimension": 0.0,
+ "dimensions_for_angle": 0,
+ "continuous_dimensions_for_angle": 0.0,
+ "retained_energy": 0.0,
+ "resulting_angle_degrees": 0.0,
+ }
+
+ probabilities = energy / total
+ square_sum = float(np.sum(probabilities * probabilities))
+ effective_dimension = 1.0 / max(square_sum, eps)
+ sorted_probabilities = np.sort(probabilities)[::-1]
+ cumulative = np.cumsum(sorted_probabilities)
+ angle = math.radians(float(max_angle_degrees))
+ required_energy = math.cos(angle) ** 2
+ k = int(np.searchsorted(cumulative, required_energy, side="left")) + 1
+ k = min(max(k, 1), values.size)
+ previous_energy = float(cumulative[k - 2]) if k > 1 else 0.0
+ current_coordinate_energy = float(sorted_probabilities[k - 1])
+ fractional_k = (k - 1) + (required_energy - previous_energy) / max(current_coordinate_energy, eps)
+ fractional_k = max(1.0, min(float(k), float(fractional_k)))
+ retained_energy = float(cumulative[k - 1])
+ resulting_angle = math.degrees(math.acos(math.sqrt(min(1.0, retained_energy))))
+ return {
+ "effective_dimension": float(effective_dimension),
+ "dimensions_for_angle": int(k),
+ "continuous_dimensions_for_angle": float(fractional_k),
+ "retained_energy": retained_energy,
+ "resulting_angle_degrees": float(resulting_angle),
+ }
+
+
+ModelFunction = Callable[[str, list[float | str | np.ndarray]], np.ndarray]
+
+
class _SafeVectorEvaluator:
- def __init__(self, aliases: dict[str, np.ndarray]):
+ def __init__(self, aliases: dict[str, np.ndarray], model_function: ModelFunction | None = None):
self.aliases = {name.upper(): np.asarray(vector, dtype=np.float64) for name, vector in aliases.items()}
+ self.model_function = model_function
self.referenced: set[str] = set()
def evaluate(self, expression: str) -> np.ndarray:
@@ -138,6 +173,12 @@ def _visit(self, node: ast.AST) -> _Value:
raise ValueError("Scalar constants must be finite and reasonably sized.")
return _Value("scalar", scalar)
+ if isinstance(node, ast.Constant) and isinstance(node.value, str):
+ text = node.value.strip().casefold()
+ if not text or len(text) > 40:
+ raise ValueError("String arguments must be short, non-blank selectors.")
+ return _Value("string", text)
+
if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.UAdd, ast.USub)):
operand = self._visit(node.operand)
sign = 1.0 if isinstance(node.op, ast.UAdd) else -1.0
@@ -145,7 +186,7 @@ def _visit(self, node: ast.AST) -> _Value:
if isinstance(node, ast.Call):
if not isinstance(node.func, ast.Name):
- raise ValueError("Only mean(...) and slerp(A, B, t) vector functions are supported.")
+ raise ValueError("Only named vector functions are supported.")
function_name = node.func.id.casefold()
if node.keywords:
raise ValueError("Vector functions accept positional arguments only.")
@@ -159,15 +200,28 @@ def _visit(self, node: ast.AST) -> _Value:
"vector",
np.mean(np.stack([np.asarray(value.value) for value in values], axis=0), axis=0),
)
+ if function_name in {"norm", "invnorm", "inverse_norm", "wov", "wqk", "wkq"}:
+ if self.model_function is None:
+ raise ValueError(f"{node.func.id} requires model auxiliary tensors loaded by the projection API.")
+ values = [self._visit(argument) for argument in node.args]
+ args: list[float | str | np.ndarray] = [
+ np.asarray(value.value)
+ if value.kind == "vector"
+ else (str(value.value) if value.kind == "string" else float(value.value))
+ for value in values
+ ]
+ return _Value("vector", self.model_function(function_name, args))
if function_name != "slerp":
- raise ValueError("Only mean(...) and slerp(A, B, t) vector functions are supported.")
+ raise ValueError(
+ "Only mean(...), slerp(A, B, t), norm(...), invnorm(...), wov(...), wqk(...), and wkq(...) vector functions are supported."
+ )
if len(node.args) != 3:
raise ValueError("slerp requires exactly three positional arguments: slerp(A, B, t).")
first = self._visit(node.args[0])
second = self._visit(node.args[1])
fraction = self._visit(node.args[2])
if first.kind != "vector" or second.kind != "vector" or fraction.kind != "scalar":
- raise ValueError("slerp requires two vectors followed by a scalar t in [0, 1].")
+ raise ValueError("slerp requires two vectors followed by a scalar t.")
return _Value(
"vector",
spherical_linear_interpolation(
@@ -206,12 +260,14 @@ def _visit(self, node: ast.AST) -> _Value:
return _Value("vector", np.asarray(left.value) / denominator)
raise ValueError(
- "Unsupported expression syntax. Use vector aliases, numeric scalars, parentheses, +, -, *, /, mean(...), and slerp(A, B, t)."
+ "Unsupported expression syntax. Use vector aliases, numeric scalars, parentheses, +, -, *, /, mean(...), slerp(A, B, t), norm(...), invnorm(...), wov(...), wqk(...), and wkq(...)."
)
-def evaluate_vector_expression(expression: str, aliases: dict[str, np.ndarray]) -> tuple[np.ndarray, tuple[str, ...]]:
- evaluator = _SafeVectorEvaluator(aliases)
+def evaluate_vector_expression(
+ expression: str, aliases: dict[str, np.ndarray], model_function: ModelFunction | None = None
+) -> tuple[np.ndarray, tuple[str, ...]]:
+ evaluator = _SafeVectorEvaluator(aliases, model_function=model_function)
vector = evaluator.evaluate(expression)
return vector, tuple(sorted(evaluator.referenced, key=lambda name: (len(name), name)))
@@ -219,6 +275,7 @@ def evaluate_vector_expression(expression: str, aliases: dict[str, np.ndarray])
def evaluate_vector_expressions(
vectors: np.ndarray,
expressions: Iterable[object],
+ model_function: ModelFunction | None = None,
) -> list[VectorExpressionResult]:
"""Evaluate request-like objects with ``expression`` and optional ``label`` attributes."""
aliases = alias_map(vectors)
@@ -226,14 +283,12 @@ def evaluate_vector_expressions(
for index, item in enumerate(expressions):
expression = str(getattr(item, "expression", "") or "").strip()
requested_label = str(getattr(item, "label", "") or "").strip()
- vector, referenced = evaluate_vector_expression(expression, aliases)
+ vector, referenced = evaluate_vector_expression(expression, aliases, model_function=model_function)
magnitude = float(np.linalg.norm(vector))
if not math.isfinite(magnitude):
raise ValueError(f"Resultant {index + 1} has a non-finite magnitude.")
if magnitude <= 1e-12:
- raise ValueError(
- f"Resultant {index + 1} is zero or near-zero and has no direction to project."
- )
+ raise ValueError(f"Resultant {index + 1} is zero or near-zero and has no direction to project.")
alias = f"R{index + 1}"
results.append(
VectorExpressionResult(
diff --git a/report/threejs/hf-vocab-sphere/tests/test_frontend_contract.py b/report/threejs/hf-vocab-sphere/tests/test_frontend_contract.py
index 88bff03c26..245962a25e 100644
--- a/report/threejs/hf-vocab-sphere/tests/test_frontend_contract.py
+++ b/report/threejs/hf-vocab-sphere/tests/test_frontend_contract.py
@@ -4,7 +4,6 @@
from html.parser import HTMLParser
from pathlib import Path
-
ROOT = Path(__file__).resolve().parents[1]
HTML_PATH = ROOT / "app" / "templates" / "index.html"
JS_PATH = ROOT / "app" / "static" / "app.js"
@@ -42,6 +41,8 @@ def test_frontend_control_contract_and_unique_ids() -> None:
"slerpFractionInput",
"showLabelAliasesInput",
"showLabelIdsInput",
+ "showLabelMagnitudesInput",
+ "showLabelParticipationInput",
"saveSettingsBtn",
"loadSettingsBtn",
"settingsFileInput",
@@ -74,9 +75,15 @@ def test_label_component_and_settings_contract() -> None:
assert "showLabelAliases: true" in javascript
assert "showLabelIds: true" in javascript
+ assert "showLabelMagnitudes: false" in javascript
+ assert "showLabelParticipation: false" in javascript
+ assert "show_label_magnitudes: state.showLabelMagnitudes" in javascript
+ assert "show_label_participation: state.showLabelParticipation" in javascript
assert "filter((value) => value.length).join(' · ')" in javascript
assert "event.key.toLowerCase() === 't'" in javascript
assert "event.key.toLowerCase() === 'i'" in javascript
+ assert "event.key.toLowerCase() === 'm'" in javascript
+ assert "event.key.toLowerCase() === 'p'" in javascript
assert "const SETTINGS_SCHEMA = 'hf-vocab-sphere/settings'" in javascript
assert "function buildSettingsSnapshot()" in javascript
diff --git a/report/threejs/hf-vocab-sphere/tests/test_vector_math.py b/report/threejs/hf-vocab-sphere/tests/test_vector_math.py
index 1b23682b56..4d56cd4021 100644
--- a/report/threejs/hf-vocab-sphere/tests/test_vector_math.py
+++ b/report/threejs/hf-vocab-sphere/tests/test_vector_math.py
@@ -5,7 +5,12 @@
import numpy as np
import pytest
-from app.vector_math import alias_for_index, evaluate_vector_expression, evaluate_vector_expressions
+from app.vector_math import (
+ alias_for_index,
+ evaluate_vector_expression,
+ evaluate_vector_expressions,
+ vector_dimension_metrics,
+)
def test_aliases_are_excel_style() -> None:
@@ -48,7 +53,7 @@ def test_slerp_interpolates_direction_and_magnitude() -> None:
"B": np.array([0.0, 4.0, 0.0]),
}
midpoint, references = evaluate_vector_expression("slerp(A, B, 0.5)", aliases)
- expected_direction = np.array([2 ** -0.5, 2 ** -0.5, 0.0])
+ expected_direction = np.array([2**-0.5, 2**-0.5, 0.0])
np.testing.assert_allclose(midpoint, expected_direction * 3.0, atol=1e-8)
assert references == ("A", "B")
@@ -58,10 +63,11 @@ def test_slerp_interpolates_direction_and_magnitude() -> None:
np.testing.assert_allclose(end, aliases["B"])
-def test_slerp_rejects_invalid_fraction() -> None:
+def test_slerp_allows_finite_extrapolation() -> None:
aliases = {"A": np.array([1.0, 0.0]), "B": np.array([0.0, 1.0])}
- with pytest.raises(ValueError, match="interval"):
- evaluate_vector_expression("slerp(A, B, 1.5)", aliases)
+ result, references = evaluate_vector_expression("slerp(A, B, -0.5)", aliases)
+ assert references == ("A", "B")
+ assert np.all(np.isfinite(result))
def test_unsafe_vector_math_is_rejected() -> None:
@@ -78,3 +84,43 @@ def test_zero_resultant_is_rejected() -> None:
vectors,
[SimpleNamespace(expression="A - A", label="zero")],
)
+
+
+def test_model_functions_can_be_nested() -> None:
+ aliases = {"A": np.array([1.0, 2.0])}
+ calls: list[str] = []
+
+ def model_function(name, args):
+ calls.append(name)
+ return np.asarray(args[0]) * (2.0 if name == "norm" else 0.5)
+
+ result, references = evaluate_vector_expression(
+ "invnorm(norm(A, 0, 'attn', 'input'), 'final')", aliases, model_function=model_function
+ )
+ np.testing.assert_allclose(result, aliases["A"])
+ assert references == ("A",)
+ assert calls == ["norm", "invnorm"]
+
+
+def test_vector_dimension_metrics_participation_ratio() -> None:
+ assert vector_dimension_metrics(np.array([1.0, 0.0, 0.0]))["effective_dimension"] == pytest.approx(1.0)
+ assert vector_dimension_metrics(np.array([1.0, 1.0, 1.0]))["effective_dimension"] == pytest.approx(3.0)
+ skewed = vector_dimension_metrics(np.array([10.0, 1.0, 0.0]))
+ assert skewed["effective_dimension"] == pytest.approx((101.0**2) / 10001.0)
+ assert skewed["dimensions_for_angle"] >= 1
+
+
+def test_qk_model_functions_are_dispatched() -> None:
+ aliases = {"A": np.array([1.0, 2.0])}
+ calls: list[str] = []
+
+ def model_function(name, args):
+ calls.append(name)
+ return np.asarray(args[0]) + (1.0 if name == "wqk" else 2.0)
+
+ result, references = evaluate_vector_expression(
+ "wkq(wqk(A, 0, 1, 2), 0, 1, 2)", aliases, model_function=model_function
+ )
+ np.testing.assert_allclose(result, np.array([4.0, 5.0]))
+ assert references == ("A",)
+ assert calls == ["wqk", "wkq"]
From acf623f5229d312a6e250defe6538376d7ac5ebf Mon Sep 17 00:00:00 2001
From: Kauna <16511995+klei22@users.noreply.github.com>
Date: Thu, 23 Jul 2026 06:46:14 -0700
Subject: [PATCH 2/2] Add layer helper vector functions
---
.../hf-vocab-sphere/app/model_store.py | 147 ++++++++++++++++++
.../threejs/hf-vocab-sphere/app/static/app.js | 37 +++++
.../hf-vocab-sphere/app/templates/index.html | 6 +-
.../hf-vocab-sphere/app/vector_math.py | 6 +-
.../tests/test_frontend_contract.py | 4 +
.../hf-vocab-sphere/tests/test_vector_math.py | 16 ++
6 files changed, 212 insertions(+), 4 deletions(-)
diff --git a/report/threejs/hf-vocab-sphere/app/model_store.py b/report/threejs/hf-vocab-sphere/app/model_store.py
index 86410efcc0..9a3b95f936 100644
--- a/report/threejs/hf-vocab-sphere/app/model_store.py
+++ b/report/threejs/hf-vocab-sphere/app/model_store.py
@@ -1021,6 +1021,13 @@ def _all_safetensor_tensors(model_name: str, *, revision: str, allow_download: b
or ".value.weight" in lname
or ".o_proj.weight" in lname
or ".out_proj.weight" in lname
+ or ".gate_proj.weight" in lname
+ or ".up_proj.weight" in lname
+ or ".down_proj.weight" in lname
+ or ".fc1.weight" in lname
+ or ".fc2.weight" in lname
+ or ".c_fc.weight" in lname
+ or ".c_proj.weight" in lname
):
tensors[str(name)] = (
handle.get_tensor(str(name)).detach().cpu().to(dtype=torch.float32).contiguous()
@@ -1199,6 +1206,133 @@ def _qk_head_slices(
return slice(q_start, q_start + q_head_dim), slice(k_start, k_start + q_head_dim)
+def _optional_norm(tensors: dict[str, torch.Tensor], args: list[float | str | np.ndarray]) -> np.ndarray:
+ try:
+ vector, weight = _norm_weight_for_args(tensors, args)
+ except ValueError:
+ return np.asarray(args[0], dtype=np.float64)
+ return _apply_rms_norm(vector, weight)
+
+
+def _softmax(values: np.ndarray) -> np.ndarray:
+ shifted = values - float(np.max(values))
+ exp = np.exp(shifted)
+ return exp / max(float(np.sum(exp)), 1e-15)
+
+
+def _apply_attention_layer(
+ assets: ModelAssets,
+ tensors: dict[str, torch.Tensor],
+ latest_vector: np.ndarray,
+ layer: int,
+ kv_vectors: list[np.ndarray],
+) -> np.ndarray:
+ _qn, q_weight = _find_tensor(
+ tensors,
+ ("self_attn.q_proj.weight", "attention.q_proj.weight", "attn.q_proj.weight", "query.weight"),
+ layer=layer,
+ ndim=2,
+ )
+ _kn, k_weight = _find_tensor(
+ tensors,
+ ("self_attn.k_proj.weight", "attention.k_proj.weight", "attn.k_proj.weight", "key.weight"),
+ layer=layer,
+ ndim=2,
+ )
+ _vn, v_weight = _find_tensor(
+ tensors,
+ ("self_attn.v_proj.weight", "attention.v_proj.weight", "attn.v_proj.weight", "value.weight"),
+ layer=layer,
+ ndim=2,
+ )
+ _on, o_weight = _find_tensor(
+ tensors,
+ ("self_attn.o_proj.weight", "attention.o_proj.weight", "attn.o_proj.weight", "out_proj.weight"),
+ layer=layer,
+ ndim=2,
+ )
+ q = q_weight.numpy().astype(np.float64)
+ k = k_weight.numpy().astype(np.float64)
+ v = v_weight.numpy().astype(np.float64)
+ o = o_weight.numpy().astype(np.float64)
+ width = latest_vector.shape[0]
+ if q.shape[1] != width or k.shape[1] != width or v.shape[1] != width or o.shape[0] != width:
+ raise ValueError("Attention layer projection widths do not match the selected vector space.")
+ latest_norm = _apply_rms_norm(*_norm_weight_for_args(tensors, [latest_vector, layer, "attn", "input"]))
+ kv_norms = [
+ _apply_rms_norm(*_norm_weight_for_args(tensors, [vector, layer, "attn", "input"])) for vector in kv_vectors
+ ]
+ # The latest/query token also contributes its own K/V entry after the supplied cache order.
+ kv_norms.append(latest_norm)
+ attention_heads, key_value_heads = _attention_head_counts(assets)
+ output = np.zeros(width, dtype=np.float64)
+ for query_head in range(attention_heads):
+ kv_head = min((query_head * key_value_heads) // attention_heads, key_value_heads - 1)
+ q_slice, k_slice = _qk_head_slices(
+ q_output_width=q.shape[0],
+ k_output_width=k.shape[0],
+ attention_heads=attention_heads,
+ key_value_heads=key_value_heads,
+ query_head=query_head,
+ key_head=kv_head,
+ )
+ v_slice, o_slice, _head_count = _wov_head_slices(
+ v_output_width=v.shape[0],
+ o_input_width=o.shape[1],
+ attention_heads=attention_heads,
+ key_value_heads=key_value_heads,
+ head=query_head,
+ )
+ query = latest_norm @ q[q_slice, :].T
+ keys = np.stack([vector @ k[k_slice, :].T for vector in kv_norms], axis=0)
+ values = np.stack([vector @ v[v_slice, :].T for vector in kv_norms], axis=0)
+ scores = keys @ query / math.sqrt(max(query.shape[0], 1))
+ context = _softmax(scores) @ values
+ output += context @ o[:, o_slice].T
+ return _optional_norm(tensors, [output, layer, "attn", "output"])
+
+
+def _activation(name: str, values: np.ndarray) -> np.ndarray:
+ key = name.casefold()
+ if key in {"gelu", "gelu_new"}:
+ return 0.5 * values * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (values + 0.044715 * values**3)))
+ return values / (1.0 + np.exp(-values))
+
+
+def _apply_mlp_layer(
+ assets: ModelAssets, tensors: dict[str, torch.Tensor], vector: np.ndarray, layer: int
+) -> np.ndarray:
+ hidden = _apply_rms_norm(*_norm_weight_for_args(tensors, [vector, layer, "ffn", "input"]))
+ try:
+ _gn, gate_weight = _find_tensor(
+ tensors, ("mlp.gate_proj.weight", "feed_forward.gate_proj.weight"), layer=layer, ndim=2
+ )
+ _un, up_weight = _find_tensor(
+ tensors, ("mlp.up_proj.weight", "feed_forward.up_proj.weight"), layer=layer, ndim=2
+ )
+ _dn, down_weight = _find_tensor(
+ tensors, ("mlp.down_proj.weight", "feed_forward.down_proj.weight"), layer=layer, ndim=2
+ )
+ gate = gate_weight.numpy().astype(np.float64)
+ up = up_weight.numpy().astype(np.float64)
+ down = down_weight.numpy().astype(np.float64)
+ output = (_activation("silu", hidden @ gate.T) * (hidden @ up.T)) @ down.T
+ except ValueError:
+ _fn, fc_weight = _find_tensor(
+ tensors, ("mlp.fc1.weight", "mlp.c_fc.weight", "feed_forward.fc1.weight"), layer=layer, ndim=2
+ )
+ _pn, proj_weight = _find_tensor(
+ tensors, ("mlp.fc2.weight", "mlp.c_proj.weight", "feed_forward.fc2.weight"), layer=layer, ndim=2
+ )
+ fc = fc_weight.numpy().astype(np.float64)
+ proj = proj_weight.numpy().astype(np.float64)
+ if fc.shape[1] == hidden.shape[0]:
+ output = _activation("gelu", hidden @ fc.T) @ proj.T
+ else:
+ output = _activation("gelu", hidden @ fc) @ proj
+ return _optional_norm(tensors, [output, layer, "ffn", "output"])
+
+
def model_vector_function(assets: ModelAssets, name: str, args: list[float | str | np.ndarray]) -> np.ndarray:
tensors = _all_safetensor_tensors(assets.model_name, revision=assets.revision, allow_download=True)
if name in {"norm", "invnorm", "inverse_norm"}:
@@ -1240,6 +1374,19 @@ def model_vector_function(assets: ModelAssets, name: str, args: list[float | str
# selected Wv-head output with zeros before/after its head slot and then
# applying the full Wo matrix, but avoids materializing the padded vector.
return (vector @ v[v_slice, :].T) @ o[:, o_slice].T
+ if name == "attn_layer":
+ if len(args) < 2:
+ raise ValueError("attn_layer expects (latest_vector, layer, optional_cache_vector, ...).")
+ latest_vector = np.asarray(args[0], dtype=np.float64)
+ layer = int(args[1])
+ kv_vectors = [np.asarray(value, dtype=np.float64) for value in args[2:]]
+ return _apply_attention_layer(assets, tensors, latest_vector, layer, kv_vectors)
+ if name == "mlp_layer":
+ if len(args) != 2:
+ raise ValueError("mlp_layer expects (vector, layer).")
+ vector = np.asarray(args[0], dtype=np.float64)
+ layer = int(args[1])
+ return _apply_mlp_layer(assets, tensors, vector, layer)
if name in {"wqk", "wkq"}:
if len(args) != 4:
raise ValueError("wqk/wkq expect (vector, layer, query_head, key_head).")
diff --git a/report/threejs/hf-vocab-sphere/app/static/app.js b/report/threejs/hf-vocab-sphere/app/static/app.js
index 75c6e93f0a..5d8491ad59 100644
--- a/report/threejs/hf-vocab-sphere/app/static/app.js
+++ b/report/threejs/hf-vocab-sphere/app/static/app.js
@@ -1371,6 +1371,41 @@ function renderArithmeticPanel() {
}
}
+
+function insertAttentionLayerExpression() {
+ const rows = selectedRows();
+ if (!rows.length) return showToast('Select tokens before inserting an attention layer helper.', 'error');
+ const defaultLatest = aliasForIndex(rows.length - 1);
+ const latest = (window.prompt('Latest token alias for Wq plus its final Wk/Wv cache entry:', defaultLatest) || '').trim().toUpperCase();
+ if (!latest) return;
+ const layerText = (window.prompt('Attention layer index:', '0') || '').trim();
+ const layer = Number(layerText);
+ if (!Number.isInteger(layer) || layer < 0) return showToast('Layer must be a non-negative integer.', 'error');
+ const defaultKv = rows.slice(0, Math.max(0, rows.length - 1)).map((_, index) => aliasForIndex(index)).join(',');
+ const kvText = (window.prompt('Comma-separated prior KV-cache aliases in order; latest is appended automatically:', defaultKv) || '').trim();
+ const kvAliases = kvText.split(',').map((value) => value.trim().toUpperCase()).filter(Boolean);
+ const cacheSuffix = kvAliases.length ? `, ${kvAliases.join(', ')}` : '';
+ $('arithmeticExpressionInput').value = `attn_layer(${latest}, ${layer}${cacheSuffix})`;
+ $('arithmeticModeSelect').value = 'expression';
+ syncArithmeticMode();
+ $('arithmeticExpressionInput').focus();
+}
+
+function insertMlpLayerExpression() {
+ const rows = selectedRows();
+ if (!rows.length) return showToast('Select tokens before inserting an MLP layer helper.', 'error');
+ const defaultTarget = aliasForIndex(rows.length - 1);
+ const target = (window.prompt('Target token alias for the MLP layer:', defaultTarget) || '').trim().toUpperCase();
+ if (!target) return;
+ const layerText = (window.prompt('MLP layer index:', '0') || '').trim();
+ const layer = Number(layerText);
+ if (!Number.isInteger(layer) || layer < 0) return showToast('Layer must be a non-negative integer.', 'error');
+ $('arithmeticExpressionInput').value = `mlp_layer(${target}, ${layer})`;
+ $('arithmeticModeSelect').value = 'expression';
+ syncArithmeticMode();
+ $('arithmeticExpressionInput').focus();
+}
+
async function addArithmeticResult() {
const rows = selectedRows();
const slerpMode = $('arithmeticModeSelect').value === 'slerp';
@@ -2184,6 +2219,8 @@ $('selectedTokens').addEventListener('click', (event) => {
});
$('anchorIdInput').addEventListener('change', () => { renderSelected(); markProjectionStale(); });
$('arithmeticModeSelect').addEventListener('change', syncArithmeticMode);
+$('insertAttentionLayerBtn').addEventListener('click', insertAttentionLayerExpression);
+$('insertMlpLayerBtn').addEventListener('click', insertMlpLayerExpression);
$('slerpFractionInput').addEventListener('input', (event) => {
$('slerpFractionLabel').textContent = `t = ${Number(event.target.value).toFixed(2)}`;
});
diff --git a/report/threejs/hf-vocab-sphere/app/templates/index.html b/report/threejs/hf-vocab-sphere/app/templates/index.html
index 12e2df6d14..5df159aa13 100644
--- a/report/threejs/hf-vocab-sphere/app/templates/index.html
+++ b/report/threejs/hf-vocab-sphere/app/templates/index.html
@@ -248,6 +248,10 @@
Vector arithmetic
+
+
+
+
@@ -269,7 +273,7 @@
Vector arithmetic
-
Expressions support aliases, scalars, parentheses, + - * /, mean(...), slerp(A,B,t), norm(A,0,'attn','input')/norm(A,'final'), invnorm(A,...), wov(A,0,0), wqk(A,0,0,0), and wkq(A,0,0,0). Functions can be nested. Example average: (A+B+C)/3 + D. Resultants are projected from the original model space, never by adding screen coordinates.
+
Expressions support aliases, scalars, parentheses, + - * /, mean(...), slerp(A,B,t), norm(A,0,'attn','input')/norm(A,'final'), invnorm(A,...), wov(A,0,0), wqk(A,0,0,0), wkq(A,0,0,0), attn_layer(A,0,B,C), and mlp_layer(A,0). Functions can be nested. Example average: (A+B+C)/3 + D. Resultants are projected from the original model space, never by adding screen coordinates.