Source-linked AI summary
Expectation-Maximization Attention Networks for Semantic Segmentation
Xia Li, Zhisheng Zhong, Jianlong Wu, Yibo Yang, Zhouchen Lin, Hong Liu
TL;DR
Self-attention captures long-range relations but computes attention against all positions, making it costly. The paper reformulates attention as EM iterations over compact bases, yielding a lightweight EMA module; experiments report superiority across three segmentation benchmarks, including 39.9% mIoU on COCO Stuff.
Problem
Self-attention captures long-range relations but uses all positions as bases, causing heavy computation and memory costs.
Method
EMA reformulates self-attention as EM iterations that estimate compact bases and reconstruct a low-rank representation, with maintenance and normalization for training stability.
Results
Extensive experiments on PASCAL VOC, PASCAL Context, and COCO Stuff demonstrate superiority over other state-of-the-art methods.
Takeaways & Limitations
EMA is a lightweight module that can be inserted into existing CNNs with little overhead while producing representations robust to input variance.
Takeaways & Limitations
Training can collapse when recurrent EM updates become unstable, constraining basis maintenance and normalization.
Abstract
from arXiv · showhide
Self-attention mechanism has been widely used for various tasks. It is designed to compute the representation of each position by a weighted sum of the features at all positions. Thus, it can capture long-range relations for computer vision tasks. However, it is computationally consuming. Since the attention maps are computed w.r.t all other positions. In this paper, we formulate the attention mechanism into an expectation-maximization manner and iteratively estimate a much more compact set of bases upon which the attention maps are computed. By a weighted summation upon these bases, the resulting representation is low-rank and deprecates noisy information from the input. The proposed Expectation-Maximization Attention (EMA) module is robust to the variance of input and is also friendly in memory and computation. Moreover, we set up the bases maintenance and normalization methods to stabilize its training procedure. We conduct extensive experiments on popular semantic segmentation benchmarks including PASCAL VOC, PASCAL Context and COCO Stuff, on which we set new records.
1. Introduction
Semantic segmentation requires assigning a category to every pixel while resolving confusing appearances and long-range context. The paper proposes EMA, an EM-based attention module that learns compact bases, reduces complexity, and is evaluated across three challenging datasets.
- Motivation: Semantic segmentation assigns a semantic category to each image pixel and must distinguish visually confusing categories across varying object appearances.The task is important for autonomous driving, image editing, and robot sensing.
- Motivation: FCN-based methods are limited by local receptive fields, motivating multi-scale context fusion and encoder-decoder structures for broader context and detail.
- Method: EMA reformulates self-attention as EM iterations that learn a compact basis set and largely reduce computational complexity.The paper presents this as the first introduction of EM iterations into attention.
- Method: EMA is implemented as a lightweight neural-network module with bases-maintenance and normalization methods intended to stabilize training.The module uses common operators and can be embedded into existing neural networks.
- Experiments: Extensive experiments on PASCAL VOC, PASCAL Context, and COCO Stuff demonstrate superiority over other state-of-the-art methods.
2. Related Works
Related segmentation methods enhance contextual aggregation or output detail through multi-scale modules, large receptive fields, and encoder-decoder designs. The paper builds on attention by computing attention maps iteratively through the EM algorithm.
- Context aggregation: ASPP, pyramid pooling, large-kernel convolution, and related methods enhance contextual aggregation or capture long-range information.
- Detailed prediction: U-Net-based variants combine high-level and mid-level features to produce more detailed segmentation outputs, especially around object boundaries.
- Attention-based approach: The proposed approach is motivated by attention and computes attention maps iteratively in the manner of the EM algorithm.
3. Preliminaries
The preliminaries introduce EM, GMM, and non-local attention as related formulations. They establish the contrast between compact manually selected bases in GMM and the data-dependent, computationally heavy bases of non-local attention.
- Expectation-Maximization Algorithm: EM maximizes likelihood for latent-variable models by alternating expectation and maximization steps until convergence.The E step estimates latent-variable posteriors and the M step updates parameters by maximizing the expected complete-data likelihood.
- Gaussian Mixture Model: GMM models data as a linear superposition of Gaussian bases whose means and covariances are estimated through EM.Responsibilities quantify how much each basis accounts for an observation.
- Non-local: The non-local module functions like self-attention, aggregating transformed features with a kernel function and normalization factor.The feature vector at location i interacts with features at other locations.
- Non-local: Unlike GMM, which typically uses K ≪ N manually selected bases, non-local attention uses the data themselves as bases, giving K = N.
- Non-local: Using all data points as non-local bases makes the representation over-complete and increases computation overhead and memory cost.
4. Expectation-Maximization Attention
EMA reformulates self-attention as alternating expectation and maximization operations that learn a compact basis set instead of using all pixels as bases. Re-estimation from these bases produces a low-rank representation while reducing complexity.
- EMA formulation: EMA uses EM iterations to find a compact basis set rather than selecting all input pixels as reconstruction bases.Responsibility estimation computes latent assignments, while likelihood maximization updates the bases.
- EMA formulation: The AE and AM steps alternate for a fixed number of iterations, after which AR reconstructs and outputs the input feature map.AE estimates Z, AM updates µ, and AR reconstructs X using the converged quantities.
- Representation: The reconstructed representation is low-rank because it is constructed from K bases with K ≪ N, placing it in a subspace of the input.The paper reports that this compactness removes unnecessary noise and makes pixel classification more tractable.
- Efficiency: EMA reduces attention complexity from O(N^2) to O(NKT), and with T treated as a small constant, the complexity becomes O(NK).The experiments report promising results with only three AE and AM iterations.
- Interpretation: In the λ →∞ limit, the alternating AE and AM updates become one-hot assignments and average assigned pixels, corresponding to soft K-means clustering.This connects EMA’s soft responsibility assignments with the hard assignments of K-means.
5. EMA Unit
EMAU embeds EMA into a lightweight residual-like neural-network unit for semantic segmentation. It uses basis maintenance and normalization strategies to address training instability and reduce computation relative to 3 × 3 convolutions.
- Module design: EMAU embeds EMA into a neural-network module for semantic segmentation using common operators and lightweight computation.The module is designed to be easily embedded into existing neural networks.
- Efficiency: Each AE, AM, and AR step has O(NKC) complexity, and the full EMAU requires about one-third the FLOPs of a same-width 3 × 3 convolution module.EMA maintains KC parameters.
- Bases maintenance: EM convergence does not guarantee a global maximum, making the initial basis values important for the resulting solution.The paper therefore treats basis initialization and maintenance as key design issues.
- Training stability: Back-propagation updates of the initial bases can suffer vanishing or exploding gradients through the unrolled iterations, potentially collapsing EMAU training.The paper uses moving averaging instead for training updates and keeps the initial bases fixed during inference.
- Training stability: Stable internal basis updates remain constrained because moving averaging requires the mini-batch average of the final bases to remain close to the initial bases.Normalization is introduced to constrain the value range while preserving basis directions.
- Bases normalization: L2Norm outperforms layer normalization and no normalization when T ≥3, while layer normalization performs worse than no normalization.L2Norm preserves each basis direction by dividing it by its length.
6. Experiments
Experiments evaluate EMANet on PASCAL VOC, PASCAL Context, and COCO Stuff using specified training configurations. The default EMANet setup uses K = 64, λ = 1, and T = 3.
- Experiments cover PASCAL VOC, PASCAL Context, and COCO Stuff semantic segmentation datasets.
- ResNet pretrained on ImageNet serves as the backbone, with a poly learning-rate policy and standard scale, crop, and flip augmentation.
- Training output stride is 16 for PASCAL VOC and Context, and 8 for COCO Stuff; ablations use ResNet-50 with batch size 12.
- The default EMANet configuration uses K = 64 bases, λ = 1, and T = 3 iterations.
6.2. Results on the PASCAL VOC Dataset
On PASCAL VOC, ablations examine iteration count, basis normalization, and maintenance, while comparisons evaluate EMANet against established baselines and state-of-the-art methods. EMANet achieves strong segmentation performance with lower computation and memory costs.
- Bases Maintenance and Normalization: Moving-average basis maintenance performs best across evaluated iteration counts and surpasses other maintenance strategies by at least 0.9 in mIoU.
- Bases Maintenance and Normalization: L2Norm outperforms layer normalization and no normalization when T ≥3, while layer normalization performs worse than no normalization.
- Iteration Number: mIoU increases monotonically with evaluation iterations, but training performance peaks at T_train = 3 and decreases thereafter.Moving Average and L2Norm alleviate the RNN-like training issue only to a certain degree.
- State-of-the-art Comparisons: EMANet outperforms DeeplabV3, DeeplabV3+, and PSANet while using less computation and memory.
- State-of-the-art Comparisons: 2.0% mIoU improvement over DeeplabV3 with the same backbone establishes a new PASCAL VOC record for EMANet.
6.3. Results on the PASCAL Context Dataset
On PASCAL Context, EMANet with a ResNet-101 backbone achieves the highest reported performance among the compared methods.
- EMANet with ResNet-101 achieves the highest performance on the PASCAL Context dataset.Even SGR+ pretrained on additional COCO Stuff data remains inferior to EMANet.
6.4. Results on the COCO Stuff Dataset
On COCO Stuff, EMANet achieves 39.9% mIoU and outperforms previous methods by a large margin.
- 39.9% mIoU is achieved by EMANet on COCO Stuff, exceeding previous methods by a large margin.
6.5. Visualization of Bases Responsibilities
The visualization shows that EMAU bases represent abstract image concepts, becoming more compact and clearer across EM iterations. Different bases specialize in distinct semantic content.
- Each selected basis corresponds to an abstract concept in the image.
- The abstract concepts become more compact and clear as the AE and AM iterations progress.
- The first two rows focus on human, wine glass, cutlery, and profile semantics.
- The last two rows focus on sailboat, mountain, airplane, and lane semantics.
7. Conclusion
The paper introduces EMA, an attention mechanism that iteratively computes a compact basis set through EM-style updates. Its low-rank reconstruction is robust to input variance, efficient, and readily integrated into existing CNNs, with experiments demonstrating effectiveness and efficiency.
- EMA computes a compact basis set by iteratively executing the EM algorithm.
- EMA reconstructs low-rank representations that are robust to input variance.
- EMAU is a lightweight module that integrates into existing CNNs with little overhead.
- Extensive benchmark experiments demonstrate the effectiveness and efficiency of EMAU.