[Paper Review] Expectation-Maximization Attention Networks for Semantic Segmentation

Xia Li, Zhisheng Zhong+4 more|arXiv (Cornell University)|Jul 31, 2019

Advanced Neural Network ApplicationsComputer Science38 references113 citations

TL;DR

EMA analyzes attention through EM iterations to learn a compact basis for pixel-wise representations, yielding a light-weight, robust semantic segmentation module (EMAU) that improves performance on standard benchmarks with reduced computation and memory.

ABSTRACT

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.

Motivation & Objective

• Motivate semantic segmentation to leverage long-range dependencies while reducing computational burden of attention.
• Reformulate self-attention as an EM process to learn a compact basis set for attention maps.
• Develop a lightweight EMAU module that can be easily integrated into CNN backbones.

Proposed method

• Reformulate attention as an Expectation-Maximization (EM) process where attention maps are latent variables and bases are parameters to learn.
• Use EMA to iteratively estimate responsibilities (E step) and update bases (M step) to obtain a low-rank reconstruction of input features.
• Introduce data re-estimation to produce a compact, noise-robust representation from the learned bases.
• Embed EMA into a neural network module (EMAU) with two 1x1 convolutions around the EMA core and a residual connection.
• Implement bases maintenance via moving-average updating of initial bases across mini-batches and apply Euclidean normalization to bases to stabilize training.
• Provide an ablation study on iteration count, maintenance strategy, and normalization to validate design choices.

Experimental results

Research questions

• RQ1 Can EM-style iterative attention learn a compact, robust basis set that reduces computation compared to standard self-attention or Non-local blocks?
• RQ2 Does the EMAU module improve segmentation accuracy on standard benchmarks while lowering FLOPs and memory usage?
• RQ3 How do initialization, maintenance (moving average), and normalization (L2Norm) of bases affect training stability and performance?

Key findings

• EMAU achieves competitive or state-of-the-art mean IoU on PASCAL VOC, PASCAL Context, and COCO Stuff with lower computational cost than some baselines.
• EM-style attention reduces complexity from O(N^2) to O(NK) with K << N and typically converges within a few iterations (T ~ 3).
• Moving-average bases maintenance and L2 normalization improve training stability and performance more than alternative strategies (e.g., gradient-based updates or LN).
• Compared to Non-local and A^2 blocks, EMAU provides similar or better performance with reduced memory and FLOPs.
• Visualization shows learned bases correspond to meaningful semantic concepts beyond simple foreground/background separation.