💻 computer science

HSMLA: Hierarchical Softmax Multi-scale Linear Attention for Efficient Vision Transformers

이 논문은 표준 비전 트랜스포머의 이차 복잡도를 극복하기 위해 ReLU 기반 선형 어텐션, 선택적 소프트맥스 정교화, 그리고 다중 스케일 심층 합성곱을 결합하여 의료 영상 분할 및 병리 분석과 같은 밀집 예측 작업에서 높은 정확도를 유지하면서도 상당한 추론 속도 향상을 달성하는 새로운 어텐션 메커니즘인 HSMLA를 소개한다.

원저자: Dong Liu, Yanxuan Yu, Renata Borovica-Gajic, Tong Geng, Ying Nian Wu

게시일 2026-08-12

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Dong Liu, Yanxuan Yu, Renata Borovica-Gajic, Tong Geng, Ying Nian Wu

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 거대하고 믿기지 않을 정도로 정교한 도시 벽화를 그리려고 한다고 상상해 보세요. 하지만 당신에게는 아주 작은 페인트 한 양동이와 작업을 마칠 수 있는 매우 짧은 시간뿐입니다. 만약 어떤 색을 사용할지 결정하기 위해 모든 건물의 모든 벽돌 하나하나를 동시에 다 살펴보려고 한다면, 당신은 시작하기도 전에 시간이 다 되어버릴 것입니다. 이것이 바로 '비전 트랜스포머(Vision Transformer)'라고 불리는 일종의 컴퓨터 두뇌가 직면한 문제입니다. 이 디지털 화가들은 이미지를 이해하는 데 매우 뛰어나지만, 고해definition 의료 스캔이나 자율주행 자동차를 위한 거리 풍경처럼 이미지가 거대해지면 압도당하고 맙니다. 그들은 모든 픽셀을 다른 모든 픽셀과 비교하려고 시도하는데, 이는 마치 경기장에 있는 모든 사람을 한꺼번에 서로에게 소개하려는 것과 같습니다. 그것은 너무 느리고 에너지를 너무 많이 소비합니다.

이를 해결하기 위해 과학자들은 '선형 어텐션(linear attention)'이라는 지름길을 시도했습니다. 모든 사람을 서로에게 소개하는 대신, 그냥 모든 사람이 방 전체에 일반적인 요약본을 외치게 하는 방식입니다. 이것은 매우 빠르지만, 단점은 디테일이 흐릿해진다는 것입니다. 마치 군중의 환호성은 들리지만, 개개인이 구체적으로 어떤 말을 외치고 있는지는 들리지 않는 것과 같습니다. 종양의 정확한 경계를 찾거나 얼굴의 미세한 선을 찾아내는 작업처럼 정밀함이 요구되는 업무에서, 이러한 흐릿함은 치명적인 결함이 됩니다. 큰 질문은 이것이었습니다: 우리는 지름길의 속도와 정밀한 관찰의 선명함을 동시에 가질 수 있을까?

여기에 연구진이 만든 HSMLA(Hierarchical Softmax Multi-scale Linear Attention)라는 새로운 발명품이 등장했습니다. HSMLA를 규칙을 언제 어겨야 할지 정확히 아는 매우 똑똑한 아트 디렉터라고 생각해보세요. HSMLA는 벽화의 전체를 동일하게 취급하는 대신, 이미지의 지루하고 텅 빈 부분(예: 하늘이나 평범한 벽)에는 효율적인 방법을 사용합니다. 이 영역들에는 시간을 낭비하지 않도록 빠른 일반 요약법을 사용합니다. 하지만 삐죽삐죽한 나뭇가지, 복잡한 건물 모서리, 혹은 의료 영상의 의심스러운 지점처럼 까다로운 것을 발견하는 즉시, 시스템은 곧바로 '슈퍼 모드'로 전환합니다. 그리고 그 작고 중요한 부분만을 위해 느리지만 정교한 픽셀 단위 비교를 수행합니다.

