M4.3 — Vision Encoders and Feature Extraction

Learning objectives

After this lesson, you should be able to:

• Explain what a vision encoder produces.
• Distinguish a global feature vector from a grid of local features.
• Explain why global average pooling loses position.
• Describe patch embedding and the structure of a Vision Transformer.
• Estimate the number of tokens and the relative attention cost for an image size.
• Explain frozen versus fine-tuned encoders.

The one-sentence idea

A vision encoder maps an image to feature vectors, and for robot control the useful form keeps one vector per region so that position is not lost.

Pool too early and you forget where things are.

What a vision encoder does

A vision encoder is a network that turns an image into features: numbers that describe what is in it in a form other models can use. Where M4.1 showed an image as raw pixels and M4.2 showed how convolution finds structure, an encoder is the full pipeline that does this and ends with a compact representation. It is the “vision” half of a vision-language-action model.

An encoder can return its result in two shapes, and the difference matters a lot for robots:

One global vector

A single embedding for the whole image, for example 512 numbers. Good for “what is this image of?”, such as classification or matching an image to a caption.

A grid of local features

One vector per region, so a (h × w) map of vectors, or a sequence of tokens. Needed for “where is it?” and for handing a detailed scene to a transformer.

Why spatial information matters

The usual way to go from a grid of features to a single vector is global average pooling: average every channel over all positions. It is simple, and it throws away where things were. The example below has a feature map with two channels, where channel 0 detects a “cup” and channel 1 detects a “bowl”. In image A the cup is on the left and the bowl is on the right. In image B they are swapped:

import numpy as np

# feature maps with shape (channels, 1 row, 4 columns)
A = np.array([[[1, 0, 0, 0]],       # channel 0 "cup" fires on the left
              [[0, 0, 0, 1]]])      # channel 1 "bowl" fires on the right
B = np.array([[[0, 0, 0, 1]],       # cup on the right
              [[1, 0, 0, 0]]])      # bowl on the left

print(A.mean(axis=(1, 2)), B.mean(axis=(1, 2)))
print(A.reshape(-1).tolist() == B.reshape(-1).tolist())

[0.25 0.25] [0.25 0.25]
False

After pooling the two scenes produce the identical vector, yet they are different arrangements, and a robot asked to “pick the cup on the left” needs exactly that difference. The unpooled maps still differ (the last line). For this reason, robot-learning models usually keep the grid of local features, as a sequence of tokens, and let later layers use position, rather than collapsing the image to one vector.

The Vision Transformer: patches as tokens

The Vision Transformer (ViT) applies the transformer of M3.4 to images by treating image patches as tokens, as M3.2 previewed. It cuts the image into a grid of non-overlapping squares, flattens each square into a vector, and projects it linearly to the model dimension:

