Teaching a Neural Network to “See” by Hiding Things From It

How I Built a Masked Autoencoder (MAE) from Scratch on Tiny ImageNet

What if the best way to teach a model to understand images… is to hide most of the image from it?

That’s the surprisingly elegant idea behind Masked Autoencoders (MAE) — a self-supervised learning technique introduced by He et al. at Meta AI in 2021. In this post, I’ll walk you through how I implemented one from scratch in PyTorch, trained it on Tiny ImageNet, and what the results looked like.

🎭 The Big Idea: Learning by Unmasking

Imagine giving someone a jigsaw puzzle where 75% of the pieces are missing and asking them to reconstruct the full picture. Sounds hard — but that’s exactly what forces deep understanding. You can’t just rely on memorizing patterns; you have to genuinely understand the structure, texture, and context of what you’re looking at.

MAE applies this exact principle to images:

1. Divide the image into a grid of patches (16×16 pixels each)
2. Randomly mask 75% of those patches — hide them completely
3. Show the model only the visible 25%
4. Ask it to reconstruct the full image, including all the missing patches

The model is forced to develop a rich internal understanding of visual content to do this well. No labels needed — it’s pure self-supervision.

Architecture: A Vision Transformer Under the Hood

The MAE architecture has two main parts: an Encoder and a Decoder, both built on the Vision Transformer (ViT) backbone.

The Building Blocks

Before assembling the full model, I built the core Transformer components from scratch:

Multi-Head Self-Attention:

class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False):
        self.scale = head_dim ** -0.5
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.proj = nn.Linear(dim, dim)

The attention mechanism computes Query, Key, Value matrices from a single linear projection, then uses scaled dot-product attention to let each patch “look at” every other patch.

MLP Block + Transformer Block: Each Transformer Block combines attention with a feed-forward MLP (hidden dim = 4×), layer normalization, and residual connections — the standard ViT recipe.

The Full MAE Model

Input Image (224×224)
       ↓
Patch Embedding (Conv2d, 16×16 patches → 196 tokens)
       ↓
Positional Encoding
       ↓
Random Masking (keep only 25% = ~49 patches)
       ↓
Encoder (12 Transformer Blocks, dim=768, 12 heads)
       ↓
Decoder Embedding (768 → 384)
       ↓
Insert Mask Tokens for missing patches
       ↓
Decoder (12 Transformer Blocks, dim=384, 6 heads)
       ↓
Prediction Head (384 → 768 = 16×16×3)
       ↓
MSE Loss on masked patches only

A key design choice: the encoder only ever sees the unmasked patches. This makes it extremely efficient — you’re processing 49 tokens instead of 196. The decoder then reconstructs everything from the latent representation plus learnable mask tokens.

The Random Masking Trick

def random_masking(self, x, mask_ratio):
    noise = torch.rand(N, L, device=x.device)
    ids_shuffle = torch.argsort(noise, dim=1)
    ids_restore = torch.argsort(ids_shuffle, dim=1)

    ids_keep = ids_shuffle[:, :len_keep]
    x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
    ...

Uniform random masking is simple but effective. By storing ids_restore, the decoder can correctly reassemble patches back into their original spatial positions.

Dataset: Tiny ImageNet

The model was trained on Tiny ImageNet — a scaled-down version of the classic ImageNet benchmark:

Split Images Classes Image Size Train 100,000 200 64×64 → resized to 224×224 Val 10,000 200 64×64 → resized to 224×224

Data augmentation was kept minimal for self-supervised pretraining — just random horizontal flip and ImageNet normalization. The goal isn’t to classify; it’s to reconstruct.

Training Setup

Hyperparameter Value Optimizer AdamW Learning Rate 1.5e-4 Weight Decay 0.05 Epochs 30 Batch Size 64 LR Schedule Cosine Annealing Mask Ratio 75% Precision Mixed (AMP) Hardware NVIDIA Tesla T4 (Kaggle)

A few training decisions worth calling out:

Mixed Precision (AMP): Using torch.amp.autocast and GradScaler gave a significant speedup by computing in float16 where possible while maintaining float32 precision where it matters.

Gradient Clipping: clip_grad_norm_(max_norm=1.0) prevents exploding gradients — critical for Transformers.

