Mixture of Depths (MoD)

Mixture of Depths (MoD) là kỹ thuật transformer, trong đó router đã học quyết định ở từng block token nào được xử lý và token nào bỏ qua block qua kết nối dư, nhờ đó mỗi token nhận lượng tính toán khác nhau. David Raposo, Adam Santoro và các cộng sự tại Google DeepMind giới thiệu kỹ thuật này vào tháng 4 năm 2024 trong bài báo Mixture-of-Depths: Phân bổ tính toán động trong mô hình ngôn ngữ dựa trên Transformer. Dung lượng cố định k xác định số token được xử lý trong mỗi block định tuyến, vì vậy ngân sách tính toán và hình dạng tensor đều được biết trước dù lựa chọn token phụ thuộc vào đầu vào. Mô hình ngôn ngữ lớn (LLM) MoD đạt chất lượng ngang baseline với cùng lượng huấn luyện FLOPs và thời gian thực, trong khi chỉ dùng một phần FLOPs cho mỗi lượt forward, đồng thời tạo token nhanh hơn hơn 50% khi lấy mẫu. Với ngân sách huấn luyện cố định, model MoD tốt nhất lớn hơn baseline tốt nhất và đạt loss thấp hơn.

Cách định tuyến MoD hoạt động#

Trước một block được định tuyến, router tuyến tính gán cho mỗi token một điểm số vô hướng. k token có điểm số cao nhất đi qua cơ chế tự chú ý và MLP của block; đầu ra của chúng được nhân với điểm số router tương ứng rồi cộng trở lại vào luồng dư. Việc nhân với điểm số đưa router vào đường truyền gradient, nhờ đó lan truyền ngược có thể giúp router học được token nào hưởng lợi. Các token còn lại đi qua mà không thay đổi và bị loại khỏi attention của block đó, vì vậy các token phía sau không thể attend tới chúng trong block đó; chi phí attention giảm theo bình phương của dung lượng: một nửa số token chỉ cần một phần tư số phép tính query-key. Thiết lập tốt nhất trong bài báo định tuyến một block trong mỗi hai block, với dung lượng bằng 12,5% độ dài chuỗi. Một phiên bản tối giản của một block được định tuyến, ở đây chỉ gồm MLP, được viết bằng PyTorch:

import torch
from torch import nn

torch.manual_seed(0)
batch, seq_len, dim = 2, 64, 128
capacity = int(0.125 * seq_len)  # 12.5% of tokens, the paper's best setting
x = torch.randn(batch, seq_len, dim)
router = nn.Linear(dim, 1)
block = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

# The router scores every token, and only the top-k tokens enter the block
weights, idx = router(x).squeeze(-1).topk(capacity, dim=1)
idx = idx[..., None].expand(-1, -1, dim)
update = block(x.gather(1, idx)) * weights[..., None]  # router weight keeps the router trainable

# Selected tokens get the residual update; the rest skip the block unchanged
out = x.scatter_add(1, idx, update)
unchanged = (out == x).all(dim=-1).sum().item()
print(f"Block ran on {batch * capacity} of {batch * seq_len} tokens; {unchanged} skipped it unchanged")

Lệnh này in Block ran on 16 of 128 tokens; 112 skipped it unchanged: MLP của block chỉ thực hiện một phần tám lượng tính toán thông thường, còn tensor đầu ra giữ nguyên hình dạng chuỗi đầy đủ để đưa vào block tiếp theo.

Định tuyến trong quá trình sinh#

Việc chọn top-k không có tính nhân quả: một token có nằm trong top k của chuỗi hay không phụ thuộc vào điểm số của các token xuất hiện sau nó, trong khi các token đó chưa tồn tại khi model sinh văn bản từng token một. Các tác giả khiến việc định tuyến có tính nhân quả theo hai cách: dùng một loss entropy chéo nhị phân phụ trợ để dạy đầu ra router dự đoán token có thuộc top-k hay không, chỉ làm thay đổi loss mô hình hóa ngôn ngữ 0,2–0,3%; hoặc dùng một MLP dự đoán nhỏ riêng biệt trên đầu vào có gradient dừng, nhanh chóng đạt độ chính xác 97–99%. Dù dùng cách nào, lấy mẫu tự hồi quy cũng chỉ giảm chất lượng đôi chút. Các router được huấn luyện còn thể hiện một quy luật: một số token kích hoạt mọi block được định tuyến trong toàn bộ độ sâu model, trong khi các token khác bỏ qua block bất cứ khi nào có thể.

MoD và các kỹ thuật liên quan#

• Hỗn hợp chuyên gia (MoE): Định tuyến mỗi token đến một hoặc vài MLP chuyên gia trong một lớp, bổ sung tham số nhưng giữ nguyên lượng tính toán cho mỗi token. Bài báo cũng kết hợp hai phương pháp thành MoDE, trong đó token có thể được định tuyến đến một chuyên gia hoặc thao tác không làm gì; hiệu quả được cộng dồn.
• Thoát sớm: Token dừng sau một lớp và bỏ qua toàn bộ các lớp còn lại. MoD có thể bỏ qua một block nhưng vẫn sử dụng các block về sau.
• Gộp token (ToMe): Gộp các token tương tự trong transformer thị giác để các lớp sau xử lý ít token hơn, không cần router đã học.
• Cắt tỉa model: Loại bỏ vĩnh viễn các trọng số để mọi đầu vào đều chạy cùng một mạng nhỏ hơn.
• Định tuyến model: Gửi toàn bộ một yêu cầu đến một trong nhiều model, áp dụng cùng ý tưởng ở cấp độ model hoàn chỉnh.

Hầu hết mạng đều cố định độ sâu khi được xây dựng. Ultralytics YOLO26 thiết lập độ sâu theo từng kích thước model: hệ số độ sâu trong YAML model là 0.50 cho các kích thước n, s và m, và 1.00 cho l và x; vì vậy, một stage C3k2 được liệt kê với 2 lần lặp sẽ chạy 1 block trong YOLO26n và 2 block trong YOLO26x, và mọi ảnh đều đi qua toàn bộ các block đó.