Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

ad-trait: A Fast and Flexible Automatic Differentiation Library in Rust

Chen Liang, Qian Wang, Andy Xu, Daniel Rakita

TL;DR

ad-trait presents a Rust-based automatic differentiation library that overhauls floating-point types via a flexible AD trait, enabling both forward-mode and reverse-mode differentiation with SIMD acceleration. The work demonstrates competitive cross-language performance for derivative computations and validates practical use through a robotics optimization case in apollo-rust. By providing an extensible trait-based design and open-source implementation, ad-trait aims to accelerate differentiable robotics and numerical optimization in Rust while highlighting current limitations and avenues for future improvements.

Abstract

The Rust programming language is an attractive choice for robotics and related fields, offering highly efficient and memory-safe code. However, a key limitation preventing its broader adoption in these domains is the lack of high-quality, well-supported Automatic Differentiation (AD)-a fundamental technique that enables convenient derivative computation by systematically accumulating data during function evaluation. In this work, we introduce ad-trait, a new Rust-based AD library. Our implementation overloads Rust's standard floating-point type with a flexible trait that can efficiently accumulate necessary information for derivative computation. The library supports both forward-mode and reverse-mode automatic differentiation, making it the first operator-overloading AD implementation in Rust to offer both options. Additionally, ad-trait leverages Rust's performance-oriented features, such as Single Instruction, Multiple Data acceleration in forward-mode AD, to enhance efficiency. Through benchmarking experiments, we show that our library is among the fastest AD implementations across several programming languages for computing derivatives. Moreover, it is already integrated into a Rust-based robotics library, where we showcase its ability to facilitate fast optimization procedures. We conclude with a discussion of the limitations and broader implications of our work.

ad-trait: A Fast and Flexible Automatic Differentiation Library in Rust

TL;DR

ad-trait presents a Rust-based automatic differentiation library that overhauls floating-point types via a flexible AD trait, enabling both forward-mode and reverse-mode differentiation with SIMD acceleration. The work demonstrates competitive cross-language performance for derivative computations and validates practical use through a robotics optimization case in apollo-rust. By providing an extensible trait-based design and open-source implementation, ad-trait aims to accelerate differentiable robotics and numerical optimization in Rust while highlighting current limitations and avenues for future improvements.

Abstract

The Rust programming language is an attractive choice for robotics and related fields, offering highly efficient and memory-safe code. However, a key limitation preventing its broader adoption in these domains is the lack of high-quality, well-supported Automatic Differentiation (AD)-a fundamental technique that enables convenient derivative computation by systematically accumulating data during function evaluation. In this work, we introduce ad-trait, a new Rust-based AD library. Our implementation overloads Rust's standard floating-point type with a flexible trait that can efficiently accumulate necessary information for derivative computation. The library supports both forward-mode and reverse-mode automatic differentiation, making it the first operator-overloading AD implementation in Rust to offer both options. Additionally, ad-trait leverages Rust's performance-oriented features, such as Single Instruction, Multiple Data acceleration in forward-mode AD, to enhance efficiency. Through benchmarking experiments, we show that our library is among the fastest AD implementations across several programming languages for computing derivatives. Moreover, it is already integrated into a Rust-based robotics library, where we showcase its ability to facilitate fast optimization procedures. We conclude with a discussion of the limitations and broader implications of our work.

Paper Structure

This paper contains 22 sections, 2 figures, 2 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Automatic Differentiation Primer
  4. Related Works
  5. Implementation Overview
  6. AD Trait
  7. Types that implement AD Trait
  8. Differentiable Function and Derivative Method Traits
  9. Differentiable Block Struct
  10. Integration with Robotics Library
  11. Evaluation 1: Benchmark Function
  12. Procedure
  13. Conditions
  14. Sub-experiment 1: Gradient Calculations
  15. Sub-experiment 2: Square Jacobian Calculations
  16. ...and 7 more sections

Figures (2)

  • Figure 1: Evaluation 2 involves assessing performance in a robotic root-finding procedure, where the goal is to determine a robot pose that positions its feet and end-effector at predefined locations or orientations.
  • Figure 2: Results for Evaluation 1, Sub-experiment 1 (Top) and Sub-experiment 2 (Bottom). Both Sub-experiments are split into forward mode AD conditions (left) and reverse mode AD conditions (right).