一、前言

【理想汽車智駕方案介紹專題 -1】端到端+VLM 方案介紹

【理想汽車智駕方案介紹專題 -2】MindVLA 方案詳解

在上述兩篇系列帖子中，筆者已對理想汽車 VLM 和 VLA 方案的框架進行了全面介紹，但對于其中的前沿技術僅做了初步探討，未進行更深入的剖析。因此，筆者計劃繼續以系列文章的形式，介紹其中涉及的相關前沿技術。

首先，將介紹 MindVLA 中采用的“MoE + Sparse Attention”高效結構。

二、MoE + Sparse Attention 高效結構

稀疏注意力（Sparse Attention）與混合專家架構（MoE）是提升大模型效率與性能的核心技術。MindGPT 針對大語言模型（LLM）進行了重新設計與預訓練，使其具備了 3D 空間理解和 3D 推理能力，不過這在一定程度上增加了大語言模型（LLM）的參數量。即便采用英偉達（NV）大力宣傳的 Thor - U 智駕芯片，在端側部署時仍有可能會出現推理效率低下的問題。 所以模型采用了 MOE 架構+SparseAttention，實現模型容量擴容的同時不會大幅度增加推理負擔。

由上圖可以看到，MindGPT 使用混合專家（MoE）架構（有 E1-E8 個專家），由 Router 動態選擇激活部分專家（而非全部），只在需要時調用相關子模型。稀疏注意力（Sparse Attention）則限制注意力機制的復雜度，只關注關鍵輸入部分（如相關物體或動作），而不是全局計算。在這張圖中，它的實現過程是：

每一層的 token 并不全互相注意，而是：

• 利用 Router 分配 token 給不同的 expert
• 每個 expert 內部只在局部或關鍵 token 上建立注意力連接

這種結構既節省計算，又保持了輸入的局部結構性，尤其適合 3D 場景中的空間約束。

2.1 Sparse Attention

參考論文 Generating Long Sequences with Sparse Transformers

傳統 transformer 的全連接注意力復雜度為$$O(n^2)$$，對于高維空間場景（例如 3D Token 很多）會顯著拖慢速度、增加計算。

稀疏注意力的核心目標是：

• 減少 token-to-token 的注意力連接（只對一部分 token 建立 attention）
• 降低計算復雜度，同時保持關鍵 token 的交互質量

那么它是如何實現目標的呢？下面將結合Sparse Transformer代碼進行介紹。

Sparse Transformer 是由OpenAI提出的稀疏注意力的最早實現之一，其核心思想是用 規則稀疏模式（Regular Sparse Patterns） 替代標準全連接注意力（full self-attention）。在保持表示能力的同時，將注意力復雜度從

$$O(n^2)$$降低為$$O（n\sqrt{n}）$$.

2.1.1 Sparse Attention 圖結構解析

Strided attention（跳躍連接）

每個 token 關注步長為 s 的若干 token，例如：

位置:      0   1   2   3   4   5   6   7   8
strided:   ↑       ↑       ↑       ↑

token 8 會關注 0， 2， 4， 6 這些步長為 2 的 token。

Local attention（局部窗口）

每個 token 也關注自己附近的窗口內 token，比如前后 2 個：

位置:       5
局部窗口:   3   4   5   6   7

組合結構

將兩者合并后，注意力圖是 局部窗口 + 跳躍連接，形成稀疏但有覆蓋力的結構：

Token 位置:         8
局部連接:         6 7 8
跳躍連接:     0 2 4

這種結構可以保證每個 token 都能快速傳播到遠端（通過跳躍），并同時保留局部建模能力。

2.1.2 代碼實現解析

由于 OpenAI Sparse Transformer 沒開源，這里參考了其思想的一個流行實現：[OpenNMT-py 或 custom PyTorch 實現]，也可參考 Reformer 等模型實現方式。

稀疏注意力 mask 構造（關鍵）

