Axial Attention: Efficient Multidimensional Self-Attention

• Axial attention is a factorized form of self-attention that computes along individual axes, reducing computational complexity compared to full attention.
• It enables global context preservation by sequentially applying attention along dimensions such as height, width, and depth in images and volumetric data.
• Axial attention has been successfully integrated in architectures for image classification, dense prediction, medical segmentation, and other domains.

Axial attention is a factorized form of self-attention for multidimensional tensors in which attention is computed along one axis at a time rather than over all tokens in a fully flattened representation. In images this usually means row-wise and column-wise attention; in volumetric data it extends to depth; in other domains the axes may be time and frequency, sequence position and pair index, or other structured coordinates. The original multidimensional Transformer formulation was motivated by the prohibitive $O(N^2)$ cost of full attention over $N=\prod_i D_i$ tensor positions, and by the observation that axis-aligned factorizations preserve long-range interaction while remaining compatible with standard deep learning frameworks (Ho et al., 2019). Subsequent work turned axial attention into a general architectural primitive for image classification, dense prediction, medical segmentation, video segmentation, wireless receivers, and scientific operator learning (Wang et al., 2020, Sun et al., 2024, Yellapragada et al., 14 Oct 2025, Cheng et al., 9 May 2026).

1. Definition and historical formulation

Axial attention was introduced for multidimensional Transformers as self-attention that naturally aligns with tensor axes. Given a tensor $$X \in \mathbb{R}^{D_1 \times D_2 \times \cdots \times D_k \times d}$$, one chooses an axis $a$, batches all remaining axes together, and applies standard 1D attention to each length-$D_a$ slice independently. Repeating this across axes yields a global receptive field without flattening the tensor into a single length-$N$ sequence (Ho et al., 2019).

In 2D vision, the canonical case is a feature map $X \in \mathbb{R}^{H \times W \times C}$. Rather than letting every pixel attend to all $HW$ positions in one step, axial attention performs one pass along columns and one pass along rows. A later medical-imaging formulation makes the connectivity intuition explicit: for any two pixels $(x_0,y_0)$ and $(x_1,y_1)$, two axial passes suffice for information exchange through two “cross” interactions, so global context is preserved even though each attention operation is one-dimensional (He et al., 2024).

This factorization was rapidly adopted in vision backbones. Axial-DeepLab replaced the spatial $N=\prod_i D_i$0 convolution in the ResNet bottleneck with height- and width-axis attention, while maintaining the bottleneck structure and Batch Normalization (Wang et al., 2020). The same principle was then adapted to medical imaging, where high resolution and small datasets make full 2D or 3D attention especially difficult to train and deploy (Valanarasu et al., 2021, Hu et al., 2022).

2. Mathematical structure and complexity

The axis-wise attention operator retains the standard scaled dot-product form. A representative 1D formulation is

$N=\prod_i D_i$1

where $N=\prod_i D_i$2, $N=\prod_i D_i$3, and $N=\prod_i D_i$4 are computed from the current axis slice and $N=\prod_i D_i$5 denotes positional bias terms when used (He et al., 2024).

For a 2D feature map, row-wise attention treats each fixed row $N=\prod_i D_i$6 as a length-$N=\prod_i D_i$7 sequence,

$N=\prod_i D_i$8

while column-wise attention treats each fixed column $N=\prod_i D_i$9 as a length-$$X \in \mathbb{R}^{D_1 \times D_2 \times \cdots \times D_k \times d}$$0 sequence,

$$X \in \mathbb{R}^{D_1 \times D_2 \times \cdots \times D_k \times d}$$1

A common ordering is width-axis attention followed by height-axis attention, or the reverse; after the pair of passes, information has propagated globally (He et al., 2024).

