PyTorchで関数を近似し、そのヤコビアンを計算する

# はじめに

Deep Learningにおいてはback propagationによってネットワークのパラメータを更新するため、ヤコビアンの計算をする必要があります。入力$x=(x_0, x_1, ..., x_n)$、 出力$y=(y_0, y_1, ..., y_m)$とするニューラルネットワークは、パラメータを$w=(w_0, w_1, ..., w_k)$として
```math
y = f_w(x)
```
と書けるわけですが、このとき
```math
\frac{\partial y_i}{\partial w_j} 
```
の計算をすることによってネットワークの学習を進めていくわけです。
ですが、一般に（数学などでは）関数$y = f_w(x)$のヤコビアンと言われたら出力を入力で微分したもの、すなわち
```math
\frac{\partial y_i}{\partial x_j}
```
を指しますよね。というわけで、この記事では$y$を$x$で微分したヤコビアンをPyTorchで計算してみたいと思います。PyTorchの`torch.autograd.functional.jacobian`を使用し、1入力1出力の場合と2入力2出力の場合、それぞれについて解析的な関数とニューラルネットによるその近似のヤコビアンを計算します。
### 注意点
* もちろん$x$, $y$, $w$は一般には多次元ベクトル$\bf{x}$, $\bf{y}$, $\bf{w}$ですが、以後あえてボールド体にしていなくても多変数ベクトルのことを示していることがあります。
* 「ヤコビアン」は上で示したヤコビ行列の**行列式**のことを指すのが正しいかもしれませんが、以下ではヤコビ行列そのもののことをヤコビアンと呼んでいます。
* 以下のコードはPyTorch1.12.0での実験結果であり、コードブロックごとにJupyter Notebook上で実行しています。


# 1入力1出力関数のヤコビアンを求める

たとえば関数
```math
y = x^2
```
という1変数入力1変数出力の関数を考えてみましょう。当然このときのヤコビアンは$y=2x$というただの導関数になります。
これをPython・PyTorchで計算してみましょう。まず必要なライブラリをインストールします。

```python
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.optim as optim
import torch.nn as nn
import torch.autograd.functional as F
```
関数を設計し、入力として10個入れてみましょう。
```python
def calc_square(x):
    return x**2

x_input = torch.tensor([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])
y_output = calc_square(x_input)
y_output
>> tensor([ 0., 1., 4., 9., 16., 25., 36., 49., 64., 81.])
```
ここからが本題で、関数`calc_square`のヤコビアンを求めます。`torch.autograd.functional.jacobian(関数名, 入力変数名)`とすればOKです。

```python
F.jacobian(calc_square, x_input)
>> tensor([[ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  2.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  4.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  6.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  8.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0., 10.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  0., 12.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  0.,  0., 14.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0., 16.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0., 18.]])
```
入力変数の長さ（≈バッチ数、今は10）の大きさを持つ対角行列で返されました。体格成分を見るときちんと$y=2x$に一致した値が得られており、`x_input`に対するヤコビアンが計算できていますね。

# 1入力1出力関数をニューラルネットで近似し、ヤコビアンを求める
同じことを今度はニューラルネットワークに対して行います。ニューラルネットワークによってまず関数$y=x^2$を近似し、そのネットワークのヤコビアンを求めるという流れです。
ネットワークのクラスを定義します。
```python
class Network(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(Network, self).__init__()
        self.net = nn.Sequential(nn.Linear(input_dim, 32),
                                 nn.Softplus(), 
                                 nn.Linear(32, 32),
                                 nn.Softplus(),
                                 nn.Linear(32, 32),
                                 nn.Softplus(),
                                 nn.Linear(32, output_dim)
                                )
        
    def forward(self, x):
        """
        x: input, (batch)
        output: (batch)
        """
        x = x.reshape(-1, 1)
        return self.net(x).squeeze()
```
### 注意点
* Deep Learning系だと、入力は`(batch_size × input_dim)`というような2次元以上のテンソルになることが多いです。例えばRGB画像が入力なら`(batch_size × Channel × Height × Width)`の4次元ですね。
* しかし今回は`(batch_size)`という1次元であり、`(batch_size × 1)`の2次元でないことに気を付けてください。`calc_square`の入力も同じ形でした。とはいえ見て分かるように、与えられた入力を`forward`メソッドの内部で`(batch_size × 1)`にreshapeしているのですが。
* 活性化関数にReLUを使うと微分が階段関数のようになってしまい、なめらかな導関数・ヤコビアンを求めることができません。Sigmoid, あるいはSoftplusなどを用いましょう。

