Token Merging (ToMe)

Hợp nhất token (ToMe) là một kỹ thuật tiên tiến được thiết kế để tối ưu hiệu năng 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 các lượt forward. Ban đầu được phát triển để tăng tốc các model Transformer Thị giác (ViT), ToMe hoạt động bằng cách xác định và kết hợp có hệ thống các token dư thừa trong mạng mà không cần thêm bất kỳ quá trình training nào. Vì độ 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 đồng giúp giảm đáng kể tổng số phép toán dấu phẩy động (FLOPs), từ đó cho phép suy luận theo thời gian thực nhanh hơn đáng kể.

Tìm hiểu quy trình hợp nhất token#

Về bản chất, ToMe khác hoàn toàn với quá trình tokenization, vốn là bước tiền xử lý ban đầu nhằm phân tách hình ảnh hoặc văn bản thành các token riêng lẻ. Trong khi tokenization tạo ra các phần tử rời rạc, Token Merging hoạt động như một cơ chế downsampling động trong quá trình thực thi forward của model.

Thuật toán thường sử dụng phép ghép cặp hai phía để đánh giá độ tương đồng giữa các token, thường tính độ tương đồng cosine giữa các key của token trong các lớp attention. Các token có thông tin trực quan hoặc ngữ nghĩa tương đồng cao 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 thông tin không gian hoặc ngữ cảnh thiết yếu được bảo toàn, đồng thời loại bỏ tải tính toán không cần thiết, cho phép các framework như PyTorch xử lý các model thị giác phức tạp nhanh hơn nhiều.

Các ứng dụng thực tế của việc hợp nhất token#

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 nặng dựa trên attention trong những môi trường bị giới hạn về tài nguyên tính toán.

1. AI tạo sinh và tổng hợp hình ảnh: Trong các model diffusion chuyển 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 token có ít chi tiết, quá trình tạo ảnh cần í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 sử dụng các model tạo sinh. Bạn có thể tìm hiểu thêm về các quy trình diffusion trong các nghiên cứu nền tảng trên arXiv.

2. Triển khai AI biên: Việc triển khai các model lớn như Mô hình Segment Anything (SAM) trên thiết bị di động vốn nổi tiếng là khó khăn do giới hạn bộ nhớ. ToMe giúp thu nhỏ footprint bộ nhớ một cách động, cho phép các tác vụ phân đoạn hình ảnh phức tạp chạy trên phần cứng biên. Trong các tình huống yêu cầu tốc độ tuyệt đối, các kỹ sư thường chuyển sang những kiến trúc không dùng attention được tối ưu hóa native như Ultralytics YOLO26 để đạt tốc độ suy luận biên end-to-end cao hơn.

Ví dụ Python: Tính độ tương đồng giữa các 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 block attention, khái niệm cốt lõi dựa trên việc tìm các token tương đồng. Đoạn mã PyTorch sau đây minh họa cách tính độ tương đồng cosine giữa một tập hợp token để xác định những token nào là ứng viên cho việc 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 machine learning hiện đại đòi hỏi sự cân bằng cẩn trọng giữa độ chính xác và tốc độ. Dù bạn sử dụng Token Merging để tối ưu một ViT tùy chỉnh hay tận dụng hiệu năng tiên tiến của YOLO26, việc quản lý các workflow dữ liệu phức tạp này được đơn giản hóa đáng kể nhờ Ultralytics Platform. Platform cung cấp một hệ sinh thái trực quan cho gán nhãn dữ liệu tự động, training trên cloud liền mạch và triển khai model mạnh mẽ trên nhiều môi trường phần cứng điện toán biên khác nhau. Các tổ chức mở rộng sáng kiến thị giác máy tính của mình dựa vào những công cụ này để đưa các model tiên tiến vào production một cách đáng tin cậy và hiệu quả.