효율적인 Video VLM을 위한 통합 Spatio-Temporal Token Scoring (STTS)

# STTS Scorer 핵심 구조 (PyTorch 의사코드)
import torch
import torch.nn as nn

class STTSScorer(nn.Module):
    def __init__(self, dim, pool_width=3):
        super().__init__()
        self.pool_width = pool_width
        # Token Pooler: Self-Attention 기반
        self.token_pooler = nn.MultiheadAttention(dim, num_heads=8, batch_first=True)
        # 3-Layer MLP for scoring (입력: 현재+이전 프레임 concat -> 2*dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim * 2, dim),
            nn.GELU(),
            nn.Linear(dim, dim // 2),
            nn.GELU(),
            nn.Linear(dim // 2, 1),
            nn.Sigmoid()
        )
    
    def forward(self, x_curr, x_prev):
        # x_curr, x_prev: (B, N, D) — 현재/이전 프레임 features
        # 1. Spatial pooling (w x w)
        x_curr_pooled = self.spatial_pool(x_curr)  # (B, N/w^2, D)
        x_prev_pooled = self.spatial_pool(x_prev)  # (B, N/w^2, D)
        
        # 2. Concat neighboring frames for temporal context
        x_concat = torch.cat([x_curr_pooled, x_prev_pooled], dim=-1)  # (B, N/w^2, 2D)
        
        # 3. Score each pooled patch
        scores = self.mlp(x_concat).squeeze(-1)  # (B, N/w^2)
        
        # 4. Expand scores back to original resolution
        scores_expanded = scores.repeat_interleave(self.pool_width**2, dim=1)  # (B, N)
        return scores_expanded

# Auxiliary Temporal Loss
def temporal_aux_loss(scores, features_l, t, pool_width=3):
    """인접 프레임 cosine similarity와 score를 정렬하는 MSE loss"""
    if t == 0:  # 첫 프레임은 항상 유지
        return torch.tensor(0.0)
    
    feat_curr = spatial_pool(features_l[t], pool_width)    # (B, N/w^2, D)
    feat_prev = spatial_pool(features_l[t-1], pool_width)  # (B, N/w^2, D)
    
    # 패치별 cosine similarity 계산
    feat_curr_norm = F.normalize(feat_curr, dim=-1)
    feat_prev_norm = F.normalize(feat_prev, dim=-1)
    cos_sim = (feat_curr_norm * feat_prev_norm).sum(dim=-1)  # (B, N/w^2)
    
    # 높은 similarity = 중복 = 낮은 중요도 -> 1 - cos_sim을 target으로
    target = 1 - cos_sim
    loss = F.mse_loss(scores, target)
    return loss

# Token Pruning: bottom-k% 제거
def prune_tokens(x, scores, k_percent):
    B, N, D = x.shape
    keep_n = int(N * (1 - k_percent / 100))
    _, top_indices = scores.topk(keep_n, dim=1)  # 중요한 토큰 유지
    top_indices_sorted = top_indices.sort(dim=1).values
    pruned = torch.gather(x, 1, top_indices_sorted.unsqueeze(-1).expand(-1, -1, D))
    return pruned  # (B, keep_n, D)