Vision Transformers for Dense Prediction
저자: René Ranftl, Alexey Bochkovskiy, Vladlen Koltun
발행년도: 2021년
인용수: 2500회
논문 링크: https://arxiv.org/abs/2103.13413
arXiv ID: 2103.13413
Vision Transformer를 Dense Prediction에 적용하기: DPT 논문 리뷰
문제 정의 (Problem Definition)
Vision Transformer(ViT)가 이미지 분류에서 CNN을 능가하는 성능을 보여준 이후, 많은 연구자들이 자연스럽게 던진 질문이 있었다. "ViT를 semantic segmentation이나 depth estimation 같은 픽셀 단위 예측 작업에도 적용할 수 있을까?" 이 논문은 바로 이 질문에 답하려 했다.
문제는 생각보다 까다로웠다. ViT는 이미지를 16×16 패치로 나누고, 각 패치를 토큰으로 변환해 처리한다. 이 과정에서 공간적 해상도가 크게 감소하는데, 224×224 이미지가 14×14 토큰 시퀀스가 되는 식이다. 분류 작업에서는 최종적으로 하나의 class 토큰만 있으면 되지만, dense prediction에서는 원본 이미지와 같은 해상도의 출력이 필요하다.
더 근본적인 문제도 있었다. CNN은 계층적 구조를 통해 다양한 스케일의 특징을 자연스럽게 추출한다. 초기 레이어는 edge나 texture 같은 저수준 특징을, 후기 레이어는 의미적 정보를 담는다. 반면 ViT의 모든 레이어는 동일한 해상도를 유지하며, 전역적 정보와 지역적 정보가 섞여 있다. 이런 구조적 차이로 인해 단순히 ViT 위에 디코더를 붙이는 접근은 한계가 있었다.
실제 환경에서의 어려움도 고려해야 했다. Depth estimation은 객체의 경계에서 급격한 깊이 변화를 정확히 포착해야 하고, semantic segmentation은 작은 객체나 얇은 구조물도 놓치지 않아야 한다. 이런 세밀한 예측을 위해서는 고해상도 특징과 의미적 정보를 효과적으로 결합하는 메커니즘이 필수적이었다.
기존 방법의 한계 (Motivation)
Dense prediction을 위한 기존 접근법들을 살펴보면, 크게 세 가지 흐름이 있었다. 첫째는 FCN(Fully Convolutional Networks) 계열로, CNN의 마지막 fully connected layer를 제거하고 deconvolution으로 해상도를 복원하는 방식이었다. 하지만 이 방법은 다운샘플링 과정에서 손실된 공간 정보를 제대로 복원하지 못했고, 특히 작은 객체나 세밀한 경계를 놓치는 경우가 많았다.
둘째는 U-Net처럼 인코더-디코더 구조에 skip connection을 추가한 방식이었다. 인코더의 각 레벨에서 나온 특징을 디코더의 대응하는 레벨과 연결해 정보 손실을 줄였다. 이 접근은 효과적이었지만, CNN의 제한된 수용 영역(receptive field) 때문에 전역적 문맥을 충분히 활용하지 못했다. 큰 객체의 전체적인 형태를 파악하거나, 멀리 떨어진 영역 간의 관계를 모델링하는 데 한계가 있었다.
셋째는 DeepLab 계열의 dilated convolution을 활용한 방법이었다. Convolution의 수용 영역을 늘려 더 넓은 문맥을 볼 수 있게 했지만, 여전히 지역적 연산에 머물렀다. ASPP(Atrous Spatial Pyramid Pooling) 같은 모듈로 다양한 스케일을 다루려 했으나, 계산 비용이 크고 설계가 복잡했다.
출처: ar5iv (Figure 1)
저자들은 CNN 기반 방법의 한계를 "제한된 수용 영역"으로 단순화했는데, 사실 최근 CNN들은 이미 상당히 큰 수용 영역을 가지고 있다. 진짜 문제는 수용 영역의 크기보다 "전역적 정보를 어떻게 효율적으로 모델링하느냐"에 있었다고 본다.
이런 상황에서 ViT의 등장은 새로운 가능성을 열었다. Self-attention 메커니즘은 이미지의 모든 위치 간 관계를 직접 모델링할 수 있었다. 문제는 이 강력한 전역적 모델링 능력을 dense prediction에 어떻게 활용할 것인가였다.
제안 방법의 핵심 아이디어 (Key Idea)
DPT의 핵심 아이디어는 의외로 단순했다. "ViT의 토큰들을 다시 이미지 형태로 재조립(reassemble)하되, 여러 레이어의 정보를 효과적으로 융합하자"는 것이었다.
구체적으로 설명하면 이렇다. ViT는 이미지를 패치로 나누고 각 패치를 토큰으로 변환해 처리한다. 이 과정에서 공간 정보가 1차원 시퀀스로 평탄화되는데, DPT는 이를 다시 2차원 특징 맵으로 재구성했다. 단순히 reshape하는 것이 아니라, 서로 다른 레이어의 토큰들을 가져와 다양한 해상도의 특징 맵을 만들었다.
이를 건물에 비유하면, ViT가 건물의 모든 방을 일렬로 늘어놓고 각 방의 관계를 파악했다면, DPT는 이 정보를 바탕으로 다시 층별 평면도를 그리는 것과 같다. 더 나아가 각 층의 정보를 적절히 조합해 건물 전체의 상세한 3D 구조를 복원했다.
기존 CNN 기반 방법과의 차별점은 명확했다. CNN은 처음부터 계층적으로 특징을 추출하지만 전역 정보 활용에 한계가 있었다. 반면 DPT는 ViT를 통해 먼저 전역적 관계를 충분히 학습한 후, 이를 다시 계층적 구조로 재구성했다. 마치 숲 전체를 본 다음에 나무 하나하나를 그리는 접근이었다.
이 아이디어가 완전히 새로운 것은 아니다. SETR(SEgmentation TRansformer) 같은 선행 연구도 ViT를 dense prediction에 적용했다. DPT의 진짜 기여는 토큰을 재조립하는 "방법"에 있었다.
아키텍처 설명 (Architecture)
DPT의 전체 구조는 크게 세 부분으로 나뉜다: ViT 인코더, Reassemble 모듈, 그리고 Fusion 디코더다.
먼저 ViT 인코더는 사전학습된 ViT-Base 또는 ViT-Large를 그대로 사용했다. 224×224 입력 이미지는 16×16 패치로 분할되어 196개의 토큰이 되고, 여기에 위치 임베딩과 class 토큰이 더해져 총 197개 토큰이 12개 또는 24개의 transformer 블록을 통과한다.
Reassemble 모듈이 DPT의 핵심이었다. 이 모듈은 ViT의 중간 레이어들(보통 4개 레이어)에서 토큰을 가져와 다양한 해상도의 특징 맵으로 변환했다. 과정은 다음과 같다:
def reassemble_tokens(tokens, image_size, patch_size):
# tokens: [B, N+1, D] (class token 포함)
# 1. class token 제거
tokens = tokens[:, 1:, :] # [B, N, D]
# 2. 2D 형태로 reshape
h, w = image_size // patch_size
tokens = tokens.reshape(B, h, w, D)
# 3. 채널 차원 조정 (1x1 conv)
features = conv1x1(tokens, out_channels)
# 4. 해상도 조정 (bilinear interpolation)
features = F.interpolate(features, scale_factor)
return features여기서 주목할 점은 서로 다른 레이어에서 추출한 특징들을 서로 다른 해상도로 변환한다는 것이다. 얕은 레이어는 고해상도로, 깊은 레이어는 저해상도로 변환해 계층적 구조를 만들었다.
Fusion 디코더는 이렇게 만들어진 다중 스케일 특징들을 점진적으로 융합했다. RefineNet 스타일의 구조를 사용했는데, residual connection과 fusion block을 통해 서로 다른 해상도의 특징을 효과적으로 결합했다. 최종 출력은 원본 이미지와 같은 해상도의 예측 맵이었다.
Reassemble 과정에서 단순 interpolation을 사용한 것은 아쉬운 부분이다. Deformable convolution이나 학습 가능한 upsampling을 사용했다면 더 정교한 재구성이 가능했을 것이다.
접근 방법의 특징 및 설계 의도 (Design Choices)
DPT의 설계에는 몇 가지 흥미로운 선택이 있었다. 첫째, 사전학습된 ViT를 그대로 사용하되 fine-tuning은 최소화했다. 이는 ImageNet에서 학습된 풍부한 시각적 표현을 최대한 활용하려는 의도였다. 실제로 ViT를 처음부터 학습하는 것보다 훨씬 효율적이었고, 적은 데이터셋에서도 좋은 성능을 보였다.
둘째, 4개의 중간 레이어를 선택적으로 사용했다. 모든 레이어를 쓰지 않은 이유는 계산 효율성과 정보의 중복성 때문이었다. 인접한 레이어들은 비슷한 정보를 담고 있어, 일정 간격으로 떨어진 레이어들을 선택하는 것이 더 효과적이었다. 마치 동영상에서 키프레임만 추출하는 것과 비슷한 원리였다.
셋째, skip connection 대신 progressive fusion을 선택했다. U-Net 스타일의 직접적인 skip connection은 ViT의 특성과 맞지 않았다. ViT의 모든 레이어가 동일한 해상도를 가지므로, 단순 concatenation보다는 점진적으로 해상도를 높이며 융합하는 방식이 더 자연스러웠다.
특히 주목할 점은 loss 설계였다. Depth estimation에는 scale-invariant loss를, segmentation에는 cross-entropy loss를 사용했는데, 별도의 edge loss나 smoothness term 없이도 좋은 결과를 얻었다. 이는 ViT의 강력한 표현력 덕분으로 보인다.
저자들은 이런 설계 선택들에 대한 ablation study를 충분히 제공하지 않았다. 특히 레이어 선택 전략이나 fusion 방식에 대한 체계적인 비교가 부족했다. "왜 4개 레이어인가?", "다른 fusion 방식은 시도해봤는가?" 같은 질문에 명확한 답이 없었다.
실험 결과 분석
DPT는 monocular depth estimation과 semantic segmentation 두 작업에서 평가되었다. Depth estimation에서는 NYU-Depth V2와 KITTI 데이터셋을, segmentation에서는 ADE20K를 사용했다.
가장 인상적인 결과는 NYU-Depth V2에서 나왔다. DPT-Large는 relative error 0.075를 기록해, 기존 최고 성능인 0.093을 크게 앞섰다. 이는 약 20%의 오류 감소에 해당하는 상당한 개선이었다. 특히 물체의 경계나 복잡한 구조에서 더 선명한 깊이 맵을 생성했다.
하지만 실험 설계에는 몇 가지 의문점이 있었다. 첫째, 비교 대상으로 선택한 baseline들이 다소 오래된 모델들이었다. 2021년 기준으로도 더 강력한 CNN 기반 모델들이 있었는데, 이들과의 공정한 비교가 부족했다. 둘째, 계산 비용에 대한 분석이 빈약했다. ViT-Large는 파라미터 수가 307M에 달하는데, 이는 대부분의 CNN 모델보다 훨씬 크다.
Ablation study도 아쉬웠다. Reassemble 모듈의 효과는 보여줬지만, 각 설계 요소의 개별적인 기여도는 불분명했다. 예를 들어, "4개 레이어를 사용한 것과 2개 또는 6개를 사용한 것의 차이는?", "다른 interpolation 방법을 사용하면?" 같은 질문에 대한 답이 없었다.
실험 결과는 인상적이지만, 재현성에 대한 우려가 있다. 학습 세부사항이나 하이퍼파라미터 설정이 충분히 공개되지 않았고, 특히 사전학습된 ViT 모델의 정확한 버전이 명시되지 않았다.
총평: 개인적 해석과 후속 연구 방향
DPT는 ViT를 dense prediction에 성공적으로 적용한 의미 있는 연구였다. 가장 큰 기여는 "토큰을 어떻게 이미지로 재구성할 것인가"라는 문제에 대한 실용적인 해법을 제시한 것이다. Reassemble과 progressive fusion이라는 간단하면서도 효과적인 방법으로 ViT의 강력한 표현력을 픽셀 단위 예측에 활용할 수 있음을 보여줬다.
하지만 몇 가지 한계도 명확했다. 첫째, 계산 비용이 너무 크다. ViT-Large 기반 DPT는 실시간 응용에는 부적합하며, 모바일이나 엣지 디바이스에서는 사용이 어렵다. 둘째, 아키텍처가 다소 복잡하고 ad-hoc한 설계가 많다. 더 원칙적이고 간결한 구조가 가능할 것 같다.
실제 프로젝트에 적용한다면, 정확도가 중요하고 계산 자원이 충분한 경우에는 좋은 선택이 될 수 있다. 하지만 실시간성이 필요하거나 리소스가 제한된 환경에서는 여전히 효율적인 CNN 기반 모델이 더 실용적일 것이다.
내가 이 연구를 이어간다면 세 가지 방향을 탐구하고 싶다. 첫째, 더 효율적인 ViT 변형(예: Swin Transformer)을 백본으로 사용해 계산 비용을 줄이는 것이다. 둘째, Reassemble 과정을 학습 가능한 모듈로 만들어 더 정교한 재구성이 가능하게 하는 것이다. 셋째, 다양한 dense prediction 작업들(instance segmentation, panoptic segmentation 등)로 확장해 DPT의 일반화 능력을 검증하는 것이다.
특히 흥미로운 가설은 "ViT의 어떤 레이어가 어떤 종류의 정보를 담고 있는가"이다. 단순히 균등하게 레이어를 선택하는 대신, 각 레이어의 특성을 분석해 작업별로 최적의 레이어 조합을 찾을 수 있지 않을까? 이를 위해 각 레이어의 attention pattern이나 feature similarity를 분석하는 실험을 설계해보고 싶다.