MAGE: MAsked Generative Encoder to Unify Representation Learning and Image Synthesis

In plain words

Image generators and image encoders for recognition are usually trained as two separate models. MAGE trains one model that can do both. Each picture is first turned into a small grid of codes from a learned image tokenizer, so the model works on symbols rather than raw pixels. During training, part of that grid is hidden and the model learns to guess the hidden part from the rest.

How much is hidden matters. If almost everything is hidden, the task becomes drawing an image from nothing, which is what generation needs. If only a little is hidden, the model learns the kind of features that a classifier can use. MAGE picks the hidden share at random from a range each time, so one training run covers both jobs.

To make a new image, the model starts from a fully hidden grid and fills it in over several rounds, keeping the guesses it is most confident about in each round. Its encoder, used without any hiding, gives features that can be tested with a simple linear classifier.

TL;DR

MAGE is masked image modeling on discrete VQGAN tokens. A ViT encoder-decoder predicts masked tokens with a masking ratio drawn from a truncated Gaussian between 0.5 and 1.0, so the same network supports class-unconditional image generation by iterative decoding and linear probing. An optional contrastive term (MAGE-C) improves the probe. Headline numbers are author-reported on ImageNet-1K at 256x256: 9.10 FID for a ViT-L generator and 78.9% linear-probe top-1 for a ViT-L encoder. The official repository is MIT-licensed, but it does not include the contrastive training code.

Problem

Generative models and self-supervised representation learners are typically trained independently, which duplicates training and maintenance §1. Masked image modeling (MIM) gives strong representations, often with very high masking ratios §1. The intro frames the gap this way: generation is reconstruction of a fully masked image, and representation learning is encoding an unmasked one §1. A fixed high masking ratio therefore serves representations but not generation. The paper reports that a variable masking ratio is necessary for generation §4.4. Pixel-level MIM targets also give blurry reconstructions, which the authors see as a barrier to using MIM for generation §1. MAGE addresses this by using discrete semantic tokens as both input and target §1, §3.1.

Method

MAGE has two parts. Pre-training (§3.1) and post-training generation (§3.2) share one network. The framework overview is fig. 3.

Pre-training (§3.1).