Checkpointing: The training loop saves a checkpoint every epoch and separately saves the best model (lowest validation loss), so you can resume if training is interrupted.

Results

Loss Curves

Training was stable and smooth across all 30 epochs:

Epoch Train Loss Val Loss 1 0.9370 0.9151 5 0.4047 0.3703 10 0.2600 0.2600 20 0.2300 0.2300 30 0.2200 0.2200

The loss dropped sharply in the first 10 epochs (from ~0.94 to ~0.26), then gradually plateaued around 0.22 MSE — classic behavior for self-supervised pretraining.

Quantitative Metrics (PSNR & SSIM)

Beyond raw MSE loss, two standard image quality metrics were evaluated on the validation set:

PSNR (Peak Signal-to-Noise Ratio): Measures reconstruction fidelity in decibels. Higher is better. Values above 20 dB indicate reasonable reconstruction quality.

SSIM (Structural Similarity Index): Measures perceptual similarity accounting for luminance, contrast, and structure. Ranges from 0 to 1, where 1 means perfect reconstruction.

The model achieved solid scores on both metrics, confirming that the reconstructions are perceptually meaningful — not just numerically close.

Visual Results: The “Magic” Moment

The most exciting output is the side-by-side visualization:

Column 1: Original Ground Truth
Column 2: Masked Input (75% hidden — shown as gray)
Column 3: Model's Reconstruction

Even with only 25% of the image visible, the model successfully fills in the missing regions with plausible textures, shapes, and colors. It’s genuinely impressive to see a model hallucinate the correct feathers of a bird or the texture of a tree trunk just from a few visible patches.

The reconstruction isn’t pixel-perfect (it doesn’t need to be), but it demonstrates that the encoder has learned a powerful representation of visual content.

Why This Matters: MAE as a Pretraining Strategy

MAE is not just a reconstruction trick — it’s a pretraining technique. The real workflow looks like this:

Phase 1: Pretrain with MAE (self-supervised, no labels needed)
              ↓
    Encoder learns rich visual representations
              ↓
Phase 2: Fine-tune the encoder on a downstream task
         (image classification, object detection, etc.)

This is analogous to how BERT works in NLP — mask some tokens, predict them, and the model learns deep language understanding. MAE does the same for vision.

The benefit: you can pretrain on millions of unlabeled images (cheap), then fine-tune on a small labeled dataset (expensive but small). This is why Meta AI’s results showed MAE-pretrained ViT models beating supervised baselines on ImageNet with far less labeled data.

Engineering Notes

A few implementation details that made a real difference:

Asymmetric Encoder-Decoder: The encoder (dim=768) is much heavier than the decoder (dim=384). This is intentional — the encoder carries the learned representation, while the decoder is a lightweight reconstruction head that gets discarded after pretraining.

Patch Reconstruction Target: The model predicts raw pixel values (patchified), not features. Simple but effective.

**unpatchify Function:** Converting the flat patch predictions back to a 2D image requires a careful reshape + einsum operation. Easy to get wrong, critical to get right.

What’s Next?

This implementation covers the core MAE pretraining loop. Natural next steps would include:

• Linear Probing: Freeze the encoder, train a linear classifier on top — measures representation quality
• Fine-tuning: End-to-end training on a classification task
• Attention Map Visualization: ViTs learn interpretable attention — visualizing where the model “looks” is fascinating
• Scale Up: Try ViT-Large or ViT-Huge on full ImageNet

Key Takeaways

• MAE is a self-supervised pretraining technique: no labels needed during pretraining
• The core idea is simple: mask 75% of image patches, reconstruct them
• Built on Vision Transformer (ViT) — the encoder only processes visible patches (efficient!)
• Trained on Tiny ImageNet for 30 epochs with AdamW + Cosine Annealing
• Loss converged from 0.93 → 0.22, with strong PSNR and SSIM scores
• The real power of MAE is as a pretraining backbone for downstream vision tasks

The full code is available as a Kaggle notebook. If you found this useful, follow for more deep learning implementation deep-dives. 🙌

Tags: Deep Learning Computer Vision PyTorch Self-Supervised Learning Vision Transformer MAE