{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4e625e7a",
   "metadata": {},
   "source": "# Lesson 52: Video Transformers and Pretraining\n\nEvery lesson so far in this part has treated an image as a bag of patch tokens (Lesson {{lesson:vision_transformers}}'s ViT) and pretrained on it with contrastive learning (Lesson {{lesson:self_supervised_learning}}), self-distillation (Lesson {{lesson:dino_self_distillation}}), or masked reconstruction (Lesson {{lesson:masked_autoencoders}}). 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 &mdash; 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 (<a href=\"../references.html#tong-2022-videomae\">Tong et al., 2022</a>, extending Lesson {{lesson:masked_autoencoders}}'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."
  },
  {
   "cell_type": "code",
   "id": "ff8ab679",
   "source": "import numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport matplotlib.pyplot as plt",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "f2577377",
   "source": "## A tiny synthetic \"video\" dataset\n\nSame `plus`/`circle` shapes as Lesson {{lesson:masked_autoencoders}}, 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** &mdash; a genuinely temporal property that doesn't exist in any individual frame, only in how frames relate to each other.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "f1793eb2",
   "source": "IMG = 16\nPATCH = 4\nT_FRAMES = 4\nT_TUBE = 2                                    # temporal tubelet depth\nN_SPATIAL = (IMG // PATCH) ** 2                # 16 spatial patches per frame\nN_TOKENS = (T_FRAMES // T_TUBE) * N_SPATIAL    # 2 * 16 = 32 space-time tokens\nTUBE_DIM = T_TUBE * PATCH * PATCH              # 2*16 = 32 pixels per tubelet\n\ndef make_clip(shape_type, cx0, cy0, vx, n_frames=T_FRAMES, size=IMG):\n    frames = []\n    for t in range(n_frames):\n        cx, cy = cx0 + vx * t, cy0\n        img = np.zeros((size, size), dtype=np.float32)\n        if shape_type == 'plus':\n            img[cy-1:cy+2, cx-3:cx+4] = 1.0\n            img[cy-3:cy+4, cx-1:cx+2] = 1.0\n        else:\n            yy, xx = np.mgrid[0:size, 0:size]\n            img[((xx-cx)**2 + (yy-cy)**2) <= 9] = 1.0\n        frames.append(img)\n    return np.stack(frames)  # (T, H, W)\n\ndef make_video_dataset(rng_local, n):\n    clips, dir_labels = [], []\n    for _ in range(n):\n        shape_type = rng_local.choice(['plus', 'circle'])\n        direction = rng_local.choice([0, 1])  # 0 = moving left, 1 = moving right\n        vx = -2 if direction == 0 else 2\n        cx0 = rng_local.integers(9, 13) if direction == 0 else rng_local.integers(3, 7)\n        cy0 = rng_local.integers(4, 12)\n        clips.append(make_clip(shape_type, cx0, cy0, vx))\n        dir_labels.append(direction)\n    return np.array(clips, dtype=np.float32), np.array(dir_labels, dtype=np.int64)\n\nrng = np.random.default_rng(7)\nX_train, dir_train = make_video_dataset(rng, 400)\nX_test, dir_test = make_video_dataset(rng, 150)\nXt = torch.tensor(X_train).unsqueeze(1)   # (B, 1, T, H, W)\nXte = torch.tensor(X_test).unsqueeze(1)\n\nfig, axes = plt.subplots(2, T_FRAMES, figsize=(7, 3.2))\nfor row, (direction_name, idx) in enumerate([('moving left', 0), ('moving right', 1)]):\n    j = int(np.argmax(dir_test == idx))\n    for t in range(T_FRAMES):\n        axes[row, t].imshow(X_test[j, t], cmap='gray', vmin=0, vmax=1); axes[row, t].axis('off')\n    axes[row, 0].set_title(direction_name, fontsize=8, loc='left')\nplt.show()",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "0bfb006f",
   "source": "## Tokenizing video: space-time patches (\"tubelets\")\n\nLesson 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.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "ccd5e48c",
   "source": "def tubelet_patchify(clip, patch_size=PATCH, tube=T_TUBE):\n    # clip: (B, C, T, H, W) -> (B, N_TOKENS, TUBE_DIM). Token order is (time_group, spatial_row, spatial_col):\n    # token s and token N_SPATIAL+s are the SAME spatial location at two different time groups.\n    B, C, T, H, W = clip.shape\n    P = patch_size\n    x = clip.unfold(2, tube, tube)            # (B, C, T//tube, H, W, tube)\n    x = x.unfold(3, P, P).unfold(4, P, P)      # (B, C, T//tube, H//P, W//P, tube, P, P)\n    x = x.permute(0, 2, 3, 4, 1, 5, 6, 7).contiguous()\n    return x.view(B, -1, C * tube * P * P)\n\ndef tubelet_unpatchify(tokens, patch_size=PATCH, tube=T_TUBE, img_size=IMG, n_frames=T_FRAMES):\n    B, N, D = tokens.shape\n    P = patch_size\n    n_side = img_size // P\n    n_time = n_frames // tube\n    x = tokens.view(B, n_time, n_side, n_side, tube, P, P)\n    x = x.permute(0, 1, 4, 2, 5, 3, 6).contiguous()\n    return x.view(B, n_frames, img_size, img_size)\n\npatches = tubelet_patchify(Xt)\nrecon_check = tubelet_unpatchify(patches)\nprint(f'{N_TOKENS} tokens per clip, {TUBE_DIM} pixels each')\nprint(f'patchify/unpatchify roundtrip max abs diff: {(recon_check - Xt.squeeze(1)).abs().max().item():.2e}')",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "a8a135ef",
   "source": "## A space-time Transformer encoder, and a masked-reconstruction decoder\n\nExactly Lesson {{lesson:masked_autoencoders}}'s MAE encoder/decoder, with one change: the positional encoding (Lesson {{lesson:transformer_architecture}}) 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.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "175ae697",
   "source": "def positional_encoding(T, D):\n    pos = torch.arange(T).unsqueeze(1).float()\n    i = torch.arange(D).unsqueeze(0).float()\n    angle_rates = 1.0 / (10000 ** (2 * (i // 2) / D))\n    angles = pos * angle_rates\n    pe = torch.zeros(T, D)\n    pe[:, 0::2] = torch.sin(angles[:, 0::2])\n    pe[:, 1::2] = torch.cos(angles[:, 1::2])\n    return pe\n\ndef random_masking(x, mask_ratio, generator=None):\n    B, N, D = x.shape\n    n_visible = max(1, int(N * (1 - mask_ratio)))\n    noise = torch.rand(B, N, generator=generator)\n    ids_shuffle = torch.argsort(noise, dim=1)\n    ids_keep = ids_shuffle[:, :n_visible]\n    ids_mask = ids_shuffle[:, n_visible:]\n    x_visible = torch.gather(x, 1, ids_keep.unsqueeze(-1).expand(-1, -1, D))\n    return x_visible, ids_keep, ids_mask\n\nclass VideoMAEEncoder(nn.Module):\n    def __init__(self, tube_dim=TUBE_DIM, embed_dim=32, n_heads=4, n_layers=2):\n        super().__init__()\n        self.embed = nn.Linear(tube_dim, embed_dim)\n        self.register_buffer('pe', positional_encoding(N_TOKENS, embed_dim))\n        layer = nn.TransformerEncoderLayer(embed_dim, n_heads, dim_feedforward=embed_dim * 2,\n                                            batch_first=True, dropout=0.0)\n        self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)\n\n    def forward(self, x_visible, ids_keep):\n        tok = self.embed(x_visible) + self.pe[ids_keep]  # positional encoding at each token's TRUE space-time index\n        return self.encoder(tok)\n\nclass VideoMAEDecoder(nn.Module):\n    def __init__(self, embed_dim=32, decoder_dim=16, tube_dim=TUBE_DIM, n_heads=4, n_layers=1):\n        super().__init__()\n        self.embed = nn.Linear(embed_dim, decoder_dim)\n        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_dim))\n        self.register_buffer('pe', positional_encoding(N_TOKENS, decoder_dim))\n        layer = nn.TransformerEncoderLayer(decoder_dim, n_heads, dim_feedforward=decoder_dim * 2,\n                                            batch_first=True, dropout=0.0)\n        self.decoder = nn.TransformerEncoder(layer, num_layers=n_layers)\n        self.pred = nn.Linear(decoder_dim, tube_dim)\n        self.decoder_dim = decoder_dim\n\n    def forward(self, encoded_visible, ids_keep, ids_mask):\n        B, n_mask = ids_mask.shape\n        vis_tok = self.embed(encoded_visible) + self.pe[ids_keep]\n        mask_tok = self.mask_token.expand(B, n_mask, -1) + self.pe[ids_mask]\n        full = torch.zeros(B, N_TOKENS, self.decoder_dim)\n        full = full.scatter(1, ids_keep.unsqueeze(-1).expand(-1, -1, self.decoder_dim), vis_tok)\n        full = full.scatter(1, ids_mask.unsqueeze(-1).expand(-1, -1, self.decoder_dim), mask_tok)\n        return self.pred(self.decoder(full))\n\nclass VideoMAE(nn.Module):\n    def __init__(self, mask_ratio=0.5, embed_dim=32, decoder_dim=16):\n        super().__init__()\n        self.mask_ratio = mask_ratio\n        self.encoder = VideoMAEEncoder(embed_dim=embed_dim)\n        self.decoder = VideoMAEDecoder(embed_dim=embed_dim, decoder_dim=decoder_dim)\n\n    def forward(self, clip, generator=None):\n        tokens = tubelet_patchify(clip)\n        x_visible, ids_keep, ids_mask = random_masking(tokens, self.mask_ratio, generator=generator)\n        encoded = self.encoder(x_visible, ids_keep)\n        pred = self.decoder(encoded, ids_keep, ids_mask)\n        D = tokens.shape[-1]\n        pred_masked = torch.gather(pred, 1, ids_mask.unsqueeze(-1).expand(-1, -1, D))\n        target_masked = torch.gather(tokens, 1, ids_mask.unsqueeze(-1).expand(-1, -1, D))\n        loss = F.mse_loss(pred_masked, target_masked)\n        return loss, pred, ids_mask, ids_keep, tokens\n\ndef train_video_mae(mask_ratio, epochs=600, lr=0.001, seed=0):\n    torch.manual_seed(seed)\n    model = VideoMAE(mask_ratio=mask_ratio)\n    opt = torch.optim.Adam(model.parameters(), lr=lr)\n    for _ in range(epochs):\n        opt.zero_grad()\n        loss, *_ = model(Xt)\n        loss.backward()\n        opt.step()\n    return model",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "f869b3c4",
   "source": "## Video tolerates extreme masking\n\nLesson 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 {{lesson:masked_autoencoders}}'s single-image MAE could tolerate — and compare against a mean-pixel baseline.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "04f2df7a",
   "source": "train_tokens = tubelet_patchify(Xt)\nmean_pixel = train_tokens.mean()\n\nprint(f'{\"mask ratio\":>10} {\"VideoMAE test MSE\":>18} {\"mean-pixel baseline\":>20}')\nvideo_mae_models = {}\nfor mr in [0.5, 0.75, 0.9]:\n    model = train_video_mae(mask_ratio=mr)\n    video_mae_models[mr] = model\n    with torch.no_grad():\n        test_loss, pred, ids_mask, ids_keep, tokens = model(Xte, generator=torch.Generator().manual_seed(123))\n        D = tokens.shape[-1]\n        target_masked = torch.gather(tokens, 1, ids_mask.unsqueeze(-1).expand(-1, -1, D))\n        baseline_mse = F.mse_loss(torch.full_like(target_masked, mean_pixel), target_masked).item()\n    print(f'{mr:>10} {test_loss.item():>18.4f} {baseline_mse:>20.4f}')",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "0cda6836",
   "source": "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.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "7987a626",
   "source": "model90 = video_mae_models[0.9]\nwith torch.no_grad():\n    _, pred, ids_mask, ids_keep, tokens = model90(Xte, generator=torch.Generator().manual_seed(42))\n\nk = 3\nmasked_view = tokens.clone()\nmasked_view[torch.arange(len(tokens)).unsqueeze(1), ids_mask] = 0.5  # gray out hidden tubelets\nrecon_view = tokens.clone()\nrecon_view[torch.arange(len(tokens)).unsqueeze(1), ids_mask] = pred[torch.arange(len(tokens)).unsqueeze(1), ids_mask]\n\norig_clip = tubelet_unpatchify(tokens)[k].numpy()\nmasked_clip = tubelet_unpatchify(masked_view)[k].numpy()\nrecon_clip = tubelet_unpatchify(recon_view)[k].numpy()\n\nfig, axes = plt.subplots(3, T_FRAMES, figsize=(7, 4.5))\nfor row, (clip, label) in enumerate(zip([orig_clip, masked_clip, recon_clip],\n                                          ['original', 'seen by encoder\\n(90% hidden)', 'reconstructed'])):\n    for t in range(T_FRAMES):\n        axes[row, t].imshow(clip[t], cmap='gray', vmin=0, vmax=1); axes[row, t].axis('off')\n    axes[row, 0].set_ylabel(label, fontsize=7)\nplt.show()",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "da0dff72",
   "source": "## A linear probe for a temporal property: direction of motion\n\nLessons 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.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "54fb2504",
   "source": "N_TIME = T_FRAMES // T_TUBE\n\ndef encode_motion_signal(encoder, clip):\n    tokens = tubelet_patchify(clip)\n    B, N, _ = tokens.shape\n    ids_all = torch.arange(N).unsqueeze(0).expand(B, -1)\n    with torch.no_grad():\n        encoded = encoder(tokens, ids_all)                       # (B, N_TOKENS, embed_dim)\n    per_time = encoded.view(B, N_TIME, N_SPATIAL, -1).mean(dim=2)  # (B, N_TIME, embed_dim): pool over SPACE only\n    return per_time[:, -1] - per_time[:, 0]                        # last time group minus first: a motion signal\n\ndef make_noisy_video_dataset(rng_local, n, noise=0.15):\n    clips, dir_labels = [], []\n    for _ in range(n):\n        shape_type = rng_local.choice(['plus', 'circle'])\n        direction = rng_local.choice([0, 1])\n        vx = -2 if direction == 0 else 2\n        cx0 = rng_local.integers(9, 13) if direction == 0 else rng_local.integers(3, 7)\n        cy0 = rng_local.integers(4, 12)\n        clip = make_clip(shape_type, cx0, cy0, vx)\n        clip = np.clip(clip + rng_local.normal(0, noise, clip.shape), 0, 1).astype(np.float32)\n        clips.append(clip)\n        dir_labels.append(direction)\n    return np.array(clips, dtype=np.float32), np.array(dir_labels, dtype=np.int64)\n\nprobe_rng = np.random.default_rng(13)\nX_probe_train, dir_probe_train = make_noisy_video_dataset(probe_rng, 8)   # only 8 labeled clips\nX_probe_test, dir_probe_test = make_noisy_video_dataset(probe_rng, 150)\nXpt = torch.tensor(X_probe_train).unsqueeze(1)\nXpte = torch.tensor(X_probe_test).unsqueeze(1)\n\ndef linear_probe_acc(encoder, seed):\n    torch.manual_seed(seed)\n    feat_train = encode_motion_signal(encoder, Xpt)\n    feat_test = encode_motion_signal(encoder, Xpte)\n    probe = nn.Linear(feat_train.shape[1], 2)\n    opt = torch.optim.Adam(probe.parameters(), lr=0.05)\n    ytr = torch.tensor(dir_probe_train)\n    for _ in range(300):\n        opt.zero_grad()\n        loss = F.cross_entropy(probe(feat_train), ytr)\n        loss.backward()\n        opt.step()\n    with torch.no_grad():\n        preds = probe(feat_test).argmax(1).numpy()\n    return (preds == dir_probe_test).mean()\n\npretrained_encoder = video_mae_models[0.75].encoder\npretrained_accs, random_accs = [], []\nfor seed in range(5):\n    pretrained_accs.append(linear_probe_acc(pretrained_encoder, seed=seed + 50))\n    torch.manual_seed(seed)\n    random_encoder = VideoMAEEncoder()\n    random_accs.append(linear_probe_acc(random_encoder, seed=seed + 50))\n\nprint(f'motion-direction probe on RANDOM (untrained) encoder:    {np.mean(random_accs):.1%}  (+/- {np.std(random_accs):.1%})')\nprint(f'motion-direction probe on VideoMAE-pretrained encoder:   {np.mean(pretrained_accs):.1%}  (+/- {np.std(pretrained_accs):.1%})')",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "753ab9f7",
   "source": "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 {{lesson:vision_transformers}}'s ViT, or any of Lessons {{lesson:masked_autoencoders}}-{{lesson:dino_self_distillation}}'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 {{lesson:masked_autoencoders}}, just applied to a richer input that has temporal structure to learn from.\n\n### Exercises\n\n1. Set `T_TUBE=1` (no temporal grouping — the encoder sees 4 independent per-frame patch grids, `N_TOKENS=64`, instead of 2 tubelet groups of `N_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?\n2. 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 (try `T_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.\n3. Increase the masking ratio sweep to include `0.95` and `0.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?",
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.x"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}