このネットワークによって1次関数を近似します。メインループ内部で毎回$[-10, 10]$の区間から1000個のサンプルを取って入力とし、それに対する正解出力$y$とネットワーク出力$y_{est}$の誤差が小さくなるように学習します。
```python
model = Network(1, 1)
optimizer = optim.Adam(model.parameters())

model.train()
loss_array = []

dist = torch.distributions.uniform.Uniform(torch.tensor([-10.]), torch.tensor([10.]))

for epoch in range(3000): 
    array = dist.sample([1000])
    x = array[:, 0]
    y = x**2
    
    y_est = model.forward(x) 
    loss = torch.linalg.norm(y_est - y)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    loss_array.append(loss.item())

    print("\r epoch:{}, loss:{}".format(epoch, loss_array[-1]), end=" ")
    
model.eval()
y_est_test = model.forward(x)
plt.scatter(x, y, label="true")
plt.scatter(x, y_est_test.detach(), label="estimated")
plt.legend()
plt.show()
plt.plot(loss_array)
plt.show()
```
![1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/604714/9f1ac561-e807-641f-1e50-a85a48a2985b.png)

学習はうまくいっています。このネットワークに対してヤコビアンを求めてみます。入力量としてはさきほど定義した`x_input`を使いましょう。
```python
F.jacobian(model, x_input)
>> tensor([[2.7725e-03, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
         0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
        [0.0000e+00, 1.8527e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
         0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
        [0.0000e+00, 0.0000e+00, 4.0002e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
         0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
        [0.0000e+00, 0.0000e+00, 0.0000e+00, 6.0969e+00, 0.0000e+00, 0.0000e+00,
         0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
        [0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 7.9527e+00, 0.0000e+00,
         0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
        [0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 9.9255e+00,
         0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
        [0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
         1.2001e+01, 0.0000e+00, 0.0000e+00, 0.0000e+00],
        [0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
         0.0000e+00, 1.3937e+01, 0.0000e+00, 0.0000e+00],
        [0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
         0.0000e+00, 0.0000e+00, 1.5920e+01, 0.0000e+00],
        [0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
         0.0000e+00, 0.0000e+00, 0.0000e+00, 1.8052e+01]])
```
やや見づらいですが、先ほどと同じ結果が（誤差が乗っていますが）だいたい得られています。グラフにすると分かりやすくなります。
```python
j_true = 2*x_input
j_est = torch.diag(F.jacobian(model, x_input))
plt.plot(range(batch), j_true, label="true")
plt.plot(range(batch), j_est, label="est")
plt.legend()
plt.show()
```
![2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/604714/fef48a35-3c00-d430-c544-ceb11aa80e24.png)
ニューラルネットワークに対してもヤコビアンを計算することができました。

