File size: 4,874 Bytes
196d16f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """A small real attention computation for CPU examples, not a pretrained VLM.
Image-token vectors are supplied explicitly; this fixture has no image encoder.
"""
import math
import torch
from torch import Tensor, nn
class TinyAttention(nn.Module):
def __init__(self, width: int):
super().__init__()
self.query = nn.Linear(width, width, bias=False)
self.key = nn.Linear(width, width, bias=False)
self.value = nn.Linear(width, width, bias=False)
self.scores = nn.Identity()
self.probs = nn.Softmax(dim=-1)
self.output = nn.Linear(width, width, bias=False)
def forward(self, x: Tensor) -> Tensor:
q, k, v = self.query(x), self.key(x), self.value(x)
scores = (q @ k.transpose(-2, -1))[:, None] / math.sqrt(x.shape[-1])
future = torch.ones(x.shape[1], x.shape[1], dtype=torch.bool, device=x.device).triu(1)
scores = self.scores(scores.masked_fill(future, -torch.inf))
probs = self.probs(scores)
return self.output((probs @ v[:, None]).squeeze(1))
class TinyBlock(nn.Module):
def __init__(self, width: int):
super().__init__()
self.attention = TinyAttention(width)
self.mlp = nn.Sequential(nn.Linear(width, width), nn.GELU(), nn.Linear(width, width))
def forward(self, x: Tensor) -> Tensor:
x = x + self.attention(x)
return x + self.mlp(x)
class TinyModel(nn.Module):
def __init__(self, width: int = 6, vocab_size: int = 12):
super().__init__()
self.embedding = nn.Embedding(vocab_size, width)
self.layers = nn.ModuleList([TinyBlock(width), TinyBlock(width)])
self.norm = nn.LayerNorm(width)
self.lm_head = nn.Linear(width, vocab_size, bias=False)
def forward(self, input_ids: Tensor, image_tokens: Tensor) -> dict[str, Tensor]:
hidden = torch.cat([image_tokens, self.embedding(input_ids)], dim=1)
for layer in self.layers:
hidden = layer(hidden)
return {"logits": self.lm_head(self.norm(hidden))}
def probing_adapter(self):
"""An explicit model protocol: Prober(model) uses this verified contract."""
from vlm_probing import HeadSite, HookPoint, ModelReadout, ModelSpec, TokenLayout, TorchModelAdapter
sites, residuals, attentions, scores, heads = {}, {}, {}, {}, {}
path_heads, path_qkv = {}, {}
for i in range(len(self.layers)):
residuals[i], attentions[i] = f"residual.{i}", f"attention.{i}"
scores[i], heads[i] = f"scores.{i}", f"heads.{i}"
sites[residuals[i]] = HookPoint(f"layers.{i}")
sites[attentions[i]] = HookPoint(f"layers.{i}.attention.probs")
sites[scores[i]] = HookPoint(f"layers.{i}.attention.scores")
sites[heads[i]] = HookPoint(f"layers.{i}.attention.output")
path_heads[i] = HeadSite(heads[i], 1, self.norm.normalized_shape[0])
path_qkv[i] = {}
for kind, component in (("q", "query"), ("k", "key"), ("v", "value")):
site = f"{kind}.{i}"
sites[site] = HookPoint(f"layers.{i}.attention.{component}")
path_qkv[i][kind] = HeadSite(site, 1, self.norm.normalized_shape[0])
sites["embeddings"] = HookPoint("layers.0", kind="input", selector=0)
sites["last_mlp"] = HookPoint(f"layers.{len(self.layers)-1}.mlp")
def layout(inputs):
ids = inputs["input_ids"]
prefix = torch.zeros(inputs["image_tokens"].shape[:2], dtype=ids.dtype, device=ids.device)
expanded = torch.cat([prefix, ids], dim=1)
valid = torch.ones_like(expanded, dtype=torch.bool)
visual = torch.zeros_like(valid)
visual[:, :prefix.shape[1]] = True
return TokenLayout(valid, visual, token_ids=expanded)
def linear_readout(full):
scale = (full.var(-1, keepdim=True, unbiased=False) + self.norm.eps).rsqrt()
return self.lm_head.weight * self.norm.weight, scale, True
spec = ModelSpec(
residuals, attentions=attentions, attention_scores=scores, heads=heads,
embeddings="embeddings", input_embeddings=self.embedding.weight,
layout=layout, project_heads=lambda _, value: value.unsqueeze(-2),
linear_readout=linear_readout, editable_attention=set(attentions),
path_heads=path_heads, path_qkv=path_qkv, path_final=residuals[len(self.layers)-1],
edges={"last_attention_to_residual": heads[len(self.layers)-1],
"last_mlp_to_residual": "last_mlp"},
notes=("Random CPU demonstration model; image embeddings are supplied directly.",),
)
return TorchModelAdapter(self, sites, readout=ModelReadout(self.lm_head, norm=self.norm),
model_id="tiny-demo", spec=spec)
|