Source-linked AI summary
MAT: Mask-Aware Transformer for Large Hole Image Inpainting
Wenbo Li, Zhe Lin, Kun Zhou, Lu Qi, Yi Wang, Jiaya Jia
TL;DR
Large-hole inpainting needs non-local context, while prior transformer methods often restrict inference to low resolution. MAT combines convolutions with a mask-aware transformer that aggregates partial valid tokens, achieving state-of-the-art performance on multiple benchmarks and supporting pluralistic completion. The method remains limited for objects with varied shapes and requires image dimensions to be multiples of 512.
Problem
Existing transformer-based inpainting methods commonly infer low-resolution predictions because of computational complexity, which can produce coarse structures for large masks.
Method
MAT combines convolutions with a customized transformer block and multi-head contextual attention that models long-range relations using dynamically selected valid tokens.
Results
MAT achieves state-of-the-art performance on multiple benchmark datasets, including Places and CelebA-HQ, and enables pluralistic completion.
Takeaways & Limitations
MAT directly processes high-resolution images while supporting efficient long-range dependency modeling and diverse generation.
Takeaways & Limitations
MAT struggles with objects having varied shapes without semantic annotations and requires padding or resizing images to dimensions that are multiples of 512.
Abstract
from arXiv · showhide
Recent studies have shown the importance of modeling long-range interactions in the inpainting problem. To achieve this goal, existing approaches exploit either standalone attention techniques or transformers, but usually under a low resolution in consideration of computational cost. In this paper, we present a novel transformer-based model for large hole inpainting, which unifies the merits of transformers and convolutions to efficiently process high-resolution images. We carefully design each component of our framework to guarantee the high fidelity and diversity of recovered images. Specifically, we customize an inpainting-oriented transformer block, where the attention module aggregates non-local information only from partial valid tokens, indicated by a dynamic mask. Extensive experiments demonstrate the state-of-the-art performance of the new model on multiple benchmark datasets. Code is released at https://github.com/fenglinglwb/MAT.
1. Introduction
Large-hole inpainting requires non-local contextual understanding, but existing transformer approaches typically operate at low resolution. MAT combines transformers and convolutions to process high-resolution images while targeting realistic, diverse completions.
- Large masks require distant contextual information to reconstruct reasonable structures and textures, which fully convolutional networks struggle to model for complicated structures.
- Existing transformer-based inpainting methods generally infer low-resolution predictions because of computational complexity, producing coarse structures that can compromise final quality.
- MAT customizes transformer blocks by removing conventional layer normalization and replacing residual learning with feature-concatenation fusion to improve large-mask training stability.
- Multi-head contextual attention models long-range dependencies efficiently using valid tokens selected by a dynamically updated mask.
- MAT combines a convolutional head, transformer body, convolutional tail, and refinement network for high-resolution pluralistic inpainting.
- MAT introduces style manipulation to support pluralistic generation and reports state-of-the-art performance on Places and CelebA-HQ benchmarks.
2. Related Work
Image completion research spans classical propagation, patch-based, statistical, and deep-learning approaches, with recent work emphasizing pluralistic generation and large-hole filling.
- Classical inpainting methods propagate neighboring information, retrieve similar patches, accelerate patch search, or use differential equations and image statistics.
- Deep-learning approaches introduced adversarial training and expanded encoder-decoder and U-Net architectures for image completion.
- Deep inpainting research also developed global and local discrimination, contextual attention, partial convolution, gated convolution, and multi-stage generation.
- Figure 2 places MAT among transformer-convolution architectures that include a convolutional head, transformer body, convolutional tail, and Conv-U-Net refinement.
- Recent work focuses on pluralistic generation and large-hole filling, including probabilistic, manifold-based, bidirectional-attention, and autoregressive approaches.
3. Method
MAT combines convolutional processing with a mask-aware transformer body to model long-range context for large-hole inpainting while supporting diverse outputs.
- Architecture: MAT uses a convolutional head, five transformer stages, a convolutional reconstruction tail, and a style manipulation module.The head extracts tokens, the transformer body models long-range interactions, and the reconstruction module upsamples output tokens.
- Convolutional Head: The convolutional head downsamples the masked image to 1/8-sized feature maps and provides local inductive priors while reducing computational cost.It contains four convolutional layers, with three used for downsampling.
- Adjusted Transformer Block: MAT’s adjusted transformer block removes layer normalization and replaces residual learning with feature-concatenation fusion learning.Attention and input features are concatenated and passed through a fully connected layer before the MLP.
- Multi-Head Contextual Attention: Multi-head contextual attention models long-range dependencies using shifted windows and a dynamic mask that excludes invalid tokens from aggregation.Attention outputs are weighted sums of valid tokens, while window shifts enable cross-window connections.
- Mask Updating Strategy: The mask is initialized from the input and updated after attention: a window becomes fully valid when it contains at least one valid token.Windows with no valid tokens remain invalid; repeated updates and shifts propagate validity.
- Style Manipulation Module: The style manipulation module fuses image-conditional and noise-unconditional styles, modulates convolution weights, and injects noise to enhance generation diversity.The modulation of different style representations leads to pluralistic outputs.
4. Experiments
Experiments evaluate MAT on 512 × 512 Places365-Standard and CelebA-HQ images using perceptual metrics and component ablations. MAT outperforms variants and competing methods while producing diverse completions, but struggles with semantically varied objects and requires size adjustment for attention.
- Datasets and Metrics: MAT is evaluated on Places365-Standard and CelebA-HQ at 512 × 512 resolution using FID, P-IDS, and U-IDS as primary large-hole metrics.The study argues that pixel-wise metrics inadequately reflect human perception for ill-posed large-hole completion.
- Ablation Study: Replacing transformer blocks with convolutions causes an obvious performance drop, especially on P-IDS and U-IDS, while distant context improves structural and texture reconstruction.The ablation compares MAT with a fully convolutional variant on a reduced 256 × 256 setting.
- Ablation Study: MAT’s adjusted transformer block improves FID by 0.39 over the original transformer block and produces more visually appealing, semantically consistent completions.The qualitative comparison includes an extremely large missing region containing an indoor scene.
- Ablation Study: Removing partial aggregation from valid tokens lowers FID by 0.1, while qualitative results show incorrect colors and blur, supporting the usefulness of multi-head contextual attention.The authors associate contextual attention with color consistency and reduced blurriness.
- Ablation Study: Removing noise style manipulation creates a large P-IDS and U-IDS gap, indicating that stochastic noise styles improve the naturalness of completed images.Style variants also produce different facial expressions, contours, appearances, and scene structures.
- Limitations and Failure Cases: MAT struggles with objects having varied shapes when trained without semantic annotations and requires padding or resizing images to multiples of 512 because of attention windows.The reported failure cases include cats and cars.
5. Conclusion
The paper presents MAT for pluralistic large-hole image inpainting, combining an adjusted transformer architecture, partial attention, and style modulation. It reports state-of-the-art performance on multiple benchmarks alongside improved image quality and generation diversity.
- Conclusion: MAT combines an adjusted transformer architecture and partial attention mechanism for pluralistic large-hole image inpainting.The framework also includes style modulation to improve generation diversity.
- Conclusion: MAT achieves state-of-the-art performance on multiple benchmarks and demonstrates strong image quality and diversity in qualitative comparisons.The conclusion attributes the framework to both high-quality and diverse completions.
- Conclusion: Style modulation improves the diversity of generated completions within the proposed framework.The conclusion identifies diversity as a central outcome of the method.
(Supplementary Material)
The supplementary material illustrates free-form masks for 512 × 512 images, distinguishing visible from invisible pixels by color.
- Free-Form Masks: Free-form masks are shown at 512 × 512 resolution, with visible pixels in white and invisible pixels in black.The figure provides examples of the mask patterns used for image completion.
A. Network Architecture
MAT uses a two-stage architecture that combines a convolutional head, multi-stage transformer body, reconstruction tail, and Conv-U-Net refinement. Style modulation supports pluralistic generation, while training uses randomly generated free-form masks.
- Network Architecture: MAT is a two-stage framework with a convolutional head, transformer body, convolutional reconstruction tail, and second-stage Conv-U-Net.The discriminator follows the design of CoModGAN.
- Network Architecture: The first stage transforms input image-and-mask features into tokens, processes them through five transformer stages, and reconstructs a full-resolution image.The transformer stages use varying resolutions and block counts {2, 3, 4, 3, 2}.
- Network Architecture: Style modulation is applied throughout reconstruction and supports pluralistic generation.The reconstruction tail and decoding layers use style representations for image-conditioned or noise-unconditional modulation.
- Network Architecture: The second-stage Conv-U-Net refines the coarse prediction using the input mask, with shortcut connections across resolutions and style-modulated decoding layers.Its encoder downsamples to H/32 × W/32 before decoding back to H × W.
- Network Architecture: Training uses free-form masks composed of randomly sized rectangles and brush strokes, with a large-mask sampling strategy.Rectangle counts, stroke counts, brush widths, and vertex numbers are randomly sampled within specified ranges.
C. Tokenization
MAT uses stacked convolutions rather than a linear projection to tokenize masked images. This design gradually fills holes, supplies multi-scale decoder shortcuts, and benefits from broader attention range and model capacity.
- Tokenization: Stacked convolutions replace ViT-style linear projection for token extraction in MAT.The comparison uses the same training setting as the ablation study.
- Tokenization: Stacked convolutions gradually fill holes and produce more effective tokens for the inpainting transformer body.This design is tailored to the inpainting problem.
- Tokenization: Multi-scale features from the convolutional tokenizer pass to the decoder through shortcut connections, improving optimization.
- Tokenization: Ablation results show performance is positively correlated with model capacity and attention range.The variants change feature width, block number, and transformer window size while keeping Conv-U-Net unchanged.
E. CelebA-HQ 256 × 256 Results
On 256×256 CelebA-HQ, MAT achieves significant improvements over other methods across FID, P-IDS, and U-IDS metrics.
- MAT is evaluated quantitatively on 256×256 CelebA-HQ.
- MAT significantly improves FID results over other methods.
- MAT also significantly improves P-IDS and U-IDS results over other methods.
F. LPIPS Results
LPIPS is treated as a reference measure because it is unsuitable for large-mask pluralistic inpainting with multiple plausible completions. MAT nevertheless achieves superior or comparable LPIPS performance and favorable fidelity-diversity results across the reported comparisons.
- LPIPS Results: LPIPS is not an appropriate primary measure for large-mask pluralistic inpainting because multiple plausible hole completions may exist.The paper reports LPIPS results only for reference.
- LPIPS Results: MAT achieves superior or comparable LPIPS performance on CelebA-HQ and Places.The comparison covers 512 × 512 datasets.
- LPIPS Results: MAT uses only 22.5% of the full Places data for training its Places model.The paper states this alongside the LPIPS comparison.
- LPIPS Results: At 1024 × 1024 on Places large-mask inpainting, MAT reports FID↓5.83 and P-IDS↑9.51%, compared with Big LaMa's FID↓6.31 and P-IDS↑4.98%.Both models were trained at 512 × 512 and transferred to 1024 × 1024.
- LPIPS Results: Compared with CoModGAN on Places, MAT obtains better FID and higher recall but slightly lower precision.The paper notes that MAT uses much less training data.
I. Additional Qualitative Results
Additional Places comparisons show MAT producing more photo-realistic inpainting results with fewer artifacts than other state-of-the-art methods. The paper omits CelebA-HQ visual comparisons because of potential copyright issues.
- Qualitative comparisons: MAT generates more photo-realistic Places results with few artifacts in qualitative comparisons against state-of-the-art methods.The comparisons are presented at 512 × 512 resolution in Figures J.4 and J.5.
- Qualitative comparisons: The additional qualitative comparisons are conducted on the Places dataset.
- Dataset coverage: Visual comparisons for CelebA-HQ are not provided because of potential copyright issues.The authors direct readers to the provided code and model if such processing is necessary.
- Image sources and licensing: The paper identifies FFHQ face images used in the paper and supplementary material and documents their sources and licenses.The listed sources include Flickr images with CC BY or CC BY-NC licenses.