[论文解读] Theseus: A Library for Differentiable Nonlinear Optimization

Luis Villaseñor-Pineda, Taosha Fan+11 more|arXiv (Cornell University)|Jul 19, 2022

Sparse and Compressive Sensing Techniques被引用 46

一句话总结

Theseus 是一个基于 PyTorch 的开源库，用于可微分非线性最小二乘（DNLS），它提供应用无关接口、稀疏求解器、批处理、GPU 加速，以及用于端到端结构化学习在机器人和视觉中的隐式微分。

ABSTRACT

We present Theseus, an efficient application-agnostic open source library for differentiable nonlinear least squares (DNLS) optimization built on PyTorch, providing a common framework for end-to-end structured learning in robotics and vision. Existing DNLS implementations are application specific and do not always incorporate many ingredients important for efficiency. Theseus is application-agnostic, as we illustrate with several example applications that are built using the same underlying differentiable components, such as second-order optimizers, standard costs functions, and Lie groups. For efficiency, Theseus incorporates support for sparse solvers, automatic vectorization, batching, GPU acceleration, and gradient computation with implicit differentiation and direct loss minimization. We do extensive performance evaluation in a set of applications, demonstrating significant efficiency gains and better scalability when these features are incorporated. Project page: https://sites.google.com/view/theseus-ai

研究动机与目标

• 为 PyTorch 上的可微分非线性最小二乘（DNLS）提供应用无关接口。
• 通过整合二阶优化器、标准成本和可学习成本，以及可微分李群和运动学，支持高效的端到端结构化学习。
• 利用稀疏性、批处理和 GPU 加速，提升机器人与视觉中 DNLS 的性能和可扩展性。

提出的方法

• 提供一个界面，用于从二阶优化器（高斯-牛顿、带自适应阻尼的 Levenberg–Marquardt、Dogleg）和一个非线性最小二乘目标函数构造优化层。
• 提供可微分李群和可微分运动学，在优化中处理 2D/3D 位姿和机器人模型。
• 通过 PyTorch 的自动微分，结合可微分的标准成本函数和鲁棒损失。
• 实现可微分的稀疏线性求解器（CPU CHOLMOD、基于 cuSolverRF 的 cudaLU、BaSpaCho），并支持批处理和 GPU 执行。
• 应用隐式微分和直接损失最小化，以实现内层 DNLS 优化的端到端梯度流。
• 启用反向传播模式，包括展开、TBPTT、隐式微分和直接损失最小化。

实验结果

研究问题

• RQ1 如何使 DNLS 在机器人和视觉任务中高效且可作为应用无关的库？
• RQ2 在 DNLS 中整合稀疏求解器、批处理和 GPU 加速有哪些好处，以及不同的向后微分模式如何比较？
• RQ3 在基于 PyTorch 的 DNLS 框架中，是否能够有效支持可微分李群和可微分运动学？
• RQ4 Theseus 在端到端学习设置和不同问题规模下，与如 Ceres 等最先进求解器相比，表现如何？

主要发现

• Theseus 通过使用稀疏求解器和跨问题的批处理，显示出比密集求解器显著更高的效率。
• 在大规模或批量 DNLS 中，Theseus 比 Ceres 快得多，在最大的测试配置中报道的加速可达约 23x。
• 通过隐式微分进行向后传播，在以更多优化迭代时仍保持恒定的计算/内存，优于展开在梯度质量上的表现。
• BaSpaCho，一个批量稀疏 Cholesky 求解器，在大型问题上提供了评估求解器中的顶级性能。
• 隐式微分与直接损失最小化的结合，使 DNLS 应用的端到端学习更具鲁棒性和可扩展性。