✅ This article is a compilation of research reviews for study purposes. It was not written for commercial purposes. If there are any issues, it will be deleted.
✅ It's written in Korean, so if you need other language, click the translate button and read it.
논문 소개
Attention U-Net은 “U-Net을 더 크게/깊게” 만든 것이 아니라, 어디를 봐야 하는지(Region of Interest; ROI)를 네트워크 내부에서 스스로 선택하게 만든 U-Net의 변형입니다.
기존 U-Net은 skip connection으로 저수준 특징을 그대로 끌어와 복원력을 얻었지만 작은 장기(pancreas 등)처럼 모양/크기 변이가 크고 대비가 낮은 대상에서는 배경 특징까지 함께 전달되어 오탐/미검이 늘기 쉽습니다. Attention U-Net은 각 skip connection 앞에 Attention Gate(AG)를 달아서 encoder의 특징과 더 거친 스케일에서 온 gating signal을 함께 보고 픽셀/영역별 가중치(α)를 만들어 해당 특징을 α로 스케일링해 “관련 있는 정보만” decoder로 보냅니다.
그 결과, 별도의 cascaded ROI localizer(외부 장기 위치추정 모델) 없이도 입력에서 불필요한 배경 반응을 억제하며 중요한 구조를 강조할 수 있고, 파라미터/연산 오버헤드는 성능이 향상된 데에 비해 크게 늘지 않음을 볼 수 있습니다.
실험적으로도 이 “선별적 skip”은 성능을 단순히 모델 용량을 늘린 것 이상으로 끌어올립니다. 예를 들어 CT-150에서 췌장 분할은 Dice가 0.814→0.840으로, Recall이 0.806→0.841로 올라가며, 파라미터는 5.88M→6.40M, 추론 시간은 0.167s→0.179s로 소폭 증가한 수준입니다.
즉, Attention U-Net의 차별점은 더 많은 특징이 아니라 더 필요한 특징만 통과시키는 게이트로 민감도와 정확도를 같이 올린다는 점에 있습니다.
배경
의료영상 분할은 임상에서 정량 지표를 빠르게 뽑아 workflow를 줄이기 위해 중요하지만, 픽셀 단위 라벨링은 매우 번거롭고 오류도 잦다. 이를 해결하기 위해 CNN 기반의 FCN이나 U-Net이 표준처럼 쓰이게 되었지만, 췌장(Pancreas)처럼 크기/형태의 개인차가 크고 대비가 낮은 장기는 한 번에 잡기 어려워 보통 아래와 같은 다단계(cascade) 구조를 사용해왔다.
(1) Region of Interest(이하 ROI)를 먼저 찾는다.
(2) 그 ROI에서 정밀 분할한다.
문제는 이런 cascade가 (같은 저수준 특징을 여러 번 추출하기 때문에) 연산량을 과하게 늘리고 외부 장기 위치 추정 모듈에 의존하게 만든다는 문제점이 있었다.
이와 관련된 연구 흐름을 보면, 초기에는 통계적인 shape model이나 multi-atlas 방식이 주류였고, 이후에는 성능을 끌어올리기 위해 cascade CNN(거친 ROI → 정밀 분할)이나 2D-FCN+RNN 등으로 발전해 온 것을 알 수 있다.
이 논문은 이런 “별도 ROI 모델 의존”을 줄이기 위해, 자연영상/NLP에서 쓰이던 attention을 의료 분할에 맞게 가져와 Attention Gate(이하 AG)를 U-Net의 skip connection에 끼워 넣는 방법을 제안한다. 특히 전역 벡터로 한 번에 게이팅하는 대신, 공간 격자(grid) 단위로 지역별 주의도를 학습해 배경 반응을 억제하고 중요한 부위만 통과시키며, 추가 파라미터/연산 오버헤드는 크게 늘리지 않겠다는 것이 핵심 문제의식과 기여이다.
아이디어
리뷰에 앞서 언급하지만 본 문단을 반드시 100% 이해하고 넘어갈 필요 없다. (딥러닝에 익숙하지 않거나 Attention 개념이 익숙하지 않은 독자라면) 다음 문단 시스템 구조 및 학습 및 추론, 실험을 읽는다면 이해에 훨씬 도움이 될 것이다. 개념이 어렵거나 익숙하지 않다면 천천히 여러 번 읽는 것을 추천한다.
Attention U-Net의 핵심 아이디어는 “U-Net이 skip connection으로 가져오는 특징을 반드시 전부 사용하는 것이 아니라, 현재 단계에서 필요한 영역(췌장 같은 작은 타겟)에 해당하는 특징만 선별적으로 통과시키자”이다.
이를 위해 각 skip connection 앞에 Attention Gate(AG)를 삽입하고, encoder의 특징 x^l과 더 coarse한 scale에서 온 gating signal g(문맥/위치 정보를 담은 query)을 함께 분석하여 픽셀(혹은 그리드별) attention 계수 αi∈[0,1]를 만든다. 그리고 skip 특징을 원소곱으로 스케일링하여 배경/잡음 반응은 줄이고 타깃 관련 활성만 decoder로 보낸다. 이렇게 하면 기존 mutli-stage/cascade처럼 ROI를 따로 찾고 crop해서 넘기는 절차 없이도, 네트워크 내부에서 점진적으로 배경을 억제하며 false positive를 줄일 수 있다.
AG의 계산은 (곱셈형보다 비용은 들지만 정확도가 더 좋다고 보고) additivate attention으로 구현한다. 또한 단일 스칼라 α만 두는 기본형을 넘어 다중 클래스/구조를 위해 multi-dimensional attention을 학습하여 “각 게이트가 특정 구조에 더 집중”하도록 확장한다. 여기서 중요한 디테일로 α정규화에 softmax 대신에 sigmoid를 써서 게이트 학습 수렴을 개선했고, gating signal을 전역 벡터로 하나만 쓰는 방식이 아니라 공간 정보를 유지한 grid-attention으로 만들어 성능을 끌어올렸다. grid-attention은 스케일들을 합쳐 query 해상도를 높이는 방식으로 성능을 올릴 수 있었다. 더 나아가 AG는 forward에서 특징을 거르는 것뿐만 아니라, backward에서도 배경에서 오는 gradient를 down-weight하기 때문에 얕은 층이 “정답에 중요한 위치”를 중심으로 업데이트되게 만든다. 계산량/파라미터 증가는 1x1x1 conv 기반 선형변환과 다운샘플링으로 억제하고, deep supervision으로 각 스케일의 특징이 충분히 의미를 갖도록 강제하여 게이트들이 다양한 전경 정보를 안정적으로 다루게 한다.
시스템 구조
Figure 1.은 표준 U-Net의 encoder-decoder 뼈대 위에 Attention Gate(AG)를 skip connection 경로에 끼운 전체 파이프라인을 보여준다. 입력은 각 스케일에 대해 2배씩 다운샘플링(예: H4 = H1 / 8)되며 decoder는 업샘플링으로 해상도를 복원하여 최종 Nc 클래스 분할 맵을 만든다. 이때, 핵심은 coarse scale에서 만든 gating signal(query)로 skip 특징을 먼저 필터링 한 뒤 concat한다는 점이다.
Figure 2.는 Attention gate 내부를 도식화하고 있다. 밑의 설명을 읽기 전 참고하자.
Fully Convolutional Network(FCN)
FCN은 분류기(FC layer) 대신 합성곱만으로 픽셀/복셀 단위 예측을 수행하는 구조로, 의료 영상에서 그래프컷/multi-atlas 같은 전통 기법보다 빠르고 정확한 경향이 있다. 그 이유는 아래와 같다.
(1) SGD로 도메인 특화 특징을 직접 학습한다.
(2) 학습된 커널이 모든 위치에 공유되어 파라미터 효율이 높다.
(3) 합성곱 연산이 영상의 공간 구조를 잘 활용한다.
네트워크는 층을 거치며 표현 xl을 점진적으로 고차원 특징으로 변환해 픽셀을 의미적으로 분리한다. 이때 각 층은 “선형 변환(합성곱) + 비선형 활성화”의 연쇄로 볼 수 있고, 논문은 ReLU σ1(x)=max(0, x)를 전형적인 선택으로 둔다. 이렇게 누적된 연산은 receptive field를 키워 넓은 문맥 정보에 조건화된 예측을 가능하게 한다.
Attention Gates for Image Analysis
일반 CNN/FCN은 충분한 문맥(semantic context)을 얻기 위해 feature-map을 점점 다운샘플링한다. 하지만 작은 장기/병변처럼 크기가 작고 형태 변이가 큰 대상에서는, 문맥이 있어도 배경에서의 오탐(false positive)을 줄이기 어렵다. 그래서 기존 프레임워크는 ROI를 먼저 찾는 별도의 localization 모델을 앞에 두고 “위치추정 → 분할”의 cascade pipeline을 쓰곤 했다. 논문에서는 이 기능을 AG를 네트워크 내부로 삽입하여 대체한다. AG는 αi∈[0,1]로 중요한 위치를 표시하고, 출력은 ^xi,cl=xi,cl⋅αi 처럼 입력 feature를 위치 별로 억제/강조한다. 즉 ROI를 잘라내지 않아도, 학습 과정에서 배경 반응을 점진적으로 suppress한다. 또한 grating vector gi는 더 coarse한 스케일에서 온 문맥을 담아 “어디를 볼지”를 결정하고, 논문에서는 성능상 유리한 additive attention 형태(아래 식 (1)(2))로 α를 계산한다.
식(1)은 AG가 현재 위치 i에서 skip connection에서 넘어온 특징 x가 얼마나 중요한가를 점수로 만드는 단계이다. x와 coarse 문맥을 담은 gating g를 각각 Wx, Wg로 같은 중간 차원(Fint)에 투영하여 더한 뒤 ReLU(σ1)을 취한다. ψ로 다시 1차원으로 눌러 compatibility score q를 만든다. 즉, 로컬 특징+전역 문맥을 합쳐 “볼만한 정도”를 계산한다.
조금 더 자세한 설명을 원한다면 아래 문단을 읽어보자.
💡 💡 💡
(I)의 Wx xi는 현재 스케일에서 로컬 특징(local feature) x를 중간 차원 Wx로 투영(projection)하여 의미 있는 형태로 바꾸는 과정을 뜻한다. 이 과정을 통해 엣지, 구조, 조직 경계 등 저수준의 local 정보를 담게 된다.
(II)의 Wg gi는 더 coarse한 scale(더 깊은 층)에서 가져온 global context “g”를 중간 차원 Wg로 투영하여 의미 있는 형태로 바꾸는 과정을 뜻한다.
이때 W란 encoder에서 온 skip feature x의 채널 수 와 coarse scale에서 온 gating feature g의 채널 수가 달라 add 연산을 하는데 어려움이 있으므로 중간 차원(F_int)으로 투영하여 더할 수 있게 만드는 윈도우다. 보통 F_int는 skip 채널보다 작게 둬서 파라미터 수를 줄이고 attention이 너무 복잡해지지 않도록 제약하는 하이퍼파라미터이다.
(I)과 (II)의 두 정보를 더하는 것이 바로 [로컬 + 전역 정보]를 같은 표현 공간에 두고 합치는 것을 의미하고 “어떤 위치가 전체 목표에서 중요한 역할을 하는지”를 비교하는 형태로 변하게 된다.
그 후, ReLU(σ1)를 통해 중요하지 않은 위치는 활성값이 0으로 수렴하도록 하고 중요한 위치는 더 양수 값을 부각한다.
이렇게 얻어진 결과가 바로 compatibility score, 즉 “이 위치의 특징이 현재 segmentation 목표와 얼마나 궁합이 잘 맞는지”를 수치로 표현한 실수 벡터가 된다.
이를 1x1x1 conv 역할을 하는 ψ로 하나의 스칼라 score로 눌러주면 중요도 점수 하나 또는 벡터가 만들어지게 된다.
class Attention_block(nn.Module):
def __init__(self, F_g, F_l, F_int):
super(Attention_block, self).__init__()
self.W_g = nn.Sequential(
nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True),
nn.BatchNorm2d(F_int)
)
self.W_x = nn.Sequential(
nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True),
nn.BatchNorm2d(F_int)
)
self.psi = nn.Sequential(
nn.Conv2d(F_int, 1, kernel_size=1, stride=1, dding=0, bias=True),
nn.BatchNorm2d(1),
nn.Sigmoid()
)
self.relu = nn.ReLU(inplace=True)
def forward(self,g,x):
g1 = self.W_g(g)
x1 = self.W_x(x)
psi = self.relu(g1+x1)
psi = self.psi(psi)
return x*psi식(2)은 (1)에서 얻은 compatibility score q를 attention coefficient α로 바꾸는 정규화 단계이다. 논문에서는 softmax 대신 sigmoid σ2를 각 위치에 독립적으로 사용하여 게이트가 너무 희소해지는 걸 막고 학습 수렴을 개선했다. 이렇게 얻은 α로 skip feature를 곱해, 전경에 해당하는 위치는 통과시키고 배경 반응은 억제한다.
조금 더 자세한 설명을 원한다면 아래 문단을 읽어보자.
💡 💡 💡
σ2는 softmax가 아니라 sigmoid이다.
softmax는 모든 위치에서 확률의 총합을 1로 만들어 소수의 위치만 강하게 선택하는 winner-takes-all 구조가 된다. 반면 의료영상 분야에서는 여러 픽셀들이 동시에 중요할 수도 있고, 여러 슬라이스가 함께 의미를 줄 수 있기 때문에 독립적인 이진 게이트처럼 작동하는 sigmoid가 훨씬 안정적이고 자연스럽다.
Figure 3(a)에서 볼 수 있듯이 특정 작업과 관련된 활성화만 보존하도록 특징 응답을 정리할 수 있다.
이때 논문은 attention coefficient를 1차원(스칼라)으로만 두지 않고 multi-dimensional AG을 통해 벡터 형태로 만들 수 있다(multi-dimension attention).
CT에서는 혈관/주변 지방 조직/췌장 실질/경계면 등의 구조들이 모두 서로 다른 appearance를 가지지만 모든 픽셀이 하나의 scalar attention으로 표현되면 정보가 너무 압축된다. 따라서 각 attention dimension을 두고 각 차원이 서로 다른 패턴에 반응하도록 학습할 수 있다.
Figure 2는 저자들이 제안한 additive attention gate의 내부를 도식화한다. skip 입력 특징 xl과 더 coarse한 스케일에서 온 gating signal g를 함께 분석하여 위치별 attention 계수 α∈[0,1]를 계산하고, 이를 xl에 곱하여 스케일링한 후 관련 activation만 통과시킨다. 계산된 α는 해상도 정합을 위해 trilinear interpolation으로 resampling된다.
이처럼 skip feature x와 gating signal g의 해상도 차이 이슈를 coarse-scale 정보에서 중요한 픽셀이 어디인지 찾도록 attention을 계산하고 픽셀 단위 대응(attention at pixel-level)하는 위 구조를 Grid-based Attention이라고 한다.
Attention Gates in U-Net Model
Attention U-Net에서 AG는 U-Net의 장점인 multi-scale feature + skip connection을 유지하면서, skip 경로를 무조건 전달하지 않고 concat 직전에 선별적으로 통과시키는 역할을 한다. 더 구체적으로 설명하자면, 더 coarse scale에서 추출된 정보(=gating)가 skip의 저수준/잡음 반응을 구분해 불필요한 활성화를 줄이고, 최종적으로 디코더가 결합하는 정보의 질을 높인다. 중요한 점은 AG가 forward에서만 작동하는 게 아니라, backward에서도 gradient를 조절한다는 것이다. 배경 영역에서 발생한 기울기는 α에 의해 down-weight 되어, 얕은 층의 파라미터가 “배경 설명”보다 과제에 관련된 공간 중심으로 업데이트되도록 유도한다. 논문은 이를 아래 업데이트 식(3)의 형태로 제시하며 첫 항이 α로 스케일되는 형태가 직관적으로 “배경 기여를 줄인다”라는 의미가 된다. 결과적으로 추가 모델 없이도 민감도/재현율을 높이는 방향으로 학습이 정렬된다.
식 (3)은 AG가 역전파에서도 학습을 “어디에 집중할지” 조절함을 보여준다. 이전 층 파라미터 Φ(l-1)에 대한 기울기는 α가 곱해져 기본 gradient가 스케일링되고, α자체가 Φ(l-1)에 의존적이므로 ∂α/∂Φ(l-1)에서 오는 추가항이 더해진다.
더 자세한 설명을 원한다면 아래 문단을 읽어보자.
💡 💡 💡
▶️ f(⋅; Φ(l-1))은 layer (l-1)의 conv+nonlinearity function으로 이 function의 output이 x_i^l이다.
▶️ a_i^l은 attention gate에서 나오는 게이트 값으로 0~1 사이의 범위를 가진다.
▶️ ^x_i^l은 gate를 통과한 skip feature이다.
논문이 증명하고 싶은 것은 “더 아래 레이어 파라미터 Φ(l-1)에 대해 gradient가 어떻게 바뀌는가” 이므로 이를 보기 위해 (I) 항으로 시작하여 (II)까지 전개한다.
(II) 항을 직관적으로 나눈다면 아래와 같다.
a. 기본 gradient x gate(= α가 곱해진 항)
b. gate 자체가 변함으로 인해 생기는 추가 gradient 항
a. 기본 gradient와 gate(α)의 곱 일반적인 conv layer의 gradient에 그 위치의 gate 값 α를 곱한 꼴이다. 이 항으로 인해 AG는 shallow layer 입장에서 “공간적으로 가중치가 다른 learning rate mask”를 씌운 것처럼 동작한다. 논문에서는 이를 Gradients originating from background regions are down weighted during the backward pass라고 표현하고 있다.
b. 게이트도 x에 의존하니까 생기는 추가 경로 α는 위 식(2)에서 밝혔듯이 이미 x에 의존하고 있고 x는 이전 레이어인 Φ(l−1)에서 왔다. chain rule 상 “α가 변함으로 인해 ^x가 바뀌는 효과”가 한 번 더 일어나게 된다. 즉, conv 파라미터 변화가 게이트 모양 자체를 조금씩 조정하게 만드는 보조 경로를 의미하게 된다
결국 핵심은 배경 위치의 α가 작으면 그쪽 gradient도 작아져, 얕은 층이 전경 중심으로 업데이트된다는 점이다.
학습 및 추론, 실험
Evaluation Datasets
실험은 2개의 복부 CT 3D 데이터셋으로 구성된다. (1) CT-150: 위암 환자에게서 획득한 150건의 3D CT로, 췌장/간/비장 경계를 3명의 연구자가 반자동으로 그린 뒤 임상의가 검증했다. (2) CT-82 (NIh-TCIA Pancreas-CT): 82건의 조영증강 3D CT이며, 췌장 마스크를 슬라이스 단위로 수동 주석한 공개 벤치마크이다.
큰 볼륨 크기와 GPU 한계 때문에 두 데이터 모두 등방성 2.0mm 해상도로 다운샘플링하여 학습 및 평가했다.
Implementation Details
저자들은 기존 SOTA(주로 2D+후처리/cascade 구조)와 달리 3D 모델로 충분한 semantic context를 확보하는 방향을 택한다. 학습은 배치 크기 2-4의 작은 배치사이즈를 선택했고 큰 네트워크는 여러 번의 forward/backward를 누적해 gradient averaging을 사용한다. 최적화는 Adam, BatchNorm, Deep supervision, 그리고 affine/축 방향 filp/random crop 등의 표준적인 data augmentation 기법을 사용했다. 입력 intensity는 선형 스케일링으로 N(0,1) 분포가 되게 정규화하고, 클래스 불균형에 덜 민감한 Sorensen-Dice loss로 학습했다. 또한 AG 파라미터는 초기에는 모든 위치의 특징이 “통과”되도록 초기화하며, hard-attention과 같은 다단계 학습은 불필요하다고 강조하고 있다.
Attention Map Analysis
학습이 진행되면서 attention coefficient(α)가 어떻게 변하는지, 테스트 이미지에서 epoch 별로 시각화한다(Figure 4.).
과정을 보면 AG는 초기에는 거의 균일 분포라서 모든 위치의 특징을 통과시키지만, 학습이 진행될수록 타깃 장기(췌장 등) 경계로 집중되며 점점 국소화된다. 또한, coarse scale의 AG는 장기의 대략적인 윤곽을 먼저 잡고, 더 fine scale로 갈수록 경계가 정교해지는 양상을 보인다. 같은 스케일에 여러 AG를 두면, 각 게이트가 서로 다른 장기/구조 subset에 특화되어 분업하듯이 attention을 학습한다는 점도 보인다.
Segmentation Experiments
기본 비교는 CT-150에서 train/test =120/30으로 abdominal CT segmentation task에서 다중 클래스 표준 U-Net과 Attention U-Net으로 진행한다. 평가는 Dice(DSC)와 표면 거리 (S2S)로 제시한다(Table 1.).
결과적으로 AG가 췌장 recall을 유의미하게(p=.005) 증가시켜 성능을 끌어올린다고 해석한다. 또 학습 데이터를 30장으로 줄인 설정에서도 개선이 일관되고 유의미(p=.01)함을 보여 데이터가 적을수록 AG의 이점이 두드러진다는 것을 알 수 있었다.
뿐만 아니라 정성 비교로는 U-Net이 놓치는 예측을 Attention U-Net이 보완하는 사례(Figure 3(b))도 제시한다.
Comparison to State-of-the-Art
SOTA 비교는 공개 벤치마크 CT-82에서 수행된다. 먼저 CT-150으로 학습한 모델을 서로 다른 데이터 세트에서의 적용 가능성을 관찰하기 위해 CT-82에 그대로 적용(BFT)했고 그 결과가 Table 3. 에 나와있다.
CT-82 데이터셋을 61/21(train/test)로 나누고 Fine-tuning 한 결과(AFT)도 수행하였다(단, 이때 출력에서는 비장/신장 노드를 제외하고 background+pancreas로 학습하였다). 결과적으로 여러 세팅에서 Attention U-Net이 U-Net 대비 일관되게 더 나은 췌장 분할/표면거리를 보여준다. 또한 CT-82에서 from-scratch 학습(SCR)과 문헌 SOTA(2D FCN, RNN 결합, multi-model 등) 요약(Table 4.)을 통해, cascade/multi-model 없이도 경쟁력 있는 성능(예: 5-fold CV에서 pancreas DSC 81.48±6.23)을 주장한다.
결과 해석
논문에서 저자들은 Attention U-Net이 외부 ROI 모듈 없이도 작은 장기(췌장 등)를 안정적으로 찾고 분할할 수 있다는 점을 실험으로 보여준 뒤, 그 의미를 아래와 같이 정리하고 있다.
(1) AG를 추가하면 단순히 필터 수를 늘린 것보다 효율적으로 성능을 올릴 수 있다. U-Net의 파라미터를 8% 정도만 늘려도 췌장 Dice가 2-3% 개선되는데, 동일한 용량을 전 층에 균일하게 뿌린 단순 확장 U-Net보다 통계적으로 유의미하게 성능이 좋다(p=0.007). 즉, 용량을 무작정 키우기보다 어디를 볼지에 투자하는 편이 더 효율적이라는 것이 핵심이다.
(2) CT-150/CT-82 양쪽에서 Attention U-Net은 U-Net 대비 일관되게 pancreas recall과 표면 거리(S2S)를 개선하면서도, 추론 속도와 파라미터 수는 비슷한 수준에 머무른다. 특히 CT-82 공개 벤치마크에서는 Attention U-Net이 단일 모델 기준 SOTA와 비슷한 수준의 Dice(5-fold에서 81.48±6.23%)를 달성하면서도 cascade/multi-model/CRF를 쓰지 않는 간단한 구조라는 점이 강조된다. 이는 AG가 중복적인 cascade/ensemble을 대체할 수 있는 실용적인 옵션임을 시사한다.
(3) 이 프레임워크는 특정 장기/데이터셋에 한정된 트릭이 아니라 일반적인 모듈형 attention 블록이다. 분할뿐 아니라 classification, regression, natural video, machine translation 등으로도 확장 가능하며, 향후에는 사전학습된 U-Net에서 attention만 파인튜닝하는 transfer learning, highway/residual gate, 더 큰 3D 모델도 고해상도 입력 등으로 성능을 끌어올릴 여지가 많다고 제안한다. 핵심적으로, 이 논문은 “의료 영상에서 ROI 탐색과 분할을 한 몸처럼 학습하는 구조”를 제안했다는 점에서 이후 다양한 attention 기반 segmentation 모델의 기반이 되는 작업으로 해석할 수 있다.
댓글