linalg 教程

本教程计算作用于 linear-algebra 向量的函数的梯度、雅可比矩阵和雅可比-向量积。你将编写作用于 Vector[Dual[Double]] 的函数，把它们传给 linalg 驱动函数，并在一个小型优化循环中使用结果。驱动函数为何如此工作见 linalg 设计。

快速上手

moon add Luna-Flow/autodiff@0.2.0
moon add Luna-Flow/linear-algebra

import {
  "Luna-Flow/autodiff",
  "Luna-Flow/autodiff/linalg",
  "Luna-Flow/linear-algebra/immut" @la,
}

fn main {
  let x = @la.Vector::from_array([2.0, 3.0])
  let g = @linalg.gradient(v => v[0] * v[0] + v[0] * v[1], x)
  println(g)
}

|7, 2|

对于 $f(x_0, x_1) = x_0^2 + x_0 x_1$ ，梯度为 $(2x_0 + x_1,\ x_0) = (7, 2)$ 。

日常任务

具名函数的梯度

当函数不是就地编写的闭包时，请为它给出显式类型。常数通过 Dual::constant 参与计算：

fn rosenbrock(v : @la.Vector[@autodiff.Dual[Double]]) -> @autodiff.Dual[Double] {
  let one = @autodiff.Dual::constant(1.0)
  let hundred = @autodiff.Dual::constant(100.0)
  let a = one - v[0]
  let b = v[1] - v[0] * v[0]
  a * a + hundred * b * b
}

fn main {
  let (value, grad) = @linalg.value_and_gradient(
    rosenbrock,
    @la.Vector::from_array([-1.0, 2.0]),
  )
  println("f    = \{value}")
  println("grad = \{grad}")
}

f    = 104
grad = |396, 200|

向量函数的雅可比矩阵

jacobian 返回一个 $m \times n$ 矩阵，其第 $i$ 行是第 $i$ 个输出的梯度。这里以极坐标到笛卡儿坐标的映射为例：

fn polar(v : @la.Vector[@autodiff.Dual[Double]]) -> @la.Vector[@autodiff.Dual[Double]] {
  let r = v[0]
  let theta = v[1]
  @la.Vector::from_array([r * theta.cos(), r * theta.sin()])
}

fn main {
  let x = @la.Vector::from_array([2.0, 0.0])
  let j = @linalg.jacobian(polar, x)
  println(j)
  println("det = \{j[0][0] * j[1][1] - j[0][1] * j[1][0]}")
}

|1, 0|
|0, 2|
det = 2

这个雅可比矩阵的行列式为 $r = 2$ ，即熟悉的极坐标面积因子。

计算雅可比-向量积

要对单个方向 $v$ 得到 $J_f(x)\,v$ ，请用方向分量为每个输入设定种子，然后对 f 求值一次：

fn polar(v : @la.Vector[@autodiff.Dual[Double]]) -> @la.Vector[@autodiff.Dual[Double]] {
  @la.Vector::from_array([v[0] * v[1].cos(), v[0] * v[1].sin()])
}

fn main {
  let x = @la.Vector::from_array([2.0, 0.0])
  let dir = @la.Vector::from_array([1.0, 0.5])
  let seeded = @la.Vector::makei(x.length(), i => @autodiff.Dual::new(x[i], dir[i]))
  let jv = polar(seeded).map(y => y.tangent())
  println(jv)
}

|1, 1|

这只需一次求值，而 jacobian 需要 $n + 1$ 次。

梯度下降

value_and_gradient 给出下降一步所需的内容：

fn bowl(v : @la.Vector[@autodiff.Dual[Double]]) -> @autodiff.Dual[Double] {
  let two = @autodiff.Dual::constant(2.0)
  let a = v[0] - @autodiff.Dual::constant(1.0)
  let b = v[1] + @autodiff.Dual::constant(0.5)
  a * a + two * b * b
}

fn main {
  let mut x = @la.Vector::from_array([3.0, 3.0])
  for _ in 0..<50 {
    let (_, g) = @linalg.value_and_gradient(bowl, x)
    x = @la.Vector::makei(x.length(), i => x[i] - 0.2 * g[i])
  }
  let (value, _) = @linalg.value_and_gradient(bowl, x)
  println("minimum near (\{x[0]}, \{x[1]}), f = \{value}")
}

minimum near (1.0000000000161655, -0.5), f = 2.6132382259185474e-22

深入学习

对偶数向量上的向量运算

Vector[Dual[Double]] 支持只需要环结构的向量运算，因此你可以用向量运算和一次折叠写出 $f(x) = \tfrac12 x^{\mathsf T} x + a^{\mathsf T}x$ ：

fn energy(v : @la.Vector[@autodiff.Dual[Double]]) -> @autodiff.Dual[Double] {
  let a = @la.Vector::from_array([1.0, -2.0, 0.5]).map(@autodiff.Dual::constant)
  let half = @autodiff.Dual::constant(0.5)
  let sq = (v * v).iter().fold(init=@autodiff.Dual::constant(0.0), (s, t) => s + t)
  let lin = (a * v).iter().fold(init=@autodiff.Dual::constant(0.0), (s, t) => s + t)
  half * sq + lin
}

fn main {
  let g = @linalg.gradient(energy, @la.Vector::from_array([1.0, 1.0, 1.0]))
  println(g)
}

|2, -1, 1.5|

梯度为 $x + a$ 。v * v 是 linear-algebra 向量的逐元素乘积。

泛型函数

针对 trait 编写的函数既可以求导，也可以在普通数上运行。给它传入一个 T 的向量：

fn[T : @autodiff.Ring] product(v : @la.Vector[T]) -> T {
  v.iter().fold(init=@autodiff.One::one(), (s, t) => s * t)
}

fn main {
  let x = @la.Vector::from_array([2.0, 3.0, 4.0])
  println("product = \{product(x)}")
  println("gradient = \{@linalg.gradient(product, x)}")
}

product = 24
gradient = |12, 8, 6|

常见陷阱

保持输出长度固定。 jacobian 从第一次调用中得知 $m$ ；之后若长度不同，会中止或丢失元素。
只读取给定范围内的下标。 驱动函数传入长度为 x.length() 的向量，且不检查访问。
输入很多时。 每个输入都要对 f 求值一次。如果有数百个输入，而别处又有廉价的反向模式可用，那么前向模式就是较慢的选择。
捕获的向量是常数。 在与输入组合之前，先用 Dual::constant 对它们做映射。

后续阅读

linalg API 列出了驱动函数及其形状和代价。
linalg 设计推导了逐列计算的方法，并将其与反向模式进行了比较。
dual 教程展示了驱动函数所做的种子设定。
linear-algebra 记录了向量和矩阵类型。