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 は、第 $i$ 行が出力 $i$ の勾配である $m \times n$ 行列を返します。ここでは極座標からデカルト座標への写像を扱います：

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$ で、極座標でおなじみの面積要素です。

ヤコビ行列ベクトル積を計算する

1 つの方向 $v$ について $J_f(x)\,v$ を得るには、各入力をその方向成分でシードし、f を 1 回評価します：

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$ 回ではなく、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 のベクトルの要素ごとの積です。

ジェネリック関数

トレイトに対して書かれた関数は、微分することも通常の数で実行することもできます。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 の評価が 1 回かかります。入力が数百個あり、他に安価なリバースモードの代替手段がある場合、フォワードモードは遅い選択です。
キャプチャしたベクトルは定数です。 入力と組み合わせる前に Dual::constant で写してください。

次のステップ

linalg API にはドライバー、その形状とコストが記載されています。
linalg の設計では列ごとの手法を導出し、リバースモードと比較しています。
dual チュートリアルではドライバーが行うシードを示しています。
linear-algebra ではベクトル型と行列型を文書化しています。