Patch-Based Stochastic Attention for Image Editing

TL;DR

The paper addresses the prohibitive memory cost of standard full attention in vision tasks by introducing PSAL, a Patch-based Stochastic Attention Layer that leverages PatchMatch to approximate nearest neighbors with linear memory $O(n)$, enabling high-resolution processing. Differentiability is restored via two strategies: PSAL-$k$ (multiple NN candidates) and patch aggregation, enabling end-to-end training in networks containing PSAL. Empirical results across image reconstruction, colorization, inpainting, and super-resolution show PSAL achieves competitive or superior quality while drastically reducing memory and compute, with high-resolution inpainting demonstrated up to $3300\times3300$ images. The approach paves the way for integrating attention into broader image-editing architectures and suggests extensions to videos and diffusion-model frameworks.

Abstract

Attention mechanisms have become of crucial importance in deep learning in recent years. These non-local operations, which are similar to traditional patch-based methods in image processing, complement local convolutions. However, computing the full attention matrix is an expensive step with heavy memory and computational loads. These limitations curb network architectures and performances, in particular for the case of high resolution images. We propose an efficient attention layer based on the stochastic algorithm PatchMatch, which is used for determining approximate nearest neighbors. We refer to our proposed layer as a "Patch-based Stochastic Attention Layer" (PSAL). Furthermore, we propose different approaches, based on patch aggregation, to ensure the differentiability of PSAL, thus allowing end-to-end training of any network containing our layer. PSAL has a small memory footprint and can therefore scale to high resolution images. It maintains this footprint without sacrificing spatial precision and globality of the nearest neighbors, which means that it can be easily inserted in any level of a deep architecture, even in shallower levels. We demonstrate the usefulness of PSAL on several image editing tasks, such as image inpainting, guided image colorization, and single-image super-resolution. Our code is available at: https://github.com/ncherel/psal

Figures (13)

• Figure 1: Illustration of patch NN search. (a) Full Attention computes a complete attention matrix but many elements have negligible weight. (b) Patch-Based Stochastic Attention only probes randomly a few elements. (c) Good matches are propagated to neighbors
• Figure 2: Full Attention (left) and PSAL (right) reconstruction using another frame of the same video. The memory constraints and the subsequent subsampling step make it impossible for the Full Attention to capture all details
• Figure 3: Architecture for our colorization network. We have used as simple an architecture as possible to isolate the contribution of the attention layer.
• Figure 4: Results on the colorization task. PSAL 3 and PSAL Aggreg. have good results despite some wrong matches (gray skin).
• Figure 5: Performance vs computation constraints (memory and GFLOPs) in the colorization task. Full Attention performs well at the cost of high memory and many GFLOPs. Local Attention is an efficient approximation of Full Attention with a limited drop in performance. Performer does not perform well. PSAL 3 and PSAL Aggreg. have better performance, largely reduced memory usage, and require less FLOPs than alternatives. $^*$Full Attention still requires a subsampling step to fit into GPU memory.