def patchify(img, p):
    """img: (C, H, W) -> (number of patches, C * p * p)"""
    C, H, W = img.shape
    x = img.reshape(C, H // p, p, W // p, p)
    return x.transpose(1, 3, 0, 2, 4).reshape((H // p) * (W // p), C * p * p)

img = np.arange(16).reshape(1, 4, 4)       # one channel, 4 x 4
print(img[0])
print(patchify(img, 2))

[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]
 [12 13 14 15]]
[[ 0  1  4  5]
 [ 2  3  6  7]
 [ 8  9 12 13]
 [10 11 14 15]]

Each row of the output is one 2 × 2 patch, flattened, read left to right and top to bottom: the first row [0 1 4 5] is the top-left square. The linear projection that follows is a matrix multiplication (M1.1), and it can be done with a single convolution whose kernel size and stride both equal the patch size. The two are the same computation:

import torch
from torch import nn

torch.manual_seed(0)
p, d = 4, 6
conv = nn.Conv2d(3, d, kernel_size=p, stride=p)            # patch embedding as a convolution
x = torch.randn(1, 3, 8, 8)
via_conv = conv(x).flatten(2).transpose(1, 2)              # (1, tokens, d)

patches = x.unfold(2, p, p).unfold(3, p, p)                # (1, 3, 2, 2, p, p)
patches = patches.permute(0, 2, 3, 1, 4, 5).reshape(1, 4, 3 * p * p)
linear = nn.Linear(3 * p * p, d)
with torch.no_grad():
    linear.weight.copy_(conv.weight.reshape(d, -1))
    linear.bias.copy_(conv.bias)
print(via_conv.shape, torch.allclose(via_conv, linear(patches), atol=1e-5))

torch.Size([1, 4, 6]) True

The patch tokens then get position information added (M3.4), and often one extra summary token, and go through a stack of transformer blocks. The output is one feature vector per patch, which is exactly the “grid of local features” a robot wants, and the extra token can serve as a global vector. Putting these pieces together:

class TinyViT(nn.Module):
    def __init__(self, patch=16, dim=32, tokens=196):
        super().__init__()
        self.embed = nn.Conv2d(3, dim, kernel_size=patch, stride=patch)
        self.cls = nn.Parameter(torch.zeros(1, 1, dim))
        self.pos = nn.Parameter(torch.zeros(1, tokens + 1, dim))
        layer = nn.TransformerEncoderLayer(dim, 4, 64, dropout=0.0, batch_first=True, norm_first=True)
        self.blocks = nn.TransformerEncoder(layer, 2, enable_nested_tensor=False)

    def forward(self, x):
        t = self.embed(x).flatten(2).transpose(1, 2)           # (B, 196, dim) patch tokens
        t = torch.cat([self.cls.expand(len(t), -1, -1), t], dim=1) + self.pos
        return self.blocks(t)

out = TinyViT()(torch.randn(2, 3, 224, 224))
print(out.shape)

torch.Size([2, 197, 32])

A batch of two 224 × 224 images becomes 197 tokens each: 14 × 14 = 196 patches plus one summary token, with 32 features per token. A real encoder is the same structure with a much larger dimension and many more blocks.

Resolution and cost

The number of tokens depends on the image size and the patch size, and the attention table grows with the square of the token count (M3.3). Doubling the resolution quadruples the tokens, and attention becomes sixteen times larger:

for size in (224, 448):
    n = (size // 16) ** 2
    print(size, n, n * n)
print((784 / 196) ** 2)

224 196 38416
448 784 614656
16.0

This is one reason robot systems often use modest image sizes, and why a robot with several cameras has to budget tokens: two cameras at 224 × 224 already feed 392 tokens, before any text, into the model that follows.

Using and training encoders

Encoders are usually pretrained (M3.5) on large image collections, with objectives such as classification, contrastive image-text matching (the subject of M4.5) or self-supervised prediction, and then reused. In a VLA, the encoder can be kept frozen, so its features stay fixed, or fine-tuned along with the rest of the model. Some systems combine more than one encoder, so that one supplies features good for semantics and another for fine spatial detail.

Features are useful because similar images produce nearby vectors, so a measure such as cosine similarity (M3.2) can compare scenes. The vectors below stand for the features of four images, made by hand to show the idea. Retrieval by similarity finds the closest one to a query:

names = ["red cup", "blue cup", "red bowl", "wrench"]
feats = np.array([[0.9, 0.1, 0.4],
                  [0.8, 0.0, 0.5],
                  [0.1, 0.9, 0.3],
                  [0.0, 0.1, 0.9]])
query = np.array([0.85, 0.05, 0.45])           # features of a new image

sims = feats @ query / (np.linalg.norm(feats, axis=1) * np.linalg.norm(query))
order = np.argsort(-sims)
print([names[i] for i in order[:2]], sims[order[0]].round(3))

['red cup', 'blue cup'] 0.996

The query is closest to the two cups. Real feature vectors have hundreds or thousands of dimensions, and what the dimensions mean is learned, not designed. What carries over is the principle: a good encoder maps images with similar content to nearby points.

Common misconceptions

Engineering takeaways

1. An encoder turns an image into features, either one global vector or a grid of local features.
2. Global average pooling discards position, so robot models usually keep local features.
3. A ViT cuts the image into patches, projects each linearly, adds positions and applies transformer blocks.
4. Patch embedding is a convolution with kernel size equal to stride equal to patch size.
5. A 224 × 224 image with 16 × 16 patches gives 196 patch tokens.
6. Doubling resolution quadruples the tokens and makes attention sixteen times larger.
7. Encoders are normally pretrained, and then frozen or fine-tuned.
8. Similar images give nearby feature vectors, which allows comparison by cosine similarity.

The encoder decides what the rest of the model gets to know about the scene.

Knowledge check

Try each question first, then reveal the explanation.

What’s next

You can now describe how a model sees an image. Next is the language side: what large language models are and how they generate text.

Back to the VLA roadmap prerequisites