Back to Home
Artificial IntelligenceJuly 8, 2026

Training A Tiny Vision Transformer From Scratch With Patchdropout

N

Written by

Nova Neural

Why I got obsessed with PatchDropout

I wanted to understand how regularization changes the internal “token flow” in a transformer for vision. Most deep learning recipes talk about dropout in generic terms, but I kept running into a very specific problem:

When you split an image into patches (small square regions) and feed them as tokens, the model can over-rely on a subset of patch positions early in training. I wanted a regularization trick that would directly perturb patch tokens, not just randomly zero individual activations.

That led me to PatchDropout—a technique where I randomly drop patch tokens during training, forcing the model to be robust to missing spatial evidence.

Below is a complete, working implementation that trains a tiny Vision Transformer from scratch on a simple synthetic dataset. The code runs end-to-end in PyTorch.

What PatchDropout means (in plain language)

A Vision Transformer (ViT) turns an image into a sequence of tokens:

  1. Split image into patches (e.g., 8×8 blocks)
  2. Flatten each patch into a vector
  3. Map vectors to an embedding dimension
  4. Add a “class token” (usually called [CLS]) used for classification
  5. Run transformer layers over the token sequence

PatchDropout modifies step (4):

  • During training, I randomly remove a fraction of patch tokens (not the [CLS] token).
  • This changes the set of evidence the transformer sees at each step.
  • During evaluation, I keep everything (no dropping).

A tiny end-to-end demo: ViT + PatchDropout

Step 1: Imports and reproducibility

import math
import random
import torch
import torch.nn as nn
import torch.nn.functional as F

from torch.utils.data import Dataset, DataLoader
def seed_everything(seed=0):
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

seed_everything(42)
device = "cuda" if torch.cuda.is_available() else "cpu"
device

Step 2: A synthetic vision dataset (fast and deterministic)

This dataset creates two classes of 32×32 grayscale images:

  • Class 0: a vertical line at a random x position
  • Class 1: a horizontal line at a random y position

This is intentionally simple—so training dynamics are easy to observe.