The computational advantage is the defining property of axial attention. Full 2D attention over $$X \in \mathbb{R}^{D_1 \times D_2 \times \cdots \times D_k \times d}$$2 tokens forms an $$X \in \mathbb{R}^{D_1 \times D_2 \times \cdots \times D_k \times d}$$3 similarity matrix, with time $$X \in \mathbb{R}^{D_1 \times D_2 \times \cdots \times D_k \times d}$$4 and memory $$X \in \mathbb{R}^{D_1 \times D_2 \times \cdots \times D_k \times d}$$5. Axial attention decomposes this into $$X \in \mathbb{R}^{D_1 \times D_2 \times \cdots \times D_k \times d}$$6 sequences of length $$X \in \mathbb{R}^{D_1 \times D_2 \times \cdots \times D_k \times d}$$7 and $$X \in \mathbb{R}^{D_1 \times D_2 \times \cdots \times D_k \times d}$$8 sequences of length $$X \in \mathbb{R}^{D_1 \times D_2 \times \cdots \times D_k \times d}$$9, for total complexity

$a$0

with attention matrices of size $a$1 per head instead of $a$2. When $a$3, full 2D attention scales as $a$4 while axial attention scales as $a$5 (He et al., 2024, Ho et al., 2019).

The same logic extends to 3D. For a volume of size $a$6, full 3D Non-Local attention scales as $a$7, whereas 3D axial attention over height, width, and depth scales as

$a$8

On a $a$9 volume, one reported example gives $D_a$0 operations for full Non-Local attention versus $D_a$1 for 3D axial attention, a reduction of more than $D_a$2 (Al-Shabi et al., 2020).

3. Positional modeling and gated variants

Because factorized attention weakens some of the inductive bias that convolution supplies implicitly, much of the axial-attention literature focuses on positional modeling. Axial-DeepLab introduced a position-sensitive formulation in which the attention logits incorporate both query-dependent and key-dependent relative terms, and the value path receives an additive relative positional contribution:

$D_a$3

Relative embeddings are shared across heads, which keeps the parameter increase modest while preserving precise spatial information for dense prediction (Wang et al., 2020).

A separate line of work introduced gating to control these positional terms on small medical datasets. Medical Transformer replaces the positional contributions by gated versions,

$D_a$4

with learnable scalar gates that regulate reliance on learned positional biases. The stated motivation is that inaccurate positional encodings can be harmful when data are scarce, so the model should be able to shrink their contribution automatically (Valanarasu et al., 2021).

The same principle appears in other medical variants. The breast-ultrasound Axial Transformer augments axial attention with three learned relative position terms per axis, $D_a$5, $D_a$6, and $D_a$7, and then introduces gates $D_a$8, $D_a$9, and $N$0 to attenuate or amplify them:

$N$1

The paper states that this is intended to stabilize training and down-weight inaccurate positional signals on small datasets (He et al., 2024).

Position design is empirically consequential. In the parameter-efficient Hybrid Axial-Attention model, absolute sinusoidal position encoding and learned absolute position encoding are combined as adaptive position embedding, whereas learnable relative position encoding is not used in the final design because the authors found it fragile on small medical datasets. On BraTS, the reported Dice scores were $N$2 for APE, $N$3 for RPE, $N$4 for LPE alone, and $N$5 for the combined LPE+APE setting (Hu et al., 2022).

4. Architectural patterns in vision and medical imaging

Axial attention has been instantiated in several distinct architectural roles. In Axial-DeepLab, it is the primary spatial operator: two sequential axial-attention layers replace the ResNet bottleneck’s $N$6 convolution, allowing the model to operate with global or very large receptive fields while preserving the bottleneck scaffold (Wang et al., 2020). In contrast, AEWin uses axial attention as one branch of a hybrid module: attention heads are split into horizontal axial, vertical axial, and local-window groups, so local and long-range dependencies are modeled in parallel rather than sequentially (Zhang et al., 2022).

Medical architectures often combine axial attention with convolutional pyramids or decoder refinements rather than replacing the backbone wholesale. Medical Transformer uses a Local-Global strategy in which a shallow global branch with gated axial attention is paired with a deeper local patch branch; their outputs are summed before the final prediction (Valanarasu et al., 2021). Hybrid Axial-Attention places a local Conv-BN-ReLU block before axial passes and uses gated residual aggregation, while Channelized Axial Attention inserts channel-relation modeling inside the axial computation itself, producing spatially varying channel weights with negligible extra FLOPs (Hu et al., 2022, Huang et al., 2021).