def build_sparse_attention_mask(seq_len, block_size=64, num_local_blocks=1, stride=2):"""構造稀疏注意力 mask：局部 + 跳躍"""mask = torch.zeros(seq_len, seq_len, dtype=torch.bool)for i in range(seq_len):# 添加局部窗口（例如前后1個block）for j in range(-num_local_blocks, num_local_blocks + 1):idx = i + j * block_sizeif 0 <= idx < seq_len:mask[i, idx] = True# 添加跳躍連接（stride）for j in range(0, seq_len, stride):mask[i, j] = Truereturn mask

這個 mask 會用于注意力權重：

attn_scores = torch.matmul(query, key.transpose(-2, -1))  # [B, H, N, N]
attn_scores[~mask] = -inf  # 掩蔽非連接位置
attn_probs = softmax(attn_scores)

應用在 Attention 模塊

核心 attention 模塊不變，只是加了 mask：

def sparse_attention(query, key, value, mask):attn_scores = torch.matmul(query, key.transpose(-2, -1))attn_scores = attn_scores.masked_fill(~mask, float('-inf'))attn_probs = torch.softmax(attn_scores, dim=-1)return torch.matmul(attn_probs, value)

這樣就可以在不修改 Transformer 主體結構的前提下使用稀疏連接。

以下是一個 9×9 attention matrix 的稀疏連接可視化（● 表示有連接）：

Token ID →0 1 2 3 4 5 6 7 8┌──────────────────
0 │●       ●       ●
1 │  ●     ●     ●
2 │    ●   ●   ●
3 │      ● ● ●
4 │      ● ● ●
5 │      ● ● ●
6 │    ●   ●   ●
7 │  ●     ●     ●
8 │●       ●       ●

中間對角線為 **局部 attention，**分布的條紋為 跳躍 attention。

2.1.3 復雜度分析

標準全連接 Self-Attention 的公式是：

$$Q{K}^T$$：產生 n×n 的權重矩陣 → 計算復雜度是：$$O(n^{2}d)$$;

Softmax ()V :$$O(n^{2}d)$$;

整體復雜度為：$$O(n^{2}d)$$，這意味著當序列變長時（如 n=8192），計算量和內存都會劇增。

Sparse Transformer 通過限制每個 token 只關注$$O(\sqrt{n})$$個位置，從而稀疏 attention 矩陣。

每個 token 的 attention 鏈接數為$$k=O(\sqrt{n})$$**，**所有 token 的總鏈接數就是$$n?k=n?\sqrt{n}$$ 。

由于只需要對是$$n?\sqrt{n}$$ 個位置計算 QK? dot product，而不是$$n^2$$ ，每個位置依然是**O(d)的向量操作，**所以 Sparse Transformer 的復雜度是是$$O(n?\sqrt{n})$$ 。

2.2 MoE

MoE（專家混合體）是神經網絡的一種架構模式，該模式將一個層或操作（例如線性層、多層感知機或注意力投影）的計算拆分為多個“專家”子網絡。這些子網絡各自獨立地進行計算，其計算結果會被整合，以生成 MoE 層的最終輸出。MoE 架構可分為密集型和稀疏型，前者意味著在處理每個輸入時都會啟用所有專家，后者則表示每個輸入僅使用部分專家。因此，MoE 架構具備以下優勢：

• 提升計算效率
• 在不顯著增加推理成本的前提下，擴大模型容量（參數量）

這些優勢與大模型的應用難題完美匹配，故而在該領域得到了較為廣泛的應用。

2.2.1 MoE 結構解析

本節首先描述 MoE 的核心組件，然后以 DeepSpeed-MoE 源代碼為例進行詳細解析。

以 PyTorch 為例，MoE 的核心組成部分通常包括以下幾個模塊：

1.Gate 模塊（路由器）

負責為每個輸入樣本選擇合適的專家（通常是 top-k 策略）：

