PaperlayerView exact revision on arXiv

Source-linked AI summary

Theseus: A Library for Differentiable Nonlinear Optimization

Luis Pineda, Taosha Fan, Maurizio Monge, Shobha Venkataraman, Paloma Sodhi, Ricky T. Q. Chen, Joseph Ortiz, Daniel DeTone, Austin Wang, Stuart Anderson, Jing Dong, Brandon Amos, Mustafa Mukadam

arXiv:2207.09442v3cs.ROcs.CVcs.LGmath.OC

TL;DR

Existing DNLS implementations are application-specific and may omit features important for efficiency. Theseus provides a shared PyTorch library of differentiable optimization components and efficiency techniques; evaluations report faster, lower-memory solving and improved scaling across large batches.

  • Problem

    Existing DNLS implementations are application-specific and do not consistently include efficiency features such as sparse solvers, batching, and GPU support.

    Source: §1 · p. 2

  • Method

    Theseus is an application-agnostic PyTorch library combining second-order optimizers, cost functions, Lie groups, sparse solvers, batching, vectorization, GPU acceleration, and differentiation modes.

    Source: §1 · p. 2

  • Results

    Up to 20x faster forward pass than Ceres when solving a batch of large problems, while Theseus with a sparse solver uses significantly less memory than a dense solver on a standard GPU.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    The shared interface supports multiple robotics and vision applications using the same differentiable components, with reported performance improvements from its efficiency-based design.

    Source: §6 · p. 10

    Theseus's current nonlinear solvers handle constraints softly, and online learning with frequently edited objectives may incur nontrivial overhead because only non-incremental settings were explored.

    Source: §6 · p. 10

Abstract

from arXiv · show

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

1 Introduction

Theseus addresses fragmented, application-specific DNLS implementations with an application-agnostic PyTorch library for end-to-end structured learning. Its shared optimization components and efficiency features support varied robotics and vision applications, with reported speed and memory gains.

1 Introduction

  • DNLS uses an objective partly parameterized by neural models and differentiable domain-specific models, mapping input tensors to minima of a weighted-squares objective.Differentiating through the optimizer retains end-to-end gradients for training on the final task loss.Source: §1 · pp. 1–2
  • Existing DNLS implementations are application-specific, often reimplement shared tools, and may omit efficiency features such as sparse solvers, batching, and GPU support.DNLS lets neural models train on a final task loss while using priors captured by an optimizer.Source: §1 · pp. 1–2

Application agnostic interface

  • The application-agnostic interface combines second-order optimizers, learnable or hand-specified costs, differentiable Lie groups, and robot kinematics.Users can build custom optimization layers and plug them into neural architectures.Source: §1 · p. 2

Efficiency based design

  • Theseus combines differentiable sparse solvers, batching, GPU acceleration, automatic vectorization, implicit differentiation, and direct loss minimization to improve efficiency.Implicit differentiation and direct loss minimization extend to second-order optimizers; standard unrolling can challenge compute, memory, and gradient quality.Source: §1 · p. 2

Highlights of results

  • Up to 20x faster forward pass than Ceres when solving a batch of large problems, alongside significantly lower memory use than a dense solver on a standard GPU.Theseus’s sparse solver is much faster and uses significantly less memory than a dense solver on a standard GPU.Source: §1 · p. 2

2 Background and related work

Nonlinear least squares provides a common optimization formulation for robotics and vision, while learning can parameterize costs, weights, or initializations. Prior differentiable implementations remain application-specific, with important shared features not consistently available.

Background

  • NLS minimizes a sum of squared vector-valued residuals over optimization variables, with each residual depending on a subset of those variables.Variables may be manifold objects, and residuals are represented as a matrix weight times a vector cost.Source: §2 · p. 2

Solving NLS

  • Gauss–Newton iteratively linearizes the objective, solves for a variable update, and repeats; Levenberg–Marquardt damps the system, while Dogleg adjusts step size with a trust region and line search.Sparse Hessians arise when costs do not depend on all variables, enabling efficient sparse linear solves.Source: §2 · p. 3

NLS with learning

  • Learning can parameterize costs, weights, or initializations, and DNLS differentiates the optimized solution with respect to those parameters for a downstream loss.Earlier approaches that learn features may train only on an intermediate loss before optimization at inference.Source: §2 · p. 3

Background and related work

  • Prior DNLS work is application-specific, with features such as learnable costs, Lie groups, kinematics, sparsity, vectorization, and implicit differentiation not consistently present.The related-work examples span learned costs and weights, sensor models, GPU bundle adjustment, and differentiable optimization applications.Source: §2 · p. 3

3 Application agnostic interface

