[Paper] PSA: Pyramid Sparse Attention을 이용한 효율적인 비디오 이해 및 생성

Source: arXiv - 2512.04025v1

Overview

이 논문은 Pyramid Sparse Attention (PSA) 라는 새로운 어텐션 모듈을 소개한다. PSA는 비디오 모델에서 자체 어텐션의 2차 비용을 크게 줄이면서도 대부분의 유용한 정보를 유지한다. 하드 바이너리 마스크를 다중 레벨 풀링된 키‑밸류(KV) 표현으로 대체함으로써, PSA는 속도와 정확도 사이의 보다 세밀한 절충을 제공한다. 이는 비디오 이해(예: 행동 인식)와 비디오 생성(예: 텍스트‑투‑비디오 합성) 모두에 실용적이다.

Key Contributions

• 피라미드‑스타일 KV 풀링: 전체 KV 블록을 버리는 대신, PSA는 각 블록에 대해 여러 풀링 버전(저해상도부터 고해상도까지)을 만들고 각 쿼리가 사용할 레벨을 스스로 결정한다.
• 쿼리당 동적 할당: 쿼리는 중요한 영역에 대해서는 고해상도 KV에, 덜 중요한 영역에 대해서는 저해상도 KV에 자동으로 어텐션을 수행하여, 전체 어텐션과 공격적인 프루닝 사이의 “보간”을 구현한다.
• 하드웨어‑친화적 커널: 저자들은 블록‑타일 구현을 설계해 GPU/TPU에 깔끔하게 매핑되도록 하였으며, 많은 희소 어텐션 기법에서 발생하는 불규칙 메모리 접근을 피한다.
• 이해와 생성 모두에 통합: PSA는 구분적 비디오 작업(예: Kinetics, Something‑Something)과 생성 작업(예: 텍스트‑투‑비디오 디퓨전) 모두에서 성능을 입증해 그 다재다능함을 보여준다.
• 오픈‑소스 공개: 코드, 사전 학습 가중치, 바로 실행 가능한 커널을 공개해 채택 장벽을 낮춘다.

Methodology

1. 블록‑와이즈 어텐션 기반 – 입력 비디오는 고정 크기의 쿼리, 키, 밸류 블록으로 나뉜다(일반적인 “블록‑스파스” 설정).
2. 다중 레벨 풀링 – 각 키/밸류 블록에 대해 PSA는 작은 피라미드를 만든다:
   • Level 0: 원본(전체 해상도) KV.
   • Level 1, 2, …: 점진적으로 풀링된(예: 평균 풀링) 버전으로, 공간/시간 해상도가 감소한다.
3. 쿼리‑주도 선택 – 경량 스코어링 네트워크가 특정 쿼리 블록에 대한 각 KV 블록의 관련성을 평가한다. 스코어에 따라 쿼리는 적절한 피라미드 레벨을 선택한다: “중요한” 블록에는 고해상도, “덜 중요한” 블록에는 저해상도.
4. 보간 및 집계 – 선택된 풀링 KV는 필요에 따라 업샘플링되고, 표준 스케일드‑닷‑프로덕트 어텐션 공식에 의해 쿼리와 결합된다. 풀링이 결정적이므로 모든 레벨을 통해 그래디언트가 흐르고, 엔드‑투‑엔드 학습이 가능하다.
5. 효율적인 커널 – 구현은 블록을 타일로 그룹화하고, 각 타일을 고정된 연산 예산으로 처리하며, 동적 메모리 할당을 피하는 CUDA 커널을 활용해 일반 하드웨어에서 빠르게 동작한다.

Results & Findings

• PSA는 밀집 어텐션과의 격차를 일관되게 좁히며(≤1 % 절대 손실) 2–2.5× 속도 향상과 30–40 % 메모리 절감을 제공한다.
• 정성적으로는, 생성된 비디오가 다른 희소 방법에 비해 움직임 경계가 더 선명하고 아티팩트가 적다.
• Ablation 연구는 동적 레벨 선택이 성능의 주요 원인임을 확인했으며, 정적 단일 레벨 풀링은 바이너리 마스크 기준으로 성능이 떨어진다.

Practical Implications

• 빠른 비디오 파이프라인: 개발자는 PSA를 기존 트랜스포머 기반 비디오 모델(예: ViViT, TimeSformer)에 바로 삽입해 전체 네트워크를 재설계하지 않고도 추론 지연을 감소시킬 수 있다.
• 엣지·모바일 배포: 메모리 사용량 감소로 비디오 트랜스포머를 VRAM이 제한된 디바이스에서도 실행 가능해져, 온‑디바이스 비디오 분석이나 AR/VR 경험에 문을 연다.
• 비용 효율적인 학습: 대형 비디오 디퓨전 모델의 학습이 더 저렴해지며, 각 전·후방 패스가 적은 FLOPs를 사용해 더 큰 배치 사이즈나 더 긴 시퀀스를 처리할 수 있다.
• 하이브리드 시스템: PSA의 블록‑타일 설계는 혼합 정밀도(FP16/FP8) 학습과 잘 맞으며, 현대 GPU 파이프라인과 정렬돼 PyTorch·TensorFlow와 같은 라이브러리에 원활히 통합된다.
• 연구 가속: 오픈소스 커널은 저밀도 연구(예: 저‑랭크 팩터화 또는 학습된 토큰 프루닝과 결합) 위한 기준점을 제공한다.

Limitations & Future Work

• 블록 크기에 제한된 세분화: PSA의 효율성은 선택한 블록 차원에 의존한다; 블록이 너무 크면 매우 미세한 시간적 디테일이 손실될 수 있다.
• 정적 풀링 레벨: 피라미드 레벨은 사전에 정의된다(예: 2×, 4× 풀링). 적응형 풀링 비율은 절충을 더욱 개선할 가능성이 있다.
• 짧은 클립에 한정된 벤치마크: 실험은 ≤2 초 클립에 초점을 맞추었으며, 시간당 수시간 길이 비디오나 스트리밍 시나리오로 확장하는 것은 아직 미해결이다.
• 하드웨어 의존성: 커널이 GPU 친화적이지만, CPU나 특수 가속기(TPU 등)에서의 성능은 다를 수 있어 별도 최적화가 필요하다.

향후 연구에서는 학습 가능한 풀링 연산자, 계층적 쿼리 라우팅, 그리고 토큰‑레벨 프루닝과의 통합을 탐색해 효율성을 더욱 높이고 고품질 비디오 작업에 필수적인 풍부한 시공간 신호를 보존할 수 있을 것이다.

Authors

• Xiaolong Li
• Youping Gu
• Xi Lin
• Weijie Wang
• Bohan Zhuang

Paper Information

• arXiv ID: 2512.04025v1
• Categories: cs.CV, cs.AI, cs.LG
• Published: December 3, 2025
• PDF: Download PDF