# 偽代碼
scores = gate(x)  # (batch_size, num_experts)
top_k_scores, top_k_indices = torch.topk(scores, k=2)

• scores 通常是通過一個線性層獲得，表示輸入樣本對各個專家的偏好程度。
• 路由器可能帶有噪聲或正則（如 Switch Transformer 中的 noisy gating）。

2.Experts 模塊（多個子網絡）

每個專家是一個獨立的神經網絡（比如一個 MLP）：

class Expert(nn.Module):def __init__(self, hidden_dim):self.ff = nn.Sequential(nn.Linear(hidden_dim, 4*hidden_dim),nn.ReLU(),nn.Linear(4*hidden_dim, hidden_dim))

通常通過參數共享或并行執行多個專家。

3.Dispatcher（稀疏調度器）

根據 Gate 的輸出，將輸入路由給選中的專家，并收集輸出：

 for i in range(num_experts):expert_input = input[mask[:, i]]expert_output = experts[i](expert_input)output[mask[:, i]] = expert_output * gate_scores[:, i]

這一步驟通常要做優化，否則容易成為性能瓶頸。

2.2.2 DeepSpeed-MoE 示例代碼

源碼地址：https://github.com/microsoft/DeepSpeed/tree/master/deepspeed/moe

其核心組件為deepspeed.moe.layer.MoE和deepspeed.moe.layers.gates 類。

1.MoE 類（deepspeed.moe.layer.MoE）

作為入口類，集成了路由器、專家、通信邏輯：

class MoE(nn.Module):def __init__(self, hidden_size, experts=..., ep_size=..., k=1):self.experts = Experts()self.gate = TopKGate()...

• ep_size： 每個專家組的并行數（Expert Parallelism）。
• k: top-k gating.
• 使用了通信優化如 All-to-All 分發樣本。

2.構造函數 init

根據官方文檔，MoE 的初始化簽名如下，參數豐富且功能強大：

class MoE(nn.Module):def __init__(self,hidden_size: int,expert: nn.Module,num_experts: int = 1,ep_size: int = 1,k: int = 1,capacity_factor: float = 1.0,eval_capacity_factor: float = 1.0,min_capacity: int = 4,use_residual: bool = False,noisy_gate_policy: Optional[str] = None,drop_tokens: bool = True,use_rts: bool = True,use_tutel: bool = False,enable_expert_tensor_parallelism: bool = False,top2_2nd_expert_sampling: bool = True)

• hidden_size：輸入和輸出的維度；
• expert：作為子模塊傳入的專家網絡（如 MLP）；
• num_experts：專家總數；
• ep_size：專家并行維度；
• k：選用 top?k 路由；
• capacity_factor 和 eval_capacity_factor：訓練/評估期間專家最大處理 token 數比例；
• min_capacity：每個專家至少能接收的 token 數；
• use_residual：是否啟用 Residual MoE 結構；
• noisy_gate_policy、drop_tokens、use_rts：路由噪聲、token 丟棄、隨機選擇；
• enable_expert_tensor_parallelism：專家參數 tensor 切分；
• top2_2nd_expert_sampling：top?2 第二專家采樣策略。

這些選項為 MoE 的訓練/推理提供了高度靈活性

3.前向函數 forward

forward 函數定義如下：

def forward(self, hidden_states: Tensor, used_token: Optional[Tensor] = None) -> Tuple[Tensor, Tensor, Tensor]:

返回三元組 (output, l_aux, exp_counts)，分別為輸出、auxiliary loss、各專家激活次數

forward 函數的核心步驟可以總結為以下 5 部分：

