Reverse-mode autodiff in Clojure

In order to learn Clojure, or at least get a feel for it, I recently implemented a proof-of-concept reverse-mode AD framework in it. Below, I’ll introduce the API briefly.

Code can be found on GitHub.

The API

clad represents the computation graph as an adjacency matrix (via core.matrix) and a map of nodes. The forward/backward passes of a function are two topological traversals over it: we compute -bottom-up to compute values and -top-down to accumulate adjoints:

(defn grad [f idx]
  (let [graph (expr/expression-graph f)]
    (fn [& y]
      (let [graph (-top-down (-bottom-up (-set-values graph y)))]
        (:adjoint
         (nth
          (filter
           (fn [node] (:is-variable node))
           (vals (into (sorted-map) (:nodes graph))))
          idx))))))

Each call to the returned function rebuilds the graph’s value/adjoint state from scratch rather than mutating a shared structure in place. This is not really efficient, but good enough for the sake of learning functional programming.

The published API is a single grad function:

(require '[clad.core :refer [grad]])

(defn f [x y]
  (/ (- 1.0 (Math/exp (- x)))
     (+ 1.0 (Math/exp (- y)))))

(def g ((grad f 0) 2.0 1.0))
;; => 0.0989

(grad f 0) returns the derivative of f with respect to its argument at index 0 (here, x), evaluated at the point (2.0 1.0).

Conclusion

I learned a functional language when I studied CS (Standard ML), but always found them a bit “academic” and didn’t give much thought to them. When learning Clojure, my opinionated view changed a bit. Quoting from Peter Norvig’s blog: A language that doesn’t affect the way you think about programming, is not worth knowing. In that sense, while I will never use Clojure professionally, learning it definitely gave me a new view on functional programming.