Tensorflow、行列の固有値でも自動微分できる模様

# あらすじ

Tensorflowで、勾配法とかで自動微分するときに途中で行列の対角化（固有値）を求めるところがあっても問題なく自動微分と勾配法による最小化ができるようです。

`GradientTape` を使って手動で `apply_gradient` する方法、自動的に微分して損失関数を `minimize` する方法も付けました。

# Version

* Tensorflow 2.1.0
* Python 3.7.3 (Anaconda)
* Windows 10

# 要点

適当な3行3列の行列 $A$ を用意して、その固有値 $\lambda$ を計算したとき、最小の固有値 $\lambda_0$ に対する行列 $A$ の各成分での微分 $d\lambda_0/dA$ （成分ごとなので3行3列の行列です）を計算したいとき。 

以下のコードで $d\lambda_0/dA$ を計算できます。Eagerモードですので、微分計算のための計算グラフを把握させるために `GradientTape()` を使います。

```python
import tensorflow as tf

A = tf.random.uniform(shape=(3, 3))

with tf.GradientTape() as g:
    g.watch(A)
    val, vec = tf.linalg.eigh(A)
    val0 = val[0]
grad0 = g.gradient(val0, A)

print(A)
print(val)
print(grad0)
```

```:A
tf.Tensor(
[[0.6102723  0.17637432 0.38962376]
 [0.3735156  0.6306771  0.19141042]
 [0.34370267 0.7677151  0.4024818 ]], shape=(3, 3), dtype=float32)
```

```:val
tf.Tensor([-0.25994763  0.34044334  1.5629349 ], shape=(3,), dtype=float32)
```

```:grad0
tf.Tensor(
[[ 5.3867564e-04  0.0000000e+00  0.0000000e+00]
 [ 3.0038984e-02  4.1877732e-01  0.0000000e+00]
 [-3.5372321e-02 -9.8626012e-01  5.8068389e-01]], shape=(3, 3), dtype=float32)
```

最後、下半分しか非ゼロになってないので、 `eigh` は下半分のみを使っているのではと推測されます。

実は要点といいつつ `val0 = val[0]` という、 slice 演算をしています。つまり固有値を求める演算以外もしています。固有値を１つ切り出さず `g.gradient(val, A)` というのも問題なく動作するのですが・・・数式での説明のしやすさを優先しました。

あと、`val0 = val[0]` ですが、 `GradientTape()`環境内で行っています。このように微分をつなげてほしい演算はすべてこの環境内でしないと、 `g.gradient()` が `None` を返すようになります。

# 応用

自動微分ゆえ、前後に演算を挟んでいくことも可能です。以降はそれを試してみます。

## 状況

ある6つの値を用意し、これを対称行列にならべ、固有値を計算します。これをある条件に当てはまるか、その誤差を計算して誤差関数とし、これを微分します。

* 固有値前の演算：行列の構成
* 固有値後の演算：固有値と目的の値との誤差の合計

ついでに、Tensorflow という機械学習のライブラリらしく、勾配法で元の値を更新してみます。

$r$ は３ｘ３対称行列の成分6個を保持、 $t = (t_0, t_1, t_2)$ は目指す固有値たちとします。

```math
r = \left(r_0, r_1, \cdots, r_5\right)
```

これを並べ替えて 

```math
A_h=\left(\begin{array}{ccc}
\frac{1}{2} r_0 & 0 & 0 \\
r_3 & \frac{1}{2} r_1 & 0 \\
r_5 & r_4 & \frac{1}{2} r_2
\end{array}
\right)
```

```math
A = A_h + A_h^T
```

として行列 $A$ を構成、これの固有値を求めます。　$\lambda = (\lambda_0, \lambda_1, \lambda_2)$ を $A$ の3つの固有値を並べたものとすると、

```math
L = \sum_{i=0}^2\left(t_i - \lambda_i\right)^2
```

を損失関数とします。$r$ を学習させるためには、

```math
\frac{\partial L}{\partial r_n} = \sum_{i,j,k}\frac{\partial A_{ij}}{\partial r_n}\frac{\partial \lambda_k}{\partial A_{ij}}\frac{\partial L}{\partial \lambda_k}
```

が必要な計算です。真ん中の $\frac{\partial \lambda_k}{\partial A_{ij}}$ が、行列の成分による固有値の微分です。

解析的には $\det(A - \lambda I)=0$ の解 $\lambda$ を求めたら計算できますが、正直3次元を超えると、よほどの疎行列でないともう手に負えないかと。なので、数値計算のできるtensorflowに頼ります。

## 手動による微分呼び出し

まずめんどくさいけど処理を追える方法。

### 準備

```python
    r = tf.Variable(initial_value=tf.random.uniform(shape=[6]))
    t = tf.constant([1.0, 1.0, 1.0])
```

`r` は３ｘ３対称行列の成分6個を保持、 `t` は目指す固有値たちです。ここでは初期値として `r` に一様乱数を与えました。　また `t` はすべて１です。


### 自動微分のための計算

```python
    with tf.GradientTape() as g:
        g.watch(r)
        A = 0.5*tf.linalg.diag(r[0:3])
        A = A + tf.pad(tf.linalg.diag(r[3:5]), [[0, 1], [1, 0]], "CONSTANT")
        A = A + tf.sparse.to_dense(tf.sparse.SparseTensor(indices=[[0, 2]], values=[r[5]], dense_shape=(3, 3)))
        A = A + tf.transpose(A)
        eigval, eigvec = tf.linalg.eigh(A)
        d = tf.reduce_sum((eigval - t)**2)
```

ここからの処理は自動微分の対象のため、 `GradientTape` の中でやります。

`A` の特定の成分を `r` から代入する方法として、紹介のために `diag` を使う方法と`SparseTensor` を使う2通りを使いました。

対角成分3つとその１つとなりの2つの成分を並べる関数として、`tf.linalg.diag`を使いました。
ただ、対角成分から１つとなりのところに並べるために `pad` を使っていますが、これは `k` オプションを使って対角から `k` ずれたところに並べるオプションがドキュメントでは紹介されていて、それを使うと

```python
   A = A + tf.linalg.diag(r[3:5], k=1)
```

と書けますが、なんか `k` オプションが動作しませんでした。なのでわざわざ `pad` を使っています。

また、 $r_5$ の１つだけを行列の角にセットする方法として、 `Tensor` には numpyのようなインデックス指定のassignはないので、面倒ですが疎行列を経由して足し合わせることでセットしました。

`A` を構成したあとは普通に対角化の `eigh` と損失関数の計算です。ここでは損失関数 $L$ を `d` としています。ややこしくてすみません。

### 微分の計算

```python
    grad_dr = g.gradient(d, r)
```

改めて、 `d` は損失関数 $L$ のことです。
`grad_dr` に長さ６のテンソルが代入されます。これでこの記事の目的は達成されました。

### 値の更新

適当なオプティマイザを使います。

```python
    opt = tf.keras.optimizers.Adam()
    opt.apply_gradients([(grad_dr, r), ])
```

ここで微分が使えるからと、

```python 
    opt.minimize(d, var_list=[r])
```

とすると、

```
TypeError: 'tensorflow.python.framework.ops.EagerTensor' object is not callable
```

というエラーが出ます。これは第一引数の `d` が `Tensor` ではなく、損失関数を返す引数なしの関数であることを要求されているからです。こっちの方法は後述します。

### 出力例

たとえば初期の `r` が

```
<tf.Variable 'Variable:0' shape=(6,) dtype=float32, numpy=
array([0.12108588, 0.8856114 , 0.00449729, 0.22199583, 0.8411281 ,
       0.54751956], dtype=float32)>
```

のとき、`A`は

```
tf.Tensor(
[[0.12108588 0.22199583 0.54751956]
 [0.22199583 0.8856114  0.8411281 ]
 [0.54751956 0.8411281  0.00449729]], shape=(3, 3), dtype=float32)
```

となります。すると損失関数の微分 $\frac{dL}{dr}$ は

```
tf.Tensor([-1.757829   -0.22877683 -1.991005    0.88798404  3.3645139   2.1900787 ], shape=(6,), dtype=float32)
```

と計算されました。これを `opt.apply_gradients([(grad_dr, r), ])` という感じで適用すると `r` の値は

```
<tf.Variable 'Variable:0' shape=(6,) dtype=float32, numpy=
array([0.12208588, 0.8866114 , 0.00549729, 0.22099583, 0.8401281 ,
       0.5465196 ], dtype=float32)>
```

となっていて、最初の値 `[0.12108588, 0.8856114 , 0.00449729, 0.22199583, 0.8411281, 0.54751956]`から若干変わっていることがわかります。

### 収束へ

せっかくなんで繰り返してみます。もし最適化が成功すれば `t` で設定したように固有値がすべて１の対称行列が得られるでしょう。

```python
    while d > 1e-8:
        with tf.GradientTape() as g:
            g.watch(r)
            A = 0.5*tf.linalg.diag(r[0:3])
            A = A + tf.pad(tf.linalg.diag(r[3:5]), [[0, 1], [1, 0]], "CONSTANT")
            A = A + tf.sparse.to_dense(tf.sparse.SparseTensor(indices=[[0, 2]], values=[r[5]], dense_shape=(3, 3)))
            A = A + tf.transpose(A)
            eigval, eigvec = tf.linalg.eigh(A)
            d = tf.reduce_sum((eigval - t)**2)
        grad_dr = g.gradient(d, r)
        opt.apply_gradients([(grad_dr, r), ])
        print("---------------")
        print(r)
        print(eigval)
        print(d)
```

`d` が2乗誤差で、これが一定以下になるまで `while` ループで繰り返しています。
これを実行すると

