Title: Theseus: A Library for Differentiable Nonlinear Optimization
Authors: Luis Pineda
Taosha Fan
Maurizio Monge
Shobha Venkataraman
Paloma Sodhi
Ricky T. Q. Chen
Joseph Ortiz
Daniel DeTone
Austin S Wang
Stuart Anderson
Jing Dong
Brandon Amos
Mustafa Mukadam
Venue: NeurIPS 2022
Keywords: robotics
differentiable optimization
nonlinear least squares
implicit differentiation

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/

TL;DR: An efficient application-agnostic open source library for differentiable nonlinear optimization built on PyTorch, providing a common framework for end-to-end structured learning in robotics and vision.

Review discussion:
STRENGTH:  potentially high utility for the ML community interested in differentiable programming paradigm
With this library, researchers can more easily design differentiable programming systems while spending less time on the differentiation part (or at least can quickly code up "proof of concept" systems using this general tool), esp. considering that the optimization is one of the hardest parts when making a system differentiable.
WEAKNESS:  novelty in this paper is unclear, if not low, as it is currently presented and described.
Another possible weakness is that, similarly to other methods that achieve differentiation by pytorch, the overhead could be high compared to simple forward C++ code, especially when the data size is small and the advantage of parallelization is not significant. An example is in Figure 3, where the existing Ceres (black dashed) outperforms most of the methods in this new library on small data.
Can you highlight or give more detail on the  key novelties of the design choices & algorithms in the library itself?
Not discussed clearly.
No
We thank the reviewer for their feedback. We are pleased to see that the reviewer thinks that with this library researchers “can more easily design differentiable programming systems while spending less time...”; this is one of our main goals! We are also glad the reviewer recognized some of the engineering challenges we tackled: “Cuda solver is usually closed-source. Making a differentiable version of it takes effort.”
The reviewer points out that the novelty of this paper is unclear, to which we respectfully disagree. The last three paragraphs of the introduction in the paper state our novel contributions in detail. Theseus is essentially **the first library** that offers a “full-stack” solution for differentiable NLS including sparse solvers, Lie-algebra support, a wide range of out-of-the-box cost functions for robotics and computer vision, support for differentiable robot kinematics, support for custom cost functions and an easy to use API. As other reviewers have already pointed out, there is value for the community in having such a library, which did not exist before Theseus.
We agree with the reviewer that there will be overhead compared with forward code written in C++, but we must stress that our purpose is to offer a flexible library for research in *differentiable* optimization, most of which is written in Python, and a significant portion of it in PyTorch.
As such, directly comparing computation against libraries that don’t support differentiation would be unfair. In the paper, we have included a benchmark against Ceres to both validate the correctness of our implementation, and to show that we have put a lot of emphasis in performance. We think the fact that Theseus can outperform Ceres in the batched setting for medium-to-large problems shows that performance is a major consideration.
But if the user only needs to do non-differentiable forward optimization, this is not the problem that Theseus was designed for.
We also refer the reviewer to Sec 6 in the paper where we discuss limitations in detail which is also acknowledged by the other reviews as “properly discussed” and “clearly addressed”.
**Strengths**:
1. The proposed library is task-agnostic allowing for end-to-end learning across a variety of robotics and vision problems.
2. The implementation supports a variety of backpropagation approaches for taking gradients including autodiff.
3. The paper shows either comparable or better results as compared to Ceres solver.
4. The motivation is well supported with several examples and has value to the community.
**Weaknesses**
1. Technically, there do not seem to be any flaws in the proposed approach.
2. The authors cite several papers that already implement related approaches. It will be good to compare a general purpose solver like Theseus with application specific implementations.
3. The writing needs
1. Reiterating from above, I would like to see comparisons with existing task specific implementations to see differences if any.
2.
In terms of the optimization, the approach builds on a two step optimization scheme.
However, this may lead to instabilities in optimization especially with the choice of wrong hyperparams. How much does the choice of hyperparameters affect the overall performance of the proposed approach?
The authors have been upfront about the limitations and have clearly addressed them in the settings.
No
We thank the reviewer for their thoughtful feedback. We are glad that the reviewer highlights the task-agnostic nature of our library and that its “motivation is well supported... and has value to the community” .
Regarding comparing with application specific implementations, in this paper we put a lot of effort in evaluating and demonstrating the correctness of our differentiable optimizer, by comparing our results with an established stand-alone optimizer (Ceres), performing thorough unit testing of all gradient computations, and adding diverse illustrative examples of how to do learning through a TheseusLayer. We hope that this offers enough evidence that our features are sound and useful. We thus argue that comparing with application specific implementations is outside the scope of the paper, considering that a lot of their performance depends on the neural architecture used, which can be quite complicated and might require extensive hyperparameter tuning. We have opted for focusing our efforts on evaluating features that are direct responsibility of our library.
Regarding the impact of hyperparameters for differentiable NLS optimization, this is likely to be problem dependent, and in many cases it can have a significant impact. We have thus added a lot of flexibility for users to easily change hyperparameters for the inner optimization when experimenting in their own research applications. This includes, number of iterations, choice of optimizer, choice of backward method, tolerance, turn on/off Lie-group gradient projection, among others. We hope that researchers find this useful in the future, and plan to adapt these options with evolving needs of the library’s users.
## Strengths
- The authors make a compelling case for the usefulness of the method implimented by their package, giving an extensive bibliography of uses in robotics and vision.
- They also implement two novel (in this context) methods for computing the adjoint derivatives for DNLS: Implicit differentiation and Direct Loss minimization.  Experiments in Sec. 5.1 show an improvement in efficiency and accuracy relative to unrolling or truncated differentiation.
- The library is very flexible providing two different sparse NLS solvers: CHOLMOD and cudaLU, and explaining the advantages and drawbacks of each.  Four methods for the backward mode are provided.  There is also support for optimizing variables in Lie groups such as orientations in SO(3) which are quite important in robotics and vision and tricky to optimize over.
- Support for batching and automatic vectorization gives the DNLS layer similar levels of parallelism to other pytorch layers.
- The supplementary material includes the complete code of the package including scripts and instructions to run several different examples in the paper.  The appendix includes nice explanations of the examples and proofs justifying the novel backward modes.  Given that the purpose of the paper is provide a solution which works flexibly across different applications, I found seeing these examples script and explanations helpful to understanding this.
## Weaknessess
- Some of the empirical evaluations are somewhat hard to interpret.  Figure 2 left is not a very clear plot and it is hard to find a strong conclusion of when to prefer either sparse solver (although the preference over dense is clear.)
- I understand that it is a prior assumption from the literature that the combination of deep learning and NLS afforded by this method is a worthwhile and superior approach in many applications relative to either DL or NLS independently.  It still would have been convenient if this comparison could be made explicitly for one of the examples.
- Figure 4, right shows that the choice of backward method inside Theseus have an impact on the ultimate performance of the model.  The authors note this is probably application dependent.  This seems like a critical aspect for end users and should probably be explored more.
## Minor Point
- Line 74: $\in$ should be $\subset$.  Do the subsets need to be disjoint?
- Line 76: parallelism issue.  ``Euclidean vector or a matrix Lie group element''
- Line 91: maybe a small semantic point, but this does not really seem like a second-order method since the 2nd derivatives do not need to be computed, which seems like an advantage.  I have seen the derivation and I understand that $J^T J$ is the second order term of the objective after linearizing the model, but it seems one should be careful to be precise here.
- Line 106: "the lastest"
## Questions
- Figure 1: It seems the larger batch has smaller speed up.   Why do you think that is?
## Post-Response Edit
I have read the other reviews and author responses.  These strengthen my positive assessment and I maintain my score.
The paper discusses limitations of various alternatives throughout the work, for example, the fact that the CHOLMOD spare solver is tied to the CPU.  A limitations section at the end lays out other limitations, for example, incompatibility with hard constraints.
No
We thank the reviewer for their comprehensive feedback. We are pleased to see that the reviewer found our library flexible, mentioning features such as sparse solvers and SO(3) support, that they see novelty in our adjoint differentiation methods in this context, and that they found our examples helpful in demonstrating the flexibility of the library.
The reviewer has pointed out that Figure 2 is somewhat hard to interpret and to derive a strong conclusion from. We agree that the explanatory text can be improved and plan to do so in the final revision. The main takeaway is that both solvers are comparable in terms of running time, with perhaps a slight advantage of LUCuda solver for smaller problems;
moreover, if not many CPU-threads are available for CHOLMOD, any advantage of cudaLU is likely to increase.
Thus, for most cases, starting with cudaLU is a good rule of thumb.
On the other hand, as mentioned in Appendix D.2, for very large problems and GPU-memory constraints the CHOLMOD-based solver might be a better choice, as the solver for the most part runs on CPU and is thus less memory constrained.
We agree with the reviewer that it would be convenient to see the advantage of structured end-to-end learning with DNLS.
However, we argue that this is outside the scope of the paper given that such advantages have already been illustrated in many previous work in this area (many of which we cite, as the reviewer pointed out), and it thus would distract from the focus of the paper, which is describing our flexibility- and efficiency-based design, and benchmarking our different features.
Similarly, we agree with the reviewer that exploring different backward modes across many applications would be very interesting, and it’s indeed one of the reasons why we have created the library: so that researchers can do this more easily in the future. Doing a comprehensive analysis of this across many applications is nevertheless outside the scope of this paper.
Regarding the minor questions, thanks for pointing these out, we will fix them in the final version of the paper.  The sets in line 74 do not have to be disjoint.
Regarding the smaller speed up with larger batch size (Fig. 1), note that in this experiment we only run a single batch and that GPU computation generally slows down with increasing batch size. Since vectorization has a multiplicative effect on the batch size, at some point the cost of having a very large batch starts outweighing the savings resulting from parallel processing. We will also clarify that the Hessian computed in NLS is approximate ($J^TJ$) unlike second order Newton‘s method.
**Strength**
- I am personally a big fan and a practitioner of deep structure models. Many of my previous work have combined deep neural networks with optimization (some are NLS while some are not). In order to train the whole model end to end, I had to write all these things myself, such as unrolling GN step, lie-algebra transformation, etc. It was fun, but indeed a bit cumbersome. It is thus very exciting to see that there is a flexible and modularized library that can greatly improve the speed of our future development. I can foresee the library being adopted by a wide variety of the community.
- The authors provides several case studies (which is just the tip of the iceberg) across computer vision and robotics and showcase how they can be implemented easily and cleanly with the task-agnostic interface.
- Theseus provide additional support/integration for sparse linear solvers, which can be used independently as standalone pytorch functions.
- Thorough analysis on computational cost and memory usage.
- I also like how the library integrates many (latest) techniques, such as projection operator for Lie-algebra, implicit differentiation, etc. It allows one to experiment with more crazy stuff such as unrolling tons of steps, or obtain more stable results with gradient explosion.
**(Minor) Weakness**
- While the aim of the package is to support generic nonlinear optimization, the current state seems to focus majorly on least square problem (*e.g.* GN solver). The title is thus a bit mis-leading imo.
- Following above, many state-of-the-art algorithms in vision and robotics adopt LM solver with learnable, data-driven $\lambda$. I don't see why this function is unavailable, or what's the major difficulty/block on this.
**Correctness**
- I'm not aware of any correctness issue in this paper.
**Minor**
- I personally gave it a try and use Theseus to replace my personal implementation on GN solver and lie-algebra (for 3D pose estimation).
While I didn't observe significant speedup (which aligns with the analyses in the paper), the code is indeed much cleaner ;)
Please see the above weakness section.
The limitations are properly discussed.
No
We thank the reviewer for their thoughtful feedback. We are glad that they found our library “very exciting” and with potential to be adopted by a wide variety of the community. We are also pleased that the reviewer sees value in our supported features like sparse solvers, projection operators for Lie-algebra, and implicit differentiation.
Moreover, it was exciting to see that the reviewer tried out the code for their own problem and found it helpful, even if not significantly faster for their problem.
The reviewer pointed out that Theseus doesn’t currently support learnable data-driven $\lambda$ for LM solver, which is correct. We agree with the reviewer that this feature should be relatively easy to support, and it’s indeed on our stack of coming features. For our dense solvers, adding this is straightforward, by simply allowing $\lambda$ to be a PyTorch tensor;
however, adding this is a bit tricker with our custom sparse solvers, and would require some non-trivial amount of work.
Thus, the main reason for its current omission has been limited time during which we have prioritized other features.

Meta-review:
role: AC
data: metareview: This paper presents Theseus, a software library which provides a new layer in the form of a differentiable nonlinear least squares (DNLS) solver. Forward pass solves the problem and the backward pass provides derivates for the optimum with respect to parameters. The reviewers uniformly appreciated the presentation of the paper and the usefulness and packaging of the proposed library. The library's features such as sparse solvers and Lie group algebra was also appreciated by the reviewers. Overall the paper and the library are a strong contribution to the Neurips community and hence I am happy to recommend for acceptance. I would urge the authors to make sure they follow through on feature requests and other suggestions made by the reviewers. 
role: PC

Decision: Accept