Mixture of Depths (MoD)

Trong các kiến trúc học sâu, hiệu quả tính toán là tối quan trọng, đặc biệt là 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ổ tài nguyên tính toán một cách linh hoạt bằng cách cho phép mạng quyết định phần nào của đầu vào cần xử lý đầy đủ và phần nào có thể bỏ qua các lớp nhất định một cách an toàn. Chiến lược điều hướng động này làm giảm computational complexity tổng thể mà không làm giảm khả năng dự đoán hoặc độ chính xác của model.

Link to this sectionTìm hiểu khái niệm#

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

Phương pháp này, được phổ biến bởi các DeepMind research gần đây và được ghi lại rộng rãi trong arXiv repository, giúp giảm đáng kể tổng số floating-point operations (FLOPs) cần thiết trong quá trình huấn luyện và suy luận (inference).

Link to this sectionPhân biệt với Mixture of Experts (MoE)#

Rất dễ nhầm lẫn khái niệm này với Mixture of Experts (MoE). Mặc dù cả hai đều sử dụng cơ chế điều hướng, chúng giải quyết các vấn đề khác nhau:

• MoE điều hướng các token đến các sub-networks (chuyên gia) khác nhau trong một lớp. Độ sâu tính toán vẫn giữ nguyên cho tất cả các token, nhưng số lượng tham số của model tăng lên.
• MoD điều hướng các token đến khối tính toán hoặc một kết nối tắt (skip connection). Số lượng tham số vẫn giữ nguyên, 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 inference latency.

Link to this sectionCác ứng dụng trong thực tế#

Khả năng lập ngân sách tính toán động làm cho kỹ thuật này trở nên vô cùng giá trị trên nhiều lĩnh vực của computer vision và xử lý ngôn ngữ tự nhiên.

1. Tối ưu hóa ngữ cảnh trong các mô hình ngôn ngữ: Các Large Language Models (LLMs) hiện đại từ các tổ chức như OpenAI và Anthropic xử lý các cửa sổ ngữ cảnh khổng lồ. Bằng cách sử dụng điều hướng độ sâu động, các model này có thể bỏ qua các từ cấu trúc hoặc từ phụ lặp lại, dành riêng các tính toán chuyên sâu cho các bước suy luận phức tạp và trích xuất 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ý hình ảnh lớn cho object detection và image segmentation đòi hỏi bộ nhớ khổng lồ. Điều hướng độ sâu cho phép mạng bỏ qua việc trích xuất đặc trưng trên các nền đồng nhất (như bầu trời trống hoặc tường trống), tập trung sức mạnh tính toán vào các đối tượng tiền cảnh phức tạp. Điều này rất quan trọng để triển khai model lên các phần cứng edge AI hạn chế về tài nguyên được tối ưu hóa bởi các thư viện CUDA optimization.

Link to this sectionVí dụ về triển khai#

Dưới đây là đoạn mã PyTorch khái niệm minh họa cách một cơ chế điều hướng cơ bản có thể bỏ qua tính toán cho một phần các token đầu vào, mô phỏng hành vi điều hướng độ 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ư PyTorch framework hoặc TensorFlow, các lập trình viên có thể tích hợp các khối model optimization tùy chỉnh này. Hơn nữa, các công cụ như Ultralytics Platform giúp các nhóm quản lý training data cần thiết để huấn luyện các router này một cách chính xác, đồng thời tích hợp liền mạch với các hệ sinh thái doanh nghiệp như Google Cloud AI.