Transformer——Q135 证明二阶优化器（Shampoo）的逆矩阵近似误差传播

该问题归类到Transformer架构问题集——训练与优化——优化器。请参考LLM数学推导——Transformer架构问题集。

1. 问题背景

随着大语言模型（LLM）的参数规模呈指数级增长，传统的一阶优化器，如随机梯度下降（SGD）及其变体 Adam 等，在训练效率和收敛速度上逐渐难以满足需求。一阶优化器仅利用梯度（一阶导数）信息更新参数，在面对复杂的损失函数地形时，容易陷入局部最优或在平坦区域缓慢收敛。而二阶优化器能够利用 Hessian 矩阵（二阶导数矩阵）所包含的曲率信息，理论上可以更精准地确定参数更新方向，加速模型收敛。

然而，直接计算和存储 Hessian 矩阵对于大规模的 LLM 来说是不可行的，因为其计算复杂度和内存需求都极高。以 GPT-3 为例，拥有 1750 亿个参数的模型，其 Hessian 矩阵的元素数量将达到惊人的量级。为解决这一问题，Shampoo 优化器应运而生。它通过对参数矩阵进行分解，并近似计算其逆矩阵来更新参数。但这种近似计算不可避免地会产生误差，而误差在训练过程中的传播情况，直接影响着优化器的性能和模型训练的稳定性，因此分析其逆矩阵近似误差传播具有重要意义。

2. 技术原理

Shampoo 优化器基础

Shampoo 优化器的核心思想是对参数矩阵进行分块分解，分别对每个块计算自适应的学习率。假设模型的参数矩阵为，Shampoo 将其沿着行和列维度进行分解。

• 矩阵分解：将在行维度上划分为r个块，在列维度上划分为s个块，即，其中，且