```
---------------
<tf.Variable 'Variable:0' shape=(6,) dtype=float32, numpy=
array([0.10630785, 0.18287621, 0.14753745, 0.16277793, 0.7271476 ,
       0.08771187], dtype=float32)>
tf.Tensor([-0.56813365  0.07035071  0.9315046 ], shape=(3,), dtype=float32)
tf.Tensor(3.3279824, shape=(), dtype=float32)
---------------
<tf.Variable 'Variable:0' shape=(6,) dtype=float32, numpy=
array([0.10730778, 0.18387613, 0.14853737, 0.16177836, 0.72614765,
       0.0867127 ], dtype=float32)>
tf.Tensor([-0.5661403   0.07189684  0.9309651 ], shape=(3,), dtype=float32)
tf.Tensor(3.3189366, shape=(), dtype=float32)
---------------
<tf.Variable 'Variable:0' shape=(6,) dtype=float32, numpy=
array([0.10830763, 0.18487597, 0.1495372 , 0.1607792 , 0.72514784,
       0.08571426], dtype=float32)>
tf.Tensor([-0.564147    0.07343995  0.9304282 ], shape=(3,), dtype=float32)
tf.Tensor(3.3099096, shape=(), dtype=float32)

以下略
```

という感じで、 `r` と `A` が少しずつ変わっていることが見て取れます。収束まですると、

```
<tf.Variable 'Variable:0' shape=(6,) dtype=float32, numpy=
array([ 9.9999946e-01,  9.9999988e-01,  9.9999732e-01,  6.9962436e-05,
        4.2644251e-07, -1.1688111e-14], dtype=float32)>
tf.Tensor([0.9999294 0.9999973 1.0000702], shape=(3,), dtype=float32)
tf.Tensor(9.917631e-09, shape=(), dtype=float32)
```

となりました。`A` はおおよそ単位行列のようですね。何度かやっても同様だったので、固有値が 1 だけの行列をこの方法で求めると単位行列に行き着くようです。

一応ループ回数をカウントしてみますと、3066とか2341とか3035とか、数千ステップかかるようです。

## 自動で微分と更新処理をやってもらう方法

要は `minimize()` を使う方法です。

対称行列 `A` の構成については、上の方の「自動微分のための計算」節を参考にしてください。

```python
    r = tf.Variable(initial_value=tf.random.uniform(shape=[6]))
    t = tf.constant([1.0, 1.0, 1.0])

    def calc_eigval(r):
        A = 0.5*tf.linalg.diag(r[0:3])
        A = A + tf.pad(tf.linalg.diag(r[3:5]), [[0, 1], [1, 0]], "CONSTANT")
        A = A + tf.sparse.to_dense(tf.sparse.SparseTensor(indices=[[0, 2]], values=[r[5]], dense_shape=(3, 3)))
        A = A + tf.transpose(A)
        eigval, eigvec = tf.linalg.eigh(A)
        return eigval
    def calc_loss(r):
        eigval = calc_eigval(r)
        d = tf.reduce_sum((eigval - t)**2)
        return d

    opt = tf.keras.optimizers.Adam()

    loss = lambda: calc_loss(r)

    while d > 1e-8:
        opt.minimize(loss, var_list=[r])
        print("---------------")
        print(r)
        print(calc_eigval(r))
        print(calc_loss(r))
```
`calc_eigval` は固有値を返す関数、 `calc_loss`は損失関数を計算する関数です。
 
`loss` は関数で `loss()` とするとこの時点での `r` の値に基づいて損失関数を計算し、その `Tensor` を返す関数です。
`minimize` の第１引数はこういう引数なしの関数を要求しています。だから 
 `gradientTape` で計算していた `d` を `minimize` に渡してもエラーが出ていたのでした。

`calc_eigval` は `calc_loss` の中に書いてしまってもいいのですが、ループ中に固有値もどう変化しているか見たかったので別の関数を用意しました。
`def calc_loss(r)` と `loss = lambda: calc_loss(r)` が `minimize` 利用のためのポイントです。`calc_loss`をはじめから引数なしで定義すればそのまま`minimize`に渡せた気がしないでもない。
ともかく、これで自分でを管理しなくても良くなりました。

上のコードをそのまま `main` 関数とかでくくって実行すると、収束の様子が出力されます。

# その他

おまけで固有値の収束の様子を、ランダムな初期値で始めた4パターンほど。線はそれぞれ<font color=#1f77b4>固有値１</font>、<font color=#ff7f0e>固有値２</font>、<font color=#2ca02c>固有値３</font>です。
右端は収束の目標とした １ で、ちゃんと動いていることが分かります。

![graphs.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/258019/0f89d8bd-34d8-a667-5cb7-295ad765027a.png)


`GradientTape()` を入れ子にすることで高次微分も可能だとのこと。思ったより自動微分のカバー範囲は広いようです。

なんで `GradientTape` メインの記事かというと、はじめ `minimize` を使う方法が分からず、手動で `apply_gradient` するしかないと思っていたからです。記事を大半書き終えてから `minimize` での方法が分かりました。