精读笔记
Problem Setting
[PoLoRA: A Preconditioned Orthogonalized LoRA Optimizer]（arXiv preprint / 2026）
这篇论文实际处理的是 LoRA 优化几何错配问题：训练参数是 A、B，但模型和 loss 看到的是 BA。这个错配导致两类已有方法都不理想：Adam 把 A、B 当普通向量参数，完全忽略矩阵与乘积结构；Muon 虽然是 matrix-aware，但若分别作用在 A、B 上，控制的是 factor update，不是 merged update；Product Muon/iMuon 类方法开始约束 BA 的线性化变化，但仍没有充分处理 loss curvature 和 factor step magnitude。
关键矛盾是：LoRA 希望在极低参数成本下快速移动函数，但低秩乘积分解带来非唯一性和不良尺度。一个在 factor 空间看起来合理的 step，可能在 BA 空间过小、过大或方向失真；一个在 BA 空间有 spectral bound 的 step，也可能因为 A/B 的小奇异值把 factor 本身推得很激进。PoLoRA 的目标就是把“方向”和“步长”都重新对齐到 BA 及其对 loss 的影响上。
Motivation
作者并不是单纯想把 Muon 用到 LoRA，而是观察到“matrix-aware optimizer 的收益不能自动迁移到 low-rank adapter”。原因是 LoRA 的核心对象不是一个自由矩阵 W，而是一个被 A、B 分解约束的 product。直接优化因子会丢掉 product geometry；只做 product-aware spectral LMO 又太保守/太盲，因为它只控制层输出变化，不知道不同方向对 loss 的敏感性。
真正缺的是一个同时满足三件事的 optimizer：更新方向应当看 BA 而不是单个 factor；方向度量应当反映 per-sample loss sensitivity，而不是只用普通 spectral norm；factor update 的实际大小需要被显式约束，否则乘积分解中的尺度病态会把理论上的 merged-space 控制破坏掉。PoLoRA 的设计基本就是补这三个缺口。
Core Idea
PoLoRA 的核心思想是把 LoRA step 解释为 merged weight BA 上的受约束 steepest descent，而不是 A、B 参数空间中的 adaptive gradient step。它先问：如果更新 ΔW，会不会让 batch 中任意 sample 的 loss 线性化变化过大？这个约束经过 outergradient set 和 Kronecker approximation 后，变成一个 preconditioned spectral norm constraint：||P^{1/2} ΔW Q^{1/2}||2 受限。于是 matrix sign step 不再是普通 Muon 的 msign(G)，而是对 P^{-1/2} G Q^{-1/2} 做 orthogonalized update。
然后作者把 ΔW 替换为 LoRA 的线性化 product update BΔA + ΔBA，并用 spectral norm subadditivity 拆成两个 factor-wise LMO。这一步的本质不是“多了两个公式”，而是把 LoRA 的信息流重新组织为：梯度先通过 product geometry 投影到 factor 更新，再通过 curvature metric 纠正方向，最后通过 spectral magnitude rule 保证更新不会因为 factorization 病态而失控。相比 prior，它引入的 inductive bias 是“LoRA step 应该在 loss-sensitive merged-space 中近似等距”，而不是在 factor coordinates 中等距。
Method
Product-aware spectral direction 解决的是 factor-wise 更新与 BA 更新不一致的问题。Muon 分别作用在 A、B 上时，无法保证 BΔA + ΔBA 的 spectral norm；PoLoRA 延续 Product Muon 的思想，直接把约束放到线性化 merged update 上。核心变化是 optimizer 的基本对象从 ΔA、ΔB 变成它们诱导的 ΔW。
Curvature preconditioning 解决的是 plain spectral norm 过于粗糙的问题。Muon 的 spectral bound 控制所有 unit input 的最大输出变化，但 finetuning 中真正相关的是 loss 对不同输入/输出方向的敏感性。PoLoRA 用 per-sample loss change 的线性化约束推出 P、Q 加权的 spectral norm，相当于在 batch-induced metric 下做 orthogonalized descent。这里 P、Q 是 diagonal Kronecker 近似，工程上轻，但理论角色是 loss-sensitive geometry。
Magnitude control 解决的是 Product Muon 的隐含不稳定性。Product Muon 约束 BΔA 和 ΔBA，不直接约束 ΔA、ΔB；当 A 或 B 的谱结构差时，factor step 可以很大。PoLoRA 把两个 factor direction 归一到共同 spectral norm ρ，并令 ρ = η/(||A||2 + ||B||2)，从而用三角不等式控制线性化 merged update。这个规则看起来简单，但很可能是实测稳定性的关键之一。
Momentum、diagonal EMA preconditioner、Newton-Schulz/power iteration 属于让方法可跑的实现层，不是核心思想。它们重要，但主要是把上述几何约束压到 LoRA finetuning 可接受的 overhead 内。
Key Insight / Why It Works
最重要的 insight 是：LoRA 优化的瓶颈不是缺少更强的 adaptive learning rate，而是更新几何没有对齐到 BA 及其 loss sensitivity。Adam 的 per-coordinate adaptivity 在 factor coordinates 中工作，但这些 coordinates 是任意且冗余的；Muon 的 matrix sign 在单个矩阵上合理，但 LoRA 的单个因子不是函数空间对象。PoLoRA 的有效性来自把 steepest descent 的 norm ball 放到“loss-aware merged update”上。
我认为最核心的贡献不是 Product Muon 部分，而是 curvature preconditioning + magnitude control 的组合。Product-awareness 已经在同期工作中出现，并且文中 ablation 显示 Product Muon 本身不能稳定超过 Adam。Curvature 给 direction 增加了 representation alignment：更新更倾向于避开 per-sample loss 高敏感方向，或在低敏感方向更大胆移动。Magnitude rule 则是 scaling/stabilization：它把低秩分解的尺度自由度压住，减少因子病态导致的 optimizer 噪声放大。
这不是 retrieval、curriculum、memory reuse 或 test-time compute；本质上是 better inductive bias + scaling control。更具体地说，它把 LoRA optimizer 的归纳偏置从“参数坐标自适应”换成“函数更新近似受控”。论文里“per-sample loss control”的理论推导很干净，但实际实现只用了 diagonal Kronecker outergradient approximation，所以不能过度解读为真正精确的 per-sample constraint。实测增益可能有相当一部分来自更好的 step-size normalization，而不完全来自 curvature 建模。
还有一点值得注意：PoLoRA 不追求严格 reparameterization invariance，且仍然赢过 iMuon。这说明在当前 LoRA finetuning regime 中，invariance 可能不是主要限制；稳定控制 merged update magnitude 和 loss-sensitive direction 更重要。这个判断对后续 LoRA optimizer 设计很有价值。
Relation To Prior Work
它最接近三条线：Muon/Scion/Shampoo 这类 matrix-aware optimizer，LoRA-RITE/iMuon/LoRA-Muon/Riemannion 这类 product-invariant 或 product-aware LoRA optimizer，以及 Spectron/QuacK 这类低秩乘积 magnitude control 方法。
与 Adam/LoRA+、rank-stabilized LoRA、μA 的差异在于，PoLoRA 不是调 scalar hyperparameter 或 rank scaling，而是改变 update direction 本身。LoRA+ 等方法仍在 Adam coordinate geometry 内工作；PoLoRA 直接换了 norm/metric。
与 Muon 的差异在于，Muon 的 matrix sign step假设被优化对象就是一个矩阵 W；PoLoRA 处理的是 W0 + BA，并且把 sign step 放进 preconditioned product geometry。与 Product Muon/iMuon/LoRA-Muon 的差异在于，PoLoRA 不止做 product-aware spectral LMO，还加入 loss-curvature preconditioning 和 factor magnitude balancing。论文也明确指出若移除后两者会退化到 Product Muon，而 Product Muon 并不稳定优于 Adam。
与 Shampoo/Mousse 的关系更微妙：preconditioned matrix sign 本身不是全新思想，Kronecker second moment 也不是全新思想。PoLoRA 的实质创新是把这些思想接到 LoRA factor gradients 上，并从 per-sample loss-control 角度给出一个能够导出 product-aware preconditioned spectral update 的框架。可以说它是“Muon × Shampoo-style curvature × low-rank product magnitude control”的一次有效重组，其中 product-space 约束和 magnitude rule 的组合是新增信息最多的部分。
Dataset / Evaluation
实验覆盖了多个 1B-8B 级别语言模型、code/math instruction tuning、rank sweep，以及一个 low-resource language shift。对 optimizer paper 来说，这个覆盖足以支持“在常见 LoRA SFT 设置下，PoLoRA 比 tuned Adam 更快达到相同 held-out loss”这一窄 claim。它还报告 wall-clock overhead，而不只报 step speedup，这点比较关键，因为 PoLoRA 的额外矩阵运算如果不能摊薄，实际价值会下降。
但 evaluation 没有真正证明更强的 downstream generalization。主要指标是 held-out loss，而且 eval split 来自同一数据源；这更像验证优化效率，而不是验证任务能力。低资源 Bengali 的实验暗示“离预训练分布越远收益越大”，但证据还弱，可能受数据规模、tokenization、模型预训练覆盖、loss landscape 差异共同影响。文中未充分说明不同 batch size、sequence length、adapter placement、训练步数延长后增益是否保持。
对核心 claim 来说，ablation 是有用的：curvature 和 magnitude 都贡献了 speedup gap。但增益归因仍不彻底，因为 curvature preconditioner、direction normalization、learning-rate stability 之间可能强耦合。现有实验更能说明“这个组合有效”，不能完全说明“per-sample loss-control derivation 是主要原因”。
Limitation
第一，理论约束和实际 optimizer 之间有明显近似链条：per-sample loss 被一阶线性化，per-sample gradients 不直接可得，outergradient set 用 Kronecker approximation，P/Q 又进一步取 diagonal，LoRA product update 也丢掉 ΔBΔA。每一步都合理，但最终控制的是一个 surrogate，不是真实 loss change。
第二，方法不严格保持 LoRA factorization 的 reparameterization invariance。作者观察到 ||B||2/||A||2 会自平衡，但这更像 empirical property，不是强保证。若某些初始化、rank、adapter placement 或训练目标导致 factor norm 不平衡，magnitude rule 可能变成保守 step-size limiter。
第三，scalability 上限不只由 FLOPs 决定。PoLoRA 在 r=256、batch size 16、2048 context 下 overhead 小，但更高 rank、更小 batch、更频繁 adapter 更新、非标准硬件或未优化 kernel 下，r^2(din+dout) 和 r^3 小矩阵操作可能不再免费。文中未充分说明这些 regime。
第四，泛化 claim 应当收窄为 optimization transfer，而不是 capability transfer。学习率跨 rank 稳定是很有价值的 optimizer property，但不等价于跨模型规模、数据分布、训练目标的稳定最优。下游任务、长期训练和 deployment latency/storage 影响都没有充分覆盖。
第五，增益来源不清的部分仍存在。PoLoRA 的一部分收益可能主要来自 scaling / magnitude normalization，而不是 curvature 建模本身；另一部分可能来自 Adam baseline 在 LoRA rank/lr scaling 上没有使用更强 recipe。虽然作者调了 Adam 学习率，但没有覆盖所有现代 LoRA scalar scaling 策略的组合基线。
Takeaway
- 1. LoRA optimizer 设计应优先对齐到 merged product BA，而不是把 A、B 当普通参数矩阵；factor-space adaptivity 很容易优化错对象。
- 2. 对 matrix-aware optimizer 来说，spectral direction 只是起点。
- 没有 curvature metric 和 magnitude control，Product Muon 这类方法可能只是形式上更几何，实测不一定超过 Adam。
- 3. PoLoRA 最值得迁移的 insight 是“用 per-sample loss-change control 推导 optimizer norm/metric”，这可能适用于其他 compositional parameterization，例如 attention product、adapter composition、低秩 pretraining 或部分 structured finetuning。
一句话总结
PoLoRA 是 LoRA 优化从 factor-coordinate adaptive update 走向 loss-aware product-space spectral update 的一篇代表性工作，真正贡献在于把 product-aware Muon、曲率预条件和低秩乘积步长控制组合成了一个实测稳定超过 Adam 的优化几何。