논문은 이 '믹스 앤 매치(mix-and-match)' 접근 방식이 게임 체인저임을 보여줍니다. 실제로 필요한 부분(이미지의 약 30%)에 대해서만 무겁고 느린 작업을 수행하고 나머지는 빠른 방법을 사용함으로써, HSMLA는 초해상도(blurry한 이미지를 선명하게 만드는 작업)와 같은 과업에서 표준 방식보다 최대 4.2배 더 빠르게 작동합니다. 의료 영상의 세계에서 결과는 더욱 인상적입니다. 장기를 찾는 데 사용되는 CT 스캔에서, 이 새로운 방식은 이전 표준보다 3.2배 더 빠르게 실행되면서도 87.3%의 Dice 점수(컴퓨터가 장기의 윤곽을 얼마나 완벽하게 그려내는지에 대한 척도)를 달성했습니다. 암을 탐지하는 데 사용되는 병리 슬라이드에서는 4.1배의 속도 향상과 함께 94.2%의 AUC(탐지 정확도에 대한 척도)에 도달했습니다.

연구진은 도시 거리의 자동차 식별부터 아주 작은 폐 결절을 찾아내는 것까지 모든 분야에서 이를 테스트했습니다. 결과는 일관되었습니다: 더 이상 속도와 정확도 사이에서 하나를 선택할 필요가 없다는 것입니다. 이 시스템은 언제 쉬어가고 언제 전력 질주해야 할지를 스스로 판단할 만큼 똑똑합니다. 이것은 단순히 이론적인 아이디어가 아닙니다. 팀은 자율주행 자동차나 의료 기기에 들어가는 칩과 같은 실제 하드웨어에서 이를 구축하고 테스트하여 작동함을 증명했습니다. 그들은 이미지의 어느 부분이 느리고 세심한 관찰을 필요로 하는지 결정하는 '게이팅(gating)' 시스템을 사용함으로써, 전체적인 그림은 선명하게 유지하면서도 가장 중요한 국소적 디테일을 정교하게 다듬을 수 있다는 것을 발견했습니다. 이는 마치 배경 화가들은 빛의 속도로 작업하고, 몇몇 숙련된 세부 묘사 화가들은 정교한 꽃과 얼굴을 위해서만 호출되어, 단 하나의 붓터치의 품질도 놓치지 않으면서 전체 벽화를 빠르게 완성하는 팀을 보유한 것과 같습니다.

기술 요약: HSMLA (Hierarchical Softmax Multi-scale Linear Attention)

1. 문제 정의

Vision Transformer (ViT)는 다양한 비전 작업에서 최첨단 성능을 달aged했으나, 고해상도 밀집 예측(dense prediction) 시 셀프 어텐션의 이차 복잡도( $O(N^2)$ )로 인해 상당한 계산 병목 현상에 직면해 있습니다. 선형 어텐션 메커니즘은 $O(N)$ 의 복잡도를 제공하지만, 일반적으로 평탄한 어텐션 분포를 생성하여 국소적 문맥(local context) 모델링 능력을 약화시킵니다. 이러한 한계는 의료 영상(장기 경계, 종양 마진) 및 시맨틱 세그멘테이션, 초해상도(super-resolution)와 같은 고해상도 밀집 예측 작업처럼 날카롭고 국소적인 어텐션이 필수적인 응용 분야에서 특히 치명적입니다. 기존의 효율적인 아키텍처들은 선형 어텐션을 균일하게 적용하거나 분류를 위해 설계된 토큰 프루닝(token pruning) 전략에 의존하는 경우가 많아, 효율성을 희생하지 않으면서 필요한 국소적 표현력을 회복하는 데 실패합니다.

2. 방법론

저자들은 전역적 효율성과 국소적 정밀도의 균형을 맞추도록 설계된 하이브리드 어텐션 메커니즘인 HSMLA(Hierarchical Softmax Multi-scale Linear Attention)를 제안합니다. 이 아키텍처는 비용이 많이 드는 전체 소프트맥스(softmax) 연산을 결정적인 공간 영역의 희소한 집합에만 할당하고, 나머지 영역에는 효율적인 선형 어텐션에 의존한다는 원칙에 따라 작동합니다.

