23 October 2006

Better Jacobians in Computing Derivatives

And another fix to the code computing derivatives: before I was trying to compute the Jacobian of functions with only one call to the function:

let jacobian f args = 
    let n = Array.length args in 
    let tags = Array.init n (fun _ -> new_tag ()) in 
    let dargs = Array.mapi (fun i arg -> D(tags.(i), arg, one)) args in 
    let result = f dargs in 
    Array.init n 
      (fun i -> 
        let fi = result.(i) in 
        Array.init n 
          (fun j -> 
            let tj = tags.(j) in 
              (fun res tag -> 
                if tag = tj then 
                  extract_deriv tag res
                  drop_deriv tag res)
That's efficient in the sense that it doesn't re-compute results within f. Unfortunately, it introduces n tags, where n is the dimensionality of the argument to f. Each tag labels a fundamental ``incremental value'' dx, dy, .... A particular D(...) expression is a tree which represents an expression a0 + a1*dx + a2*dx*dy + ... which potentially contains one term for each of the 2^n different combinations of the dx, dy, .... This means that, though we save on repeated invocations of f, the cost for computing the jacobian becomes exponential in the dimension of the argument because each operation (+, sin, etc) has to manipulate these trees to a depth of 2^n.

The following simple change wastes effort computing f dargs multiple times, but scales only linearly in complexity with the dimension of args. This is a huge win. I've also added a function insert_tag which inserts a new D(tag, x, one) term within the derivative-computing wrappers (but does so in way consistent with the heap property of the diff-tree). The modified functions are below:

  let rec insert_tag tag = function 
    | (C(x) as d) -> D(tag, d, one)
    | (D(tagd, x, dx) as d) when tag < tagd -> 
        D(tag, d, one)
    | D(tagd, x, dx) when tagd = tag -> 
        raise (Failure "Cannot insert tag which is already present.")
    | D(tagd, x, dx) -> 
        D(tagd, insert_tag tag x, dx)
  let d f = 
    fun x -> 
      let tag = new_tag () in 
      let res = f (insert_tag tag x) in
      extract_deriv tag res

  let jacobian f args = 
    let n = Array.length args in 
    let jac = Array.make_matrix n n zero in 
    for j = 0 to Pervasives.(-) n 1 do 
      let tag = new_tag () in 
      let dargs = Array.mapi 
          (fun i arg -> if Pervasives.(=) i j then insert_tag tag arg else arg)
          args in 
      let result = f dargs in 
      Array.iteri (fun i res -> 
        jac.(i).(j) <- extract_deriv tag res)

No comments: