Channelized Axial Attention (CAA)

• Channelized Axial Attention (CAA) is a unified mechanism that integrates spatial and channel attention within an efficient axial self-attention framework.
• It addresses conflicts in independent dual-attention schemes by applying spatially-varying channel modulation and grouped vectorization to reduce memory usage.
• Empirical evaluations on benchmarks like PASCAL Context and Cityscapes show that CAA consistently improves mIoU performance with minimal computational cost.

Channelized Axial Attention (CAA) integrates spatial and channel attention mechanisms into a unified, computationally efficient framework for semantic segmentation. Originally motivated by limitations in existing dual attention schemes, CAA applies spatially-varying channel attention within the axial self-attention paradigm, resolving conflicts arising from independent treatment of spatial and channel dimensions while maintaining tractable memory and compute requirements. Empirical results on standard semantic segmentation benchmarks consistently demonstrate CAA's effectiveness compared to state-of-the-art models, with minimal overhead due to innovations in both attention formulation and vectorization strategy (Huang et al., 2021).

1. Motivation and Background

Traditional approaches for semantic segmentation often deploy dual-attention mechanisms, in which spatial attention (identifying "where" in the image to focus) and channel attention (identifying "which" feature channels are important) are calculated independently—either in parallel or sequentially—followed by their fusion through addition. This independent computation can cause the two attention mechanisms to make conflicting decisions. For instance, spatial attention may highlight a small object, but if channel attention (computed globally) suppresses its feature channels, the representation is degraded.

Applying full 2D self-attention to an $H \times W$ feature map incurs a computational cost of $O((HW)^2 \cdot C)$ and analogous prohibitive memory requirements. Axial attention factorizes 2D attention into two sequential 1D attentions—along the height and width axes—reducing the complexity to $O(HW^2C + H^2WC)$, a tractable alternative for typical image sizes. CAA further addresses the interplay and integration of channel and spatial dependencies, providing a unified, lightweight mechanism that mitigates the shortcomings of disjointed dual-attention schemes.

2. Mathematical Formulation

Let $X \in \mathbb{R}^{C \times H \times W}$ denote the input feature map. CAA operates by integrating channelized modulation within the axial decomposition of spatial self-attention.

2.1 Spatial Axial Attention Decomposition

• Full spatial attention: For each location $(i,j)$,

$$Y_{i,j} = \sum_{m=1}^H \sum_{n=1}^W \text{Softmax}_{m,n}\left(Q_{i,j}^T K_{m,n}\right) V_{m,n}$$

where $Q, K, V$ are linear projections of $X$.

• Axial reduction: The dot product is decomposed into two 1D steps:
  • Column (height) attention

  $$A_\text{col}(i,j; m,j) = \text{Softmax}_m\big( \theta(X_{i,j})^T \theta(X_{m,j}) \big)$$

  The intermediate feature is

  $$\alpha_{i,j,m,n} = A_\text{col}(i,j; m,j) \cdot g(X_{m,n})$$ - Row (width) attention

  $O((HW)^2 \cdot C)$0

  The output is aggregated as

  $O((HW)^2 \cdot C)$1

  $O((HW)^2 \cdot C)$2

2.2 Channelization Between Axial Passes

CAA introduces channel relation modules between the two axial steps:

• Column-wise channel attention:

$O((HW)^2 \cdot C)$3

where $O((HW)^2 \cdot C)$4 applies mean pooling over $O((HW)^2 \cdot C)$5 and $O((HW)^2 \cdot C)$6 while retaining the $O((HW)^2 \cdot C)$7 index, $O((HW)^2 \cdot C)$8 are learned weights, and $O((HW)^2 \cdot C)$9 denotes the sigmoid function.

• Row-wise channel attention:

$O(HW^2C + H^2WC)$0

where $O(HW^2C + H^2WC)$1 averages over $O(HW^2C + H^2WC)$2 and $O(HW^2C + H^2WC)$3 but retains $O(HW^2C + H^2WC)$4.

• Complete CAA output:

$O(HW^2C + H^2WC)$5

This mechanism ensures that channel attention is independently optimized per spatial location, shaping how information is aggregated in both directions before final output.

3. Efficient Computation via Grouped Vectorization

The two-stage spatial attention inherently generates large intermediate tensors, leading to excessive memory usage for even moderate spatial resolutions. A naive, fully vectorized implementation is memory-prohibitive, while loop-based implementations are slow. CAA introduces grouped vectorization:

• The spatial dimension is partitioned into $O(HW^2C + H^2WC)$6 groups (typically 4).
• Attention and feature tensors are reshaped for groupwise processing.
• Each group is processed in parallel, fully vectorized within the group.
• Outputs from all groups are concatenated and unpadded as needed.