핵심 구성 요소:

Multi-Scale ReLU Linear Attention (Global Backbone):
EfficientViT에서 영감을 받은 HSMLA는 전역 문맥을 효율적으로 계산하기 위해 ReLU 기반 선형 어텐션을 사용합니다. 입력 토큰은 다양한 수용 영역(receptive fields)에서 국소 문맥을 포착하기 위해 다중 스케일 심층 컨볼루션(kernel size 3×3, 5×5, 7×7 사용)을 통해 투영되고 풍부해집니다. 선형 어텐션 메커니즘은 결합 법칙(associativity)을 활용하여 복잡도를 $O(Nd^2)$ 로 줄이며, 레이어당 한 번씩 전역 집계 항( $Z$ 및 $D$ )을 계산하고 이를 모든 쿼리에 재사용합니다.
Hierarchical Block-Level Gating:
선형 어텐션의 국소적 선명도 부족 문제를 해결하기 위해, HSMLA는 토큰 레벨이 아닌 블록 레벨(연속적인 공간 타일)에서 작동하는 학습된 게이팅 메커니즘을 도입합니다.
- Head-Level Scoring: 각 어텐션 헤드는 다중 스케일 특징을 기반으로 게이팅 점수를 계산합니다.
- Block-Level Aggregation: 이 점수들은 집계되어 각 공간 블록에 대한 공유 게이트 $M_b \in [0,1]$ 를 생성합니다.
- Routing: 게이트 값이 임계값 $\tau$ (통상적으로 약 30%의 블록을 선택)를 초과하는 블록만 로컬 소프트맥스 정밀화 경로로 라우팅됩니다. 이 블록 레벨 접근 방식은 연속적인 메모리 액세스를 보장하고 per-token 분기 발산(branch divergence)을 방지하여 하드웨어 친화적입니다.
Two-Path Residual Kernel:
선형 어텐션을 소프트맥스로 이진 교체하는 대신, HSMLA는 잔차 정밀화(residual refinement) 전략을 채택합니다:
- Global Path: 모든 토큰은 선형 어텐션 경로( $O_{linear}$ )를 통해 전역 문맥을 유지합니다.
- Local Path: 선택된 블록에 대해서는 작은 윈도우( $w \times w$ ) 내에서 로컬 소프트맥스 어텐션이 계산됩니다.
- Residual Correction: 최종 출력은 전역 선형 출력과 희소 잔차 보정( $R_i = O_{local} - O_{lin-local}$ )의 합입니다. 이 공식은 전역 문맥이 결코 버려지지 않음을 보장하며, 값비싼 소프트맥스 연산을 국소 이웃으로 제한하여 복잡도를 $O(\alpha N^2 d)$ 에서 $O(\alpha N w^2 d)$ 로 줄입니다.
Hardware-Aware Design:
학습 과정에는 하드웨어 친화적인 활성화 패턴을 강제하기 위한 두 가지 정규화 항이 포함됩니다:
1. Budget Term: 예측 가능한 FLOPs와 메모리 트래픽을 보장하기 위해 목표 소프트맥스 예산 비율( $\rho$ )로부터의 편차에 페널티를 부여합니다.
2. Spatial Smoothness Term: 인접한 블록 게이트 간의 급격한 변화에 페널티를 부여하여, 흩어진 활성화가 아닌 연속적인 타일 클러스터를 유도합니다. 이는 메모리 계층 구조(HBM $\to$ L2 $\to$ SRAM) 사용을 최적화하고 커널 런치 오버헤드를 줄입니다.
Block Composition:
HSMLA는 심층 컨볼루션 및 피드 포워드 네트워크(FFN)와 어텐션 모듈을 교차 배치하는 "샌드위치" 아키텍처(EfficientViT와 유사)에 통합됩니다.

3. 주요 기여