class LineDataset(Dataset):
    def __init__(self, n=2000, img_size=32, thickness=2):
        self.n = n
        self.img_size = img_size
        self.thickness = thickness

    def __len__(self):
        return self.n

    def __getitem__(self, idx):
        # Class label: 0 => vertical line, 1 => horizontal line
        y = idx % 2
        img = torch.zeros(1, self.img_size, self.img_size)

        if y == 0:
            x = torch.randint(0, self.img_size, (1,)).item()
            x0 = max(0, x - self.thickness // 2)
            x1 = min(self.img_size, x0 + self.thickness)
            img[:, :, x0:x1] = 1.0
        else:
            yy = torch.randint(0, self.img_size, (1,)).item()
            y0 = max(0, yy - self.thickness // 2)
            y1 = min(self.img_size, y0 + self.thickness)
            img[:, y0:y1, :] = 1.0

        # Add tiny noise so the task isn't purely thresholding
        img += 0.05 * torch.randn_like(img)
        img = img.clamp(0.0, 1.0)
        return img, torch.tensor(y, dtype=torch.long)

Step 3: Patch embedding

I split each image into non-overlapping patches using a convolution with kernel size = patch size. That yields patch tokens efficiently.

class PatchEmbed(nn.Module):
    def __init__(self, img_size=32, patch_size=8, in_chans=1, embed_dim=64):
        super().__init__()
        assert img_size % patch_size == 0
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = img_size // patch_size
        self.num_patches = self.grid_size * self.grid_size

        self.proj = nn.Conv2d(
            in_chans,
            embed_dim,
            kernel_size=patch_size,
            stride=patch_size,
            bias=True
        )

    def forward(self, x):
        # x: (B, C, H, W)
        x = self.proj(x)  # (B, E, H/ps, W/ps)
        x = x.flatten(2).transpose(1, 2)  # (B, N, E)
        return x

Step 4: PatchDropout module

This is the core. During training:

  • I compute a keep mask for patch tokens only.
  • I always keep the [CLS] token.
  • For simplicity and correctness, I keep the tensor shape fixed by:
    • Sampling a binary mask per token
    • Multiplying token embeddings by mask
    • Rescaling by 1/(1-drop_prob) (standard inverted dropout behavior)

This avoids variable-length sequences, which keeps the rest of the ViT straightforward.

class PatchDropout(nn.Module):
    def __init__(self, drop_prob=0.25):
        super().__init__()
        self.drop_prob = drop_prob

    def forward(self, patch_tokens):
        """
        patch_tokens: (B, N, E) where N excludes the CLS token
        """
        if not self.training or self.drop_prob == 0.0:
            return patch_tokens

        keep_prob = 1.0 - self.drop_prob

        # mask: (B, N, 1)
        mask = (torch.rand(patch_tokens.size(0), patch_tokens.size(1), 1, device=patch_tokens.device)
                < keep_prob).type_as(patch_tokens)

        # Inverted dropout scaling
        patch_tokens = patch_tokens * mask / keep_prob
        return patch_tokens

Step 5: The transformer encoder

I use PyTorch’s nn.TransformerEncoderLayer for the heavy lifting. It expects shape (S, B, E), so I transpose appropriately.

class TinyViT(nn.Module):
    def __init__(
        self,
        img_size=32,
        patch_size=8,
        in_chans=1,
        num_classes=2,
        embed_dim=64,
        depth=4,
        num_heads=4,
        mlp_ratio=4.0,
        patch_drop_prob=0.25
    ):
        super().__init__()

        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, 1 + num_patches, embed_dim))
        self.pos_drop = nn.Dropout(0.0)

        self.patch_drop = PatchDropout(patch_drop_prob)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim,
            nhead=num_heads,
            dim_feedforward=int(embed_dim * mlp_ratio),
            dropout=0.1,
            activation="gelu",
            batch_first=False,  # we'll provide (S, B, E)
            norm_first=True
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=depth)

        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes)

        # Initialization
        nn.init.trunc_normal_(self.cls_token, std=0.02)
        nn.init.trunc_normal_(self.pos_embed, std=0.02)

    def forward(self, x):
        # x: (B, C, H, W)
        patch_tokens = self.patch_embed(x)  # (B, N, E)

        # Apply PatchDropout to patch tokens only (not CLS)
        patch_tokens = self.patch_drop(patch_tokens)  # (B, N, E)

        B, N, E = patch_tokens.shape
        cls_tokens = self.cls_token.expand(B, -1, -1)  # (B, 1, E)

        tokens = torch.cat([cls_tokens, patch_tokens], dim=1)  # (B, 1+N, E)
        tokens = tokens + self.pos_embed[:, : tokens.size(1), :]
        tokens = self.pos_drop(tokens)

        # TransformerEncoder expects (S, B, E)
        tokens = tokens.transpose(0, 1)  # (S, B, E)
        tokens = self.encoder(tokens)    # (S, B, E)
        tokens = tokens.transpose(0, 1)  # (B, S, E)

        cls_final = tokens[:, 0]  # (B, E)
        cls_final = self.norm(cls_final)
        logits = self.head(cls_final)  # (B, num_classes)
        return logits

Step 6: Training loop

I train two versions:

  • Baseline: patch_drop_prob = 0.0
  • PatchDropout: patch_drop_prob = 0.3

This makes the effect easy to see without over-explaining.

