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)$として

y = f_w(x)

と書けるわけですが、このとき

\frac{\partial y_i}{\partial w_j} 

の計算をすることによってネットワークの学習を進めていくわけです。
ですが、一般に（数学などでは）関数$y = f_w(x)$のヤコビアンと言われたら出力を入力で微分したもの、すなわち

\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出力関数のヤコビアンを求める

たとえば関数

y = x^2

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

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個入れてみましょう。

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です。

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$を近似し、そのネットワークのヤコビアンを求めるという流れです。
ネットワークのクラスを定義します。

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}$の誤差が小さくなるように学習します。

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()

学習はうまくいっています。このネットワークに対してヤコビアンを求めてみます。入力量としてはさきほど定義したx_inputを使いましょう。

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]])

やや見づらいですが、先ほどと同じ結果が（誤差が乗っていますが）だいたい得られています。グラフにすると分かりやすくなります。

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入力2出力関数のヤコビアンを求める

変数の数を増やしてみましょう。
ヤコビアンと聞けば必ず出てくる、極座標変換を考えます。2次元空間での極座標 $(r, \theta)$から直交座標$(x, y)$への変換は、

\begin{pmatrix} x \\ y \end{pmatrix} = \begin{pmatrix} r\cos\theta \\ r\sin\theta \end{pmatrix}

で与えられますから、そのヤコビアンは

\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で計算します。

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つからなるタプルであることに気をつけましょう。
入力を与えてヤコビアンを計算します。

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です。

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つの変数のタプルで与えます。

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])

先程と同じく学習させます。

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()

先程定義した$r$, $\theta$によってヤコビアンを求めてみます。

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]$で動かします。

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()

よく一致していることがわかります。

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

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を使いたくない場合はこれをカスタマイズするのも手でしょう。