README.md

---
title: MNIST 2x2 Transformer Demo
emoji: 🔢
colorFrom: gray
colorTo: blue
sdk: gradio
sdk_version: 5.29.0  # or your current gradio version
app_file: app.py
pinned: false
---

# Transformer MNIST 2×2 — Image-to-Sequence Prediction

This project implements a minimal Transformer-based model that takes a 2×2 grid of MNIST digits as input and autoregressively predicts the corresponding 4-digit sequence. It serves as a practical deep dive into the inner workings of the Transformer architecture and basic multimodality concepts, combining vision (image patches) with language modeling (digit sequences).

## 1. Project Overview

The goal is to understand how a vanilla Transformer encoder-decoder can be applied to a simple multimodal task: mapping an image input to a discrete token sequence. This project focuses on building each architectural component from scratch and wiring them together cleanly.

## 2. Task Definition

- **Input:** a 2×2 grid composed of 4 random MNIST digits, forming a 56×56 grayscale image.
- **Output:** the 4-digit sequence corresponding to the digits in the grid (top-left → bottom-right).
- **Modeling approach:** sequence-to-sequence using an autoregressive decoder with special `<start>` and `<finish>` tokens.

## 3. Model Architecture

The model follows a clean encoder-decoder Transformer architecture:

- **Feature Extractor:** splits the 56×56 image into 16 non-overlapping patches of 14×14 pixels and projects each to a 64-dimensional embedding.
- **Transformer Encoder:** processes the 16 patch embeddings using standard multi-head self-attention, positional embeddings, and MLP blocks.
- **Transformer Decoder:** autoregressively predicts the digit sequence:
  - Uses masked self-attention over token embeddings
  - Attends to encoder output via cross-attention
  - Outputs a sequence of logits over a vocabulary of 13 tokens (digits 0–9, `<start>`, `<finish>`)
- **Tokenizer:** handles token ↔ digit conversions and input preparation.

## 4. Training Setup

- **Dataset:** MNIST, wrapped into a custom `MNIST_2x2` PyTorch dataset that returns the stitched image and 4-digit target.
- **Batch size:** 64
- **Epochs:** 10
- **Loss:** `CrossEntropyLoss` over vocabulary tokens
- **Optimizer:** Adam
- **Hardware:** Apple M4 with `mps` acceleration
- **Logging:** `tqdm` per-batch loss tracking for clear training progress

## 5. Evaluation

Evaluation is done on the held-out MNIST test set using greedy decoding:

- Starts with <start> token
- Predicts one token at a time (no teacher forcing)
- Stops after 4 tokens or if <finish> is predicted

### Evaluation Metrics

- **Sequence accuracy:** % of samples where all 4 digits are predicted correctly
- **Per-digit accuracy:** % of individual digits predicted correctly across all positions

### final results after 10 epochs of training

- **training loss at epoch 10:** 0.0101
- **Sequence accuracy:** 93.77% on held-out test set
- **Per digit accuracy:** 98.43% on held-out test set