DivPrune: 다양성 기반 Visual Token Pruning for Large Multimodal Models

import torch
import torch.nn.functional as F

def divprune(visual_tokens: torch.Tensor, keep_ratio: float = 0.1) -> torch.Tensor:
    """
    DivPrune: Max-Min Diversity 기반 visual token 선택
    Args:
        visual_tokens: [M, D] shape의 visual token 텐서
        keep_ratio: 유지할 토큰 비율 (0.1 = 10% 유지)
    Returns:
        선택된 토큰 [M_kept, D]
    """
    M, D = visual_tokens.shape
    M_keep = max(1, int(M * keep_ratio))

    # 코사인 거리 행렬 사전 계산 (1 - cosine_similarity)
    normed = F.normalize(visual_tokens, dim=-1)
    sim_matrix = normed @ normed.T          # [M, M]
    dist_matrix = 1.0 - sim_matrix          # cosine distance
    dist_matrix.fill_diagonal_(float('inf'))  # 자기 자신 제외

    selected_idx = []
    remaining = list(range(M))

    # Stage 1: 첫 번째 토큰 선택 (다른 모든 토큰과의 최소 거리가 가장 큰 것)
    min_dists = dist_matrix[remaining][:, remaining].min(dim=1).values
    first = remaining[min_dists.argmax().item()]
    selected_idx.append(first)
    remaining.remove(first)

    # Stage 2: 선택된 집합과의 최소 거리가 가장 큰 토큰을 반복 추가
    while len(selected_idx) < M_keep:
        sel_tensor = torch.tensor(selected_idx)
        rem_tensor = torch.tensor(remaining)
        # 각 remaining 토큰 → 선택된 토큰들과의 최소 거리
        dists_to_sel = dist_matrix[rem_tensor][:, sel_tensor].min(dim=1).values
        next_idx = remaining[dists_to_sel.argmax().item()]
        selected_idx.append(next_idx)
        remaining.remove(next_idx)

    return visual_tokens[torch.tensor(selected_idx)]

# 사용 예시 (LLaVA-style)
# visual_tokens: vision encoder 출력 [576, 4096]
# pruned = divprune(visual_tokens, keep_ratio=0.10)  # 576 -> ~58개 토큰
# llm_input = torch.cat([text_tokens, pruned], dim=0)