1. Tokenize. Each 256x256 image is mapped by a frozen VQGAN-style tokenizer to 16x16 = 256 discrete codes from a codebook of 1024 entries [§3.1 Tokenization] app. C [config/vqgan.yaml#L4-L14].
2. Sample a masking ratio m_r from a truncated Gaussian with mode 0.55, standard deviation 0.25, and bounds 0.5 and 1.0 §3.1 Masking Strategy tab. 15. The text describes the ratio as sampled per image; the code draws one ratio per batch (see Pitfalls).
3. Mask. Hide m_r * 256 tokens with a learnable mask token [M] §3.1 Masking Strategy fig. 3.
4. Drop. Because m_r >= 0.5, drop a further 0.5 * 256 = 128 tokens from the masked set. Dropped tokens are removed from the encoder input §3.1 Masking Strategy.
5. Encode. Prepend a learnable class-like token [C0] and run a standard ViT encoder on the remaining tokens §3.1 Encoder-Decoder Design.
6. Pad and decode. Pad the encoder output back to full length using the [C] feature of the same image instead of a shared mask token; the decoder reconstructs all positions §3.1 Encoder-Decoder Design app. B.
7. Loss. Cross-entropy between the decoder's codebook prediction and the ground-truth token, averaged over masked positions only [§3.1 Reconstructive Training] eq. 1.

# MAGE pre-training loss, reconstruction part  [§3.1] [eq. 1]
# Input:  images (B, 3, 256, 256)
# Output: scalar loss

def reconstruction_loss(images):
    # 1. Tokenize with the frozen VQGAN-style tokenizer  [§3.1 Tokenization]
    targets = tokenizer(images)                   # (B, 256), ids < 1024
    # 2. Draw a masking ratio from the truncated Gaussian  [§3.1] [tab. 15]
    ratio = truncated_normal(mode=0.55, std=0.25, low=0.5, high=1.0)
    # 3. Rank random noise per image; the lowest ranks are masked
    #    [§3.1 Masking Strategy]
    rank = argsort(argsort(uniform(B, 256)))      # (B, 256)
    masked = rank < ceil(256 * ratio)             # (B, 256) bool
    # 4. Drop half the sequence, chosen among the masked tokens  [§3.1]
    dropped = rank < ceil(256 * 0.5)              # subset of masked
    # 5. Encoder input: prepend [C0]; drop tokens; masked-but-kept -> [M]
    inputs = concat(C0_ID, where(masked, MASK_ID, targets))   # (B, 257)
    keep = concat(True, not dropped)              # (B, 257) bool
    feats = encoder(token_embed(inputs)[keep])    # (B, 129, d)  [§3.1]
    # 6. Project to decoder width; pad every non-visible position with the
    #    [C] feature of the same image  [§3.1] [app. C]
    feats = decoder_embed(feats)                  # (B, 129, d_dec)
    visible = not masked, with [C0] also visible  # (B, 257) bool
    padded = repeat(feats[:, 0], 257)             # (B, 257, d_dec)
    padded[visible] = feats[visible[keep]]        # restore visible tokens
    decoded = decoder(padded + learned_pos)       # 8 blocks  [app. C]
    # 7. Codebook logits for positions 1..256; smoothed CE on masked ones
    logits = head(decoded)[:, 1:, :1024]          # (B, 256, 1024)
    loss = smoothed_ce(logits[masked[:, 1:]],     # label smoothing 0.1
                       targets[masked[:, 1:]])    # [app. C] [tab. 15]
    return mean(loss)                             # mean, not sum (see §Pitfalls)

Equation 1 is the reconstruction term, written as a sum over masked positions of the negative log-probability of the true token given the unmasked tokens eq. 1. Following MAE, the loss is applied only to masked tokens; the authors report that adding the loss on all tokens reduces both generation and representation quality §3.1 Reconstructive Training.

Contrastive co-training (MAGE-C, §3.1). A two-layer MLP is placed on the encoder output after global average pooling. An InfoNCE loss is computed between two augmented views of the same image, with other images in the batch as negatives eq. 2. The paper uses no color jitter, grayscale or noise augmentations, arguing that the reconstruction term prevents shortcut solutions §3.1. The total loss is the reconstruction term plus lambda times the contrastive term, with lambda = 0.1 eq. 3 tab. 15. Appendix B says the contrastive term is applied only when the masking ratio is below 0.6 app. B.

# MAGE-C contrastive term  [§3.1] [eq. 2] [tab. 15]
# Input:  two masked views of a batch, each with masking ratio < 0.6 [app. B]
# Output: scalar InfoNCE loss; total loss = rec + 0.1 * contrastive

def contrastive_term(views_a, views_b):
    # 1. Average-pool the encoder output over the tokens  [§3.1]
    pooled_a = mean(encoder(views_a), axis=1)     # (B, d)
    pooled_b = mean(encoder(views_b), axis=1)     # (B, d)
    # 2. Two-layer MLP, then L2 normalization; activation not stated  [§3.1]
    z_a = normalize(mlp(pooled_a))                # (B, d_z)
    z_b = normalize(mlp(pooled_b))                # (B, d_z)
    # 3. InfoNCE: the matching view is the positive, the rest negatives
    #    [eq. 2] with temperature 0.2 [tab. 15]
    scores = z_a @ z_b.T / 0.2                    # (B, B)
    return cross_entropy(scores, arange(B))       # mean over the batch

Generation by iterative decoding (§3.2, app. C). Generation starts from a blank canvas where every token is [M]. The procedure is MaskGIT-style and runs for T = 20 steps with tau = 6.0 for class-unconditional generation app. C.

# Class-unconditional generation by iterative decoding  [§3.2] [app. C]
# Input:  steps T = 20, temperature tau = 6.0 [app. C]
# Output: 16x16 token grid, decoded to a 256x256 image

def generate(model):
    # 1. Start from a blank canvas: all 256 tokens are [M]  [app. C]
    tokens = full(256, MASK_ID)
    for t in range(1, T + 1):
        # 2. Predict codebook probabilities for every position  [app. C]
        probs = softmax(model(tokens)[:, :1024])    # (256, 1024)
        # 3. Sample one token per position  [app. C, step 2]
        sampled = categorical(probs)                # (256,)
        # 4. Confidence = log p + tau * Gumbel noise; known tokens = +inf
        #    [app. C, step 2]
        gumbel = gumbel_noise(256)                  # (256,)
        confidence = log(prob_of(sampled)) + tau * gumbel
        confidence[known(tokens)] = inf             # keep unmasked tokens
        # 5. Number still masked after round t: cosine schedule [app. C, step 3]
        n_masked = floor(256 * cos(pi * t / (2 * T)))
        # 6. Re-mask the n_masked least-confident positions  [app. C, step 3]
        tokens = sampled
        tokens[lowest_k(confidence, n_masked)] = MASK_ID
    return vqgan_decode(tokens.reshape(16, 16))    # [§3.2]

Linear probing and fine-tuning (§3.2, §4.3). For classification the encoder output is average-pooled without the class token and fed to a linear head §3.2 app. C. Probing and fine-tuning follow MAE's setup app. C; the hyperparameters are in the table below.

Hyperparameters

Pre-training values are from tab. 15. "Code" entries cite the official repository at the commit in the frontmatter.

Results

Author-reported. ImageNet-1K, 256x256 unless noted. Baselines and their parameter counts are as the paper reports them tab. 1.

Class-unconditional generation tab. 1. FID and IS, lower FID is better. "w.a." is the weak-augmentation variant.

Linear probing, ImageNet-1K top-1 tab. 2. Parameter counts for MAGE include the tokenizer.

Few-shot linear classification, ImageNet-1K tab. 3. Top-1 accuracy with 5, 10, 13 or 25 labeled images per class.

Fine-tuning, ImageNet-1K top-1 tab. 4. The paper says this is at par with DINO and slightly below MoCo v3 §4.3.

Class-conditional generation app. A, tab. 7. Frozen MAGE ViT-B encoder plus a 113M class-conditional decoder: FID 6.93, IS 195.8, against MaskGIT at FID 6.18 and IS 182.1 app. A, tab. 7.

Transfer. With 25 samples per class on eight datasets, MAGE outperforms SimCLR and MAE on 6 of 8 §4.3, fig. 5. Appendix A gives 1, 5 and 10 shot results on the same datasets, where the paper reports MAGE better than SimCLR and MAE on most datasets app. A.

Ablations (appendix B and §4.4; epochs as stated per table).

Decoder choice: 8 blocks at width 768 is the chosen ViT-B decoder, with 11.6 FID and 73.3% linear probe at 800 epochs. Wider or deeper variants in the same table change FID by at most 0.3 and linear probe by at most 0.2 points app. B, tab. 12. The ablations use ViT-B for 800 epochs unless noted app. B.

Paper vs code

Official code: https://github.com/LTH14/mage at commit 1becb14475354fc40df35ba7c7c7bf418a137cd3, MIT License, copyright Tianhong Li (first author; the README lists the author's contact). The repository holds the pre-training model, generation, and linear-probe and fine-tune scripts. The README says the MAGE-C training script is "coming soon" README.md.

• Contrastive term not in code. The pre-training model returns only the reconstruction loss models_mage.py. Its checkpoint is provided, but its temperature (0.2), weight (0.1), MLP head and the 0.6 cap on the masking ratio cannot be checked against code tab. 15 README.md.
• Positional embeddings. The appendix says the model uses two learnable positional embeddings, one for the encoder and one for the decoder app. C. In code, the encoder's sin-cos pos_embed is never used in the encoder path; positions come from the learned position embeddings inside the token-embedding module models_mage.py models_mage.py. The decoder adds a learned position embedding models_mage.py. Its fixed sin-cos decoder_pos_embed is set up but not used [code: models_mage.py#L205].
• Warmup. §4.1 says 80 warmup epochs §4.1. The appendix table and the code both use 40 tab. 15 main_pretrain.py. The README command also uses 40 README.md.
• Code defaults differ from paper. main_pretrain.py defaults to 400 epochs and base LR 1e-3 main_pretrain.py main_pretrain.py. The README command gives 1600 epochs and 1.5e-4, matching the paper README.md.
• Loss normalisation. Equation 1 is written as a sum over masked positions eq. 1. The code divides by the number of masked tokens in the batch models_mage.py.
• Label smoothing. Equation 1 has no smoothing term eq. 1, but the appendix and the code use 0.1 tab. 15 models_mage.py.
• Generation temperature. The paper uses a constant tau = 6.0 on the Gumbel noise app. C. The code multiplies tau by (1 - t/T), so the noise decays linearly gen_img_uncond.py. The code also clamps the number of re-masked tokens to at least 1 and at most (unknown - 1) [code: gen_img_uncond.py#L82-L83]. Neither clamp nor decay is in the paper.
• Output head. The output projection reuses the input token-embedding matrix, but through a detached copy, so the output loss sends no gradient into that matrix models_mage.py. The paper does not specify its output head.
• Masking ratio per batch. The code draws one ratio per call, so every image in a batch shares it models_mage.py. The paper's text reads as per-sample §3.1 Masking Strategy.
• Weak augmentation. The 0.8-1.0 crop used for the "w.a." rows §4.2 is not a command-line option in main_pretrain.py. The default is scale (0.2, 1.0) main_pretrain.py.
• Tokenizer. The paper says the tokenizer follows the VQGAN first stage app. C. The code loads vqgan_jax_strongaug.ckpt models_mage.py. This checkpoint is not in the repository and must be downloaded [code: README.md#L48]. Its training recipe is not in the repository.

Pitfalls

• Masking ratio is sampled per batch in code. The text reads as per-image, but the released code shares one ratio across the batch [code: models_mage.py#L270]. Decide which one you implement and record it in your notes.
• [C] padding, not a learned mask token, on the decoder side. The decoder sees the image's encoder [C] feature at every position that was not visible to the encoder, masked or dropped models_mage.py. The learned [M] token only appears in the encoder input [code: models_mage.py#L288]. The cls_token parameter is initialised but not used on the forward path models_mage.py. The [C0] input is the token id 1100 from the class-label range of the vocabulary, not a learned vector [code: models_mage.py#L169].
• Vocabulary layout. The vocabulary has 1024 codebook ids, 1000 class-label ids and one mask id, 2025 in total models_mage.py. Output logits are restricted to the first 1024 ids in both training [code: models_mage.py#L347-L353] and generation gen_img_uncond.py.
• Mean, not sum. Equation 1 is a sum over masked positions, but the code averages over them in the batch models_mage.py. The README learning rate is set for the code's average, so a literal sum needs a rescaled learning rate.
• Generation cost. Each of the 20 steps runs the full encoder and decoder on all 257 positions, with no token dropping gen_img_uncond.py. Generation is therefore 20 full forward passes per batch of images.
• Tokenizer shortcuts. The paper reports that feeding unquantized VQGAN features gives much easier reconstruction and worse linear probes §4.4 Tokenization tab. 6. Quantization is part of the method, not an optimisation.
• FID depends on crop scale. The authors attribute part of the gap between strong and weak augmentation to zoom and off-centre crops in generated images, which FID penalises §4.2. Compare FID only within the same augmentation.
• Compare epochs. Some ablations run at 400 or 800 epochs tab. 5 [app. B, tab. 11]. Their linear-probe numbers are not comparable with the 1600-epoch headline (74.7) unless the table says so.
• Fine-tuning is not a strength. The paper reports fine-tuning below MoCo v3 and only at par with DINO §4.3. The authors suspect the quantized tokens and leave the question open §4.3.
• Ratio range matters for generation. A fixed masking ratio yields FID above 50, so the variable range is needed for generation tab. 5.
• Unstated values. The paper does not state ViT-L encoder heads, the decoder heads, the MLP activation for MAGE-C, or the tokenizer's training recipe app. C §3.1. Use the code for these (see Paper vs code) or mark them "not stated" in your implementation notes.

Lineage

• Builds on masked autoencoders (MAE) for the encoder-decoder, the [C] padding and the drop-before-encode step §3.1 app. B. Compared against MAE, MAE with GAN loss, iBOT, MoCo v3, DINO, MSN and SimCLR §4.3 app. B.
• Uses a VQGAN first stage as the tokenizer §3.1. Uses MaskGIT's iterative decoding for generation §3.2. Compared against MaskGIT, ADM and BigGAN for generation tab. 1.
• Uses an InfoNCE contrastive term similar to SimCLR for MAGE-C §3.1.
• Superseded by: not stated in this card.