Theseus provides an application-agnostic interface for building differentiable nonlinear least-squares layers from optimizers, objectives, and variables. Its examples show how the same interface supports varied robotics and vision tasks, including optimization with Lie groups.

3 Application agnostic interface

  • The core API combines an Objective of cost functions and weights with an Optimizer inside TheseusLayer, which connects the optimization block to other PyTorch modules.Users provide named tensors and receive optimized-variable tensors; the output can feed downstream modules while preserving the differentiable computation graph.Source: §3 · pp. 4–5
  • Theseus supplies Gauss–Newton, Levenberg–Marquardt, and Dogleg optimizers, plus standard costs, robust losses, and custom costs with automatic Jacobians.The library also provides analytical Lie-group operators and derivatives, tangent-space gradient projection, and a heuristic extension for first-order optimizers on manifolds.Source: §3 · p. 5
  • Differentiable forward kinematics and an analytical manipulator Jacobian can be used within any Theseus CostFunction.The implementation wraps Differentiable Robot Model and builds a differentiable kinematics function from a standard robot model file.Source: §3 · p. 5

3.1 Example applications

  • Examples use the shared interface for pose graph optimization, tactile state estimation, bundle adjustment, motion planning, and homography estimation.Application-specific work is mainly defining data handling, neural models, or custom costs before assembling a TheseusLayer and outer loop.Source: §3.1 · p. 5

4 Efficiency based design

Theseus targets efficiency through batching and vectorization, sparse linear solvers, and multiple ways to differentiate through DNLS. The design combines parallel computation and sparsity-aware solving with backward methods suited to different optimization settings.

4.1 Batching and vectorization

  • Theseus supports parallel batches of DNLS problems and automatically vectorizes same-type operations, reducing computation overhead in forward and backward passes.For the PGO example, automatic vectorization increases memory use by up to ~82% for forward and ~55% for backward.Source: §4.1 · p. 6, Figure 2 · p. 6

4.2 Handling sparsity with linear solvers beyond PyTorch

  • Differentiable sparse solvers exploit the sparsity of NLS linear systems, with Theseus assembling objectives into solver-ready sparse structures.The available solvers are CPU-based CHOLMOD, GPU cudaLU, and BaSpaCho, a batched sparse Cholesky solver with GPU support.Source: §4.2 · p. 6
  • BaSpaCho combines supernodal Cholesky with sparse elimination, avoiding an external Schur-complement workaround for very sparse problems such as bundle adjustment.Its elimination tree groups column blocks with similar sparsity patterns so dense operations can use BLAS or cuBLAS.Source: §4.2 · p. 7

4.3 Backward modes for DNLS

  • Implicit differentiation uses the first-order optimality condition and implicit function theorem to compute derivatives at an exact solution; Theseus derives them by differentiating a Newton step.Unrolling is a standard alternative, but often uses only a few steps because of compute, memory, and vanishing-gradient challenges; truncation yields biased derivatives.Source: §4.3 · p. 7
  • Direct loss minimization augments the inner-loop problem to form a finite-difference scheme for the outer-loss gradient.Theseus’s version uses only the outer-loss gradient vector and adds regularization so the modified objective remains a sum of squares without changing the limit.Source: §4.3 · p. 8

5 Evaluation

The evaluation profiles Theseus on pose graph optimization (PGO) and compares backward modes on tactile state estimation. It reports sparse-solver scalability and mode-specific trade-offs in runtime, memory, and validation loss.

5.1 Profiling forward and backward pass of Theseus for DNLS

  • The PGO evaluation measures full forward-and-backward time across problem scales and batch sizes, using implicit differentiation and sparse or dense solvers.It uses 10 inner-loop iterations, 20 outer-loop epochs, and an Nvidia V100 GPU for Python computation.Source: §5.1 · p. 8
  • The reported total times were 20.81s for dense, 10.96s for CHOLMOD, 2.86s for cudaLU, and 2.25s for BaSpaCho, making BaSpaCho fastest among the four.At batch size 128, dense solves at most 256 poses, while BaSpaCho scales to 2048; with 2048 poses, dense cannot solve at any batch size.Source: §5.1 · pp. 8–9

5.2 Profiling Theseus as stand-alone NLS optimizer

  • The stand-alone optimizer evaluation compares Theseus with Ceres on 256 PGO problems over 10 iterations, while Figure 4 charts forward-pass speedup across problem scales.The supplied passages identify the comparison and figure scope but do not report Figure 4's speedup values.Source: §5.2 · p. 9, Figure 4 · p. 9

