MobileLLM:面向设备端的小型语言模型优化
一、核心创新
MobileLLM通过深度优化架构设计,在小型模型上取得显著突破:
- 125M参数模型相比现有SOTA提升2.7%
- 350M参数模型提升4.3%
- 创新的即时块级权重共享(MobileLLM-LS)进一步提升0.7-0.8%
二、技术方法详解
1. 前馈网络优化
- 深入研究激活函数选择
- 验证SwiGLU在小型模型中的有效性
- 打破传统认知,证明架构设计对小模型性能影响显著
2. 深度vs宽度权衡
- 传统观点认为模型性能主要取决于:
- 参数数量
- 训练数据集大小
- 训练迭代次数
- 研究发现:
- 对于小型模型,增加深度比增加宽度更重要
- 这一发现为后续层共享策略奠定基础
3. 嵌入层优化
- 参数占比分析:
- 125M模型中嵌入层占比>20%(512维度,32k词表)
- LLaMA-70B中仅占0.7%
- 优化策略:
- 输入嵌入权重复用作输出层
- 通过重分配节省参数来增加层数
- 最终实现性能提升
4. 注意力机制优化
- 查询头数量:
- 实验验证16个查询头最优
- Key-Value头从16减至4
- 125M模型性能基本不变
- 350M模型仅降低0.2个百分点,但模型大小减少10%
- GQA优化:
- 结合更大的嵌入维度
- 125M模型额外提升0.4个百分点
5. 层共享创新
研究了三种权重共享策略:
- 即时块级重复
- 整体重复
- 反向共享
最终选择即时块级共享因为:
- 性能表现最佳
- 适应硬件内存层级
- 共享权重可以留在缓存中
- 支持立即两次计算
三、实验验证
1. 训练细节
- 快速探索:12万次迭代,0.25T tokens
- 完整训练:48万次迭代,1T tokens
2. 评估任务
零样本常识推理任务评估:
- ARC-easy
- ARC-challenge
- BoolQ
- PIQA
- SIQA
- HellaSwag
- OBQA
- WinoGrande
- TQA
- RACE
3. 性能表现
-
基础模型性能:
- 125M超越OPT、GPT-Neo、Galactica
- 优于Pythia-160M和RWKV169M
- TQA提升4.3个百分点
- RACE阅读理解显著提升
-
下游任务:
-
聊天能力:
- AlpacaEval评测
- MT-Bench基准测试
- 350M版本达到48.2%胜率(vs GPT-3)
-
API调用能力:
- 意图匹配与LLaMAv2 7B相当
- 结构精确匹配表现优秀
- Rouge分数略低但不影响实用性
-
4. 量化兼容性
- W8A8 PTQ量化测试
- 精度损失<0.5个百分点
- 与层共享机制完全兼容
四、扩展验证
扩展原理到更大模型:
- 600M版本
- 1B版本
- 1.5B版本(达到59.4%准确率,超越Qwen1.5-1.8B)
五、技术价值
- 建立小型模型新标准
- 提供系统优化方法论
- 适用于设备端部署
- 为大模型瘦身提供思路