Transformers Debugging

Question

# ML Debugging Interview: Transformers KV-Cache & Autoregressive Generation

## Context

You are implementing a **mini GPT-style autoregressive Transformer** to generate small MNIST-like digit images, where each pixel (0 – 255) is treated as a token.  
Your task is to **debug and optimize** the implementation below.

The model should:

- Train a GPT-like model to predict pixel intensities (sequence modeling task)
- Achieve training loss < 2.8 on MNIST
- Produce recognizable generated digits
- Implement **KV-cache optimization** for faster generation

There are **intentional bugs** hidden in the Transformer implementation — including errors in positional encoding, attention masking, scaling, and inf signs.

---

## Part 1 – Dataset

```python
import os
from torchvision.datasets import MNIST
from torchvision import transforms
from torch.utils.data import Dataset
import torch

class CustomMNIST(MNIST):
    """MNIST dataset using a custom GitHub mirror."""
    mirrors = ["https://raw.githubusercontent.com/mkolod/MNIST/master/"]
    resources = [
        ("train-images-idx3-ubyte.gz", "f68b3c2dcbeaaa9fbdd348bbdeb94873"),
        ("train-labels-idx1-ubyte.gz", "d53e105ee54ea40749a09fcbcd

Accepted Answer

A detailed reference solution is available to Crack ML Interview members.