Token Merging (ToMe) – Gộp Token Để Tăng Tốc Vision Transformer

Mục tiêu kiến thức

Token Merging (ToMe) giảm số lượng token bằng cách gộp các token giống nhau.

Bipartite Soft Matching:

• Chia tokens thành 2 set A và B (cách nhau 1 bước).
• Tính cosine similarity giữa mỗi token trong A và B.
• Với mỗi token trong A, tìm token gần nhất trong B (best match).
• Gộp r cặp token có similarity cao nhất.
• Merge: average hai token lại (hoặc weighted average theo importance).

Yêu cầu

def bipartite_soft_matching(
    x: torch.Tensor,  # [B, N, C]
    r: int            # số cặp cần gộp
) -> tuple[torch.Tensor, torch.Tensor]:
    # Returns (merged_x [B, N-r, C], merge_indices [B, N])
    pass

Ràng buộc

• Thời gian chạy tối đa: 5000ms.
• Giới hạn bộ nhớ: 512MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.