Masked Autoencoder (MAE) — From Scratch in PyTorch

A complete from-scratch implementation of Masked Autoencoders Are Scalable Vision Learners (He et al., 2021) built entirely in PyTorch.

"If we are not able to rebuild it, we don't understand it."

---

Overview

This repository reimplements the Masked Autoencoder (MAE) architecture from first principles to deeply understand self-supervised learning and Vision Transformers.

The implementation covers the entire training pipeline:

• Patch embedding
• Random masking
• Transformer encoder and decoder
• Learnable positional embeddings
• Reconstruction objective
• End-to-end pretraining
• Unit testing for every component

No pretrained weights, no external MAE libraries, and no high-level wrappers were used.

---

What is MAE?

A Masked Autoencoder is a self-supervised vision model that learns rich image representations by solving a simple pretext task:

1. Split an image into patches.
2. Randomly mask 75% of them.
3. Encode only the remaining 25% visible patches.
4. Reconstruct the missing patches.

The model never sees labels and learns purely from the structure of images.

---

Why MAE?

Previous vision methods process every image patch, making training expensive.

MAE introduces a simple but powerful idea:

• Mask most of the image.
• Run the expensive Transformer encoder only on visible patches.

Since self-attention scales quadratically with sequence length (O(N²)), processing only 25% of the patches dramatically reduces computation and memory requirements.

This efficiency allows MAE to:

• train larger Vision Transformers,
• scale to larger datasets,
• learn strong representations from unlabeled images.

This is why the paper is called:

Masked Autoencoders Are Scalable Vision Learners

---

Architecture

Image (224×224)
    │
    ▼
PatchEmbed
    │
    ├── Split into 196 patches (16×16)
    └── Linear projection to embeddings
    │
    ▼
Random Masking
    │
    ├── Keep 25% visible patches
    └── Mask 75% of patches
    │
    ▼
MAE Encoder
    │
    ├── Transformer blocks
    ├── Learnable positional embeddings
    └── Processes visible patches only
    │
    ▼
Latent Representations
    │
    ▼
MAE Decoder
    │
    ├── Insert learnable mask tokens
    ├── Restore original patch ordering
    ├── Add decoder positional embeddings
    └── Reconstruct all patches
    │
    ▼
Prediction Head
    │
    ▼
Pixel Reconstruction

---

Random Masking and Sequence Restoration

During masking, patches are randomly shuffled and only the visible patches are passed to the encoder.

To reconstruct the image, the decoder:

1. Projects encoder latents.
2. Inserts learnable mask tokens.
3. Uses ids_restore to recover the original patch ordering.
4. Adds positional embeddings and reconstructs the complete image.

This asymmetric encoder-decoder design is one of the key innovations of MAE.

---

Reconstruction Objective

The decoder predicts pixel values for all patches, including both visible and masked patches.

However, the reconstruction loss is computed only on the masked patches:

$$\mathcal{L} = \frac{1}{|\mathcal{M}|}\sum_{i \in \mathcal{M}} |x_i - \hat{x}_i|^2$$

where $\mathcal{M}$ denotes the set of masked patches.

This prevents the model from learning a trivial identity mapping and encourages the encoder to learn meaningful semantic representations.

---

Results

Trained on Oxford Flowers102:

• 8,189 images
• Resolution: 224×224
• 200 epochs of self-supervised pretraining

The model successfully reconstructs the global structure and color of masked flowers using only 25% of visible patches as context.

Reconstruction Example 1

Reconstruction Example 2

Reconstruction Example 3

---

Project Structure

mae-from-scratch/
├── README.md
├── assets/
│   ├── reconstructions_flowers_01.png
│   ├── reconstructions_flowers_02.jpg
│   └── reconstructions_flowers_03.png
├── checkpoints/
├── notebooks/
│   └── mae_from_scratch.ipynb
└── src/
    ├── __init__.py
    ├── components/
    │   ├── __init__.py
    │   ├── masking.py
    │   └── patch_embed.py
    ├── data/
    │   └── dataset.py
    ├── engine/
    │   ├── __init__.py
    │   ├── loss.py
    │   └── train.py
    ├── layers/
    │   ├── __init__.py
    │   ├── attention.py
    │   ├── block.py
    │   └── mlp.py
    ├── models/
    │   ├── __init__.py
    │   ├── decoder.py
    │   ├── encoder.py
    │   └── mae.py
    ├── positionembedding/
    │   ├── __init__.py
    │   └── learnable.py
    └── test/
        ├── __init__.py
        ├── 01_test_patch_embed.py
        ├── 02_test_pos_embd.py
        ├── 03_test_random_masking.py
        ├── 04_test_attention.py
        ├── 05_test_mlp.py
        ├── 06_test_block.py
        ├── 07_test_encoder.py
        ├── 08_test_decoder.py
        ├── 09_test_mae.py
        └── 10_test_losses.py

---

Implementation Details

Every component was implemented from scratch.

Component	File
Patch Embedding	src/components/patch_embed.py
Random Masking	src/components/masking.py
Learnable Positional Embedding	src/positionembedding/learnable.py
Multi-Head Self-Attention	src/layers/attention.py
MLP Block	src/layers/mlp.py
Transformer Block	src/layers/block.py
MAE Encoder	src/models/encoder.py
MAE Decoder	src/models/decoder.py
Full MAE Model	src/models/mae.py
Loss Function	src/engine/loss.py
Training Loop	src/engine/train.py
Dataloader	src/data/dataset.py

Every module includes unit tests covering:

• shape correctness,
• gradient flow,
• optimizer updates,
• end-to-end integration.

---

Training Configuration

Hyperparameter	Value
Image Size	224×224
Patch Size	16×16
Mask Ratio	75%
Encoder Depth	12
Encoder Embedding Dimension	512
Decoder Depth	4
Decoder Embedding Dimension	256
Optimizer	AdamW
Learning Rate	1.5e-4
Weight Decay	0.05
Betas	(0.9, 0.95)
Scheduler	Cosine Annealing
Epochs	200
Batch Size	32

---

Model Size

• Total Parameters: ~41.9M
• Trainable Parameters: ~41.9M

---

How to Run

Install dependencies:

pip install torch torchvision matplotlib

Run locally:

jupyter notebook notebooks/mae_from_scratch.ipynb

Or open directly on Kaggle:

---

Future Work

• Fine-tune the encoder on downstream classification tasks.
• Linear probing of learned representations.
• Experiment with Rotary Positional Embeddings (RoPE).
• Study different masking ratios.
• Train on larger datasets (Tiny ImageNet, ImageNet).
• Scale to ViT-Base and ViT-Large configurations.
• Add stronger data augmentation.

---

References

• He, K., Chen, X., Xie, S., Li, Y., Dollár, P., & Girshick, R. (2021).

  Masked Autoencoders Are Scalable Vision Learners

  https://arxiv.org/abs/2111.06377

---

Author

Wajdi Zanteni

Second-Year Engineering Student at Sup'Com Training Manager at MLS Machine Learning Sup'Com Member of RLS Research Labs Sup'Com