Token Merging (ToMe)

Gộp token (ToMe) là phương pháp tăng tốc Transformer, đặc biệt là Vision Transformer (ViT), bằng cách kết hợp các token tương tự bên trong mạng để mỗi block xử lý ít token hơn. Ảnh chứa nhiều phần dư thừa, chẳng hạn các patch bầu trời hoặc mặt đường gần như giống nhau; lấy trung bình các token như vậy sẽ giảm phép tính mà vẫn giữ thông tin. ToMe hoạt động trên model đã huấn luyện mà không cần huấn luyện lại và cũng có thể được áp dụng trong quá trình huấn luyện. Daniel Bolya cùng các cộng sự tại Georgia Tech và Meta AI giới thiệu phương pháp này trong bài báo Token Merging: Your ViT But Faster, được trình bày miệng tại ICLR 2023. Khi dùng trực tiếp, phương pháp này tăng gấp đôi thông lượng của ViT-L ở 512 pixel và ViT-H ở 518 pixel trên ảnh, với mức giảm độ chính xác chỉ 0,2% đến 0,3%; đồng thời tăng tốc model video ViT-L lên 2,2 lần với mức giảm tương tự.

Cách gộp token hoạt động#

ViT chia ảnh thành các patch, mỗi patch tương ứng một token; mỗi block chạy self-attention, có chi phí tăng theo bình phương số token, rồi đến MLP, có chi phí tăng tuyến tính. ToMe chèn bước gộp giữa attention và MLP của mỗi block, mỗi lần loại bỏ một số token cố định r; vì vậy model 24 block với r = 8 kết thúc với ít hơn 192 token so với ban đầu. Phép gộp sử dụng ghép mềm hai phía, được chọn vì nhanh và không cần phân cụm lặp:

1. Chia: Chia các token luân phiên thành hai tập A và B.
2. Ghép: Kết nối mỗi token trong A với token tương tự nhất trong B, đo bằng độ tương đồng cosine của các key attention được lấy trung bình qua các head.
3. Giữ các cạnh tốt nhất: Giữ lại r cặp tương tự nhất.
4. Gộp: Lấy trung bình từng cặp được giữ lại để tạo thành một token, gán trọng số theo số patch ban đầu mà mỗi token đại diện, rồi nối lại hai tập.

Vì token đã gộp đại diện cho nhiều patch, ToMe cộng logarit kích thước của mỗi token vào điểm attention, gọi là attention tỷ lệ, để token được tạo từ nhiều patch nhận được phần attention tương ứng. Toàn bộ quy trình có thể vi phân, cho phép huấn luyện model bằng ToMe để tăng tốc hơn nữa.

Ví dụ PyTorch này áp dụng ghép mềm hai phía cho ViT DINOv2-small của Meta từ Transformers của Hugging Face, trên ảnh xe buýt kích thước 448 x 448; ví dụ gộp r token sau mỗi block trong 12 block và so sánh embedding ảnh thu được với embedding của model không gộp token:

from urllib.request import urlopen

import torch
import torch.nn.functional as F
from PIL import Image
from transformers import AutoImageProcessor, AutoModel

processor = AutoImageProcessor.from_pretrained("facebook/dinov2-small")
model = AutoModel.from_pretrained("facebook/dinov2-small")
image = Image.open(urlopen("https://ultralytics.com/images/bus.jpg")).resize((448, 448))
pixels = processor(images=image, do_resize=False, do_center_crop=False, return_tensors="pt").pixel_values


def merge(x, r):
    """Bipartite soft matching: merge the r most similar pairs between alternating token sets A and B."""
    cls, a, b = x[:, :1], x[:, 1::2], x[:, 2::2]
    similarity = F.normalize(a[0], dim=-1) @ F.normalize(b[0], dim=-1).T
    best, match = similarity.max(dim=-1)  # each A token's most similar B token
    order = best.argsort(descending=True)
    merged, kept = order[:r], order[r:]
    index = match[merged][None, :, None].expand(1, r, x.shape[-1])
    b = b.scatter_reduce(1, index, a[:, merged], reduce="mean")  # average merged A tokens into their B matches
    return torch.cat([cls, a[:, kept], b], dim=1)


