Mixture of Depths (MoD)

Trong các kiến trúc deep learning, hiệu quả tính toán là yếu tố tối quan trọng, đặc biệt khi xử lý các chuỗi dài hoặc đầu vào có độ phân giải cao. Một phương pháp mới phân bổ động tài nguyên tính toán bằng cách cho phép mạng quyết định phần nào của đầu vào cần được xử lý đầy đủ và phần nào có thể bỏ qua một cách an toàn một số layer nhất định. Chiến lược định tuyến động này làm giảm độ phức tạp tính toán tổng thể mà không làm giảm năng lực dự đoán hoặc độ chính xác của model.

Tìm hiểu khái niệm#

Hỗn hợp độ sâu (MoD) là một kỹ thuật kiến trúc chủ yếu được áp dụng cho các kiến trúc Transformer, trong đó model học cách động bỏ qua việc tính toán cho các token cụ thể tại nhiều layer khác nhau. Các transformer truyền thống xử lý mọi token qua mọi layer, bất kể đó là thông tin quan trọng hay nội dung bổ sung. Ngược lại, model MoD sử dụng một cơ chế router để đánh giá các token và gán điểm cho chúng. Chỉ những token có điểm cao nhất—lên đến giới hạn capacity được xác định trước—mới được truyền qua các block tính toán chuyên sâu, chẳng hạn như cơ chế attention hoặc các layer feed-forward dense. Những token còn lại bỏ qua block thông qua các kết nối residual, qua đó tạo ra một "hỗn hợp độ sâu" trong đó các token khác nhau trải qua mức độ xử lý với độ sâu khác nhau.

Phương pháp này, được phổ biến nhờ nghiên cứu của DeepMind gần đây và được trình bày chi tiết trong kho lưu trữ arXiv, giúp giảm đáng kể tổng số phép toán dấu phẩy động (FLOPs) cần thiết trong cả quá trình training và inference.

Phân biệt với Hỗn hợp chuyên gia (MoE)#

Khái niệm này rất dễ bị nhầm lẫn với một Hỗn hợp chuyên gia (MoE). Mặc dù cả hai đều sử dụng cơ chế định tuyến, chúng giải quyết các vấn đề khác nhau:

• MoE định tuyến các token đến những mạng con (expert) khác nhau trong một layer. Độ sâu tính toán vẫn giống nhau đối với mọi token, nhưng số lượng tham số của model tăng lên.
• MoD định tuyến các token đến block tính toán hoặc một kết nối skip. Số lượng tham số hoàn toàn không đổi, nhưng độ sâu tính toán giảm đối với các token ít quan trọng hơn, trực tiếp cải thiện độ trễ inference.

Các ứng dụng trong thực tế#

Khả năng phân bổ ngân sách tính toán một cách động khiến kỹ thuật này trở nên đặc biệt hữu ích trong nhiều lĩnh vực của thị giác máy tính và xử lý ngôn ngữ tự nhiên.

1. Tối ưu hóa context trong các model ngôn ngữ: Các mô hình ngôn ngữ lớn (LLMs) hiện đại từ những tổ chức như OpenAI và Anthropic xử lý các cửa sổ context khổng lồ. Bằng cách sử dụng định tuyến độ sâu động, các model này có thể bỏ qua những từ bổ sung mang tính cấu trúc hoặc lặp lại, dành năng lực tính toán chuyên sâu cho các bước suy luận phức tạp và việc trích xuất thông tin thực tế.

2. AI thị giác độ phân giải cao: Trong các hệ thống thị giác tiên tiến như model Ultralytics YOLO26, việc xử lý các ảnh lớn cho phát hiện đối tượng và phân đoạn ảnh đòi hỏi lượng memory rất lớn. Định tuyến độ sâu cho phép mạng bỏ qua việc trích xuất feature trên các nền đồng nhất (chẳng hạn bầu trời trống hoặc những bức tường trống), tập trung năng lực tính toán vào các đối tượng tiền cảnh phức tạp. Đây là yếu tố thiết yếu khi triển khai model trên phần cứng AI biên bị giới hạn tài nguyên và được tối ưu bằng các thư viện tối ưu hóa CUDA.

Ví dụ triển khai#

Dưới đây là một đoạn mã PyTorch mang tính khái niệm, minh họa cách một cơ chế định tuyến cơ bản có thể bỏ qua việc tính toán cho một phần token đầu vào, mô phỏng hành vi định tuyến độ sâu.

import torch
import torch.nn as nn


class MixtureOfDepthsBlock(nn.Module):
    def __init__(self, d_model, capacity_factor=0.5):
        super().__init__()
        self.capacity_factor = capacity_factor
        self.router = nn.Linear(d_model, 1)
        self.heavy_compute = nn.Sequential(nn.Linear(d_model, d_model * 4), nn.GELU(), nn.Linear(d_model * 4, d_model))

    def forward(self, x):
        # x shape: (batch_size, seq_len, d_model)
        seq_len = x.size(1)
        capacity = int(seq_len * self.capacity_factor)

        # 1. Compute routing scores
        scores = self.router(x).squeeze(-1)  # Shape: (batch_size, seq_len)

        # 2. Identify top-k tokens to process
        topk_indices = torch.topk(scores, capacity, dim=1).indices

        # 3. Create an output tensor mirroring the input (residual baseline)
        output = x.clone()

        # 4. Apply heavy computation only to the selected tokens
        for b in range(x.size(0)):
            selected_tokens = x[b, topk_indices[b]]
            processed_tokens = self.heavy_compute(selected_tokens)
            output[b, topk_indices[b]] += processed_tokens

        return output


# Example usage
dummy_input = torch.randn(2, 64, 128)  # Batch=2, Seq=64, Dim=128
mod_block = MixtureOfDepthsBlock(d_model=128, capacity_factor=0.5)
output = mod_block(dummy_input)
print(f"Output shape: {output.shape}")  # Expect (2, 64, 128)

Bằng cách tận dụng các framework như framework PyTorch hoặc TensorFlow, developer có thể tích hợp các block [tối ưu hóa model](https://ultralytics-translation-2.invalid tùy chỉnh này. Ngoài ra, các công cụ như Ultralytics Platform giúp các team quản lý dữ liệu training cần thiết để training chính xác các router này, đồng thời tích hợp liền mạch với các hệ sinh thái enterprise như Google Cloud AI.