FlashSwin이 Swin의 윈도 메모리 한계를 없애 큰 창과 촘촘한 토큰을 가능하게 한다
FLASHSWIN: Unlocking Large Windows and Dense Tokens in Swin Vision Transformers with Memory Efficient Attention
무엇인가
Swin과 SwinV2 같은 계층적 윈도 트랜스포머는 어텐션을 국소 창 안으로 제한해 비용을 토큰 수에 선형으로 만들었지만, 실제로는 창 크기 M=8~16, 패치 크기 4×4라는 거친 설정에 묶여 있었다. 이유는 메모리다. 표준 윈도 어텐션은 창마다 M²×M² 점수 행렬을 실제로 만들기 때문에 창당 O(M⁴), 특징맵 전체로는 O(HW·M²) 메모리를 쓴다. 여기에 Swin은 학습된 상대 위치 바이어스를 점수 행렬에 원소별로 더하는데, 이 항 역시 점수 행렬과 그 기울기를 전부 물질화해야 한다. 결과적으로 토큰 밀도, 창 크기, 입력 해상도 세 축이 모두 작게 유지되고, 고해상도 파인튜닝은 정확도 이득을 다 쓰기도 전에 하드웨어 메모리 한계에 부딪힌다.
어떻게 동작하나
FlashSwin은 Swin의 계층 구조와 어텐션 패턴은 그대로 두고, 물질화된 행렬이나 거친 그리드를 전제하는 세 부품만 교체한다. 첫째, 윈도 내 어텐션을 FlashAttention 기반 IO 인지 융합 커널로 계산한다. 타일 단위로 온라인 소프트맥스를 누적해 N×N 행렬을 저장하지 않으므로 창당 보조 메모리가 O(M⁴)에서 O(M²)로, 스테이지 단위로는 대략 O(HW)로 줄어든다. 근사가 아니라 동일한 정확 소프트맥스 어텐션이다. 둘째, 융합 커널이 만들지 않는 점수 행렬에 더해지던 상대 위치 바이어스를 버리고, 쿼리와 키에 적용하는 윈도 로컬 학습형 2D RoPE로 대체한다. 헤드 차원 d를 절반으로 나눠 앞 절반은 세로 좌표 y, 뒤 절반은 가로 좌표 x를 인코딩하고, 두 축이 하나의 학습 주파수 집합 θ를 공유한다. 회전이 R(p)ᵀR(p′)=R(p′−p)로 합성되므로 내적이 변위 (x′−x, y′−y)에만 의존해, 덧셈 항 없이도 상대 위치 구조가 보존된다. 주파수는 θᵢ=exp(ωᵢ)로 표준 기하 RoPE 스케줄에서 초기화해 네트워크와 함께 최적화하며, 헤드당 학습 파라미터는 d/2개뿐이다. 좌표는 각 창의 원점 기준으로 국소 정의되고, 시프트 윈도에서는 순환 이동 후 창별로 다시 계산하며 래핑된 창의 마스킹은 커널에 덧셈 마스크로 전달한다. Swin의 이산 바이어스 테이블이 M=32에서 헤드당 3969개 항목을 저장하고 학습 중 본 오프셋에만 정의되는 것과 달리, 이 방식은 창 크기와 무관하게 임의 해상도에서 보간 없이 평가된다. 셋째, 겹치지 않는 패치 임베딩을 겹치는 스템으로 바꾼다. 커널 2p−1, 스트라이드 p, 패딩 p−1의 합성곱으로 스트라이드는 유지해 토큰 수는 같지만 인접 토큰의 수용 영역이 겹쳐 패치 경계를 넘는 고주파 디테일이 살아남는다.
무엇과 다른가
효율 결과는 창 크기 의존성의 제거로 요약된다. 패치 4에서 FlashSwin-T는 M을 8에서 32로 키워도 12.4GB로 변하지 않는 반면, Swin-T는 15.7→90.1GB(5.7배), SwinV2-T는 15.0→70.4GB로 늘어난다. 처리량도 Swin-T는 1173→442 img/s, SwinV2-T는 871→371 img/s로 떨어지지만 FlashSwin-T는 약 1K img/s를 유지하며, M=32에서 Swin-T보다 2.4배 빠르다. M=8에서는 융합 커널 오버헤드를 상각하지 못해 Swin-T가 약간 빠르지만 M=16부터 역전된다. 패치 2의 밀집 그리드에서는 격차가 더 커진다. FlashSwin-T는 창 크기와 무관하게 47.7GB, 약 290 img/s를 유지하는데, Swin 두 변형은 M=16에서 이미 120GB를 넘고 M=32에서는 Swin-T가 체크포인팅으로도 학습되지 않으며 SwinV2-T는 체크포인팅으로 138.9GB·57 img/s에 그친다. FlashSwin은 같은 조건에서 체크포인팅 없이 47.7GB·272 img/s로 메모리 2.9배 적고 처리량 4.8배 높다. 패치 4→2로 토큰이 4배 늘 때 메모리는 3.85배(12.4→47.7GB)만 늘고 창 확대는 사실상 비용이 없다.
어떻게 쓰나
정확도에서도 이득이 확인된다. 동일 연산(5.9 GFLOPs, M=8·p=4)에서 FlashSwin-T는 ImageNet-1K 82.3%로 SwinV2-T(81.9%)를 앞선다. M=16에서 82.9%, ImageNet-V2는 70.8→71.9%로 오르지만 M=32는 FLOPs를 27% 더 쓰고도 IN-1K 개선이 없고 V2에서 +0.1에 그친다. 가장 큰 이득은 토큰 밀도에서 나온다. M=32·p=2에서 같은 약 28M 모델이 IN-1K 84.1%, IN-V2 73.4%로 M=8·p=4 대비 각각 1.8, 2.6포인트 오르며, 이는 약 28M 파라미터로 Base급(약 84.5%) 정확도에 근접한 것이다. 분포 이동에서는 격차가 더 벌어진다. M=32·p=2에서 IN-A 42.8%, IN-R 46.6%, IN-Sketch 33.0%, IN-C mCE 50.7로, SwinV2-T 대비 IN-A 개선이 동일 구성의 +3.4포인트에서 +18.1포인트로 커지는 반면 클린 IN-1K에서는 +2.2포인트에 불과하다. 융합 어텐션만으로는 이 이득이 설명되지 않는다는 통제 실험도 있다. 덧셈 위치 바이어스를 랭크 16 분해로 유지한 FlashBias-T는 M=32·p=4에서 82.7%로 FlashSwin-T와 0.2포인트 차이지만 파라미터가 29.5M으로 더 많고, IN-A 4.5, IN-R 2.6, Sketch 3.0포인트 뒤진다. 구성 요소 제거 실험은 바이어스를 그냥 버리고 융합 커널만 쓰면 80.0%(SwinV2-T보다 1.9포인트 낮음), 학습형 윈도 로컬 RoPE를 넣으면 81.8%로 SwinV2-T와 동등, 겹치는 스템이 +0.5%를 더해 82.3%가 된다.
전제와 한계
다운스트림에서도 창과 밀도가 서로 다른 일을 한다. ADE20K UPerNet에서 FlashSwin-T는 M=32에서 46.88 mIoU / 35.55 Boundary-IoU / 49.72 BF를 기록하고, 분류와 달리 M이 커질수록 계속 좋아진다. M=32 고정에서 패치를 4→2로 줄이면 mIoU는 0.40 오르는 데 그치지만 Boundary-IoU는 0.87, BF는 1.34 올라 약 2~3배 이득이 경계에 집중된다. 반대로 밀도 고정에서 창을 8→32로 넓히면 mIoU 2.49, Boundary-IoU 2.44, BF 2.36이 함께 오른다. FlashBias-T는 같은 M=32에서 mIoU 1.73, Boundary-IoU 1.78, BF 1.42 뒤진다. COCO Mask R-CNN(512×512)에서 FlashSwin-T는 M=8에서 38.5 box AP / 35.3 mask AP로 SwinV2-T(37.3/34.4)를 앞서고, M=32·p=2에서 44.1/39.9까지 오르며 Swin-T보다 7.0/5.7포인트 높다. 이득은 작은 객체에 집중돼 AP_S가 박스 20.1→27.1, 마스크 14.4→19.1로 오르는 동안 AP_L은 59.7로 변하지 않는다. 약 197M 파라미터의 FlashSwin-L은 IN-22K로 192², p=2, M=32로 90에폭 사전학습해 IN-1K 86.91%(SwinV2-L 86.93%, Swin-L 86.51%), IN-V2 78.02%(77.48%, 76.81%), ADE20K 52.58 mIoU로 Swin-L보다 +2.45, SwinV2-L보다 +4.23포인트를 얻고 Boundary-IoU +3.02, BF +2.98을 더한다.
개발자 관점에서 이 논문의 실질적 메시지는 창 크기가 더 이상 메모리 급증의 원인이 아니라는 점이다. 고해상도 의료영상, 병리 슬라이드, 항공영상처럼 다운샘플링이 과제 신호를 지워버리는 도메인에서 백본을 직접 학습시킬 때, 창과 패치를 키우는 선택지가 열린다. 다만 토큰 수 자체는 여전히 메모리를 지배하므로 p=2 구성은 FLOPs를 5.9에서 37.4 GFLOPs로 늘리는 맞바꿈이고, 체크포인팅 없이 배치 128을 돌릴 수 있는지가 실무 판단 기준이 된다. 또한 FlashAttention을 윈도 어텐션에 그대로 끼워 넣으려는 시도는 위치 정보 손실로 정확도가 1.9포인트 떨어진다는 점을 기억해야 하며, RoPE를 쿼리·키로 옮기는 변경이 사실상 필수다. 재현 시에는 융합 커널이 bf16 Q/K/V를 받고 RoPE를 fp32로 적용하며 소프트맥스를 fp32로 누적한다는 혼합 정밀도 설정을 확인해야 한다.
저자들이 밝힌 한계는 분명하다. FlashSwin은 창 크기에 대한 메모리 의존을 없앴지만 토큰 밀집화는 여전히 연산을 늘리므로 p=2는 FLOPs를 추가로 지불하는 선택이다. 현재 벤치마크는 상대적으로 낮은 해상도에서 수행돼 네이티브 고해상도 처리를 정당화하는 영역을 충분히 검증하지 못했다. 패치 4 대 2 비교 실험도 ImageNet-1K의 중앙 단변 길이가 374픽셀이고 512픽셀에 도달하는 이미지가 6.4%뿐이라, 거친 패치가 업샘플된 이미지에서 세밀한 패치와 비슷해진다는 것만 보여줄 뿐 네이티브 고해상도 입력이 불필요하다는 결론은 아니라고 저자들은 못 박는다. 기가픽셀 병리, 고해상도 의료영상, 3D 볼륨과 비디오로 확장하는 것이 향후 과제로 제시된다.