---
title: Efficient Patchwise Axial Self-Attention
url: https://www.emergentmind.com/topics/patchwise-axial-self-attention
type: topic
---

# Efficient Patchwise Axial Self-Attention

Patchwise axial self-attention is a family of self-attention mechanisms that factorize global attention into more memory- and computation-efficient operations by performing attention along single axes of high-dimensional data or within and across spatial patches. This structure enables context modeling in vision and volumetric tasks that would be infeasible for standard quadratic self-attention. Architectures including the Axial Transformer, AEWin Transformer, GASA-UNet, and multiscale Self-Attentive Convolutions (MSAC) instantiate various forms of patchwise axial self-attention in both 2D and 3D domains, with specific designs for balancing local detail and global context while controlling cost [1912.12180, 2209.08726, 2409.13146, 1912.01521].

## 1. Principles of Axial and Patchwise Axial Self-Attention

Conventional self-attention on a tensor $X \in \mathbb{R}^{N^1 \times N^2 \times \cdots \times N^d \times D}$ requires $O((\prod_{i=1}^d N^i)^2)$ operations and memory, quickly becoming infeasible for large spatial domains. **Axial self-attention** addresses this by performing attention sequentially along each axis (e.g., rows, columns, depth), reducing the cost to $O(dN^{1+1/d})$ for hypercube tensors [1912.12180]. For images, axial attention applies separate attention to each row or column in succession, factorizing global dependence.

**Patchwise axial self-attention** further reduces cost and increases flexibility by (1) partitioning input data into patches or windows and (2) performing axial or windowed attention within and/or across these regions. Typical variants combine local window attentions with global axial operations or use a two-stage local-global block structure. This is motivated by the need to capture both fine local cues and long-range dependencies without incurring prohibitive compute/memory requirements [2209.08726, 1912.12180].

## 2. Canonical Algorithms and Variants

Several approaches elaborate the concept of patchwise axial self-attention, differing in their partitioning schemes, ordering, and fusion strategies:

- **Two-Stage Patchwise Axial Attention** [1912.12180]:
  1. **Partitioning**: The input is split into non-overlapping $g \times g$ patches.
  2. **Intra-Patch (Local) Axial Attention**: For each patch, independently apply row- and column-wise self-attention.
  3. **Patch Summarization**: Each patch is reduced to a vector (e.g., by mean pooling).
  4. **Inter-Patch (Global) Axial Attention**: Row- and column-wise attention across the grid of patch summaries.
  5. **Feedback**: Broadcast or add the patch-level global representations back to the fine spatial grid.

- **Axially Expanded Window Attention (AEWin)** [2209.08726]:
  1. Tokens are divided into local, non-overlapping square windows and row/column stripes.
  2. Attention is computed in parallel on each local window (fine granularity) and on horizontal/vertical stripes (coarse granularity), with heads split among these groups.
  3. Outputs from all groups are concatenated and projected.

- **Global Axial Self-Attention for 3D Volumes** [2409.13146]:
  1. 3D feature maps are collapsed along axes using 2D convolutions to generate three sets of 1D patch sequences for width, height, and depth.
  2. The concatenated sequence undergoes multi-head self-attention.
  3. Learnable 1D positional embeddings are injected post-attention to retain spatial identity.
  4. Attention-enhanced features are “unflattened” by axis and concatenated, restoring a 3D tensor for fusion.

- **Self-Attentive Convolutions (SAC, MSAC)** [1912.01521]:
  1. Q/K/V projections become $n \times m$ convolutions.
  2. Sliding window (patchwise) attention is computed locally within each patch.
  3. Multiscale context is captured by running SAC blocks at various patch sizes in parallel, with outputs summed or concatenated.


## 3. Mathematical Formulation

Patchwise axial self-attention methods employ the standard QKV projection and scaled dot-product, but apply it on limited axes or regions.

### Generic Axial Self-Attention (single axis):

Let $X \in \mathbb{R}^{L \times D}$ denote a sequence along a given axis.

\[
Q = X W^Q,\quad K = X W^K,\quad V = X W^V
\]
\[
A = \mathrm{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right)
\]
\[
O = A V
\]

### Two-Stage Patchwise Axial Attention [1912.12180]:

- **Local (intra-patch):** For each patch, two rounds per axis (row, col), $g \times g$ attention matrices.
- **Global (inter-patch):** Attention on $P_h \times P_w$ patch grid, again per axis.

