Phase 2, Section 3: Automatic Differentiation

Phase 2: AD as an Effect — algebraic effect and handler infrastructure. Section 3 of 4: Automatic Differentiation.

1. What this section commits to

This section specifies the implementation of reverse-mode automatic differentiation (AD) as a first-class language effect in kina-ai. Specifically, it establishes:

• The theoretical mapping of reverse-mode AD to delimited continuations and CPS transformations.
• The Diff effect structure and the reverseMode handler mechanics.
• Syntactic rules for custom gradient registrations.
• Implicit differentiation rules using the Implicit Function Theorem (IFT) for fixed-point/solver loops.

2. AD as Delimited Continuation

Reverse-mode AD records a sequence of forward operations (the tape) and later walks them in reverse to accumulate adjoints. In a language with algebraic effects, the forward tape is represented implicitly by the suspended execution stack (delimited continuations), and the backward pass is the evaluation of the resumption callback.

2.1 The Diff Effect

To track and intercept differentiable variables, we define the Diff effect:

effect Diff {
  fn active(x: tensor<F32, [D]>) -> tensor<F32, [D]>
}

The active operator tags a tensor as tracking gradients. The reverseMode handler intercepts this operation, wrapping the value in a dual node.

3. The reverseMode Handler

The reverseMode handler converts a differentiable function f: fn() !{Diff} -> tensor<F32, []> into a pair of its primal result and a function to calculate gradients:

  Primal computation: evaluates f(), returning result y.
  Adjoint propagation: walks the captured fiber frames in reverse, multiplying Jacobian-vector products.

3.1 Inline Differentiation Example

Below is a complete .kina code example running a simple expression under the reverseMode handler:

fn square(x: tensor<F32, [D]>) -> tensor<F32, [D]> {
  x * x
}

fn run_differentiation(val: tensor<F32, [D]>) -> tensor<F32, [D]> {
  // Run the forward computation under the reverseMode handler
  let (out, grad_fn) = handle {
    let active_x = perform Diff.active(val);
    square(active_x)
  } with reverseMode;
  
  // Backpropagate starting with an adjoint of 1.0
  let initial_adjoint = tensor<F32, [D], @host>.broadcast([D]);
  grad_fn(initial_adjoint)
}

4. Custom Gradients

Certain operations (such as custom CUDA kernels, numerically unstable functions, or operations containing singular points) require custom derivatives. We provide a custom_grad block pattern:

fn custom_sigmoid(x: tensor<F32, [D]>) -> tensor<F32, [D]> {
  let y = 1.0 / (1.0 + exp(-x));
  
  // Custom gradient registration using handler boundary
  handle y with {
    return(result) => result,
    Diff.active(v), k => {
      let dy = k(v);
      // pullback: dx = dy * y * (1.0 - y)
      dy * result * (1.0 - result)
    }
  }
}

5. Implicit Differentiation (IFT)

For loops or iterative processes that solve a fixed-point equation $f(x^, \theta) = 0$, unrolling the solver iterations in reverse-mode AD consumes $O(N)$ memory and accumulates numerical noise. The Implicit Function Theorem (IFT) states: $$\frac{\partial x^}{\partial \theta} = -\left(\frac{\partial f}{\partial x^*}\right)^{-1} \frac{\partial f}{\partial \theta}$$

We introduce the Implicit effect to define and differentiate fixed-point equations:

effect Implicit {
  fn solve(f: fn(tensor<F32, [D]>) -> tensor<F32, [D]>, init: tensor<F32, [D]>) -> tensor<F32, [D]>
}

5.1 The ift Handler

When the ift handler intercepts Implicit.solve, it:

1. Forward Pass: Runs the solver (e.g. Newton-Raphson) in pure mode to convergence, obtaining $x^*$.
2. Backward Pass: Discards the solver tape. It solves the linear system $\left(\frac{\partial f}{\partial x^*}\right)^T dy = dw$ to propagate the adjoint without unrolling.

fn run_solver_diff(init: tensor<F32, [D]>) -> tensor<F32, [D]> {
  handle {
    perform Implicit.solve(fn(x: tensor<F32, [D]>) -> tensor<F32, [D]> {
      // equation: x * x - 2.0 = 0
      x * x - 2.0
    }, init)
  } with ift
}