This results in a significant reduction in peak memory usage by a factor of $O(HW^2C + H^2WC)$7, with minimal computational overhead and no significant impact on inference speed. Computational complexity remains $O(HW^2C + H^2WC)$8 per forward pass, and with $O(HW^2C + H^2WC)$9 groups, peak intermediate memory drops by approximately $X \in \mathbb{R}^{C \times H \times W}$0.

Implementation Variant	Memory Usage	Inference Throughput
Full vectorization	$X \in \mathbb{R}^{C \times H \times W}$1	High
Grouped vectorization	$X \in \mathbb{R}^{C \times H \times W}$2	High
Loop over rows/columns	Minimal	Low

Grouped vectorization, therefore, marries the computational efficiency of vectorized operations with memory profiles suitable for deployment on contemporary GPUs.

4. Architectural Integration and Training

CAA modules are incorporated within standard ResNet-101 backbones for semantic segmentation. After the final ResNet conv4_x block (corresponding to a feature stride of 16), two or three CAA modules are placed in series. Alternatively, an architecture may insert one CAA block at the output of both conv3_x (stride 8) and conv4_x (stride 16), fusing their outputs via bilinear upsampling and summation as in DeepLabV3+.

Key network and training characteristics:

• Backbone: ResNet-101 pretrained on ImageNet.
• Feature stride: 16 for ablation, up to 8 for benchmark runs.
• Optimizer: SGD with momentum 0.9, weight decay $X \in \mathbb{R}^{C \times H \times W}$3.
• Learning rate schedule: Poly decay, initial learning rate $X \in \mathbb{R}^{C \times H \times W}$4.
• Batch size: 16 (8 on Cityscapes), with synchronized BatchNorm.
• Data augmentation: Random horizontal flip, scale in $X \in \mathbb{R}^{C \times H \times W}$5, random crop.
• Training iterations: 40K (PASCAL, COCO-Stuff), 80K (Cityscapes).
• Additional computational overhead: Each CAA block adds $X \in \mathbb{R}^{C \times H \times W}$60.03 GFLOPs compared to $X \in \mathbb{R}^{C \times H \times W}$78.8 GFLOPs for the full ResNet-101 backbone.

5. Empirical Performance and Comparative Analysis

CAA has been evaluated across three major semantic segmentation datasets using mean Intersection-over-Union (mIoU):

Dataset	DANet (%)	SPG R (%)	EMANet (%)	CAA (%)
PASCAL Context (OS=8)	52.6	52.8	53.1	55.0
COCO-Stuff 10K (OS=8)	39.7	—	39.9	41.2
Cityscapes (OS=8)	81.5	—	81.8	82.6

Ablation studies on PASCAL Context (OS=16) reveal:

• No channelization: 50.27%
• 1 FC layer, 128 hidden: 50.75%
• 3 FC layers, 128: 50.85%
• 5 FC layers, 128 (best): 51.06%
• Applying channelization to full self-attention: +0.67% mIoU over vanilla self-attention.
• “Axial+SE (sequential)” only achieves marginal improvement over axial alone.

Grouped vectorization with $X \in \mathbb{R}^{C \times H \times W}$8 almost matches full vectorization throughput using about 1/4 the peak memory.

6. Strengths, Limitations, and Applicability

Strengths

• Seamless integration of spatial and channel attention, enabling locally adaptive channel weighting without either "overriding" the other.
• Demonstrable, consistent +1–2% mIoU improvements over established dual-attention and transformer-based segmentation backbones across benchmarks.
• The grouped vectorization method allows practical training and inference on large feature maps using modern GPU hardware.
• Minimal computational increase (less than 0.03 GFLOPs per CAA block) relative to overall network complexity.

Limitations

• The relative overhead of the inner MLPs may be counterproductive when feature maps are very small (e.g., $X \in \mathbb{R}^{C \times H \times W}$9).
• If the domain exhibits weak channel inter-dependencies or highly decorrelated feature channels, the benefits of CAA may diminish.

A plausible implication is that CAA’s primary gains manifest in high-resolution feature maps and domains where both spatial and channel dependencies are significant.

7. Implementation and Availability

CAA is implemented on top of ResNet-101 with an output stride of 16 for ablation studies and stride 8 for final comparisons. The learning regime employs standardized SGD hyperparameters, synchronized BatchNorm, and diverse data augmentation techniques suitable for semantic segmentation. The total model size is approximately 50 MB per checkpoint. Model code and trained weights are to be released upon publication.

Reference: "Channelized Axial Attention for Semantic Segmentation -- Considering Channel Relation within Spatial Attention for Semantic Segmentation" (Huang et al., 2021).