Stiefel-AdamW：用于线性分解模块的几何感知 AdamW

现代深度学习中普遍存在的结构模式是线性分解模块：$W = BA$ 形式的子模块，其中两个参数矩阵直接相乘，没有中间的非线性。此类块出现在 LoRA 适配器、低秩压缩层、自注意力的查询键产品中，并且具有共同的病理学特征：分解是非唯一的，这可能会破坏训练的稳定性并限制可用的学习率。尽管如此，因式分解块通常使用忽略底层几何的标准欧几里得方法进行优化。我们引入了 Stiefel-AdamW，它是 AdamW 的近乎直接替代品，可在出现此类块的任何地方使用。通过在 Stiefel 流形上约束一个因子，同时保留另一个欧几里得因子，Stiefel-AdamW 将完整的 $\mathrm{GL}(\mathbb{R}^r)$ 规范对称性松弛为紧凑的正交对称性，排除了因子爆炸，同时保留了赋予 AdamW 实际强度的坐标方向对角预处理。力矩估计是在环境欧几里得空间中执行的，几何形状仅通过切线空间投影和流形回缩进入。 AdamW 的实现开销很小，并且我们表明生成的优化器继承了黎曼方法的稳定性优点和标准收敛保证。我们在 GPT2、ViT 和 Mistral 7B 的 LoRA 式微调以及 OpenWebText 上对 GPT2 的全面预训练上验证了 Stiefel-AdamW，显示出在强基线上的持续改进，而与 AdamW 相比基本上没有额外成本。