Total complexity is given by:
\[
T_{\mathrm{patchwise}} \approx 2D S^2 g + 2D \frac{S^3}{g^3}
\]
where $S$ is image width/height, $g$ is patch size, and $D$ is channel dimension.

### GASA Block [2409.13146]:

Extract 1D patches along each axis using specialized 2D convolutional projections, concatenate, and apply MHSA:

\[
Q = P W^Q, \quad K = P W^K, \quad V = P W^V
\]
\[
O = \mathrm{softmax}\left(\frac{Q K^T}{\sqrt{d_k}} \right) V
\]
\[
O' = O + PE
\]

$O'$ is then rebroadcast into 3D context via broadcasting per axis.

### AEWin Head Splitting [2209.08726]:

Divide $K$ heads into $K/2$ window (local), $K/4$ horizontal, $K/4$ vertical (axial). Attention is computed in three parallel streams and outputs concatenated.

## 4. Computational Complexity

Patchwise axial self-attention dramatically reduces the quadratic cost of global self-attention on high-dimensional data. Key metrics for major variants include:

| Method                   | Time   | Memory   | Dominant Matrix Size             |
|--------------------------|--------|----------|----------------------------------|
| Global Self-Attention    | $O(N^2D)$   | $O(N^2)$   | $(N\times N)$                    |
| Axial (Row+Col)          | $O(DN^{3/2})$| $O(N^{3/2})$| $(S\times S)$ per axis           |
| Patchwise Axial (2-stage)| $O(DN g + D S^3 / g^3)$| $O(N)$ or $O(N/g^2)$| $(g\times g)$ (local), $(P_h\times P_w)$ (global) |
| GASA (3D volume)         | $O(W(H D)^2 + H(W D)^2 + D(W H)^2)$ |      | Each is $O(N^4)$ for $N=W=H=D$   |

These reductions enable use on larger spatial/volumetric data or higher resolution, while still providing global receptive fields through compounding local/global blocks [1912.12180, 2409.13146, 2209.08726].


## 5. Positional Information and Fusing Axes

Axial splitting or patch partitioning loses explicit information about a voxel/pixel’s global position. Patchwise axial attention variants re-inject positional awareness via:

- **Concatenated absolute positional embeddings** (by patch and by stripe) [2409.13146].
- **Relative positional bias tables** (by local window, row, or column) [2209.08726].
- **Conditional positional encoding** (depthwise local convolutions) [2209.08726].

Outputs aggregated from attention heads assigned to different axes or windows are concatenated or summed, followed by a final linear projection or convolution to fuse information [2209.08726, 1912.01521].

## 6. Integration and Empirical Performance

Patchwise axial self-attention has been incorporated successfully into both encoder-decoder and hierarchical transformer architectures:

- **GASA-UNet** integrates the GASA block between encoder and decoder, concatenating attention-enhanced 3D features with the encoder output [2409.13146]. The GASA-enhanced nnUNet baseline achieves Dice score improvements up to $+1.5$ and NSD increase of up to $+2.38$ on small or ambiguous structures, while adding only $\sim1$M parameters and $0.3$ GFLOPs.
- **AEWin Transformers** alternate between attention on windows and axial stripes, embedding this hybrid block into multi-stage architectures with patch merging for hierarchical context [2209.08726].
- **Multiscale SAC** (MSAC) modules run patchwise attention at different spatial scales in parallel and fuse outputs, supporting flexible backbone integration; however, large-scale benchmarking is not yet reported [1912.01521].

## 7. Architectural and Hyperparameter Trade-Offs

Design choices critically affect trade-offs between expressiveness, efficiency, and practical feasibility:

- **Patch Size ($g$):** Smaller $g$ increases local operation cost but reduces global (across-patch) cost; larger $g$ does the reverse. Optimal $g$ often scales as $S^{1/4}$, balancing computation.
- **Number of Heads ($h$):** Increasing heads distributes attention, allowing axes-local specialization at marginally higher cost.
- **Depth and Fusion:** Interleaving local and global blocks or stacking more local/global layers can increase expressivity while controlling memory.
- **Positional Embeddings/Bias:** Critical for restoring the unique spatial identity lost in axis or patch flattening.

A plausible implication is that architectures employing patchwise axial self-attention achieve near-global receptive field with close-to-linear cost, preserving both local detail and global coherence. This framework supports efficient scaling for high-resolution vision and medical imaging tasks, often outperforming pure local or axial-only baselines in empirical evaluation [1912.12180, 2409.13146, 2209.08726].

Source: https://www.emergentmind.com/topics/patchwise-axial-self-attention