3、MobileViT中的Transformer模块：局部与全局特征建模、多头自注意力机制、计算效率优化
好，咱们今天来啃一块硬骨头——MobileViT里的Transformer模块。
说实话，我第一次在移动端模型里看到Transformer时，心里是打鼓的。你想想看，Transformer那玩意儿在NLP里动辄上亿参数，放到手机上跑？这不是开玩笑吗？但苹果那帮人还真就把它塞进去了，而且效果出奇的好。今天我就带你拆开看看，他们到底是怎么做到的。
3.1 为什么要在MobileViT里用Transformer？
先聊个本质问题：CNN不是挺好的吗？为什么非要折腾Transformer？
我个人习惯把这个问题拆成两个层面来看：
- CNN的局限：卷积操作本质上是局部运算。哪怕你堆100层，每个像素能看到的范围也就那么大。这就导致了一个问题——全局依赖建模能力弱。比如你要识别一张图里远处的一只鸟和近处的一棵树之间的关系，CNN得靠层层堆叠才能“看到”，而且容易丢失信息。
- Transformer的优势：自注意力机制一上来就是全局的。每个像素直接跟所有其他像素做交互，长距离依赖一次搞定。但代价你也知道——计算量是O(n²)，n是像素数。一张224x224的图，n=50176，算力直接爆炸。
所以MobileViT的思路很聪明：我不全盘用Transformer，我只在关键地方用。具体来说，就是把图像切成一个个小patch，在patch内部用CNN做局部建模，在patch之间用Transformer做全局交互。这样既拿到了全局视野，又把计算量控制在了可接受范围内。
核心思想：局部用CNN精雕细琢，全局用Transformer宏观把控。两者互补，缺一不可。
3.2 局部与全局特征建模：MobileViT的“分而治之”
咱们直接看MobileViT block的结构。我画了一张图，帮你理清整个流程：
你看这个流程，其实就三步：
- 局部建模：先用一个3×3卷积对每个位置做局部特征提取。这一步跟普通CNN没区别，就是抓边缘、纹理这些细节。
- 展开成序列：把特征图切成一个个不重叠的patch（比如2×2或4×4），然后拉直成一维向量。这一步是关键——把图像数据变成序列数据，Transformer才能吃。
- 全局建模：把序列丢进Transformer，让每个patch跟所有其他patch做自注意力。这一步学到了全局上下文。
最后再折叠回特征图，跟原始特征做残差连接。嗯，这里要注意：残差连接不能丢，否则梯度传不回去，深层网络根本训不动。
我的经验：patch size的选择很讲究。我试过1×1（相当于每个像素一个patch），计算量直接炸了。也试过8×8，全局信息又不够。最后发现2×2或4×4在大多数任务上是个不错的折中。你可以根据你的输入分辨率调一调。
3.3 多头自注意力机制：让模型“多角度”看问题
多头自注意力（Multi-Head Self-Attention, MHSA）是Transformer的核心。说白了，就是让模型从多个不同的角度去理解输入之间的关系。
我打个比方：你看一张街景照片，可能同时关注到红绿灯、行人、车辆。每个“头”就像你大脑里的一个注意力通道——一个头专门看颜色，一个头专门看位置，一个头专门看形状。最后把所有这些信息合并起来，你才能做出“现在能不能过马路”的判断。
在MobileViT里，MHSA的数学过程是这样的：
# 伪代码：多头自注意力计算过程
def multi_head_self_attention(x, num_heads=4):
# x: [batch, num_patches, dim]
batch, n, d = x.shape
head_dim = d // num_heads # 每个头的维度
# 1. 线性投影得到 Q, K, V
Q = linear(x) # [batch, n, d]
K = linear(x)
V = linear(x)
# 2. 拆分成多个头
Q = reshape(Q, [batch, n, num_heads, head_dim])
K = reshape(K, [batch, n, num_heads, head_dim])
V = reshape(V, [batch, n, num_heads, head_dim])
# 3. 每个头独立计算注意力
# 注意力分数 = softmax(Q @ K^T / sqrt(head_dim))
attn_scores = softmax(Q @ K.transpose(-2,-1) / sqrt(head_dim))
# 4. 加权求和
out = attn_scores @ V # [batch, n, num_heads, head_dim]
# 5. 合并多头
out = reshape(out, [batch, n, d])
out = linear(out) # 最终投影
return out
这里有个细节：除以 sqrt(head_dim) 是为了防止内积过大导致softmax梯度消失。我刚开始做的时候忘了这步，结果训练loss死活不降，查了半天才发现是这个问题。
| 参数 | 典型值 | 说明 |
|---|---|---|
| num_heads | 4 | MobileViT中常用4个头，平衡性能与效率 |
| head_dim | 32 | 每个头的维度，总维度 = num_heads × head_dim |
| patch_size | 2×2 或 4×4 | 控制序列长度，影响计算量 |
避坑指南：我曾经在某个项目里把num_heads设成了8，想着“头越多效果越好”。结果模型不仅没变好，反而因为参数量暴增导致过拟合。后来我悟了——移动端模型资源有限，4个头已经够用了。别盲目堆参数。
3.4 计算效率优化：MobileViT的精打细算
好了，前面讲了原理，现在聊聊最实际的问题——怎么让Transformer在手机上跑起来。
Transformer的计算瓶颈在哪？就是自注意力那块。对于n个patch，计算复杂度是O(n²·d)。n稍微大一点，比如64×64=4096个patch，计算量就奔着千万级别去了。
MobileViT用了几个招数来优化：
- 减少patch数量：通过下采样先把特征图缩小，再进Transformer。比如输入是32×32，用4×4的patch，序列长度只有64。这比直接用像素级序列（1024个）少了16倍。
- 共享投影矩阵：Q、K、V的线性投影共享同一套参数。虽然理论上会损失一点表达能力，但参数量直接减半。在移动端，这点牺牲是值得的。
- 使用FFN的瓶颈设计：前馈网络（FFN）先降维再升维。比如输入是128维，先降到64维做非线性变换，再升回128维。这样计算量比直接做128→128少了将近一半。
我贴一段实际代码，你看看MobileViT里Transformer block是怎么实现的：
import torch
import torch.nn as nn
class MobileViTTransformerBlock(nn.Module):
def __init__(self, dim, num_heads=4, expansion_factor=2):
super().__init__()
self.norm1 = nn.LayerNorm(dim)
self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
self.norm2 = nn.LayerNorm(dim)
# 瓶颈FFN：先降维再升维
hidden_dim = dim // expansion_factor
self.ffn = nn.Sequential(
nn.Linear(dim, hidden_dim),
nn.GELU(),
nn.Linear(hidden_dim, dim)
)
def forward(self, x):
# 自注意力 + 残差
x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
# FFN + 残差
x = x + self.ffn(self.norm2(x))
return x
你看这个代码，核心就两个操作：自注意力 + 残差，FFN + 残差。简洁明了，没有花里胡哨的东西。
效率对比：同样处理32×32的特征图，标准ViT需要约1.2G FLOPs，而MobileViT的Transformer模块只需要约0.15G FLOPs。效率提升了8倍，而精度只下降了不到1%。这就是优化的价值。
3.5 小结：Transformer模块的设计哲学
回顾一下，MobileViT的Transformer模块其实就三个关键词：
- 局部+全局：CNN做细节，Transformer做宏观，各司其职
- 多头注意力：从多个角度理解特征关系，提升表达能力
- 效率优先：每个设计决策都以“能不能在手机上跑”为底线
我个人觉得，MobileViT最厉害的地方不是它用了Transformer，而是它知道什么时候该用Transformer，什么时候不该用。这种“取舍”的智慧，比单纯堆砌技术难得多。
好，这一章就到这里。下一章咱们聊聊怎么把MobileViT真正部署到手机上，包括模型转换、量化、推理优化这些实战内容。到时候我会分享一些踩过的坑，保证让你少走弯路。