---
title: Class Attention Block (CAB)
url: https://www.emergentmind.com/topics/class-attention-block-cab
type: topic
---

# Class Attention Block (CAB)

A Class Attention Block (CAB) is a neural network module that operates on either channel or class-token dimensions and is designed to improve joint feature aggregation and attention-based reweighting. CAB modules are leveraged in diverse architectures, including U-shaped models for dense prediction and Vision Transformers for class-token-based representation learning. The two primary paradigms are: Channel Attention Bridge (CAB) for multi-stage channel fusion in encoder-decoder models [2211.01784] and Class-Attention Block (CAB) for global class-token refinement in Transformer models [2211.12292]. Both variants focus on extracting and combining the most informative features across input channels or tokens and modulating the information flow to subsequent network stages.

## 1. Channel Attention Bridge (CAB) in Encoder-Decoder Architectures

The Channel Attention Bridge (CAB) is designed for efficient multi-stage feature fusion in U-shaped architectures, such as MALUNet for medical image segmentation. In a six-stage encoder-decoder model, CAB sits in the skip-connection path between encoder and decoder, aggregating encoded features from stages 1 to 5, which have increasing channel sizes $\{8, 16, 24, 32, 48\}$ [2211.01784].

Given per-stage feature maps $t_i \in \mathbb{R}^{C_i \times H_i \times W_i}$, CAB executes the following steps:
1. **Channel Descriptor Extraction**: Each $t_i$ is globally average pooled to $t_i' \in \mathbb{R}^{C_i}$.
2. **Multi-Stage Channel Fusion**: All channel descriptors are concatenated, $T = \mathrm{Concat}(t_1', ..., t_5') \in \mathbb{R}^{128}$.
3. **Local Correlation via 1D Convolution**: $T$ is passed through a 1D convolution, yielding $T' = \mathrm{Conv1D}(T)$.
4. **Stage-Wise Global Gating**: For each stage $i$, a stage-specific fully connected layer applies, yielding attention vector $Att_i = \sigma(\mathrm{FC}_i(T')) \in (0,1)^{C_i}$, with $\sigma$ the element-wise sigmoid.
5. **Channel-Wise Reweighting and Residual Addition**: The attention vector is broadcast and used for channel-wise scaling and residual addition: $Out_i = t_i + t_i \odot Att_i$.

This mechanism ensures each decoder block receives an adaptively fused and selectively reweighted collection of encoder features, improving the expressivity-to-parameter ratio in compute-constrained settings. CAB employs a lightweight 1D convolution (e.g., kernel size $K=3$, output dimension $D=128$, zero dilation) and per-stage fully connected layers ($W_{fc,i} \in \mathbb{R}^{C_i \times D}$) without normalization layers. All parameters are initialized using standard procedures (He initialization for Conv1D, Xavier for FC).

## 2. Mathematical Formulation of CAB in MALUNet

The CAB module in MALUNet [2211.01784] is defined mathematically by equations (6)-(10):

\[
\begin{align*}
(6) \quad & t_i' = \mathrm{GAP}(t_i) \\
(7) \quad & T = \mathrm{Concat}(t_1', t_2', ..., t_{s-1}') \\
(8) \quad & T' = \mathrm{Conv1D}(T) \\
(9) \quad & Att_i = \sigma(\mathrm{FC}_i(T')) \\
(10) \quad & Out_i = t_i + t_i \odot Att_i
\end{align*}
\]

where $s=6$ (number of encoder stages) and $i=1 \dots 5$. $\mathrm{GAP}$ is global average pooling; $\mathrm{Conv1D}$ and $\mathrm{FC}_i$ are learnable; $\sigma$ is sigmoid.

Local context fusion occurs via Conv1D, enabling short-range channel interaction among stages, while FC layers provide stage-specific global attention. The design avoids normalization and maintains low parameter cost (on the order of a few thousand parameters for Conv1D and $\sum C_i \times D$ for FCs).

## 3. Standard Class-Attention Block in Vision Transformers

The Class-Attention Block (CAB) in transformer-based architectures (e.g., CaiT) incorporates a learnable class token $\theta \in \mathbb{R}^D$ alongside patch tokens $b \in \mathbb{R}^{N \times D}$, forming an augmented sequence $p = [\theta; b]$ [2211.12292].

The standard attention mechanism proceeds as follows:
- Linear projections for query ($Q = W_q \theta$), key ($K = W_k p$), and value ($V = W_v p$).
- Attention weights: $A = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d'}}\right)$, where $d' = D/h$ and $h$ is the number of heads.
- Attended class-token update: $O = W_o (A V)$.
- MLP refinement and residuals: class-token is updated through an MLP and residual addition: $f = v + x'$, where $x' = O + \theta$; all projections and MLPs are learnable.

CAB thereby aggregates global image representations via direct attention between class and patch tokens, which are pivotal for classification tasks and task-specific transfer in continual learning.

## 4. Gated Class-Attention Block for Continual Learning

The Gated Class-Attention Block (GCAB) extends the standard transformer CAB to address catastrophic forgetting in exemplar-free continual learning [2211.12292]. GCAB introduces task-specific soft masks that gate both forward activations and parameter gradients.

Key mechanisms:
1. **Soft Mask Parametrization**: For each task $t$, the mask $m^t \in (0,1)^D$ is defined by $m^t = \sigma(s A t)$, where $A \in \mathbb{R}^{D \times T}$ is a learnable embedding, $s$ a scaling factor, and $t$ a one-hot task selector.
2. **Gated Attention Application**: All Q/K/V projections and MLP activations are modulated. For example, $Q^t = W_q(\theta \odot m^t_i)$, $K^t = W_k(p \odot m^t_i)$, and so on.
3. **Sparsity Regularization**: The cumulative mask $m^{<t}$ is updated by $m^{<t} = \max(m^t, m^{<t-1})$. A loss term encourages sparsity in new masks except at already-allocated capacity.
4. **Gradient Masking/Weight Protection**: Gradients are masked during backpropagation to preserve weights used by previous tasks: $W_q \gets W_q - \eta (M^{<t}_q \odot \partial \mathcal{L} / \partial W_q)$, where $M^{<t}_{q,kl} = 1 - \min(m^{<t}_{i,k}, m^{<t}_{QK,l})$.

GCAB operates at the final transformer block, enabling task-specific activation patterns and selective plasticity. At inference, all stored task-specific masks are sequentially applied, and outputs concatenated, obviating the need for task-ID during test time. This approach distinguishes GCAB from hard parameter-isolation approaches, as it leverages shared weights and soft task gating.

## 5. Hyper-parameters and Implementation Details

CAB implementations are distinguished by carefully selected hyper-parameters:
- **For MALUNet CAB** [2211.01784]:
    - Number of encoder stages fused: $s-1=5$
    - Conv1D kernel size: typically $3$ or $5$ (padding $\lfloor K/2 \rfloor$)
    - Conv1D input/output dims: input $128$, output $128$
    - No dilation, batch normalization, or layer normalization
    - Initialization: He for Conv1D, Xavier for FC
- **For GCAB** [2211.12292]:
    - Mask embedding $A \in \mathbb{R}^{D \times T}$
    - Mask scaling $s$ dynamically set during training
    - Most soft masks are shared across Q/K/V and MLP layers, minimizing additional parameters
    - Regularization parameter $\lambda_{GCAB}$ for sparsity

Both designs avoid unnecessary complexity and maintain computational efficiency, while yielding pronounced performance gains in their respective application domains.

## 6. Applications and Empirical Impact

CAB modules have demonstrated performance increases aligned with their efficient attention-driven fusion or selection:
- **MALUNet** achieves improvement over UNet on skin lesion segmentation—specifically, mIoU and DSC metrics increased by $2.39\%$ and $1.49\%$ respectively. This result is achieved alongside dramatic reductions in parameter count ($44\times$) and computational cost ($166\times$), positioning CAB as an essential primitive in lightweight segmentation models [2211.01784].
- **GCAB** enables exemplar-free class incremental training for Vision Transformers, achieving competitive results on datasets such as CIFAR-100, Tiny-ImageNet, and ImageNet100 without rehearsal. The gating mechanism facilitates plasticity towards new tasks while constraining catastrophic forgetting, with no requirement for test-time task identification and modest inference cost increase (limited to the last transformer block) [2211.12292].

These empirical results validate the architectural advantages of CAB and its variants for both dense prediction and continual learning scenarios.

## 7. Comparative and Methodological Context

CAB in encoder-decoder architectures and class-attention in transformers address different but related challenges—feature fusion and class-wise representation learning. Unlike generic additive skip connections, the Channel Attention Bridge applies coordinated per-stage attention, guided by both local (Conv1D) and global (FC) context. Transformer-based CAB leverages a dedicated class token and token-class interaction, and in GCAB, further incorporates task-specific gating for continual learning.

GCAB distinguishes itself from classical parameter-isolation techniques (PackNet, Piggyback, HAT) by employing shared weights and runtime-applied soft masks, which are learned and regularized to enforce sparsity and weight protection, respectively [2211.12292]. This design enables flexible, scalable continual learning without explicit hard routing or known task identifiers.

By integrating these modules, networks achieve targeted feature selection, inter-stage knowledge distillation, and—when appropriately extended—robust task transfer without sacrificing architectural efficiency.

Source: https://www.emergentmind.com/topics/class-attention-block-cab