Small-object segmentation has been a recurring application. CaraNet combines axial attention with reverse attention in its Axial Reverse Attention module and with a Channel-wise Feature Pyramid using dilation rates $N$7. The reverse mask is $N$8, and the axial-attended feature is reweighted as $N$9, so uncertain regions receive more refinement. The paper reports top-rank mean Dice performance on small medical objects across polyp datasets and a gain on BraTS 2018 small tumors over PraNet (Lou et al., 2023). CANet uses an Axial Attention Catching module after decoder transposed convolutions in 3D kidney parsing, where depth-, height-, and width-axis outputs are concatenated to refine boundaries; reported Dice scores on KiPA2022 are $X \in \mathbb{R}^{H \times W \times C}$0 for kidney, $X \in \mathbb{R}^{H \times W \times C}$1 for tumor, $X \in \mathbb{R}^{H \times W \times C}$2 for artery, and $X \in \mathbb{R}^{H \times W \times C}$3 for vein (Bu et al., 2022).

Three-dimensional medical segmentation has produced additional variants. GASA-UNet applies global axial self-attention separately along width, height, and depth at the U-Net bottleneck, inflates the three outputs back to 3D, adds learnable absolute positional embeddings after attention, and concatenates the result with the encoder feature. On BTCV, AMOS, and KiTS23, the reported overhead relative to nnUNet is small, from $X \in \mathbb{R}^{H \times W \times C}$4G FLOPs to $X \in \mathbb{R}^{H \times W \times C}$5G FLOPs, while Dice and NSD improve consistently, especially on KiTS23 (Sun et al., 2024).

5. Generalizations beyond 2D vision

Axial attention has also been generalized by redefining what counts as an “axis.” The following representative mappings appear across the literature.

Domain	Axes	Representative paper
3D medical volumes	height, width, depth	(Al-Shabi et al., 2020)
Video segmentation	time-height and time-width trajectories	(He et al., 2023)
OFDM neural receivers	temporal and spectral axes	(Yellapragada et al., 14 Oct 2025)
PDE operators on meshes	charted row and column axes	(Cheng et al., 9 May 2026)
Phylogenetic inference	sequence axis and pair axis	(Teska et al., 12 May 2026)

In 3D lung-nodule classification, the model applies axial attention sequentially over height, width, and depth, uses a shared embedding with learnable 3D positional encoding $X \in \mathbb{R}^{H \times W \times C}$6, and removes local filters entirely. The reported 3D Axial-Attention model reaches AUC $X \in \mathbb{R}^{H \times W \times C}$7 and Accuracy $X \in \mathbb{R}^{H \times W \times C}$8 on LIDC-IDRI, outperforming both 2D axial and prior Non-Local baselines (Al-Shabi et al., 2020).

In video segmentation, Axial-VS introduces axial-trajectory attention. Within a clip, dense features are reshaped so that attention is applied along the height axis and then the width axis, while temporal pooling follows the estimated axis-wise trajectory. The stated complexity drops from $X \in \mathbb{R}^{H \times W \times C}$9 for dense trajectory attention to $HW$0, making stacked dense tracking feasible on high-resolution video features (He et al., 2023).

In wireless communications, axial self-attention is applied to the OFDM resource grid $HW$1 by factoring attention into temporal and spectral passes. The reported complexity decreases from $HW$2 for global MHSA to $HW$3, and for $HW$4, $HW$5, $HW$6 the attention-matrix size falls from $HW$7 to $HW$8 entries, an approximately $HW$9 reduction. The corresponding receiver achieves $(x_0,y_0)$0 GFLOPs per inference versus $(x_0,y_0)$1 for a global-MHSA Transformer and $(x_0,y_0)$2 for a CNN-ResNet baseline (Yellapragada et al., 14 Oct 2025).

