Lesson 52: Video Transformers and Pretraining¶
Every lesson so far in this part has treated an image as a bag of patch tokens (Lesson 47's ViT) and pretrained on it with contrastive learning (Lesson 50), self-distillation (Lesson 51), or masked reconstruction (Lesson 49). A video is the same idea with one more axis: a sequence of frames, each itself a grid of patches, so a token now needs a spatial position and a time index. This lesson extends the self-supervised pretraining story from still images to video — spatio-temporal attention (treating space and time as one combined token sequence, or factoring attention into separate spatial and temporal passes for efficiency), and masked video pretraining in the style of VideoMAE (Tong et al., 2022, extending Lesson 49's masking idea to space-time tubes, exploiting the heavy redundancy between nearby frames with very high masking ratios). It closes out this part's self-supervised-pretraining cluster before the course turns to vision-language models.
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
A tiny synthetic "video" dataset¶
Same plus/circle shapes as Lesson 49, but now each example is a short clip: T=4 frames of a 16x16 grid, with the shape translating a fixed number of pixels per frame. The label this lesson cares about isn't the shape's identity (a single-frame property) but its direction of motion — a genuinely temporal property that doesn't exist in any individual frame, only in how frames relate to each other.
IMG = 16
PATCH = 4
T_FRAMES = 4
T_TUBE = 2 # temporal tubelet depth
N_SPATIAL = (IMG // PATCH) ** 2 # 16 spatial patches per frame
N_TOKENS = (T_FRAMES // T_TUBE) * N_SPATIAL # 2 * 16 = 32 space-time tokens
TUBE_DIM = T_TUBE * PATCH * PATCH # 2*16 = 32 pixels per tubelet
def make_clip(shape_type, cx0, cy0, vx, n_frames=T_FRAMES, size=IMG):
frames = []
for t in range(n_frames):
cx, cy = cx0 + vx * t, cy0
img = np.zeros((size, size), dtype=np.float32)
if shape_type == 'plus':
img[cy-1:cy+2, cx-3:cx+4] = 1.0
img[cy-3:cy+4, cx-1:cx+2] = 1.0
else:
yy, xx = np.mgrid[0:size, 0:size]
img[((xx-cx)**2 + (yy-cy)**2) <= 9] = 1.0
frames.append(img)
return np.stack(frames) # (T, H, W)
def make_video_dataset(rng_local, n):
clips, dir_labels = [], []
for _ in range(n):
shape_type = rng_local.choice(['plus', 'circle'])
direction = rng_local.choice([0, 1]) # 0 = moving left, 1 = moving right
vx = -2 if direction == 0 else 2
cx0 = rng_local.integers(9, 13) if direction == 0 else rng_local.integers(3, 7)
cy0 = rng_local.integers(4, 12)
clips.append(make_clip(shape_type, cx0, cy0, vx))
dir_labels.append(direction)
return np.array(clips, dtype=np.float32), np.array(dir_labels, dtype=np.int64)
rng = np.random.default_rng(7)
X_train, dir_train = make_video_dataset(rng, 400)
X_test, dir_test = make_video_dataset(rng, 150)
Xt = torch.tensor(X_train).unsqueeze(1) # (B, 1, T, H, W)
Xte = torch.tensor(X_test).unsqueeze(1)
fig, axes = plt.subplots(2, T_FRAMES, figsize=(7, 3.2))
for row, (direction_name, idx) in enumerate([('moving left', 0), ('moving right', 1)]):
j = int(np.argmax(dir_test == idx))
for t in range(T_FRAMES):
axes[row, t].imshow(X_test[j, t], cmap='gray', vmin=0, vmax=1); axes[row, t].axis('off')
axes[row, 0].set_title(direction_name, fontsize=8, loc='left')
plt.show()
Tokenizing video: space-time patches ("tubelets")¶
Lesson 47's ViT cuts a single image into non-overlapping 4x4 patches. A video adds a time axis to cut along too: group every T_TUBE=2 consecutive frames together, and within that group apply the same 4x4 spatial patchify as before — each resulting tubelet is a small space-time block (2x4x4 = 32 pixels) that becomes a single token, instead of a 4x4 = 16-pixel single-frame patch. With T_FRAMES=4 split into 2 temporal groups of 16 spatial patches each, a clip becomes N_TOKENS=32 tokens total — the same total pixel count as 4 independent ViT-style frames, just regrouped so nearby-in-time pixels share a token.
def tubelet_patchify(clip, patch_size=PATCH, tube=T_TUBE):
# clip: (B, C, T, H, W) -> (B, N_TOKENS, TUBE_DIM). Token order is (time_group, spatial_row, spatial_col):
# token s and token N_SPATIAL+s are the SAME spatial location at two different time groups.
B, C, T, H, W = clip.shape
P = patch_size
x = clip.unfold(2, tube, tube) # (B, C, T//tube, H, W, tube)
x = x.unfold(3, P, P).unfold(4, P, P) # (B, C, T//tube, H//P, W//P, tube, P, P)
x = x.permute(0, 2, 3, 4, 1, 5, 6, 7).contiguous()
return x.view(B, -1, C * tube * P * P)
def tubelet_unpatchify(tokens, patch_size=PATCH, tube=T_TUBE, img_size=IMG, n_frames=T_FRAMES):
B, N, D = tokens.shape
P = patch_size
n_side = img_size // P
n_time = n_frames // tube
x = tokens.view(B, n_time, n_side, n_side, tube, P, P)
x = x.permute(0, 1, 4, 2, 5, 3, 6).contiguous()
return x.view(B, n_frames, img_size, img_size)
patches = tubelet_patchify(Xt)
recon_check = tubelet_unpatchify(patches)
print(f'{N_TOKENS} tokens per clip, {TUBE_DIM} pixels each')
print(f'patchify/unpatchify roundtrip max abs diff: {(recon_check - Xt.squeeze(1)).abs().max().item():.2e}')
A space-time Transformer encoder, and a masked-reconstruction decoder¶
Exactly Lesson 49's MAE encoder/decoder, with one change: the positional encoding (Lesson 46) now has to disambiguate N_TOKENS=32 space-time positions instead of 16 purely-spatial ones, so attention can tell "same place, next time group" apart from "same time, next place." This lesson's toy encoder applies attention jointly across all 32 tokens at once — every token can directly attend to every other token, space or time, in a single pass. Real video Transformers often factorize this into two cheaper passes instead: a spatial attention pass within each frame (or time group), followed by a temporal attention pass across time groups at the same spatial location. Joint attention costs $O(N^2)$ over all $N$ space-time tokens; factorized attention costs roughly $O(N \cdot (T + S))$ — far cheaper once a real clip has dozens of frames and hundreds of spatial patches, at the cost of needing two passes to mix information that joint attention mixes in one.
def positional_encoding(T, D):
pos = torch.arange(T).unsqueeze(1).float()
i = torch.arange(D).unsqueeze(0).float()
angle_rates = 1.0 / (10000 ** (2 * (i // 2) / D))
angles = pos * angle_rates
pe = torch.zeros(T, D)
pe[:, 0::2] = torch.sin(angles[:, 0::2])
pe[:, 1::2] = torch.cos(angles[:, 1::2])
return pe
def random_masking(x, mask_ratio, generator=None):
B, N, D = x.shape
n_visible = max(1, int(N * (1 - mask_ratio)))
noise = torch.rand(B, N, generator=generator)
ids_shuffle = torch.argsort(noise, dim=1)
ids_keep = ids_shuffle[:, :n_visible]
ids_mask = ids_shuffle[:, n_visible:]
x_visible = torch.gather(x, 1, ids_keep.unsqueeze(-1).expand(-1, -1, D))
return x_visible, ids_keep, ids_mask
class VideoMAEEncoder(nn.Module):
def __init__(self, tube_dim=TUBE_DIM, embed_dim=32, n_heads=4, n_layers=2):
super().__init__()
self.embed = nn.Linear(tube_dim, embed_dim)
self.register_buffer('pe', positional_encoding(N_TOKENS, embed_dim))
layer = nn.TransformerEncoderLayer(embed_dim, n_heads, dim_feedforward=embed_dim * 2,
batch_first=True, dropout=0.0)
self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)
def forward(self, x_visible, ids_keep):
tok = self.embed(x_visible) + self.pe[ids_keep] # positional encoding at each token's TRUE space-time index
return self.encoder(tok)
class VideoMAEDecoder(nn.Module):
def __init__(self, embed_dim=32, decoder_dim=16, tube_dim=TUBE_DIM, n_heads=4, n_layers=1):
super().__init__()
self.embed = nn.Linear(embed_dim, decoder_dim)
self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_dim))
self.register_buffer('pe', positional_encoding(N_TOKENS, decoder_dim))
layer = nn.TransformerEncoderLayer(decoder_dim, n_heads, dim_feedforward=decoder_dim * 2,
batch_first=True, dropout=0.0)
self.decoder = nn.TransformerEncoder(layer, num_layers=n_layers)
self.pred = nn.Linear(decoder_dim, tube_dim)
self.decoder_dim = decoder_dim
def forward(self, encoded_visible, ids_keep, ids_mask):
B, n_mask = ids_mask.shape
vis_tok = self.embed(encoded_visible) + self.pe[ids_keep]
mask_tok = self.mask_token.expand(B, n_mask, -1) + self.pe[ids_mask]
full = torch.zeros(B, N_TOKENS, self.decoder_dim)
full = full.scatter(1, ids_keep.unsqueeze(-1).expand(-1, -1, self.decoder_dim), vis_tok)
full = full.scatter(1, ids_mask.unsqueeze(-1).expand(-1, -1, self.decoder_dim), mask_tok)
return self.pred(self.decoder(full))
class VideoMAE(nn.Module):
def __init__(self, mask_ratio=0.5, embed_dim=32, decoder_dim=16):
super().__init__()
self.mask_ratio = mask_ratio
self.encoder = VideoMAEEncoder(embed_dim=embed_dim)
self.decoder = VideoMAEDecoder(embed_dim=embed_dim, decoder_dim=decoder_dim)
def forward(self, clip, generator=None):
tokens = tubelet_patchify(clip)
x_visible, ids_keep, ids_mask = random_masking(tokens, self.mask_ratio, generator=generator)
encoded = self.encoder(x_visible, ids_keep)
pred = self.decoder(encoded, ids_keep, ids_mask)
D = tokens.shape[-1]
pred_masked = torch.gather(pred, 1, ids_mask.unsqueeze(-1).expand(-1, -1, D))
target_masked = torch.gather(tokens, 1, ids_mask.unsqueeze(-1).expand(-1, -1, D))
loss = F.mse_loss(pred_masked, target_masked)
return loss, pred, ids_mask, ids_keep, tokens
def train_video_mae(mask_ratio, epochs=600, lr=0.001, seed=0):
torch.manual_seed(seed)
model = VideoMAE(mask_ratio=mask_ratio)
opt = torch.optim.Adam(model.parameters(), lr=lr)
for _ in range(epochs):
opt.zero_grad()
loss, *_ = model(Xt)
loss.backward()
opt.step()
return model
Video tolerates extreme masking¶
Lesson 49 found images reconstruct well even with 75% of patches hidden, and traced that to spatial redundancy (a patch of sky predicts its neighbor). Video adds a second, much stronger source of redundancy: nearby frames look almost identical. Sweep the masking ratio up to 90% — far beyond what Lesson 49's single-image MAE could tolerate — and compare against a mean-pixel baseline.
train_tokens = tubelet_patchify(Xt)
mean_pixel = train_tokens.mean()
print(f'{"mask ratio":>10} {"VideoMAE test MSE":>18} {"mean-pixel baseline":>20}')
video_mae_models = {}
for mr in [0.5, 0.75, 0.9]:
model = train_video_mae(mask_ratio=mr)
video_mae_models[mr] = model
with torch.no_grad():
test_loss, pred, ids_mask, ids_keep, tokens = model(Xte, generator=torch.Generator().manual_seed(123))
D = tokens.shape[-1]
target_masked = torch.gather(tokens, 1, ids_mask.unsqueeze(-1).expand(-1, -1, D))
baseline_mse = F.mse_loss(torch.full_like(target_masked, mean_pixel), target_masked).item()
print(f'{mr:>10} {test_loss.item():>18.4f} {baseline_mse:>20.4f}')
Reconstruction error barely rises from 50% to 90% masking, and stays comfortably below the naive baseline even with just 3-4 of the 32 tokens ever shown to the encoder. At 90% masking, with only 3-4 tokens visible out of 32, the vast majority of at least one entire time group typically goes unseen — and the model can still often recover it, because the other time group's visible tokens pin down where the shape is and how fast it's moving, and a straight line extrapolates the rest. This is exactly why real VideoMAE pretrains with masking ratios around 90-95%, far above image MAE's ~75%: redundancy across frames stacks on top of redundancy across space, so hiding more of a clip destroys proportionally less genuinely new information than hiding the same fraction of a single image would.
model90 = video_mae_models[0.9]
with torch.no_grad():
_, pred, ids_mask, ids_keep, tokens = model90(Xte, generator=torch.Generator().manual_seed(42))
k = 3
masked_view = tokens.clone()
masked_view[torch.arange(len(tokens)).unsqueeze(1), ids_mask] = 0.5 # gray out hidden tubelets
recon_view = tokens.clone()
recon_view[torch.arange(len(tokens)).unsqueeze(1), ids_mask] = pred[torch.arange(len(tokens)).unsqueeze(1), ids_mask]
orig_clip = tubelet_unpatchify(tokens)[k].numpy()
masked_clip = tubelet_unpatchify(masked_view)[k].numpy()
recon_clip = tubelet_unpatchify(recon_view)[k].numpy()
fig, axes = plt.subplots(3, T_FRAMES, figsize=(7, 4.5))
for row, (clip, label) in enumerate(zip([orig_clip, masked_clip, recon_clip],
['original', 'seen by encoder\n(90% hidden)', 'reconstructed'])):
for t in range(T_FRAMES):
axes[row, t].imshow(clip[t], cmap='gray', vmin=0, vmax=1); axes[row, t].axis('off')
axes[row, 0].set_ylabel(label, fontsize=7)
plt.show()
A linear probe for a temporal property: direction of motion¶
Lessons 49-51 all pooled their encoder's output by averaging over every token, producing one feature vector per image. Doing the same thing here — averaging over all 32 space-time tokens at once — would throw away exactly the information a motion-direction probe needs: it would blend "where the shape was early on" and "where it ended up" into one indistinguishable blur. Instead, pool over space only, separately within each of the two time groups, giving one vector per time group; then take the simple difference between the last time group's pooled vector and the first's. That difference is a direct read-out of "how did the pooled representation change from early to late in the clip" — precisely a motion signal, and precisely the kind of information joint spatio-temporal attention (not naive whole-clip pooling) makes available.
N_TIME = T_FRAMES // T_TUBE
def encode_motion_signal(encoder, clip):
tokens = tubelet_patchify(clip)
B, N, _ = tokens.shape
ids_all = torch.arange(N).unsqueeze(0).expand(B, -1)
with torch.no_grad():
encoded = encoder(tokens, ids_all) # (B, N_TOKENS, embed_dim)
per_time = encoded.view(B, N_TIME, N_SPATIAL, -1).mean(dim=2) # (B, N_TIME, embed_dim): pool over SPACE only
return per_time[:, -1] - per_time[:, 0] # last time group minus first: a motion signal
def make_noisy_video_dataset(rng_local, n, noise=0.15):
clips, dir_labels = [], []
for _ in range(n):
shape_type = rng_local.choice(['plus', 'circle'])
direction = rng_local.choice([0, 1])
vx = -2 if direction == 0 else 2
cx0 = rng_local.integers(9, 13) if direction == 0 else rng_local.integers(3, 7)
cy0 = rng_local.integers(4, 12)
clip = make_clip(shape_type, cx0, cy0, vx)
clip = np.clip(clip + rng_local.normal(0, noise, clip.shape), 0, 1).astype(np.float32)
clips.append(clip)
dir_labels.append(direction)
return np.array(clips, dtype=np.float32), np.array(dir_labels, dtype=np.int64)
probe_rng = np.random.default_rng(13)
X_probe_train, dir_probe_train = make_noisy_video_dataset(probe_rng, 8) # only 8 labeled clips
X_probe_test, dir_probe_test = make_noisy_video_dataset(probe_rng, 150)
Xpt = torch.tensor(X_probe_train).unsqueeze(1)
Xpte = torch.tensor(X_probe_test).unsqueeze(1)
def linear_probe_acc(encoder, seed):
torch.manual_seed(seed)
feat_train = encode_motion_signal(encoder, Xpt)
feat_test = encode_motion_signal(encoder, Xpte)
probe = nn.Linear(feat_train.shape[1], 2)
opt = torch.optim.Adam(probe.parameters(), lr=0.05)
ytr = torch.tensor(dir_probe_train)
for _ in range(300):
opt.zero_grad()
loss = F.cross_entropy(probe(feat_train), ytr)
loss.backward()
opt.step()
with torch.no_grad():
preds = probe(feat_test).argmax(1).numpy()
return (preds == dir_probe_test).mean()
pretrained_encoder = video_mae_models[0.75].encoder
pretrained_accs, random_accs = [], []
for seed in range(5):
pretrained_accs.append(linear_probe_acc(pretrained_encoder, seed=seed + 50))
torch.manual_seed(seed)
random_encoder = VideoMAEEncoder()
random_accs.append(linear_probe_acc(random_encoder, seed=seed + 50))
print(f'motion-direction probe on RANDOM (untrained) encoder: {np.mean(random_accs):.1%} (+/- {np.std(random_accs):.1%})')
print(f'motion-direction probe on VideoMAE-pretrained encoder: {np.mean(pretrained_accs):.1%} (+/- {np.std(pretrained_accs):.1%})')
The pretrained encoder's probe clears the random encoder's by a wide margin, from only 8 labeled clips — the masked-reconstruction pretraining task, despite never being told anything about "direction" or even seeing a label, was forced to represent how the scene changes between time groups well enough that a linear probe can read direction straight off it. This is the video-specific payoff of this lesson's whole pipeline: a still-image encoder (Lesson 47's ViT, or any of Lessons 49-51's pretrained encoders) has no mechanism to represent motion at all, because it never sees more than one frame at a time. Giving the encoder space-time tokens and a masked-reconstruction objective that spans multiple time groups is what makes a temporal property like "which way is this moving" representable in the first place — the same self-supervised recipe as Lesson 49, just applied to a richer input that has temporal structure to learn from.
Exercises¶
- Set
T_TUBE=1(no temporal grouping — the encoder sees 4 independent per-frame patch grids,N_TOKENS=64, instead of 2 tubelet groups ofN_TOKENS=32). Retrain and re-probe. Does the motion-direction probe still work, worse, or about the same — and what does the answer say about whether tubelets themselves (versus just some positional encoding that disambiguates time) are doing the essential work here? - This lesson's motion signal is
per_time[:, -1] - per_time[:, 0], which only ever looks at the first and last time groups. With more than 2 time groups (tryT_FRAMES=8,T_TUBE=2, giving 4 time groups), sketch how you'd generalize this into a feature that uses all the time groups, not just the endpoints — and why a real video model's temporal attention makes this generalization automatic instead of hand-designed. - Increase the masking ratio sweep to include
0.95and0.97(only 1 of 32 tokens visible). At what point does reconstruction finally collapse toward the mean-pixel baseline, and does that breaking point tell you anything about how many independent pieces of information this particular synthetic dataset's clips actually contain?