5.3 Backward modes analysis

  • Implicit and DLM derivative times remain constant as inner-loop iterations increase, whereas Unroll grows linearly and is slowest; Trunc-5 and Trunc-10 fall between them.The backward-time ordering is Implicit = DLM < Trunc-5 < Trunc-10 << Unroll.Source: §5.3 · p. 9
  • Unroll memory rises from ∼34MBs to ∼262MBs with more inner-loop iterations, while Implicit and DLM remain constant at ∼28MBs and ∼29MBs.The reported validation-loss ranking is application-dependent and affected by inner and outer optimizer learning rates; Implicit achieves the best validation loss after 100 epochs in this example.Source: §5.3 · p. 10

6 Discussion

Theseus combines differentiable nonlinear least squares with an application-agnostic interface, and reports up to 20x faster forward passes than Ceres on batches of large problems. The discussion presents shared components across robotics and vision examples while noting current implementation limits.

6 Discussion

  • For batches of large problems, Theseus’s forward pass is up to 20x faster than Ceres.The authors identify sparsity and implicit differentiation, alongside autodiff and GPU acceleration, as key ingredients.Source: §6 · p. 10
  • The application-agnostic interface reuses differentiable components across pose-graph optimization, tactile estimation, bundle adjustment, motion planning, and homography estimation.Theseus supplies example applications and users primarily define application-specific data, models, or cost functions.Source: §6 · p. 10
  • Current limits include soft-only constraints, non-learnable LM damping, non-incremental settings, and no distributed training beyond PyTorch’s support.The authors also note unoptimized overhead for frequently edited objectives and potential gains from moving Python code to C++.Source: §6 · p. 10

D Differentiable Lie group details

  • For Lie groups, Theseus computes tangent-space gradients by applying a linear projection to Euclidean gradients, with the projection available in closed form.This projection supports automatic differentiation on Lie-group tangent spaces.Source: D Differentiable Lie group details · pp. 12–13

E.1 Pose graph optimization

  • Theseus provides coordinate-independent relative-pose errors for PGO, and matches Ceres’s objective values on all evaluated benchmark datasets with chordal initialization.The paper presents initial and final objective values in Table 1 and qualitative results in Figure 7.Source: §E.1 · p. 13, Table 1 · p. 13

E.2 Tactile state estimation

  • The tactile example differentiates through offline pose optimization to fine-tune a relative-pose model, using 56 training and 7 test trajectories.It uses 63 trajectories of length 25 and does not consider an incremental setting.Source: §E.2 · p. 14

E.3 Bundle adjustment; E.4 Motion planning; E.5 Homography estimation; F BaSpaCho: Batched Sparse Cholesky

  • Examples also optimize bundle-adjustment soft-kernel radius, initialize motion-planning trajectories for two-step optimization, and train CNN features through homography optimization.The BaSpaCho example implements supernodal Cholesky, exploiting dense BLAS/cuBLAS operations and architecture-aware clustering trade-offs.Source: §E.3 · p. 14, §E.4 · p. 14, §E.5 · p. 15

G Benchmark details and additional results

Additional profiling compares sparse and dense solvers on PGO problems with smaller fixed batches and pose counts. Sparse solvers become faster relative to dense as problem scale increases, with BaSpaCho fastest among the reported solvers.

G.1 Forward and backward times with smaller batch size and number of poses

  • At 512 poses and batch size 16, BaSpaCho takes 2.34s per batch versus 21.37s for dense, a roughly 9.1x speedup.At that scale, cudaLU and CHOLMOD each achieve roughly 3.7x speedup over dense.Source: §G.1 · p. 16
  • At 64 poses and batch size 16, total forward/backward time is 0.54s for BaSpaCho, 0.94s for CHOLMOD, 1.23s for dense, and 1.32s for cudaLU.At 128 poses, BaSpaCho takes 0.80s versus 1.98s for dense.Source: §G.1 · p. 16

G.2 Forward and backward times with respect to sparsity

The sparsity profiling varies PGO loop-closure probability and compares forward and backward times across pose counts and batch sizes. Lower sparsity generally increases runtime, while dense forward time changes little with sparsity.

G.2 Forward and backward times with respect to sparsity

  • Loop-closure probabilities from 0.05 to 0.40 represent Cube datasets from high to low sparsity, respectively.Higher loop-closure probability means more loop-closure edges and a less sparse Hessian.Source: §G.2 · p. 16
  • Figures 11 and 12 show Theseus’s average forward and backward times on PGO problems across different sparsity levels, with fixed problem sizes of 2048 and 256 poses, respectively.Source: §G.2 · p. 16
  • Lower sparsity generally takes more time, while dense forward-pass time stays nearly constant as sparsity changes.More loop-closure edges add cost-function terms and can increase the time needed to assemble the approximate Hessian.Source: §G.2 · p. 17

G.3 Scalability of Theseus

Theseus scales to larger pose-graph optimization (PGO) problems with sparse solvers than with a dense solver. The solvers differ in maximum solvable problem size and memory demands.

