# Patch-Based Stochastic Attention for Image Editing

arXiv: https://arxiv.org/abs/2202.03163 · HTML: https://mlsanity.com/paper/2202.03163/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 44 references that are in the map.
- First posted on arXiv 7 Feb 2022.

## Builds on (most cited first, 24 of 24)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [Non-local Neural Networks](https://mlsanity.com/paper/1711.07971/) (2017)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Context Encoders: Feature Learning by Inpainting](https://mlsanity.com/paper/1604.07379/) (2016)
- [Image Super-Resolution Using Very Deep Residual Channel Attention Networks](https://mlsanity.com/paper/1807.02758/) (2018)
- [Self-Attention Generative Adversarial Networks](https://mlsanity.com/paper/1805.08318/) (2018)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [Image Transformer](https://mlsanity.com/paper/1802.05751/) (2018)
- [Generative Image Inpainting with Contextual Attention](https://mlsanity.com/paper/1801.07892/) (2018)
- [Stand-Alone Self-Attention in Vision Models](https://mlsanity.com/paper/1906.05909/) (2019)
- [Non-Local Recurrent Network for Image Restoration](https://mlsanity.com/paper/1806.02919/) (2018)
- [Combining Markov Random Fields and Convolutional Neural Networks for Image Synthesis](https://mlsanity.com/paper/1601.04589/) (2016)
- [Neural Nearest Neighbors Networks](https://mlsanity.com/paper/1810.12575/) (2018)
- [Shift-Net: Image Inpainting via Deep Feature Rearrangement](https://mlsanity.com/paper/1801.09392/) (2018)
- [SCRAM: Spatially Coherent Randomized Attention Maps](https://mlsanity.com/paper/1905.10308/) (2019)

## Built on by (most cited first, 3 of 3)

- [VRWKV-Editor: Reducing quadratic complexity in transformer-based video editing](https://mlsanity.com/paper/2509.25998/) (2025)
- [Learned Nonlocal Feature Matching and Filtering for RAW Image Denoising](https://mlsanity.com/paper/2604.17453/) (2026)
- [A Study of the Limits of Collaborative DCT-Based Image Denoising via Interpretable Neural Networks](https://mlsanity.com/paper/2609.29334/) (2026)

## Abstract (arXiv)

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