def forward(self, hidden_states, used_token=None):# 1?? Gating 階段：計算每個 token 的專家 logits，并選出 top?k 專家gates, load, indices, expert_capacity = self.gate(hidden_states, self.training)# gates: (B, k) 專家權重，load: auxiliary balance loss，indices: 專家索引# 2?? Capacity 控制：根據 capacity_factor 限制每個專家最多處理的 token 數# 用 expert_capacity 來計算實際可接收的 token 數# 3?? Dispatch 階段：將 token 分發給對應專家dispatch_mask, combine_mask = create_masks(indices, expert_capacity)# dispatch_mask: 用于提取每個專家的 token# 重塑 hidden_states 方便通信：expert_inputs = torch.einsum("b h, b e -> e b h", hidden_states, dispatch_mask)# 4?? all_to_all 分發 token：跨 GPU 路由 token 到對應專家所在 GPUexpert_inputs = all_to_all(expert_inputs, self.expert_parallel_group)# 5?? 專家計算階段：每個專家在本地 receive 的 token 上執行 forwardexpert_outputs = self.experts(expert_inputs)# 6?? all_to_all 收集結果：各專家輸出回傳給原始 GPUexpert_outputs = all_to_all(expert_outputs, self.expert_parallel_group)# 7?? 合并輸出：將專家輸出按照 token-group 重組回 batch 維度output = torch.einsum("e b h, b e -> b h", expert_outputs, combine_mask)return output, load, indices.bincount(...)

4.TopKGate 類（deepspeed.moe.layers.gates）

實現 top-k 路由，并支持 noisy gate：

class TopKGate(nn.Module):def forward(self, input):logits = self.w_gating(input)topk_vals, topk_indices = torch.topk(logits, k)...

還包括負載均衡 loss。

5.Expert 通信部分

使用 all_to_all 通信分發數據，保證跨 GPU 的負載均衡和數據交換效率。

在多 GPU（分布式）訓練中，如果我們有多個專家網絡（Experts），這些專家可能是跨 GPU 分布的。舉個例子：

• 有 4 張 GPU，每張 GPU 上部署 2 個專家，共 8 個專家；
• 假設 batch 里的一部分 token 需要被路由到第 3 個專家（在 GPU 2 上），另一部分需要被送到第 6 個專家（在 GPU 4 上）；
• 那么就必須跨 GPU 發送這些 token —— 所以通信效率就非常關鍵。

這時候就需要 all_to_all 機制來高效完成這個通信過程。

torch.distributed.all_to_all 是一種跨 GPU 的點對點通信機制，作用是每張 GPU 都給其他所有 GPU 發送一塊數據，同時從其他 GPU 收到對應的數據塊。

具體來說：

• 每個進程（GPU）上都有自己的輸入 token；
• 每個進程根據路由器（Gate）的輸出，將 token 分為 N 份，分別屬于 N 個專家（也就是 N 個目標 GPU）；
• 然后用 all_to_all 一次性把這 N 份數據分發到對應的 GPU 上；
• 所有專家完成前向計算后，再用 all_to_all 把結果發回。

參考 deepspeed/moe/utils/all_to_all.py：

def all_to_all(input, group):# input shape: (num_local_experts, tokens_per_expert, hidden_size)output = torch.empty_like(input)torch.distributed.all_to_all_single(output, input, group=group)return output

通常輸入需要 reshape 成：

[local_experts, tokens_per_expert, hidden_dim]

之后經過兩次 all_to_all：

• tokens 分發階段：將 token 從本地發送到它所需的專家所在的 GPU；
• 結果收集階段：將處理后的輸出再收集回原始的 GPU。

參考鏈接

Applying Mixture of Experts in LLM Architectures

https://www.cnblogs.com/theseventhson/p/18247463

需要 reshape 成：

[local_experts, tokens_per_expert, hidden_dim]

之后經過兩次 all_to_all：

• tokens 分發階段：將 token 從本地發送到它所需的專家所在的 GPU；
• 結果收集階段：將處理后的輸出再收集回原始的 GPU。

參考鏈接

Applying Mixture of Experts in LLM Architectures

https://www.cnblogs.com/theseventhson/p/18247463

https://zhuanlan.zhihu.com/p/1888975857147691214