3D Axial Attention Module

• 3D Axial Attention Modules are neural components that decompose dense 3D self-attention into sequential axis-wise operations, capturing global context with reduced computational overhead.
• They implement separate 1D self-attention along depth, height, and width, and are integrated in architectures like U-Net for enhanced boundary refinement and segmentation accuracy.
• Empirical results indicate a 10–20× speedup and 8–12× memory reduction, proving their value in tasks such as medical image segmentation and volumetric classification.

Three-dimensional (3D) Axial Attention Modules are neural network mechanisms that enable efficient long-range contextual modeling in volumetric data by factorizing global self-attention into computationally tractable, axis-wise operations. Originating as a solution to the prohibitive memory and computation footprint of dense 3D self-attention, these modules sequentially apply 1D self-attention along separate spatial axes—depth, height, and width—facilitating information exchange across the volume with greatly reduced cost. Multiple architectural instantiations exist, including variants integrated into U-Net decoders for boundary refinement, cross-plane transfer architectures, and transformer-based medical image pipelines. This class of module is now foundational in a wide spectrum of medical imaging networks and volumetric transformers.

1. Mathematical Foundation and Axial Decomposition

The canonical 3D Axial Attention mechanism replaces direct attention over the flattened $N = D \times H \times W$ voxel space with three consecutive passes, each restricted to a single axis. For an input tensor $$X \in \mathbb{R}^{B \times C \times D \times H \times W}$$, separate $Q$, $K$, $V$ projections are computed for each axis via parameter-shared or axis-specific linear layers (often implemented as $1 \times 1 \times 1$ convolutions). Each attention pass reshapes $X$ so the non-attended axes are subsumed into the batch dimension: for example, height-axis attention views $F$ as $F_H \in \mathbb{R}^{(B D W) \times H \times C'}$. Per-head queries $Q_a$, keys $$X \in \mathbb{R}^{B \times C \times D \times H \times W}$$0, and values $$X \in \mathbb{R}^{B \times C \times D \times H \times W}$$1 are then computed, followed by scaled dot-product attention:

$$X \in \mathbb{R}^{B \times C \times D \times H \times W}$$2

where $$X \in \mathbb{R}^{B \times C \times D \times H \times W}$$3 and $$X \in \mathbb{R}^{B \times C \times D \times H \times W}$$4. Each axis-specific output is reshaped to the original volume and concatenated along the channel axis, followed by a projection back to $$X \in \mathbb{R}^{B \times C \times D \times H \times W}$$5 channels:

$$X \in \mathbb{R}^{B \times C \times D \times H \times W}$$6

Optionally, the block includes learned positional encoding $$X \in \mathbb{R}^{B \times C \times D \times H \times W}$$7, layer normalization, and a position-wise MLP with nonlinearity (commonly ReLU). Residual connections are typically employed, either after attention or after MLP sublayers. This design is implemented in models such as the Axial Attention Catching (AAC) block in CANet (Bu et al., 2022), GASA-UNet's global axial block (Sun et al., 2024), and more generally in multidimensional transformers (Ho et al., 2019).

2. Module Architecture, I/O Shapes, Normalization, and Activation

The essential module expects an input tensor $$X \in \mathbb{R}^{B \times C \times D \times H \times W}$$8 and produces output in the same spatial dimensions; channel size may change depending on projection. The number of attention heads $$X \in \mathbb{R}^{B \times C \times D \times H \times W}$$9 varies by implementation (e.g., $Q$0 in CANet AAC; $Q$1 in GASA-UNet), with per-head dimension $Q$2. Intermediate $Q$3 tensors have shapes $Q$4, split along the channel axis. Layer normalization is often applied per voxel and channel before each attention and MLP. Nonlinear activations include ReLU and occasionally GELU following attention or within MLP sublayers. The GASA-UNet global axial attention block omits LayerNorm and residual-MLP sublayers, relying on final channel-wise concatenation and skip connection for signal propagation (Sun et al., 2024).

3. Integration Strategies in Network Pipelines

Axial attention modules have been integrated at different positions in volumetric network architectures, with consequences for representational power and computational overhead:

• Decoder Insertion: In CANet, AAC is placed directly after each decoder upsampling/transposed-conv and before merging with encoder skip features, to refine edge information early in the upsampling path (Bu et al., 2022).
• Bottleneck Augmentation: In GASA-UNet, a global axial block sits at the bottleneck between encoder and decoder, concatenating axial-attended features with encoder features prior to decoder stages (Sun et al., 2024).
• Full Replacement: Some works, such as 3D Axial-Attention for lung nodule classification (Al-Shabi et al., 2020), replace all convolutional encoder blocks with axial attention blocks, resulting in fully attention-based feature extraction.
• Cross-Plane Conditioning: In Axial-Centric Cross-Plane Attention architectures, query-key-value allocation is explicitly assigned such that only axial-plane tokens query, with coronal/sagittal keys and values, reinforcing the clinical interpretive asymmetry across imaging planes (Park et al., 25 Feb 2026).

Integration points are chosen to maximally exploit long-range context where it is most beneficial, such as object boundaries or bottleneck compression, and to balance overhead.

4. Computational Complexity and Efficiency

The primary advantage of axial attention is the reduction in computational and memory complexity. Full self-attention over an $Q$5-voxel volume requires $Q$6 operations due to the size of the attention matrix. Axial factorization reduces this to:

$Q$7

For cubic grids $Q$8, this becomes $Q$9. Empirical measurements report 10–20$K$0 speedup and 8–12$K$1 memory reduction at $K$2 voxels (Bu et al., 2022, Al-Shabi et al., 2020). This efficiency is critical for large-scale medical volumes and enables architectural choices (e.g., stacking multiple blocks, higher channel counts) that are infeasible for dense 3D attention.

Optimizations include fusing $K$3, $K$4, $K$5 projections into a single convolution, shared projection weights across axes, FP16 training, and parameter sharing. Certain designs (e.g., CANet AAC) use half-precision arithmetic and axis-wise parameter sharing to minimize overhead (Bu et al., 2022).

5. Positional Encoding Strategies

Proper spatial context is necessary to avoid permutation invariance in non-local attention. 3D positional encoding mechanisms include:

• Additive Absolute Encoding: Learnable tensors of shape $K$6 are added to features before attention (CANet AAC, GASA, 3D Axial-Attention) (Bu et al., 2022, Sun et al., 2024, Al-Shabi et al., 2020).
• Per-Axis Vector Decomposition: Positional embeddings constructed from separate learnable vectors $K$7, $K$8, $K$9, combined at each voxel $V$0 as $V$1 (Al-Shabi et al., 2020).
• Relative Position Biases: Optionally, per-axis relative position embeddings are introduced as additive biases in the attention logits (Ho et al., 2019).

Empirical ablations indicate that placement and type of positional embedding can affect segmentation metrics; e.g., post-attention positional embedding yielded superior results on BTCV, AMOS, and KiTS23 (Sun et al., 2024).

6. Empirical Benefits and Practical Performance

A consistent outcome across studies is that 3D axial attention modules yield measurable improvements in volumetric segmentation and classification, particularly in boundary delineation and small structure sensitivity:

• Segmentation Accuracy: CANet AAC boosted Dice for kidney, tumor, artery, and vein by 0.3–0.7% over channel-extended nn-UNet; refinement of vessel boundaries was confirmed by reduction in Hausdorff distance (by up to 8 mm) and average surface distance (by up to 0.8 mm) (Bu et al., 2022).
• Classification Metrics: 3D Axial-Attention achieved higher AUC and accuracy over 2D axial and non-local baselines on LIDC-IDRI (AUC 96.17% vs. 94.74–95.62%), with robust gains across all confusion matrix metrics (Al-Shabi et al., 2020).
• Comparative Ablations: GASA-UNet's axial block consistently improved Dice and Normalized Surface Dice (NSD) over standard nn-UNet, CBAM, and alternative axial attention variants. The impact was most pronounced for fuzzy and small-volume structures (Sun et al., 2024).
• Design Sensitivity: Axial-centric cross-plane attention models validate the importance of directional QKV assignment and positional encoding—reversing the allocation degrades classification accuracy by up to 6.2% on organ tasks (Park et al., 25 Feb 2026).

These findings indicate axial attention's superior ability to aggregate global context and sharpen fine structural details, while remaining computationally feasible in large 3D volumes.

7. Limitations, Challenges, and Extensions

While axial attention offers reductions in resource requirements and empirical benefits, certain challenges persist:

• Edge Fuzziness: Tissue boundaries with weak contrasts remain difficult; boundary losses or edge-aware objective functions may further enhance delineation (Sun et al., 2024).
• Limited Penetration: Single-location GASA insertion does not propagate global context throughout the entire network; multi-stage insertion is a plausible extension (Sun et al., 2024).
• Positional Encoding Variants: Most designs employ absolute learned embeddings; exploration of rotary or relative embeddings is largely unaddressed.
• Annotation Variability: Segmentation accuracy is upper-bounded by ground-truth labeling consistency, especially at ambiguous boundaries.

Potential directions include combining GASA with boundary-aware losses, inserting blocks at multiple stages, unsupervised or semi-supervised pretraining of attention parameters, and generalizing the architecture to structure-sensitive 3D tasks such as registration or anomaly detection (Sun et al., 2024). A plausible implication is that cross-modal volumetric transformers may leverage hybrid axial/global blocks for improved data efficiency as medical imaging datasets continue to expand.

---

Key References:

• Axial Attention Catching in CANet (Bu et al., 2022)
• GASA-UNet Global Axial Attention Block (Sun et al., 2024)
• Axial Attention in Multidimensional Transformers (Ho et al., 2019)
• 3D Axial-Centric Cross-Plane Attention (Park et al., 25 Feb 2026)
• 3D Axial-Attention for Lung Nodule Classification (Al-Shabi et al., 2020)