[RFC] Add forward mode for autodiff

[erizmr]

In this issue, we would like to share a draft implementation plan for the forward mode autodiff.

Background

In general, there are two modes for autodiff: reverse mode and forward mode. The two modes have their advantage in different scenarios. The reverse mode is more efficient when the number of inputs is much more than the outputs (e.g., machine learning cases, thousands of trainable parameters and one scalar loss). On the contrary, the forward mode is more efficient. In addtion, the second-order derivatives can be efficiently computed by combining both the forward and reverse mode.

For a roadmap for the autodiff feature in Taichi, please check out #5050.

Goals

• Implement forward mode autodiff.
• Design python interface for forward and reverse mode.
• Make it possible to apply both forward/reverse mode iteratively (e.g., forward(reverse())), preparing for computing second-order derivatives.

Implementation Roadmap

• Implement forward mode autodiff.

  • Allocate dual snodes for fields [autodiff] Allocate dual and adjoint snode #5083
  • Implement forward mode routine [autodiff] Add forward mode pipeline for autodiff pass #5098
  • Distinguish three kinds of kernels: primal, forward ad and reverse ad [autodiff] Add forward mode pipeline for autodiff pass #5098
  • Implement forward ad transformation for operators currently supported by reverse mode [autodiff] Support basic operations for forward mode autodiff #5160
  • Handle mutation (e.g., for-loops) [autodiff] Support for-loop and mutation for forward mode #5225
  • Handle control flow [autodiff] Support control flow for forward mode #5231
  • Unify the shared components for reverse and forward mode [autodiff] Extract shared components for reverse and forward mode #5088
  • Handle the stop gradients for dual snodes
  • Handle the lazy gradient for dual snodes (allocate dual by forward mode context manager) [autodiff] Add a context manager for forward mode autodiff #5146 [autodiff] Refactor dual allocation for forward mode autodiff #5224
  • Support all Unary operators [autodiff] Support unary ops for forward mode #5366
  • Support all Binary operators [autodiff] Support binary operators for forward mode #5389
  • Support all Ternary operators [autodiff] Add ternary operators for forward mode #5405

• Design python interface for forward and reverse mode.
  - [x] Decouple the adjoint and grad, make the grad including both adjoint and dual [autodiff] Allocate dual and adjoint snode #5083

  • A context manager to trigger forward mode [autodiff] Add a context manager for forward mode autodiff #5146
  • Keep the grad indicate adjoint, expose dual for forward mode [autodiff] Refactor dual allocation for forward mode autodiff #5224

• Python test cases

  • Test cases for all the operators [autodiff] Support unary ops for forward mode #5366 [autodiff] Support binary operators for forward mode #5389 [autodiff] Add ternary operators for forward mode #5405
  • Test cases for context manager [autodiff] Refactor dual allocation for forward mode autodiff #5224
  • Test cases for for-loop [autodiff] [test] Add more complex for loop test cases for forward mode #5592
  • Test cases for if [autodiff] Fix alloca block and add control flow test case for forward mode #5301
  • Test cases for atomic [autodiff] [test] Add atomic test for forward autodiff #5286

• Second-order derivative

  • Proof of concept for forward on reverse second order derivative [autodiff] Proof of concept for second-order derivative #5117

Discussions

• How many kernels we need to compile for forward mode autodiff?

Currently in reverse mode, two kernels (original kernel and grad kernel) for evaluating function values and compute the gradients respectively. However, in forward mode autodiff, the derivatives are computed eagerly during the function evaluating process, i.e., the functions values and gradients can be computed using only kernel. This raise the question whether need to compile one or two kernels.`

Update: three kinds of kernels are generated: primal, forward ad and reverse ad according to different autodiff modes, see #5098.

[victoriacity]

I wonder if explicitly differentiating a function as in JAX will be supported, for example,

@ti.func
def f(x): return x**3 + 2*x**2 - 3*x + 1

dfdx = forward(f)

@ti.kernel
def k() -> float:
    return dfdx(1.0)
k() # returns 4.0

[erizmr]

I think it is possible to support similar features. A naive current Taichi equivalent is:

import taichi as ti

ti.init()

x = ti.field(float, shape=(), needs_grad=True)
y = ti.field(float, shape=(), needs_grad=True)

@ti.kernel
def f(): 
    y[None] += x[None]**3 + 2*x[None]**2 - 3*x[None] + 1
    
def dfdx(_x):
    x[None] = _x
    y.grad[None] = 1.0
    f.grad()
    return x.grad[None]

print(dfdx(1.0))

For more general case, it may require to specify the input and output if we would like to generate dfdx for the users. A possible implementation might be:

import taichi as ti

ti.init()

x1 = ti.field(float, shape=(), needs_grad=True)
x2 = ti.field(float, shape=(), needs_grad=True)
x3 = ti.field(float, shape=(), needs_grad=True)
y = ti.field(float, shape=(), needs_grad=True)

@ti.kernel
def f(): 
    y[None] += x1[None]**3 + 2*x2[None]**2 - 3*x3[None] + 1

def backward(f, input_field, out_field):
    import numpy as np
    out_field.grad[None] = 1.0
    def _dfdx(inputs):
        for i, x in enumerate(inputs):
            input_field[i].from_numpy(np.array(inputs[i]))
        f.grad()
        ret = []
        for x in input_field:
            ret.append(x.grad.to_numpy())
        return ret
    return _dfdx

dfdx = backward(f, [x1, x2, x3], y)

print(dfdx([1.0, 2.0, 3.0])) # [3, 8, -3]