Token Merging (ToMe)

Token Merging (ToMe) là một kỹ thuật tiên tiến được thiết kế để tối ưu hóa hiệu suất và hiệu quả của các kiến trúc Transformer bằng cách giảm số lượng token được xử lý trong quá trình truyền tiến (forward pass). Ban đầu được phát triển để tăng tốc các mô hình Vision Transformer (ViT), ToMe hoạt động bằng cách xác định và kết hợp các token dư thừa trong mạng một cách có hệ thống mà không cần thêm bất kỳ bước huấn luyện nào. Do độ phức tạp tính toán của cơ chế self-attention tăng theo bình phương số lượng token, việc hợp nhất các token tương tự nhau giúp giảm đáng kể tổng số phép toán dấu phẩy động (FLOPs), cho phép real-time inference nhanh hơn đáng kể.

Link to this sectionHiểu về quy trình Token Merging#

ToMe khác biệt cơ bản với tokenization, vốn là bước tiền xử lý ban đầu để chia nhỏ hình ảnh hoặc văn bản thành các tokens riêng lẻ. Trong khi tokenization tạo ra các phần tử rời rạc, Token Merging đóng vai trò là một cơ chế lấy mẫu xuống (downsampling) động trong quá trình thực thi truyền tiến của mô hình.

Thuật toán này thường sử dụng khớp hai phía (bipartite matching) để đánh giá độ tương đồng của token, thường là tính độ tương đồng cosin giữa các key của token trong các lớp attention. Các token chia sẻ thông tin hình ảnh hoặc ngữ nghĩa rất giống nhau sẽ được hợp nhất với nhau—thường bằng cách lấy trung bình các đặc trưng của chúng. Điều này đảm bảo rằng thông tin không gian hoặc ngữ cảnh cần thiết được bảo toàn trong khi loại bỏ tải trọng tính toán không cần thiết, cho phép các framework như PyTorch xử lý các mô hình thị giác phức tạp nhanh hơn nhiều.

Link to this sectionCác ứng dụng thực tế của Token Merging#

Token Merging đã trở thành một chiến lược tối ưu hóa quan trọng để triển khai các kiến trúc dựa trên attention nặng trong các môi trường bị hạn chế về khả năng tính toán.

1. Generative AI and Image Synthesis: Trong các mô hình khuếch tán (diffusion) văn bản-thành-hình ảnh phổ biến, ToMe thường được sử dụng để tăng tốc quá trình tạo ảnh. Bằng cách hợp nhất các token nền hoặc các token ít chi tiết, quá trình tạo ảnh yêu cầu ít bước hơn, giúp tiết kiệm đáng kể tài nguyên GPU và giảm độ trễ cho người dùng cuối dựa vào các mô hình tạo sinh. Bạn có thể tìm hiểu thêm về các quy trình khuếch tán trong nghiên cứu nền tảng trên arXiv.

2. Edge AI Deployments: Việc triển khai các mô hình khổng lồ như Segment Anything Model (SAM) lên các thiết bị di động là rất khó khăn do hạn chế về bộ nhớ. ToMe giúp thu hẹp dung lượng bộ nhớ một cách linh hoạt, cho phép các tác vụ image segmentation phức tạp chạy trên phần cứng biên. Đối với các kịch bản mà tốc độ thuần túy là yếu tố quan trọng, các kỹ sư thường chuyển sang các kiến trúc không dùng attention được tối ưu hóa nguyên bản như Ultralytics YOLO26 để suy luận biên end-to-end nhanh hơn.

Link to this sectionVí dụ Python: Tính toán độ tương đồng của token#

Mặc dù việc tích hợp ToMe vào một kiến trúc hoàn chỉnh đòi hỏi phải sửa đổi các khối attention, nhưng khái niệm cốt lõi dựa vào việc tìm các token tương tự. Đoạn mã PyTorch sau đây minh họa cách tính độ tương đồng cosin giữa một tập hợp các token để xác định những token nào là ứng viên để hợp nhất.

import torch
import torch.nn.functional as F

# Simulate a batch of 4 image patches (tokens) with 64-dimensional features
tokens = torch.randn(1, 4, 64)

# Normalize the tokens to easily compute cosine similarity via dot product
normalized_tokens = F.normalize(tokens, p=2, dim=-1)

# Compute the similarity matrix between all tokens (1 x 4 x 4)
similarity_matrix = torch.matmul(normalized_tokens, normalized_tokens.transpose(1, 2))

# Tokens with high similarity scores (close to 1.0) off the diagonal
# are prime candidates for Token Merging.
print("Similarity Matrix:", similarity_matrix)

Các pipeline học máy hiện đại đòi hỏi sự cân bằng cẩn thận giữa độ chính xác và tốc độ. Cho dù bạn đang sử dụng Token Merging để tối ưu hóa một ViT tùy chỉnh hay dựa vào hiệu suất tiên tiến của YOLO26, việc quản lý các quy trình dữ liệu phức tạp này được đơn giản hóa đáng kể bởi Ultralytics Platform. Nền tảng này cung cấp một hệ sinh thái trực quan cho data annotation tự động, huấn luyện trên đám mây liền mạch và triển khai mô hình mạnh mẽ trên các môi trường phần cứng edge computing đa dạng. Các tổ chức mở rộng quy mô các sáng kiến computer vision của họ dựa vào các công cụ này để đưa các mô hình tiên tiến nhất vào sản xuất một cách đáng tin cậy và hiệu quả.