효율적인 비디오 VLM을 위한 통합 시공간 토큰 스코어링

비디오를 이해하는 AI 모델은 수많은 프레임을 처리해야 하므로 연산량이 매우 많고 속도가 느립니다. 이 논문은 중요하지 않은 시각 정보를 지능적으로 골라내 삭제함으로써, 성능 저하를 최소화하면서도 모델의 작동 속도를 획기적으로 높이는 STTS 기법을 제시하여 고성능 비디오 AI의 대중화 가능성을 열었습니다.

Transformer 기반의 비디오 모델은 각 프레임을 수백 개의 패치 토큰으로 분할하여 처리합니다. 프레임 수가 늘어날수록 토큰 시퀀스가 길어지며, Attention Mechanism의 연산 비용이 시퀀스 길이 N의 제곱에 비례해 증가하기 때문에 메모리 부족과 속도 저하 문제가 발생합니다.

STTS는 모든 토큰이 동일하게 중요하지 않다는 점에 착안합니다. 배경처럼 고정된 부분이나 이전 프레임과 거의 차이가 없는 중복된 정보는 과감히 삭제합니다. 이를 위해 모델의 최종 출력에서 전달되는 Gradient를 통해 공간적 중요도를 배우고, 인접 프레임 간의 유사도를 측정하는 Auxiliary Loss를 통해 시간적 중복성을 파악합니다.

이렇게 걸러진 토큰들은 '패킹' 과정을 거쳐 빈틈없이 모아집니다. 결과적으로 모델은 꼭 필요한 정보에만 Attention을 집중할 수 있게 되어, 훨씬 적은 연산량으로도 긴 비디오의 핵심 맥락을 정확하게 파악할 수 있게 됩니다.

STTS는 ViT의 특정 레이어 뒤에 삽입되는 플러그인 모듈이다. 입력 토큰 $X \in \mathbb{R}^{T \times N \times D}$를 받아 공간적 풀링을 거친 후, 3-Layer MLP를 통해 각 패치의 중요도 점수를 계산한다.

공간적 중요도 학습을 위해 Bias Injection 방식을 사용한다. [계산된 점수 $S$ 입력 → Attention Matrix에 로그 스케일로 더하는 연산 수행 → 가중치가 조정된 Attention 결과 획득] 과정을 거친다. 이 값의 의미는 텍스트 조건 없이도 하위 태스크의 Gradient만으로 중요한 영역을 스스로 식별하게 만드는 것이다.

시간적 중복성은 Neighboring-Frame Cosine Similarity를 활용한 Auxiliary Loss $\mathcal{L}_{sim}$으로 제어한다. [인접 프레임 패치 특징 벡터 입력 → L2 정규화 후 내적 연산 수행 → 0~1 사이의 유사도 값 산출] 과정을 거친다. 유사도가 1에 가까울수록 중복된 정보임을 의미하며, 이를 통해 변화가 적은 토큰에 낮은 점수를 부여하도록 유도한다.

가지치기 후 발생하는 Sparse한 토큰들을 처리하기 위해 First-Fit Descending 패킹 알고리즘을 적용한다. [프레임별 생존 토큰 수 정렬 → 가용 용량이 있는 빈(bin)에 순차적 배치 → 조밀한 텐서 생성] 과정을 통해 PyTorch 등 프레임워크에서 실제 하드웨어 연산 속도 향상을 이끌어낸다.

Molmo2-4B 모델을 기반으로 13개 비디오 QA 벤치마크에서 평가한 결과, 50% 토큰 제거 시 평균 성능 하락은 0.7%에 불과했다. 특히 NextQA와 VideoMME 같은 복잡한 추론 태스크에서는 30% 제거 시 오히려 성능이 소폭 향상되는 현상이 관찰됐다. 이는 불필요한 노이즈 토큰이 제거되어 Attention Mechanism이 핵심 정보에 더 잘 집중했기 때문이다.

효율성 측면에서 256프레임 설정 시 학습 처리량은 2.25배, 추론 속도는 2.22배 향상됐다. 이는 토큰 수가 많아질수록 STTS의 효율 개선 효과가 Quadratic하게 커짐을 의미하며, 긴 비디오 처리에 매우 유리하다.

Ablation Study를 통해 Auxiliary Loss의 중요성을 입증했다. 보조 손실이 없는 경우 성능이 Random Pruning보다 낮아졌는데, 이는 LLM 자체가 미세한 시간적 중복성을 스스로 파악하기 어렵기 때문에 명시적인 신호가 필수적임을 시사한다.

STTS는 ViT 인코더 내부에서 토큰을 영구적으로 제거(Hard Pruning)하는 방식을 취한다. 이는 ViT 출력 이후에만 토큰을 줄이는 기존 Post-ViT 방식들과 달리, ViT 후반부 레이어와 LLM 전체의 연산량을 동시에 줄이는 효과를 가져온다.

Scorer 아키텍처는 Self-Attention 기반의 Token Pooler와 3-Layer MLP로 구성된다. 현재 프레임과 이전 프레임의 특징을 Concatenate하여 입력함으로써 시간적 맥락을 확보하며, 첫 번째 프레임은 비교 대상이 없으므로 항상 보존하여 기준점으로 삼는다.

학습 시에는 Differential Learning Rate를 적용한다. LLM은 1e-5, ViT는 5e-6, STTS 모듈은 1e-4로 설정하여 스코어링 모듈이 빠르게 수렴하도록 유도한다. 이는 고정된 백본 모델의 성능을 유지하면서도 효율적인 가지치기 정책을 학습하는 데 최적화된 설정이다.

Test-time Scaling 기법을 통해 추론 시 프레임 수를 늘려 성능을 추가로 개선할 수 있다. 50% 제거 모델에서 프레임 수를 2배로 늘리면(128프레임), 원래 64프레임 모델과 동일한 연산 비용으로 0.5~1%의 성능 이득을 얻을 수 있음을 확인했다.

ViT의 매우 이른 레이어(l=0, 1)에서 가지치기를 수행할 경우, 충분한 특징 추출이 이루어지지 않아 성능이 크게 저하되는 한계가 있다. 또한, 인접 프레임 간의 코사인 유사도에 기반한 보조 손실 함수는 매우 빠른 움직임이 있는 영상에서는 최적이 아닐 수 있다.