G.3 Scalability of Theseus

  • At 8192 poses, CHOLMOD solves batch size 256, while cudaLU and BaSpaCho solve batch size 32; dense fails at 4096 poses or more, even at batch size 1.The maximum solvable batch sizes and pose counts are shown in Figure 13.Source: §G.3 · p. 17, Figure 13 · p. 17
  • Dense solves only small PGO problems because of memory limits, and cannot solve any with 4096 poses or more, even at batch size 1.The profile covers 64 to 8192 poses and batch sizes 8 to 1024.Source: §G.3 · p. 17
  • CHOLMOD runs computation on CPU, while cudaLU and BaSpaCho require extra GPU memory to solve linear systems.This helps explain why CHOLMOD solves larger DNLS problems than cudaLU and BaSpaCho in these tests.Source: §G.3 · p. 17

G.4 Comparison with Ceres on smaller batch size and number of poses

Theseus is compared with Ceres on PGO problems using a fixed batch size of 16 and 256 poses. Ceres is faster at smaller scales, while Theseus is faster as problem scale increases.

G.4 Comparison with Ceres on smaller batch size and number of poses

  • As PGO problem scale increases, Theseus becomes faster than Ceres; Ceres is faster at small batch sizes and pose counts.Figure 14 shows the speedup of Theseus over Ceres for different PGO problem scales.Source: §G.4 · p. 17, Figure 14 · p. 17
  • The comparison uses PGO problems with a fixed batch size of 16 and 256 poses.Theseus is evaluated as a stand-alone NLS optimizer.Source: §G.4 · p. 17
  • Ceres is faster than Theseus for PGO problems at small batch sizes and numbers of poses.The reported pattern is similar to Figure 4.Source: §G.4 · p. 17

H.1 Experimental details

The tactile state estimation experiments use trajectory data and specify separate training and test setups. Table 2 summarizes backward modes and their limitations.

H.1 Experimental details

  • The tactile state estimation dataset uses 56 trajectories of length 25 for training and 7 for testing.These counts follow from the dataset description across the stated training and test split.Source: §H.1 · pp. 17–18
  • Table 2 provides a summary of backward modes and their limitations.The table caption identifies its scope.Source: Table 2 · p. 18
  • Training uses batch size 8 for 100 epochs, yielding 700 batches for averaging time and memory results.The test inner loop runs for 50 iterations regardless of the number used during training.Source: §H.1 · p. 18

H.2 Additional results

Additional tactile state estimation results compare memory use and training behavior across backward modes. Implicit has the lowest reported peak memory, but its training is unstable; Unroll degrades with more inner-loop iterations.

H.2 Additional results

  • Implicit uses the least peak memory at approximately 22 MB, followed by DLM at approximately 29 MB.Figure 15 (left) shows forward-pass peak memory consumption.Source: §H.2 · p. 18, Figure 15 · p. 18
  • Implicit training oscillates between low and high values despite higher test performance, suggesting early stopping and hyperparameter tuning may be required.The training curves are shown for 10 and 50 inner-loop iterations.Source: §H.2 · p. 18, Figure 15 · p. 18
  • Unroll's performance degradation relative to other methods as inner-loop iterations increase also occurs during training, suggesting possible numerical issues from unrolling many optimization steps.Figure 15 (right) reports training loss with 50 inner-loop iterations.Source: §H.2 · p. 18, Figure 15 · p. 18

H.3 Backward modes summary and limitations

The appendix compares backward-mode tradeoffs, states the local conditions behind implicit differentiation, and describes a sum-of-squares-compatible direct-loss-minimization procedure.

H.3 Backward modes summary and limitations

  • Implicit differentiation cannot learn initial optimization values and may yield inaccurate gradients when the inner problem is not solved optimally.Unroll and Trunc can have vanishing or exploding gradients over many iterations; DLM's epsilon and Trunc's backward-iteration count may need tuning.Source: §H.3 · p. 18

H.4.1 Implicit function theorem

  • A continuously differentiable root function with a nonsingular Jacobian in θ gives a locally single-valued, continuously differentiable solution mapping.The theorem is stated around a point where the root function is zero.Source: §H.4.1 · p. 19

H.4.2 Proof of Prop. 1

  • Differentiating a Newton step with automatic differentiation yields the implicit derivative used in the proof.Source: §H.4.2 · p. 19

H.4.3 Direct loss minimization for use in Theseus

  • DLM linearizes the loss around θ*, then uses a completing-the-square reformulation to obtain an approximate vector-Jacobian product within a sum-of-squares objective.The added squared term introduces a small gradient bias when ε is small; in practice, θ_direct is obtained with one Gauss–Newton iteration from θ*.Source: §H.4.3 · p. 20

Related papers

Loading 2207.09442v3…