# 2入力2出力関数のヤコビアンを求める
変数の数を増やしてみましょう。
ヤコビアンと聞けば必ず出てくる、極座標変換を考えます。2次元空間での極座標 $(r, \theta)$から直交座標$(x, y)$への変換は、
```math
\begin{pmatrix} x \\ y \end{pmatrix} = \begin{pmatrix} r\cos\theta \\ r\sin\theta \end{pmatrix}
```
で与えられますから、そのヤコビアンは
```math
\begin{pmatrix} \frac{\partial x}{\partial r} && \frac{\partial x}{\partial \theta} \\ \frac{\partial y}{\partial r} && \frac{\partial y}{\partial \theta} \end{pmatrix} = \begin{pmatrix} \cos\theta && -r\sin\theta \\ \sin\theta && r\cos\theta \end{pmatrix} 
```
となります。これをPyTorchで計算します。
```python
def calc_xy_from_polar(r, theta):
    """
    r: (batch)
    theta: (batch)
    return: (batch) * 2
    """
    x = r * torch.cos(theta)
    y = r * torch.sin(theta)
    return (x, y)
```
入力は長さ`batch_size`の1次元のテンソルふたつ（それぞれ$r$, $\theta$）、出力は`(batch_size) * 2`という1次元テンソル2つからなるタプルであることに気をつけましょう。
入力を与えてヤコビアンを計算します。
```python 
r = torch.tensor([1., 2., 3., 4., 5.])
theta = torch.tensor([torch.pi, torch.pi/2, torch.pi/3, torch.pi/4, torch.pi/6])

calc_xy_from_polar(r, theta)

J = F.jacobian(calc_xy_from_polar, (r, theta))
J

>> ((tensor([[-1.0000e+00, -0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],
          [-0.0000e+00, -4.3711e-08,  0.0000e+00,  0.0000e+00,  0.0000e+00],
          [-0.0000e+00, -0.0000e+00,  5.0000e-01,  0.0000e+00,  0.0000e+00],
          [-0.0000e+00, -0.0000e+00,  0.0000e+00,  7.0711e-01,  0.0000e+00],
          [-0.0000e+00, -0.0000e+00,  0.0000e+00,  0.0000e+00,  8.6603e-01]]),
  tensor([[ 8.7423e-08, -0.0000e+00, -0.0000e+00, -0.0000e+00, -0.0000e+00],
          [ 0.0000e+00, -2.0000e+00, -0.0000e+00, -0.0000e+00, -0.0000e+00],
          [ 0.0000e+00, -0.0000e+00, -2.5981e+00, -0.0000e+00, -0.0000e+00],
          [ 0.0000e+00, -0.0000e+00, -0.0000e+00, -2.8284e+00, -0.0000e+00],
          [ 0.0000e+00, -0.0000e+00, -0.0000e+00, -0.0000e+00, -2.5000e+00]])),
 (tensor([[-8.7423e-08,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],
          [-0.0000e+00,  1.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],
          [-0.0000e+00,  0.0000e+00,  8.6603e-01,  0.0000e+00,  0.0000e+00],
          [-0.0000e+00,  0.0000e+00,  0.0000e+00,  7.0711e-01,  0.0000e+00],
          [-0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  5.0000e-01]]),
  tensor([[-1.0000e+00, -0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],
          [-0.0000e+00, -8.7423e-08,  0.0000e+00,  0.0000e+00,  0.0000e+00],
          [-0.0000e+00, -0.0000e+00,  1.5000e+00,  0.0000e+00,  0.0000e+00],
          [-0.0000e+00, -0.0000e+00,  0.0000e+00,  2.8284e+00,  0.0000e+00],
          [-0.0000e+00, -0.0000e+00,  0.0000e+00,  0.0000e+00,  4.3301e+00]])))

print(J[0].shape)
>> torch.Size([2, 5, 5])
```
見づらいですが、ヤコビアン$J$は`torch.Size([2, 5, 5])`という3次元テンソルを2つ連ねたタプル`(J[0], J[1])`で与えられています。それぞれのテンソルは第1行、第2行のヤコビアンを示しています（第1出力$x$についての微分と、第2出力$y$の微分を返してくれています）。そして$J[0], J[1]$それぞれの最初の2次元がヤコビアンの列数を示します。

すなわち、2×2のヤコビアンのうち
* 第1行第1列（$\cos\theta$）は$J[0][0]$であり、
* 第1行第2列（$-r\sin\theta$）が$J[0][1]$、
* 第2行第1列（$\sin\theta$）が$J[1][0]$、
* 第2行第2列（$r\cos\theta$）が$J[1][1]$
になります。そしてそれぞれのヤコビアン要素は5×5の対角行列で示されているというわけです。

例として第2行第1列を確認しておきます。対角成分が入力$(\pi, \pi/2, \pi/3, \pi/4, \pi/6)$のサインになっていればOKです。
```python
torch.diag(J[1][0])
>> tensor([-8.7423e-08,  1.0000e+00,  8.6603e-01,  7.0711e-01,  5.0000e-01])
```
誤差はありますが、おおむね一致しているのを確認できます。

# 2入力2出力関数をニューラルネットで近似し、ヤコビアンを求める
まずモデルのクラスを定義します。入力が2変数、出力は（上で確認したような形式に合わせて）2つの変数のタプルで与えます。
```python 
class Network_2(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(Network_2, self).__init__()
        self.net = nn.Sequential(nn.Linear(input_dim, 32),
                                 nn.Softplus(), 
                                 nn.Linear(32, 32),
                                 nn.Softplus(),
                                 nn.Linear(32, 32),
                                 nn.Softplus(),
                                 nn.Linear(32, output_dim)
                                )
        
    def forward(self, r, theta):
        """
        r, theta: (batch)
        output: (x, y). each size is (batch).
        """
        _input = torch.cat([r.unsqueeze(1), theta.unsqueeze(1)], axis=1) #(batch, 2)
        _output = self.net(_input) #(batch) -> (batch)
        return (_output[:, 0], _output[:, 1])
```
先程と同じく学習させます。
```python
model = Network_2(2, 2)
optimizer = optim.Adam(model.parameters())

loss_array = []

dist = torch.distributions.uniform.Uniform(torch.tensor([-10., -10.]), torch.tensor([10., 10.]))

for epoch in range(30000):
    model.train()
    array = dist.sample([1000])
    r = array[:, 0]
    theta = array[:, 1]
    
    x, y = calc_xy_from_polar_2(r, theta)
    x_est, y_est = model.forward(r, theta) #(batch) * 2 -> (batch) * 2
    loss = torch.linalg.norm(x_est - x) + torch.linalg.norm(y_est - y)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    loss_array.append(loss.item())
    
    print("\r epoch:{}, loss:{}".format(epoch, loss_array[-1]), end=" ")

model.eval()
plt.plot(loss_array)
plt.show()
```
![3.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/604714/e9f69f2f-c76b-0107-8713-8934825a8103.png)

