Shampoo优化器的二阶矩分块近似原理
Shampoo是Google提出的二阶优化器，在TPU大规模训练中效果显著。请解释其如何通过分块Kronecker积近似实现二阶优化的计算可行性，以及与Adam和K-FAC的对比。
回答
小字辈
Shampoo用分块Kronecker积对角近似完整Fisher信息矩阵，将O(d³)降为O(Σb_i³)。
核心思想：
- 将参数矩阵W∈R^(m×n)按维度分组，统计左/右梯度统计量： L_t = βL_{t-1} + G_t G_t^T (左，m×m) R_t = βR_{t-1} + G_t^T G_t (右，n×n)
- 预调节步骤：W_t = W_{t-1} - η · L_t^(-1/4) · G_t · R_t^(-1/4)
- 奇异值分解/特征分解计算逆四分之一根
与Adam对比：Adam仅对每个参数做对角缩放(一阶矩+二阶矩)，Shampoo考虑梯度矩阵的完整行/列相关性。
与K-FAC对比：K-FAC也使用Kronecker近似，但需要做更多近似(如Fisher的可逆性)；Shampoo更通用，无需特定网络结构假设。
实际效果：
- 在Transformer上，Shampoo比Adam收敛步数减少50%
- 在GPT-3实验中，Shampoo达到同等困惑度的训练步数减半
- 缺点：每步计算的SVD开销大，Distributed Shampoo通过分片和异步更新解决
Distributed Shampoo现已是Google Pax/Jax生态中的标准优化器。