with torch.no_grad():
    for r in (0, 32, 64):
        x, counts = model.embeddings(pixels), []
        for layer in model.encoder.layer:
            counts.append(x.shape[1])
            x = merge(layer(x), r)  # merge r tokens after every transformer block
        embedding = model.layernorm(x)[:, 0]  # CLS token: the image embedding
        reference = embedding if r == 0 else reference
        similarity = F.cosine_similarity(embedding, reference).item()
        print(f"r={r}: mean tokens per block {sum(counts) / 12:.0f}, CLS similarity to unmerged {similarity:.3f}")

Model không gộp xử lý 1025 token trong mỗi block, gồm 1.024 patch cộng với token CLS. Gộp 32 token mỗi block giảm mức trung bình xuống 849, trong khi embedding ảnh vẫn giữ độ tương đồng cosine 0.983 so với bản gốc; gộp 64 token mỗi block giảm xuống 673 với độ tương đồng 0.887, cho thấy r đánh đổi lượng tính toán để lấy độ trung thực.

Ứng dụng của gộp token#

Token Merging for Fast Stable Diffusion áp dụng cùng ý tưởng cho các block Transformer của mạng khử nhiễu Stable Diffusion, nơi phương pháp có thể loại bỏ tới 60% token, tạo ảnh nhanh hơn tới hai lần và dùng bộ nhớ ít hơn tới 5,6 lần mà không cần huấn luyện bổ sung. Ngoài ảnh, bài báo ToMe cho biết phương pháp tăng gấp đôi thông lượng của ViT-B âm thanh được huấn luyện với phép gộp, với mức giảm mAP 0,4%, đồng thời tăng tốc tinh chỉnh autoencoder có mặt nạ trên video lên tới 2 lần. Các nhà nghiên cứu áp dụng phép gộp và cắt tỉa tương tự cho token thị giác của model thị giác-ngôn ngữ, trong đó LLaVA-1.5 dùng 576 token trong ngữ cảnh của model ngôn ngữ cho mỗi ảnh.

Trong gói Ultralytics, Segment Anything Model (SAM) của Meta mã hóa mỗi ảnh kích thước 1024 x 1024 pixel thành 4.096 token patch; trong so sánh MobileSAM, bộ mã hóa ảnh ViT-H của model này chiếm 452 ms trong tổng 456 ms mà SAM cần cho mỗi ảnh trên GPU, chính là phần tính toán mà các phương pháp giảm token nhắm đến.

Gộp token và các khái niệm liên quan#

• Cắt tỉa token: Các phương pháp như DynamicViT và EViT loại bỏ token bị đánh giá là không quan trọng, làm mất thông tin và thường cần huấn luyện. Thay vào đó, ToMe gộp token để lấy trung bình thông tin và hoạt động trên các model chưa được huấn luyện bằng phương pháp này.
• Token hóa: Tạo các token ban đầu từ patch hoặc văn bản trước khi mạng chạy. ToMe giảm số lượng token bên trong mạng.
• Flash attention và attention thưa: Giảm chi phí hoặc tăng tốc attention với số lượng token nhất định. ToMe giảm số lượng token, nhờ đó cũng tăng tốc các lớp MLP, và có thể kết hợp với hai phương pháp này.
• Mixture of Depths (MoD): Cho phép bộ định tuyến được học quyết định token nào bỏ qua một block trong model ngôn ngữ. ToMe hợp nhất vĩnh viễn các token dư thừa bằng quy tắc tương đồng cố định.
• Giảm độ phân giải đầu vào: Giảm đồng đều số lượng token trên toàn ảnh. ToMe chủ yếu gộp token ở các vùng dư thừa như nền, đồng thời giữ các vùng chi tiết ở độ phân giải đầy đủ.