先程定義した$r$, $\theta$によってヤコビアンを求めてみます。
```python
J_nn = F.jacobian(model, (r, theta))
J_nn
>>((tensor([[-1.0118,  0.0000,  0.0000,  0.0000,  0.0000],
          [ 0.0000,  0.0238,  0.0000,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  0.5466,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  0.0000,  0.6485,  0.0000],
          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.7945]]),
  tensor([[-0.0120,  0.0000,  0.0000,  0.0000,  0.0000],
          [ 0.0000, -2.0013,  0.0000,  0.0000,  0.0000],
          [ 0.0000,  0.0000, -2.7921,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  0.0000, -2.5708,  0.0000],
          [ 0.0000,  0.0000,  0.0000,  0.0000, -2.2662]])),
 (tensor([[0.0162, 0.0000, 0.0000, 0.0000, 0.0000],
          [0.0000, 1.0233, 0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.8791, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.7105, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.4974]]),
  tensor([[-0.4103,  0.0000,  0.0000,  0.0000,  0.0000],
          [ 0.0000,  0.1548,  0.0000,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  1.3270,  0.0000,  0.0000],
          [ 0.0000,  0.0000,  0.0000,  2.7108,  0.0000],
          [ 0.0000,  0.0000,  0.0000,  0.0000,  4.3991]])))
```
ニューラルネットを用いなかった時と比べて、`J_nn`の形状に差異があります。`J_nn`は2つの要素からなるタプル`(J[0], J[1])`ではありますが、`J[0]`はまた2つの要素からなるタプル`(J[0][0], J[0][1])`です。要素のどれがどのヤコビアン要素に対応するかは先程と同じです。

最後に、ニューラルネットに対して計算したヤコビアンがどの程度正確なのかを確認します。上で求めた`calc_xy_from_polar`関数の[0][0]成分と、このネットワークの[0][0]成分を比較します。いずれも理論的には$\cos\theta$であるはずです。

簡単のため$r=1$に固定し、$\theta$のみを$[-2\pi, 2\pi]$で動かします。
```python
theta_test = torch.arange(-2*torch.pi, 2*torch.pi, 0.1)
r_test = torch.ones_like(theta_test)

J_00 = torch.diag(F.jacobian(calc_xy_from_polar, (r_test, theta_test))[0][0])
J_nn_00 = torch.diag(F.jacobian(model, (r_test, theta_test))[0][0])
plt.plot(theta_test, J_00, label="true")
plt.plot(theta_test, J_nn_00, label="est")
plt.legend()
plt.show()
```
![4.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/604714/8072243b-a3bf-ab0a-ca0e-84e3b10f2194.png)
よく一致していることがわかります。

以上、`torch.autograd.functional.jacobian`の使い方を紹介しました。この機能の内部でどのような計算が行われているかについては確認していませんが、`backward`を用いて似たような機能を実装することは可能です。例えば[このリンク](https://discuss.pytorch.org/t/derivative-of-model-outputs-w-r-t-input-features/95814/2 "torch_derivative_of_model")が参考になります。
リンク先ではMNISTの分類問題を考え、`(1, 784)`次元の入力`x_input`から(10)次元の出力`preds`への非線形関数であるニューラルネットワークを作成したときに、そのネットワークのヤコビアンを求めようとしています。回答としては次のようなものが挙げられています。

```python
J = torch.zeros ((1, 784, 10))   # loop will fill in Jacobian
x_input.requires_grad = True
preds = model(x_input)
for i in range (10):
    grd = torch.zeros ((1, 10))   # same shape as preds
    grd[0, i] = 1    # column of Jacobian to compute
    preds.backward (gradient = grd, retain_graph = True)
    J[:,:,i] = x_input.grad   # fill in one column of Jacobian
    input.grad.zero_()   # .backward() accumulates gradients, so reset to zero
```
これは出力`preds`をループ内でひとつひとつ微分していくというやり方であり、回答者によれば「基本的には`torch.autograd.functional.jacobian`も同じことやってるんじゃないかと思う」そうです。`torch.autograd.functional.jacobian`を使いたくない場合はこれをカスタマイズするのも手でしょう。