BASA가 창문 이동만으로 고해상도 생성 어텐션을 4.52배 가속한다

무엇인가

디퓨전 트랜스포머(DiT)는 이미지·영상 생성에서 강력한 성능을 내지만, 전체 어텐션의 제곱 복잡도 때문에 해상도를 올릴수록 비용이 급격히 커진다. 기존 선택지는 두 갈래였다. Swin처럼 격자를 겹치지 않는 창문으로 나누는 분할 창문 어텐션은 정규적인 계산 패턴 덕분에 이론적 복잡도만큼 실측 가속이 나오지만, 창문 사이 상호작용이 끊겨 격자 모양 아티팩트가 생긴다. 반대로 세밀한 슬라이딩 창문 어텐션은 경계를 넘는 상호작용을 복원해 화질을 살리지만, 불규칙한 계산 패턴 때문에 인덱싱·메모리 접근 오버헤드가 커서 이론 대비 실측 가속 격차가 크고, 하드웨어 백엔드마다 전용 커널을 따로 만들어야 한다. 이 논문은 창문 간 소통을 확보하면서도 정규적인 계산과 백엔드 이식성을 유지할 수 있는가를 질문으로 삼는다.

어떻게 동작하나

BASA의 첫 번째 구성 요소는 이동 창문 희소 어텐션이다. 잠재 격자를 w_h×w_w 크기의 겹치지 않는 창문으로 나누되(영상은 각 창문이 전체 시간축을 유지해 시간 방향 상호작용을 보존한다), 레이어마다 창문 경계를 다르게 어긋나게 한다. 후보 오프셋을 a_i = (iδ) mod w, δ = max(1, floor(w/L))로 만들고, 이를 순서대로 배정하지 않고 앞 절반과 뒤 절반을 교차 배치해 (b_0,…,b_{L-1}) = (a_0, a_K, a_1, a_{K+1}, …, a_{K-1}, a_{L-1})로 재배열한다. w=8, L=4라면 (0,2,4,6)이 (0,4,2,6)이 된다. 여기에 디노이징 스텝별 오프셋까지 더해 Δ_{l,s} = ((b_l^h + c_s^h) mod w_h, (b_l^w + c_s^w) mod w_w)로 최종 이동량을 정한다. 각 레이어에서 격자를 -Δ만큼 순환 이동해 창문으로 자르고, 사전학습된 Q/K/V 투영으로 표준 scaled dot-product softmax 어텐션을 수행한 뒤 다시 Δ만큼 되돌린다. 한 레이어에서 경계로 갈라졌던 토큰이 다음 레이어에서는 같은 창문에 들어오게 되어 전역 정보 교환이 이루어지고 격자 아티팩트가 사라진다. 어텐션 비용은 O(N²d)에서 O(N·n_w·d)로 줄어든다.

무엇과 다른가

두 번째 요소는 K/V 풀링 전역 메모리다. 이동 창문만으로는 수용 영역이 여전히 국소적이라 고해상도에서 레이아웃 일관성이 흔들릴 수 있다. 그래서 키와 값을 잠재 격자 형태로 되돌려 공간 영역별로 평균 풀링한 K̄, V̄를 만들고, 이를 각 지역 창문의 키·값 쪽에만 이어 붙인다(쿼리는 그대로). 쿼리 시퀀스와 출력 길이가 변하지 않으므로 모델 상태에 추가 토큰이 생기지 않고, 기존 DiT 어텐션 모듈에 그대로 끼워 넣을 수 있다. 영상에서는 공간 방향으로만 풀링하고 시간축은 유지해 시간 해상도를 뭉개지 않는다. 다만 공간 풀링은 인스턴스 수준의 세밀한 특징을 버릴 수 있어, 비슷한 객체가 다른 위치에 있을 때 이를 구분하지 못하고 반복 객체를 만들어낼 수 있다. 이를 보완하는 것이 세 번째 요소인 내용 적응형 실제 창문 라우팅이다. 각 쿼리 창문의 쿼리 토큰을 평균 내 서술자 q_i를 만들고, 후보 창문마다 키 서술자 k_j를 만들어 s_ij = q_i·k_j^T/√d로 관련도를 잰다. 무작정 전역 검색을 하지 않고 후보 집합 C_i 안에서만 TopK를 뽑아 R_i를 정하고, 선택된 창문은 압축되지 않은 원본 키·값 토큰을 그대로 제공한다. 최종적으로 각 쿼리 창문은 자기 지역 창문, 라우팅된 창문, 풀링 전역 메모리로 구성된 S_i = W_i ∪ R_i ∪ G에 어텐션한다.

