MARch\'e: Fast Masked Autoregressive Image Generation with Cache-Aware Attention

TL;DR

MARche9 introduces a cache-aware attention mechanism with selective KV refresh to significantly speed up masked autoregressive image generation, maintaining quality while reducing redundant computations.

Contribution

It presents a training-free, efficient generation framework that reduces computation in masked autoregressive models through novel attention and key/value refresh strategies.

Findings

Achieves up to 1.7x speedup in image generation.

Maintains image quality with negligible degradation.

Provides a scalable solution for masked transformer models.

Abstract

Masked autoregressive (MAR) models unify the strengths of masked and autoregressive generation by predicting tokens in a fixed order using bidirectional attention for image generation. While effective, MAR models suffer from significant computational overhead, as they recompute attention and feed-forward representations for all tokens at every decoding step, despite most tokens remaining semantically stable across steps. We propose a training-free generation framework MARch\'e to address this inefficiency through two key components: cache-aware attention and selective KV refresh. Cache-aware attention partitions tokens into active and cached sets, enabling separate computation paths that allow efficient reuse of previously computed key/value projections without compromising full-context modeling. But a cached token cannot be used indefinitely without recomputation due to the changing contextual information over multiple steps. MARch\'e recognizes this challenge and applies a technique called selective KV refresh. Selective KV refresh identifies contextually relevant tokens based on attention scores from newly generated tokens and updates only those tokens that require recomputation, while preserving image generation quality. MARch\'e significantly reduces redundant computation in MAR without modifying the underlying architecture. Empirically, MARch\'e achieves up to 1.7x speedup with negligible impact on image quality, offering a scalable and broadly applicable solution for efficient masked transformer generation.