[论文解读] McTorch, a manifold optimization library for deep learning

Mayank Meghwanshi, Pratik Jawanpuria+3 more|arXiv (Cornell University)|Oct 3, 2018

Human Pose and Action Recognition参考文献 22被引用 24

一句话总结

McTorch 是一个 PyTorch 扩展，通过将黎曼优化功能直接集成到 PyTorch 框架中，实现了深度学习中的流形约束优化。它允许用户通过模块化设计将模型参数约束到流形（如 Stiefel 流形（正交矩阵）和 PositiveDefinite 流形（对称正定矩阵）），该设计将流形与优化器解耦，从而可与标准的 PyTorch 模块和优化器（如 SGD 和 Adagrad）无缝集成。

ABSTRACT

In this paper, we introduce McTorch, a manifold optimization library for deep learning that extends PyTorch. It aims to lower the barrier for users wishing to use manifold constraints in deep learning applications, i.e., when the parameters are constrained to lie on a manifold. Such constraints include the popular orthogonality and rank constraints, and have been recently used in a number of applications in deep learning. McTorch follows PyTorch's architecture and decouples manifold definitions and optimizers, i.e., once a new manifold is added it can be used with any existing optimizer and vice-versa. McTorch is available at https://github.com/mctorch .

研究动机与目标

• 通过在 PyTorch 中原生支持流形优化，弥合深度学习框架与流形优化之间的差距。
• 降低研究人员和实践者在深度学习模型中应用流形约束（如正交性和正定性）的门槛。
• 实现在 PyTorch 的动态计算图和 autograd 系统中无缝集成流形约束层和优化器。
• 通过重用 PyTorch 的张量操作和计算图，支持高效、GPU 加速的黎曼优化。
• 提供一个模块化、可扩展的库，允许独立添加新流形和优化器，并在不同模型中复用。

提出的方法

• 扩展 PyTorch 的 Parameter 类，使其可接受流形规范，从而实现参数在流形上的初始化，并通过 rgrad 属性追踪黎曼梯度。
• 引入一个 Manifold 基类，包含关键黎曼运算的抽象方法：rand（随机点生成）、retr（切向量映射的收缩）、egrad2rgrad（欧氏梯度到黎曼梯度的转换）。
• 修改现有的 PyTorch 优化器（如 SGD、Adagrad、ConjugateGradient），使其使用黎曼梯度并应用基于收缩的更新，而非欧氏更新。
• 引入 ManifoldShapeFactory，用于生成与层权重形状兼容的流形对象，确保初始化符合流形约束。
• 通过在层定义中新增 weight_manifold 属性，集成流形约束的线性和卷积层（Linear、Conv1d、Conv2d、Conv3d）。
• 所有数值运算均使用 PyTorch 的张量 API 实现，确保对 CPU 和 GPU 计算的完整支持，并保证端到端可微分。

实验结果

研究问题

• RQ1 如何高效且无缝地将流形优化集成到 PyTorch 深度学习框架中？
• RQ2 解耦流形与优化器的模块化设计，是否能实现深度学习中黎曼优化的灵活且可扩展的使用？
• RQ3 McTorch 在保持参数在流形上严格可行的前提下，对常见深度学习操作（如反向传播和 GPU 加速）的支持程度如何？
• RQ4 McTorch 如何在标准深度学习架构中实现正交权重初始化等流形约束的实际应用？
• RQ5 在真实世界深度学习任务中，使用 McTorch 对正交性或正定性等约束进行模型训练的可行性与性能如何？

主要发现

• McTorch 通过极少的代码修改成功扩展了 PyTorch，支持流形约束优化，用户仅需一个参数属性即可将层权重约束到 Stiefel 和 PositiveDefinite 等流形。
• 该库与 PyTorch 的 autograd 系统完全兼容，确保反向传播能通过 egrad2rgrad 变换正确计算黎曼梯度。
• 优化器如 Adagrad 和 SGD 已成功适配为执行黎曼优化，通过使用收缩和黎曼梯度，保持了参数在流形上的可行性。
• 模块化设计使得任意新流形可与任意现有优化器配合使用，反之亦然，显著提升了可扩展性并降低了实现开销。
• McTorch 支持 CPU 和 GPU 计算，所有数值运算均使用 PyTorch 的张量操作实现，确保高性能并兼容现有深度学习工作流。
• 该库以 BSD-3-Clause 许可证发布，公开获取地址为 https://github.com/mctorch/mctorch，目前正开发示例和基准测试以支持社区采用。