한계 분석: 본 논문은 표준 선형 어텐션이 특히 의료 영상과 같이 밀집 예측에 필요한 날카로운 국소 구조를 모델링하는 데 실패한다는 점을 식별했습니다.
HSMLA 아키텍처: 다중 스케일 선형 어텐션과 콘텐츠 인식형 쿼리 희소 선택적 소프트맥스 게이팅을 결합한 새로운 계층적 어텐션 모듈을 제안합니다.
잔차 정밀화 전략: 전역 선형 집계와 희소 국소 보정을 분리하여, 전역 문맥을 유지하면서도 국소 정밀도를 더하는 이중 경로 공식을 제시합니다.
하드웨어 공동 설계: 학습된 게이팅 마스크가 GPU 메모리 계층 및 실행 패턴과 일치하도록 하기 위한 특정 정규화 손실과 타일 압축 추론 파이프라인을 도입했습니다.

4. 실험 결과

저자들은 네 가지 비전 설정(시맨틱 세그멘테이션, 이미지 분류, 초해상도, 생물 의학 영상)에 대해 HSMLA를 평가했습니다.

시맨틱 세그멘테이션 (Cityscapes & ADE20K):
- HSMLA-B2는 Cityscapes에서 81.8% mIoU를 달성했으며, SegFormer-B2 대비 3.7배 빠른 속도를 보였습니다.
- ADE20K에서는 동일한 3.7배 속도로 45.9% mIoU에 도달했습니다.
- 정확도(+1.5 mIoU)와 지연 시간(latency) 모두에서 EfficientViT-B2를 능가했습니다.
이미지 분류 (ImageNet-1K):
- HSMLA-B2는 82.1% Top-1 정확도를 달성하여 ConvNeXt-B와 대등한 성능을 보였으며, Swin Transformer-B 대비 2.3배 빠른 속도로 EfficientViT-B2를 앞질렀습니다(+0.3%).
생물 의학 영상 (Biomedical Imaging):
- CT 장기 세그멘테이션 (BTCV): 87.3% Dice 및 8.6 mm HD95를 달성하였으며, 정확도와 경계 묘사 모두에서 TransUNet 및 SwinUNet을 크게 능가하며 3.2배 빠른 속도를 기록했습니다.
- 병리 WSI (CAMELYON16): 94.2% AUC와 4.1배 빠른 속도에 도달했습니다.
- 수술 도구 세그멘테이션 (EndoVis 2017): Jetson AGX Orin에서 메모리 사용량을 줄이면서 42 FPS로 89.7% Dice를 달성했습니다.
- 폐 결절 탐지 (LUNA16): 95.0% 민감도와 2.9배 빠른 속도를 나타냈습니다.
초해상도 (DIV2K):
- HSMLA-B1은 Restormer 대비 4.2배 빠른 속도로 31.92 dB PSNR을 달성했습니다.

5. 의의 및 주장

본 논문은 HSMLA가 선형 어텐션의 효율성과 소프트맥스 어텐션의 국소 표현력 사이의 간극을 성공적으로 메운다고 주장합니다. 쿼리의 희소한 부분집합만을 풀 소프트맥스 어텐션으로 라우팅함으로써, HSMLA는 밀집 예측 작업 전반에서 정확도를 유지하거나 향상시키면서 최대 4.2배의 추론 시간 가속을 달립다.

저자들은 HSMLA가 특히 다음과 같은 고해상도 및 의료 영상 응용 분야에서 매우 중요하다고 강조합니다:

모델이 에지 디바이스(예: Jetson 클래스 GPU)에서 엄격한 지연 시간 및 메모리 예산 하에 작동해야 하는 경우.
미세한 구조적 세부 사항(장기 경계, 종양 마진)을 보존하는 것이 결정적인 경우.
전역 문맥을 희생하지 않으면서 진단적으로 유의미한 영역에 계산을 동적으로 할당할 수 있는 경우.

이 연구는 다중 스케일 선형 어텐션과 콘텐츠 인식형 쿼리 희소 소프트맥스 정밀화를 결합하는 것이 자원이 제한된 환경에서 표준 이차(quadratic) 어텐션 메커니즘에 대한 실질적이고 효과적인 대안이 될 수 있음을 입증합니다.