Meta人工智能团队早前开源了一个可微分非线性最小二乘（NLS）的优化技术库Theseus。据介绍，Theseus基于PyTorch，适用于机器人和计算机视觉等用例，允许研究人员能够轻松地将专家领域知识融入现代AI架构。它将knowledge表示为优化问题，并在基于梯度的学习过程中将其作为模块化“优化层”添加到架构之中。这个domain knowledge与训练数据不同，它可以帮助模型实现更为准确的预测。例如，为了确保机器人的运动平稳，研究人员可以将机器人的实现方式和运动模式的knowledge作为一个层，同时对机器人进行端到端的运动训练。

Theseus是一个为可微非线性优化提供跨应用框架的库。Theseus非常高效，可通过支持批处理、GPU加速、稀疏解算器和隐式微分来加速计算和内存。Meta宣称，它比谷歌基于C++的Ceres Solver（不支持端到端学习）快四倍。

Theseus融合了将prior knowledge注入人工智能系统的两种主流方法的最佳方面。在深度学习出现之前，研究人员使用更简单、独立的AI优化算法来解决机器人中的单个问题。通过计算手动选择的因素组合的最小值，机器人系统学会了执行命令的最佳方式。这种方法有效但不灵活。特定于应用的优化算法通常难以适应新的系统或环境。

另一方面，深度学习方法的可扩展性要高得多，但它们需要大量的数据，会产生有效但在training domain之外脆弱的解决方案。

为了训练特定应用的深度学习模型，研究人员使用精心选择的损失函数来衡量模型预测数据的效果。但要通过反向传播更新模型权重，每个层必须是可微分，允许误差信息通过网络流动。传统的优化算法不是端到端可微分，所以研究人员面临着一个权衡：他们可以放弃优化算法，转而进行专门针对特定任务的端到端深度学习，并有可能失去优化的效率及其泛化能力。

或者，他们可以离线训练深度学习模型，并在推理时将其添加到优化算法中。第二种方法的优势是将深度学习和prior knowledge相结合，但由于深度学习模型是在没有预存在信息或特定任务的错误函数的情况下训练，所以其预测可能不准确。

为了融合相关策略，减轻它们的弱点并利用它们的优势，Theseus将优化结果转换为一个可以插入任何神经网络架构的层。这样，修改可以通过优化层反向传播，从而允许研究人员对最终任务损失的特定domain knowledge进行微调，并作为端到端深度学习模型的一个组成环节。

Theseus（绿色）如上所示， objective由上游神经模型的输出张量（灰色）和prior knowledge（橙色）组成。Theseus层的输出是最小化 objective的张量。

工作原理

NLS测量一个非线性函数与它要预测的实际数据之间的差异。较小的值意味着函数非常适合数据集。从映射和估计到规划和控制，NLS在机器人和视觉问题的表述中非常普遍。例如，机器人朝向目标的路线可以表述为NLS优化问题：为了绘制最快的安全轨迹，系统会寻找一个可以最小化行程时间和不必要行为的sum-of-costs objective，比方说急转弯或与环境中的障碍物碰撞。sum-of-costs objective同时可以捕获传感器测量误差，从而优化机器人或camera的过去轨迹。

通过令NLS可微分，Theseus提供了可微非线性优化，并作为研究人员可以插入神经网络的层。输入张量定义了加权平方和objective函数，输出张量是产生objective最小值的参数。通过优化器进行微分，可以保留端到端计算梯度的能力。

这将优化器和已知先验集成到深度学习训练循环中，允许模型对domain knowledge进行编码，并学习实际的任务损失。同时，深度学习模型将从训练过程中的感知或语言教学中提取更大的objective。这样，研究人员可以在训练循环中使用已知的运动学模型来开发objective预测模型。这种将已知先验与神经组件模块化混合的技术可以提高数据效率和泛化能力。

为了提高效率，Theseus集成了对稀疏解算器、自动矢量化、批处理、GPU加速和带隐式微分的梯度计算的支持。正如autodiff和GPU加速推动了PyTorch在NumPy的发展一样，稀疏性和隐式微分推动了Theseus的发展，与通常只支持稀疏性的Ceres等解算器形成对比。

对于标准GPU，使用稀疏解算器的Theseus要比使用密集解算器快得多，所需内存同样要少得多。另外，Ceres的GPU支持有限，不支持分批或端到端学习。最后，隐式微分比标准展开产生更好的梯度。隐式差异同时具有恒定的内存和计算占用空间。

为何重要

Theseus提供了一个可以利用传统技术、视觉方法和深度学习的互补优势的通用框架。可微分优化作为归纳先验，提高数据效率和泛化能力。

现有的实现是特定于任务，通常效率低下。Theseus可以跨应用，所以AI社区可以通过训练在多个任务和环境中表现出色的精确模型来取得更快的进展。

相关论文：Theseus: A Library for Differentiable Nonlinear Optimization

Theseus目前已经开源。具体源代码请访问这里，而项目页面请访问这里。