CATO extends axial attention to neural PDE operators by learning a continuous latent chart $(x_0,y_0)$3 and conditioning row and column axial attention on continuous RoPE in chart coordinates. The operator is paired with derivative-aware losses on solution values, gradients, and an auxiliary flux-like field. Across evaluated datasets, the paper reports an average improvement of approximately $(x_0,y_0)$4 over the strongest competing baselines while reducing the number of parameters by $(x_0,y_0)$5 (Cheng et al., 9 May 2026).

Tropical Axial Attention changes the algebra rather than the domain. In phylogenetic tree inference, softmax dot-product attention is replaced by max-plus projections and tropical distances, and axial factorization alternates between the sequence axis and the unordered leaf-pair axis. This produces a piecewise-linear attention map aligned with dynamic-programming formulations and tropical geometry (Teska et al., 12 May 2026).

A distinct but terminologically related case is axial-centric cross-plane attention for 3D CT classification. Here “axial” refers to the anatomical axial plane rather than axis-wise factorization: axial tokens serve as queries, while coronal and sagittal tokens provide keys and values. The model is therefore clinically axial-centric rather than an instance of classical intra-tensor axial attention (Park et al., 25 Feb 2026).

6. Empirical behavior, trade-offs, and recurrent misconceptions

A common misconception is that axial attention is only a local approximation to full attention. In fact, many implementations are global along the active axis: Axial-DeepLab uses full-line attention when memory permits, GASA-UNet attends over the entire depth, height, and width lines of a 3D feature map, and the original Axial Transformer emphasized that axial composition preserves full expressiveness over joint distributions while reducing resource use (Ho et al., 2019, Wang et al., 2020, Sun et al., 2024).

A second misconception is that factorization necessarily destroys global context. Two-pass arguments recur across the literature: in 2D, width-then-height or height-then-width attention suffices for any pair of positions to influence each other, and in 3D the same principle extends by chaining three axis-wise operations (He et al., 2024, Al-Shabi et al., 2020). This suggests that the main compromise of axial attention is not reach but inductive structure: interactions are mediated through axis-aligned compositions rather than a single unrestricted affinity matrix.

The principal trade-off is residual cubic cost in the side length rather than full quadratic cost in the number of tokens. Even in 2D, $(x_0,y_0)$6 can remain heavy for very high resolutions; the breast-ultrasound study explicitly notes that windowed axial attention or downsampled tokens may be necessary, and that anisotropy can cause one axis to dominate compute (He et al., 2024). Window-augmented designs such as AEWin address this by parallelizing axial stripes with local windows and allocating more heads to the window branch than to the axial branches (Zhang et al., 2022).

Another recurrent limitation is sensitivity to positional strategy on small datasets. Medical Transformer and the breast Axial Transformer both introduce learnable gates precisely because positional terms can be noisy or overfit in data-scarce regimes (Valanarasu et al., 2021, He et al., 2024). Hybrid Axial-Attention reports an especially stark failure case for LPE alone, and Channelized Axial Attention argues that separating spatial and channel attentions can create conflicts that are reduced when channel relations are inserted inside the axial computation (Hu et al., 2022, Huang et al., 2021).

Finally, axial attention does not guarantee superior boundary quality simply by providing global context. On BUSI breast ultrasound, the reported Axial Transformer achieves the highest precision, $(x_0,y_0)$7, and the highest F1, $(x_0,y_0)$8, but its IoU, $(x_0,y_0)$9, and recall, $(x_1,y_1)$0, are lower than those of Res-UNet, indicating missed lesions and incomplete boundaries despite improved small-lesion localization and fewer false positives (He et al., 2024). This pattern has motivated hybrids that combine axial attention with reverse attention, channel-aware gating, boundary-sensitive decoding, or local windows rather than treating axial factorization as a complete replacement for convolutional inductive bias (Lou et al., 2023, Zhang et al., 2022).