The frozen stand-in, revisited
Every VLA you built in Phase 1 saw the world through FrozenVisionEncoder — a small,
random-init, never-trained conv stack (ch1.7). We were honest about it: a frozen random CNN
is a fixed nonlinear projection of the pixels, not perception. It preserved coarse layout and
nothing more, and --break blind proved the policy barely used it. The whole reason to reach
for a pretrained backbone — SigLIP, DINOv2 — was to close that gap.
Before you can adapt a pretrained vision tower, you have to know what one is. This chapter
builds the architecture underneath all of them, from scratch, over cached PushT frames: a
Vision Transformer. No convolutions, no transformers, no einops — a patchify reshape,
a CLS token, learned positions, and the exact pre-norm attention Block you fused with in
ch1.8, re-derived here. Open vit.py. Seven regions: setup, data, patches,
model, train, probe, report.
A ViT sees no picture
Here is the single idea the rest of the chapter defends. A ViT does not process an image the way a CNN does — sliding receptive fields over a grid. It cuts the image into patches, flattens each into a plain vector, and treats the result as a SEQUENCE — a bag of tokens. The only thing that tells the model two patches were neighbors is a learned positional embedding we add to each. Order is not baked in; it is a parameter.
# A ViT sees NO image — it sees a SEQUENCE of flat patch vectors. patchify cuts the# (B, H, W, 3) image into an 8x8 grid of 8x8x3 patches and flattens each to 192 numbers.# The grid is row-major: token t=(grid_row, grid_col). The reshape is the whole trick and# the whole trap: to keep each patch's pixels contiguous you must PERMUTE the grid axes in# front of the within-patch axes. Skip the permute and the patch-grid axis interleaves with# the pixel axis — "patches" become a global permutation of the real ones. The trap is that# this is SILENT to a coarse probe (a permutation of a bag is the same bag), so the quadrant# accuracy barely moves; the damage only shows in the attention map. That is the deliberate# failure the exercise has you write, predict, and diagnose (--break patch_interleave).def patchify(images: torch.Tensor, patch: int, interleave_bug: bool = False) -> torch.Tensor: """(B, H, W, C) -> (B, num_patches, patch*patch*C). Correct: grid axes first, then pixel axes, so each row of the output is ONE spatially-contiguous patch.""" B, H, W, C = images.shape gh, gw = H // patch, W // patch x = images.reshape(B, gh, patch, gw, patch, C) # split H->(gh,patch), W->(gw,patch) if not interleave_bug: x = x.permute(0, 1, 3, 2, 4, 5) # (B, gh, gw, patch, patch, C): patch contiguous # BUG PATH (no permute): flattening (B, gh, patch, gw, patch, C) mixes grid-row with # patch-row into the token axis and grid-col into the feature axis — patches interleave. return x.reshape(B, gh * gw, patch * patch * C)patchify is the whole trick, and the whole trap. To turn a (B, 64, 64, 3) image into
(B, 64, 192) — 64 patches, each an 8×8×3 = 192-vector — you split H → (grid_row, patch_row) and W → (grid_col, patch_col), then permute the two grid axes in front of the
two pixel axes before flattening, so each output row is one spatially-contiguous patch. Skip
the permute and the grid axis interleaves with the pixel axis: every "patch" becomes a
scramble. The shape is still (B, 64, 192), so nothing errors. Hold that thought — it is the
deliberate failure at the end.
Cache the frames, label them for free
# Cache frames ONCE from the scripted expert (ch1.7's replay-in-process pattern), at the# free-tier 64x64. Each cached frame carries a CHEAP label read straight off sim state:# which quadrant the T-block occupies. No human annotation, no learned label — the probe# target is a deterministic function of the state that produced the pixels.def quadrant(tee_x: float, tee_y: float) -> int: """0=NE, 1=NW, 2=SW, 3=SE. The target sits at the origin; the block spawns in an annulus around it, so the sign of (tee_x, tee_y) is an unambiguous scene fact.""" top, right = tee_y >= 0.0, tee_x >= 0.0 return {(True, True): 0, (True, False): 1, (False, False): 2, (False, True): 3}[(top, right)] def cache_frames(episodes: int, seed: int, stride: int): """Replay the scripted PushT expert and keep every Nth frame + its quadrant label + the episode it came from. Deterministic given seed (ch1.7's contract).""" env = PushTEnv() frames, labels, ep_index = [], [], [] for e in range(episodes): obs = env.reset(seed + e) expert = ScriptedExpert(noise=0.0, seed=seed + e) step, done = 0, False while not done: action = expert.action(env) if step % stride == 0: # obs is the CURRENT state; render_frame renders that same state frames.append(env.render_frame(IMG_HW, IMG_HW)) labels.append(quadrant(float(obs[2]), float(obs[3]))) # obs[2:4] = tee_x, tee_y ep_index.append(e) obs, _, done, _ = env.step(action) step += 1 return (np.asarray(frames, np.uint8), np.asarray(labels, np.int64), np.asarray(ep_index, np.int64)) frames, labels, ep_index = cache_frames(args.episodes, args.seed, args.frame_stride)# Split by EPISODE (not by frame) so near-duplicate frames from one episode never straddle# train/test — the honest way to measure a representation (ch1.6). Last ~25% of episodes -> test.test_ep = ep_index >= int(math.ceil(args.episodes * 0.75))train_idx = np.where(~test_ep)[0]test_idx = np.where(test_ep)[0]if len(test_idx) == 0: # tiny smoke budgets can leave no held-out episode; fall back to a frame split test_idx, train_idx = train_idx[-len(train_idx) // 3:], train_idx[: -len(train_idx) // 3]# Frames as float in [-1, 1], on device, kept whole so patchify runs on the batch (B, 64, 64, 3).images = (torch.from_numpy(frames).to(device).float() / 127.5 - 1.0)labels_t = torch.from_numpy(labels).to(device)majority_baseline = float(np.bincount(labels[test_idx], minlength=NUM_QUADRANTS).max() / len(test_idx))print(f"cached {len(frames)} frames ({len(train_idx)} train / {len(test_idx)} test), " f"quadrant balance {np.bincount(labels, minlength=NUM_QUADRANTS).tolist()}")We replay the scripted PushT expert (the ch1.7 pattern) and cache ~740 frames once, at
the free-tier floor of 64×64 — half the pixels of ch1.7's 96×96. Each frame carries a label
that costs nothing: which quadrant the block sits in, read straight off the simulator
state (obs[2:4] is the block's x, y). No annotation, no human — a deterministic function of
the state that produced the pixels. We split by episode, not by frame, so near-duplicate
frames from one episode never straddle train and test (ch1.6).
The ViT itself
# The ViT, from scratch. A transformer Block is the SAME shape ch1.8 fused vision+language# with: pre-norm multi-head self-attention (Q/K/V are nn.Linear, a scaled-dot-product# softmax mixes tokens, an output projection), then a per-token MLP, both on residuals. No# transformers, no einops. Here there is no key padding — the sequence is a fixed 1+64 tokens.class Block(nn.Module): def __init__(self, dim: int, heads: int) -> None: super().__init__() self.heads = heads self.ln1, self.ln2 = nn.LayerNorm(dim), nn.LayerNorm(dim) self.qkv = nn.Linear(dim, 3 * dim) self.proj = nn.Linear(dim, dim) self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim)) self.last_attn = None # (B, L, L) head-averaged attention, kept for the rollout viz def forward(self, x: torch.Tensor) -> torch.Tensor: B, L, dim = x.shape h, hd = self.heads, dim // self.heads qkv = self.qkv(self.ln1(x)).reshape(B, L, 3, h, hd).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] # each (B, h, L, hd) scores = (q @ k.transpose(-2, -1)) / math.sqrt(hd) # (B, h, L, L) attn = scores.softmax(dim=-1) self.last_attn = attn.mean(1).detach() # (B, L, L): average over heads x = x + self.proj((attn @ v).transpose(1, 2).reshape(B, L, dim)) return x + self.mlp(self.ln2(x)) class TinyViT(nn.Module): """Patch-embed -> [CLS, patch_0..patch_63] + learned positions -> pre-norm blocks. The CLS row after the blocks is the pooled scene representation; `head` reads the quadrant.""" def __init__(self, dim: int, depth: int, heads: int, num_classes: int) -> None: super().__init__() self.patch_embed = nn.Linear(PATCH_DIM, dim) # Conv2d(3,dim,8,stride=8) is the SAME op self.cls = nn.Parameter(torch.zeros(1, 1, dim)) self.pos = nn.Parameter(0.02 * torch.randn(1, NUM_PATCHES + 1, dim)) # LEARNED position tags self.blocks = nn.ModuleList([Block(dim, heads) for _ in range(depth)]) self.norm = nn.LayerNorm(dim) self.head = nn.Linear(dim, num_classes) def features(self, images: torch.Tensor, interleave_bug: bool = False, shuffle_pos: torch.Tensor | None = None) -> torch.Tensor: B = images.shape[0] tokens = self.patch_embed(patchify(images, PATCH, interleave_bug)) # (B, 64, dim) seq = torch.cat([self.cls.expand(B, -1, -1), tokens], dim=1) # prepend CLS -> (B, 65, dim) pos = self.pos if shuffle_pos is not None: # --break shuffle_pos: permute the 64 PATCH position tags (keep CLS at 0) pos = torch.cat([self.pos[:, :1], self.pos[:, 1:][:, shuffle_pos]], dim=1) seq = seq + pos for blk in self.blocks: seq = blk(seq) return self.norm(seq[:, 0]) # the CLS row: pooled over all 64 patches def forward(self, images: torch.Tensor, interleave_bug: bool = False) -> torch.Tensor: return self.head(self.features(images, interleave_bug))The Block is ch1.8's, re-derived: pre-norm multi-head self-attention (Q/K/V are three
nn.Linear, a scaled-dot-product softmax mixes tokens, an output projection), then a per-token
MLP, both on residuals. TinyViT embeds the patches with a single nn.Linear — a Conv2d
with stride=8 is the exact same operation, which is how the real towers write it — prepends
a learned CLS token, adds the learned position table, runs the blocks, and reads the CLS
row as the pooled scene representation. Free-tier dims are deliberately tiny: dim 96, depth 2, heads 3. This is a nanoGPT-tiny, in pixels.
Training a transformer is not free
# Train the ViT end-to-end to classify the block's quadrant. Nothing about the label needs# a ViT — a linear probe on raw pixels could do it — the POINT is what the ATTENTION learns:# to route the CLS token to the patches that actually contain the block. Cross-entropy, Adam,# and a LINEAR LR-WARMUP: a from-scratch transformer cold-starts badly, and without warmup# some seeds pin at chance for the whole run (measured — try --warmup 1). Warmup + a modest lr# are why the model trains RELIABLY across seeds. --break patch_interleave trains through the# buggy (globally-permuted) patchify — it still fits the coarse label (the trap), but its# attention never lines up with the image.INTERLEAVE = args.break_mode == "patch_interleave"torch.manual_seed(args.seed) # trained ViT initmodel = TinyViT(args.dim, args.depth, args.heads, NUM_QUADRANTS).to(device)optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)train_images, train_labels = images[train_idx], labels_t[train_idx]loss_fn = nn.CrossEntropyLoss()train_loss, step = float("nan"), 0for epoch in range(args.epochs): for group in optimizer.param_groups: # linear warmup, then hold: lr ramps 0 -> args.lr over --warmup epochs group["lr"] = args.lr * min(1.0, (epoch + 1) / args.warmup) model.train() epoch_loss, nb = 0.0, 0 for batch in torch.randperm(len(train_idx), generator=gen).split(args.batch_size): logits = model(train_images[batch], INTERLEAVE) loss = loss_fn(logits, train_labels[batch]) optimizer.zero_grad() loss.backward() optimizer.step() epoch_loss, nb = epoch_loss + loss.item(), nb + 1 if args.rerun: rr.set_time("step", sequence=step) rr.log("vit/loss/train", rr.Scalars([loss.item()])) step += 1 train_loss = epoch_loss / nb if epoch % 5 == 0 or epoch == args.epochs - 1: print(f"epoch {epoch:3d} ce {train_loss:.4f}") # A random-init ViT of the SAME shape, never trained — the reference the probe must beat.torch.manual_seed(args.seed + 1)random_vit = TinyViT(args.dim, args.depth, args.heads, NUM_QUADRANTS).to(device)# --break shuffle_pos: a FIXED permutation of the 64 patch position tags (pixels untouched).shuffle = torch.randperm(NUM_PATCHES, generator=gen).to(device) if args.break_mode == "shuffle_pos" else NoneA from-scratch transformer does not just train — it cold-starts badly. Two things keep it
honest here: a modest learning rate and a linear LR warmup. Drop the warmup
(--warmup 1) and some seeds pin at chance for the entire run — the loss never leaves
ln 4. This is not a detail to hide; it is a real, measured property of small transformers,
and it is why every production recipe warms up. Feel it before you trust the numbers.
The measurement: does the representation carry the scene?
# THE MEASUREMENT. Freeze each backbone, extract its CLS feature on train+test, and fit a# CLOSED-FORM linear probe (least-squares onto one-hot quadrants — deterministic, no SGD,# ch1.7's lstsq trick). Test accuracy of that probe is representation QUALITY: how linearly# the scene fact falls out of the pooled feature. Trained ViT > random-init ViT > majority.@torch.no_grad()def cls_features(vit: TinyViT, imgs: torch.Tensor) -> np.ndarray: vit.eval() return vit.features(imgs, INTERLEAVE, shuffle).cpu().double().numpy() def linear_probe(feat_tr: np.ndarray, y_tr: np.ndarray, feat_te: np.ndarray, y_te: np.ndarray) -> float: onehot = np.zeros((len(y_tr), NUM_QUADRANTS)) onehot[np.arange(len(y_tr)), y_tr] = 1.0 X = np.concatenate([feat_tr, np.ones((len(feat_tr), 1))], axis=1) # + intercept column W, *_ = np.linalg.lstsq(X, onehot, rcond=None) pred = (np.concatenate([feat_te, np.ones((len(feat_te), 1))], axis=1) @ W).argmax(1) return float((pred == y_te).mean()) y_tr, y_te = labels[train_idx], labels[test_idx]probe_trained = linear_probe(cls_features(model, images[train_idx]), y_tr, cls_features(model, images[test_idx]), y_te)probe_random = linear_probe(cls_features(random_vit, images[train_idx]), y_tr, cls_features(random_vit, images[test_idx]), y_te) # Attention ROLLOUT: how the CLS token's attention flows to each patch through the whole# stack. Add the residual (0.5*A + 0.5*I), renormalize, multiply the per-layer matrices, and# read the CLS row over the 64 patches -> an 8x8 map. Trained: concentrates on the block.# Random-init: washes out near-uniform. (Abnar & Zuidema's rollout, the standard ViT viz.)@torch.no_grad()def cls_rollout(vit: TinyViT, imgs: torch.Tensor) -> np.ndarray: vit.eval() vit.features(imgs, INTERLEAVE, shuffle) # populate each block's last_attn rolled = None for blk in vit.blocks: a = blk.last_attn a = 0.5 * a + 0.5 * torch.eye(a.shape[-1], device=a.device)[None] a = a / a.sum(-1, keepdim=True) rolled = a if rolled is None else a @ rolled grid = rolled[:, 0, 1:].reshape(-1, GRID, GRID) # CLS -> patches, as an 8x8 map return (grid / grid.amax(dim=(1, 2), keepdim=True)).cpu().double().numpy() # per-frame normalizedWe do not grade the ViT by its training accuracy — of course it fits its own labels. We grade
its representation with a linear probe: freeze the backbone, extract the CLS feature,
and fit a closed-form least-squares read-out of the quadrant (ch1.7's lstsq trick, no SGD).
Probe accuracy is how linearly accessible the scene fact is in the pooled vector. We run it
on the trained ViT and on a same-shape random-init ViT, and compare to the majority guess.
The measured direction, seed 0 (and it holds on every seed):
| representation | probe accuracy |
|---|---|
| trained ViT (CLS) | ~0.89 |
| random-init ViT (CLS) | ~0.69 |
| majority guess | ~0.33 |
| chance (1/4) | 0.25 |
Read this carefully, because it is more honest than "the ViT learned to see." A random-init ViT — never trained — already probes to ~0.69, far above the majority guess. That is the point: which quadrant is nearly a bag-of-patches property, and a random projection of patches preserves it. What training adds is real but layered on top: it lifts the probe another ~0.20, and it teaches the attention to concentrate — the CLS token's attention-rollout peaks ~10–16× over uniform on the patches that actually contain the block, where the random-init ViT's attention is flat. That concentration is the toy: hover a frame and watch the trained map find the block while the random map washes out. Report the order (trained > random > majority), never the exact %: these are rendered images (not bitwise across CPU arches) read through a small held-out set (ch1.6).
Break it: the bug your accuracy can't see
Now collect on the patchify trap. You wrote the reshape in the exercise; the classic mistake is
to forget the permute, globally permuting the patch set. Predict what it does to the probe,
then run --break patch_interleave.
If you predicted the probe crashes to chance, you are in good company — and wrong. The probe is silent: ~0.89 clean, ~0.92 interleaved. A coarse label is a permutation-invariant property of the bag, and a shape-preserving reshape bug is exactly a permutation — so the ViT relearns the fixed reordering and the accuracy shrugs. That silence is the danger. A patchify bug will not show up in your metric; it shows up only in the attention map, which no longer lines up with the image. The lesson of the ViT is also its footgun: patches are a bag.
The companion misconception — "a ViT sees the image as a picture" — falls to
--break shuffle_pos, which permutes the learned position embeddings with the pixels
byte-for-byte identical. The coarse probe survives (a bag barely needs order for a quadrant),
but the trained model's edge over random nearly vanishes (the gap collapses from ~0.20 toward
zero) and its attention map scrambles. The pixels never changed; the position tags did. The
model was never looking at a picture — it was looking at tokens wearing position tags, and the
spatial structure it learned lived in those tags.
Read the real thing
Everything here has a production form. The read-the-real-thing segment pairs this chapter with
a real SigLIP ViT vision tower (google-research/big_vision, models/vit.py, or the SmolVLA
vision_model in the already-pinned huggingface/lerobot). You will recognize every piece:
the patch embedding is a Conv2d(stride=patch) — our flattened-patch Linear — the class/pool
token is our CLS, the learned positional embeddings are our pos table, and the pre-norm
attention stack is our Block. What the real tower adds is scale and, above all, pretraining
on image–text pairs, which is what makes its features object- and language-aligned and
transferable. That is the Scale Lab, and the reason to prefer adapt-pretrained when performance
matters.
What's next
You built the architecture a vision backbone is made of and measured, honestly, what a tiny from-scratch one carries: a coarse scene fact, more accessible after training, with attention that localizes the object — and a random-init baseline closer than you'd like, because the task is easy. The gap between this and a policy-ready visual representation is exactly the gap pretraining fills. The next chapters put a real, pretrained tower to work.