어떻게 쓰나

네 번째 요소는 dense 교사 기반 희소 적응이다. Q/K/V 투영과 softmax 형식은 그대로 두지만 수용 영역이 바뀌면 사전학습된 dense 모델과의 괴리가 생기므로, dense DiT를 얼려둔 교사로 두고 희소 학생을 LoRA로 적응시킨다. 손실은 세 가지다. 플로우 매칭 목표와의 L_fm, 학생과 교사 예측 사이의 L_pred, 그리고 각 레이어의 어텐션 출력을 맞추는 L_attn(MSE)이다. 어텐션 맵이 아니라 어텐션 출력을 정렬하기 때문에 dense 어텐션 행렬을 명시적으로 물질화하지 않아도 된다. 이미지는 L_img = L_fm + αL_pred + βL_attn, 영상은 여기에 플로우 매칭 스케줄러의 타임스텝 가중 w(t)를 곱한 L_vid를 쓴다. 적응이 끝나면 교사는 버리고 희소 학생만으로 추론한다.

전제와 한계

실험은 단일 A100에서 FLUX.1-dev와 Wan2.1-T2V-1.3B로 수행했다. 이미지 해상도 외삽은 COCO2014 검증셋에서 캡션 1,000개를 뽑아 FLUX로 1024×1024 원본을 만든 뒤 SDEdit 파이프라인으로 2048×2048까지 외삽했고, 20 디노이징 스텝, SDEdit 가이던스 7.5, 10K FLUX 생성 샘플로 어텐션 레이어 증류를 했다. 비교 대상은 CLEAR, Swin, STA이며 모두 동일한 dense 교사 증류를 적용했다. BASA는 w=32에서 CLIP-I 97.60, DINO 95.95를 유지하면서 지연을 2.443배 줄였고, CLIP-T와 IS는 dense FLUX와 비슷하거나 약간 앞섰다. w=8은 화질 저하가 더 뚜렷한데 지연 감소는 크지 않았고, w=16이 가장 균형 잡힌 지점, w=32는 참조 보존이 중요할 때의 선택지다. 같은 창문 크기의 경직된 지역 어텐션과 비교해 BASA-32는 Swin과 증류된 STA-32보다 모든 참조 유사도 지표에서 앞섰다. CLEAR는 비슷한 화질을 내지만 어텐션 연산량이 더 필요하다. FLUX에서 실측 가속이 이론 추정치의 90%를 넘었고, 어텐션 지연은 약 9ms까지 내려갔다.

