FastAV: Audio-Visual Large Language Model 추론을 위한 효율적인 Token Pruning

# FastAV 핵심 로직 의사코드 (PyTorch 기반)
import torch
import torch.nn.functional as F

def attention_rollout(attention_matrices, alpha=0.5):
    """
    attention_matrices: list of [batch, heads, seq, seq] tensors (layer별)
    """
    rollout = None
    for attn in attention_matrices:
        # 헤드 평균
        attn_mean = attn.mean(dim=1)  # [batch, seq, seq]
        # residual connection 반영
        identity = torch.eye(attn_mean.size(-1), device=attn.device).unsqueeze(0)
        attn_mod = alpha * attn_mean + (1 - alpha) * identity
        rollout = attn_mod if rollout is None else torch.bmm(attn_mod, rollout)
    return rollout  # [batch, seq, seq]

def global_pruning(token_indices, rollout, keep_position=750):
    """중간 레이어에서 position 750 이후 low-informative 토큰 제거"""
    # rollout 기준으로 중요도 낮은 뒤쪽 토큰 제거
    importance = rollout.sum(dim=1)  # [batch, seq]
    mask = (token_indices < keep_position) | (importance > importance.median())
    return mask

def fine_pruning(query_last, keys, prune_ratio=0.2):
    """
    query_last: [batch, heads, 1, dim] — 마지막 query 토큰
    keys: [batch, heads, seq, dim]
    """
    scores = torch.softmax(
        torch.matmul(query_last, keys.transpose(-2, -1)) / (keys.size(-1) ** 0.5),
        dim=-1
    ).mean(dim=1).squeeze(1)  # [batch, seq]
    
    k = int(scores.size(-1) * (1 - prune_ratio))
    _, keep_indices = scores.topk(k, dim=-1)
    return keep_indices