MaskGIT

The problem, from first principles

MaskGIT is an iterative generative model using masked token prediction. During training, tokens are randomly masked and a transformer predicts them. During generation, tokens with highest uncertainty are iteratively refined in parallel.

Training loss is standard cross-entropy on masked positions. The key is the masking schedule and iterative refinement strategy.

From theory to code

Implement:

maskgit_loss(logits, targets, mask)

Computes cross-entropy loss only on masked token positions.

Constraints

• logits shape: (N, L, V) where V is vocab size.
• targets shape: (N, L).
• mask shape: (N, L), binary.
• Return scalar loss.