[논문 리뷰] Expectation-Maximization Attention Networks for Semantic Segmentation

Xia Li, Zhisheng Zhong외 4명|arXiv (Cornell University)|2019. 07. 31.

Advanced Neural Network Applications참고 문헌 38인용 수 113

한 줄 요약

EMA는 EM 반복을 통해 주의(attention)를 분석하여 픽셀 단위 표현의 Compact 기초를 학습하고, 가볍고 강건한 의미 분할 모듈(EMAU)을 도출하여 표준 벤치마크에서 계산 및 메모리 감소로 성능을 향상시킵니다.

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.

연구 동기 및 목표

• 장거리 의존성과 함께 의미 분할을 촉진하면서 주의의 계산 부담을 줄인다.
• 셀프 어텐션을 EM 프로세스로 재구성하여 어텐션 맵에 대한 컴팩트한 기저 집합을 학습한다.
• CNN 백본에 쉽게 통합될 수 있는 경량의 EMAU 모듈을 개발한다.

제안 방법

• 주의를 예상-최대화(EM) 프로세스로 재구성하고 주의 맵을 잠재 변수로, 기저를 학습해야 할 매개변수로 삼는다.
• EMA를 사용하여 책임도(담당도)를 반복적으로 추정(E 단계)하고 기저를 업데이트(M 단계)하여 입력 특징의 저랭크 재구성을 얻는다.
• 학습된 기저로부터 컴팩트하고 노이즈에 강한 표현을 생성하기 위한 데이터 재추정(data re-estimation)을 도입한다.
• EMA 코어를 둘러싼 두 개의 1x1 컨볼루션과 잔차 연결을 갖는 신경망 모듈(EMAU)에 EMA를 삽입한다.
• 미니배치 간 초기 기저의 이동평균 업데이트를 통해 기저를 유지하고, 학습의 안정화를 위해 기저에 유클리드 정규화(Euclidean normalization)을 적용한다.
• 반복 횟수, 유지 전략, 정규화에 대한 소거 연구(ablation study)를 통해 설계 선택의 타당성을 검증한다.

실험 결과

연구 질문

• RQ1 EM 스타일의 반복 주의가 표준 셀프 어텐션이나 Non-local 블록에 비해 계산을 줄이는 컴팩트하고 강건한 기저 집합을 학습할 수 있는가?
• RQ2 EMAU 모듈이 표준 벤치마크에서 분할 정확도를 향상시키면서 FLOPs와 메모리 사용량을 줄이는가?
• RQ3 초기화, 유지(이동 평균), 기저의 정규화(L2Norm)가 학습 안정성과 성능에 어떤 영향을 미치는가?

주요 결과

• EMAU는 PASCAL VOC, PASCAL Context, COCO Stuff에서 일부 기준선보다 낮은 계산 비용으로 경쟁력 있거나 최첨단(mean IoU)을 달성한다.
• EM 스타일의 주의는 복잡도를 O(N^2)에서 O(NK)로 감소시키고 K가 N보다 작은 경우 일반적으로 몇 차례의 반복(T ~ 3) 내에 수렴한다.
• 이동 평균 기저 유지와 L2 정규화가 학습 안정성과 성능을 대체 전략(예: 기울기 기반 업데이트나 LN)보다 더 향상시킨다.
• Non-local 및 A^2 블록과 비교하여 EMAU는 메모리 및 FLOPs를 줄이면서 비슷하거나 더 나은 성능을 제공한다.
• 시각화를 통해 학습된 기저가 단순 전경/배경 구분을 넘어 의미 있는 의미적 개념에 부합한다는 것을 보여준다.