--- title: Patchwise Stochastic Attention (PSAL) url: https://www.emergentmind.com/topics/patchwise-stochastic-attention-psal type: topic --- # Patchwise Stochastic Attention (PSAL) Patchwise Stochastic Attention (PSAL) refers to a collection of attention mechanisms in deep learning architectures that combine non-local, patch-based reasoning with stochastic components. PSAL approaches enable scalable, memory-efficient attention across high-resolution inputs and, in some variants, incorporate uncertainty quantification via stochastic representations. The term PSAL encompasses two distinct families: (1) Patch-based Stochastic Attention Layers employing stochastic nearest-neighbor search for sparse, scalable attention [2202.03163], and (2) Patchwise Stochastic Attention mechanisms that represent patches as probabilistic distributions, enabling attention computations informed by measures such as the Wasserstein distance [2311.18645]. ## 1. Mathematical Formulations ### 1.1 Patch-Based Stochastic Attention Layer (Sparse Approximate NN-Based) Given an input feature map $X\in\mathbb{R}^{H\times W\times C}$, patches of size $p\times p$ are extracted via a linear operator $P$: $$ Q = P(X) \in \mathbb{R}^{n\times d}, $$ where $n$ is the number of patches and $d=p\cdot p\cdot C$. Attention is classically computed using a dense score matrix $S_{\text{full}}\in\mathbb{R}^{n\times n}$. PSAL restricts attention to a sparse set $\psi(i)\subset\{1,\ldots,n\}$ of $k$ approximate nearest neighbors per query $i$, forming a sparse score matrix $S$: $$ S_{i,j} = \begin{cases} s(Q_i, K_j), & \text{if } j\in\psi(i)\\ -\infty, & \text{otherwise} \end{cases} $$ with $A_{i,j}$ computed as a row-wise softmax restricted to $\psi(i)$ and the output $Y=AV$. ### 1.2 Distribution-Based Patchwise Stochastic Attention Each patch $x\in\mathbb{R}^d$ is embedded as a Gaussian: $$ x \longmapsto \mathcal{N}(\mu(x), \Sigma(x)), $$ where $\mu(x) = W_\mu x + b_\mu$ and $\Sigma(x) = \operatorname{diag}\bigl(\mathrm{ELU}(W_\sigma x + b_\sigma) + 1\bigr)$. The 2-Wasserstein distance between two Gaussians $\mathcal{N}_i,\mathcal{N}_j$ is $$ W_2^2(\mathcal{N}_i, \mathcal{N}_j) = \|\mu_i - \mu_j\|_2^2 + \operatorname{Tr}\left(\Sigma_i + \Sigma_j - 2(\Sigma_j^{1/2} \Sigma_i \Sigma_j^{1/2})^{1/2}\right). $$ Attention logits are given by $s_{ij} = -\gamma W_2^2(\mathcal{N}_i, \mathcal{N}_j)$, normalized with softmax. ## 2. Stochastic and Approximate Attention Mechanisms Sparse PSAL realizes non-local attention by leveraging the PatchMatch algorithm to identify $k$ approximate nearest neighbors per query patch, reducing computational and memory requirements from $O(n^2)$ to $O(nk)$ [2202.03163]. PatchMatch proceeds iteratively via: - Random initialization of neighbors. - Alternating propagation (using jump-flooding) and random search steps. - Updates with best found candidates, maintaining a heap of top-$k$ matches. In the Gaussian PSAL variant, stochasticity is introduced at the representational level: image patches are modeled as elliptical Gaussians, thereby encoding uncertainty at the patch level. This enables explicit modeling of patch uncertainty and facilitates distance-aware attention. ## 3. Differentiability and Aggregation Strategies Differentiability in sparse PSAL is achieved for $k>1$ by applying the softmax over the restricted neighbor set $\psi(i)$. Gradients propagate through attention weights directly. Two aggregation strategies ensure effective learning: - **PSAL-k**: Top-$k$ matches are maintained per query, yielding $O(nk\log n)$ memory. - **PSAL-Agg**: Supports are augmented with neighbors of spatial neighbors (up to $p^2$ per query), improving match robustness at $O(p^2 n)$ memory. Both approaches maintain end-to-end differentiability for network training [2202.03163]. ## 4. Integration into Deep Architectures PSAL is implemented as a drop-in replacement for full dot-product or non-local attention layers in CNNs or ViTs. Integration procedure: - Project feature maps to query, key, value triplets via 1x1 convolutions. - Extract patches to form $Q, K, V$ matrices. - Employ PSAL to compute the sparse (or Wasserstein-based) attention output. - Aggregate outputs and scatter to the pixel or feature locations. - Add residual connections and proceed with convolutional or transformer blocks [2202.03163, 2311.18645]. In Wasserstein-based PSAL built on ViT-B and the data2vec self-supervised framework, transformer blocks are augmented with projection layers $(W_\mu,W_\sigma)$ for mean and variance, and the multi-head Wasserstein self-attention mechanism replaces conventional dot-product attention. There is a memory and compute overhead (e.g., $\sim$9.3h for 300 epochs on CIFAR-100 with 8xA40 GPUs versus 4h for baseline) [2311.18645]. ## 5. Regularization and Loss Formulations Distribution-based PSAL introduces Wasserstein-based regularization terms: - **Generic loss**: $\mathcal{L} = \mathcal{L}_{\rm base} + \lambda \sum_{i