영상은 VBench에서 프롬프트 20개를 뽑아 각 5개씩 총 100개 영상을 생성했다. dense Wan이 832×480, 15fps로 81프레임 원본을 만들고, 이를 공간 양축으로 2배 외삽해 1664×960으로 키웠다(50 디노이징 스텝, CFG 5.0). BASA w=(10,10), p=(4,4) 설정은 희소도 90.00%, 지연 112.84ms, 4.52배 가속으로 가장 큰 폭을 기록했고, w=(15,26), p=(4,4)는 커널 효율이 더 높으면서 4.04배를 냈다. 풀링 계수를 (4,4)에서 (2,2)로 낮추면 전역 K/V 토큰이 늘어 희소도가 90.00%에서 70.17%로, 가속이 4.52배에서 1.99배로 떨어진다. 품질 쪽에서는 품질 우선 설정인 w=(10,10), p=(2,2)가 주체 일관성과 배경 일관성에서 각각 0.12%, 0.22%만 떨어지고 운동 매끄러움과 이미징 품질은 오히려 조금 올랐다. 빠른 설정인 w=(10,10), p=(4,4)도 4.52배 가속에서 주체·배경 일관성과 미학 품질이 소폭만 줄었고, 운동 매끄러움은 거의 그대로, 이미징 품질은 65.22%에서 67.02%로 올랐다. CLEAR는 VBench 점수는 강하지만 실측 가속이 1.05배에 그쳤고, Swin은 이론 희소도는 높지만 역동성·미학·이미징 품질이 크게 떨어졌으며, STA는 트리톤 커널을 써도 가속 폭이 작았다. 이 비교가 보여주는 핵심은 FLOP 감소만으로 실측 가속이 보장되지 않는다는 점이다.

이미지 절제 실험에서는 풀링 전역 메모리를 제거하고 이동 창문만 남기면 PSNR이 24.09에서 23.33으로, CLIP-I이 97.60에서 96.86으로, DINO가 95.95에서 94.72로 떨어지고 FID는 20.97에서 24.05로 나빠졌다. 그래도 고립된 지역 창문보다는 크게 앞서므로 창문 이동 자체의 효과는 확인된다. 두 요소를 함께 쓸 때 여섯 개 유사도 지표에서 가장 좋은 결과가 나와 상보적이라는 점이 드러난다. CLIP-T와 IS는 상대적으로 안정적이어서, 이 기법이 기본적인 텍스트 정렬보다 dense 참조 일관성을 개선하는 쪽에 기여한다는 해석이 가능하다.

개발자 관점에서 이 논문의 실질적 가치는 커스텀 커널 없이 기존 어텐션 백엔드에 그대로 얹을 수 있다는 데 있다. 순환 이동과 창문 재배열, 풀링, TopK 라우팅 모두 표준 SDPA 호출로 표현되므로, 이미 FLUX나 Wan 같은 DiT 파이프라인을 돌리고 있고 고해상도에서 어텐션 지연이 병목이라면 백엔드별 커널을 새로 작성하지 않고도 시도해볼 수 있다. 다만 도입 전에 확인할 것은 세 가지다. 창문 크기 w와 풀링 계수 p가 품질·속도 절충을 좌우하므로 워크로드에 맞는 조합을 실측해야 하고, 라우팅의 후보 집합 설계가 전역 검색 비용과 지역 일관성 사이에서 추가적인 튜닝 지점이 되며, dense 교사 증류 단계가 필요하므로 사전학습된 dense 모델과 그 출력을 만들 수 있는 환경이 전제된다.

한계와 전제는 저자들이 결론과 각 절에서 밝힌다. 결론은 적절한 크기의 창문에서 품질이 안정적이라고 표현하는데, 실제로 w=8처럼 창문이 너무 작으면 화질 저하가 뚜렷한 반면 지연 감소는 크지 않아 실효 구간이 제한된다. 공간 풀링은 인스턴스 수준의 세밀한 특징을 버릴 수 있어 유사 객체가 다른 위치에 있을 때 반복 객체를 유발할 수 있고, 이를 보완하는 라우팅도 무구조 전역 검색을 피하려고 후보 집합 C_i로 제한하기 때문에 후보 설계에 따라 얻을 수 있는 장거리 정보가 달라진다. 또한 희소 학생의 성능은 dense 교사와의 증류에 의존하므로, 교사 모델과 적응 데이터(이미지의 경우 10K 샘플)를 준비해야 한다는 전제가 붙는다. 논문은 별도의 한계 절을 두지는 않았고, 위 제약들은 방법 설명과 실험 해석 과정에서 드러난다.