def train_one(model, loader, optimizer, criterion):
    model.train()
    total_loss = 0.0
    correct = 0
    total = 0

    for imgs, labels in loader:
        imgs = imgs.to(device)
        labels = labels.to(device)

        optimizer.zero_grad(set_to_none=True)
        logits = model(imgs)
        loss = criterion(logits, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item() * imgs.size(0)
        preds = logits.argmax(dim=1)
        correct += (preds == labels).sum().item()
        total += imgs.size(0)

    return total_loss / total, correct / total

@torch.no_grad()
def eval_one(model, loader, criterion):
    model.eval()
    total_loss = 0.0
    correct = 0
    total = 0

    for imgs, labels in loader:
        imgs = imgs.to(device)
        labels = labels.to(device)

        logits = model(imgs)
        loss = criterion(logits, labels)

        total_loss += loss.item() * imgs.size(0)
        preds = logits.argmax(dim=1)
        correct += (preds == labels).sum().item()
        total += imgs.size(0)

    return total_loss / total, correct / total
# Data
train_ds = LineDataset(n=3000, img_size=32, thickness=2)
test_ds  = LineDataset(n=1000, img_size=32, thickness=2)

train_loader = DataLoader(train_ds, batch_size=64, shuffle=True, num_workers=0)
test_loader  = DataLoader(test_ds, batch_size=64, shuffle=False, num_workers=0)

criterion = nn.CrossEntropyLoss()

def run_experiment(patch_drop_prob):
    model = TinyViT(
        img_size=32,
        patch_size=8,
        in_chans=1,
        num_classes=2,
        embed_dim=64,
        depth=4,
        num_heads=4,
        mlp_ratio=4.0,
        patch_drop_prob=patch_drop_prob
    ).to(device)

    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.05)

    history = []
    for epoch in range(1, 16):
        tr_loss, tr_acc = train_one(model, train_loader, optimizer, criterion)
        te_loss, te_acc = eval_one(model, test_loader, criterion)
        history.append((tr_loss, tr_acc, te_loss, te_acc))
        print(f"[patch_drop_prob={patch_drop_prob}] epoch {epoch:02d} "
              f"train acc={tr_acc:.3f} test acc={te_acc:.3f}")
    return history

history_base = run_experiment(0.0)
history_drop = run_experiment(0.3)

What I actually observed when running it

On this synthetic dataset, both models learn quickly (the task is easy), but PatchDropout typically:

  • reduces overconfidence early on (loss curve is a bit smoother)
  • improves test accuracy a little in the later epochs
  • makes training less “spiky” because the model can’t depend on the same patch positions every step

Even though the dataset is simple, the mechanism is real: PatchDropout forces the transformer to aggregate across patch evidence rather than memorizing patch/token statistics.

Visual sanity check: confirming patches are being “dropped”

I added a small debug snippet to verify that patch tokens are being masked only in training mode.

def debug_patch_dropout(model, x):
    model.train()
    with torch.no_grad():
        patch_tokens = model.patch_embed(x)
        dropped = model.patch_drop(patch_tokens)
    # Percentage of near-zero token entries (rough proxy)
    near_zero = (dropped.abs() < 1e-6).float().mean().item()
    return near_zero

model_dbg = TinyViT(patch_drop_prob=0.3).to(device)
x, _ = train_ds[0]
x = x.unsqueeze(0).to(device)

print("near-zero fraction in training (rough):", debug_patch_dropout(model_dbg, x))
model_dbg.eval()
with torch.no_grad():
    patch_tokens = model_dbg.patch_embed(x)
    dropped_eval = model_dbg.patch_drop(patch_tokens)
print("near-zero fraction in eval (should be ~0):", (dropped_eval.abs() < 1e-6).float().mean().item())

When the model is in .eval() mode, PatchDropout returns tokens untouched—exactly what I wanted.

Where PatchDropout fits in a real project

PatchDropout is most compelling when:

  • you treat spatial patches as meaningful discrete evidence
  • you see overfitting to specific patch positions (often due to dataset bias)
  • you want token-level regularization that’s more targeted than generic dropout

In production-style training runs, I’d pair it with:

  • standard weight decay (already used above)
  • careful learning rate schedules
  • potentially patch masking augmentations as a preprocessing step (mixing strategies can help)

Conclusion

I built a tiny Vision Transformer from scratch and added a dedicated PatchDropout regularizer that drops patch tokens (not the [CLS] token) during training by masking token embeddings with inverted dropout scaling. Running the full PyTorch script showed smoother training behavior and a small generalization gain on a simple line-classification dataset, and the debug checks confirmed the masking happens only in training mode.