本記事はTensorflowにおけるベクトルとヤコビアンの積の計算手法についてまとめたものです。

Tensorflowでのgradientsの基礎

Tensorflowで微分を行う関数gradientsは、スカラー値に対する微分を計算します。典型的な利用例は、cross_entropyなどで表現される損失関数$L$に対して、パラメーター$W$で微分するケースです。バッチサイズを$N$, クラス数を$D$, パラメーターの次元数を$M$とすると以下のようになります。

# logits.shape = [N, D]
logits = network(images)
# losses.shape = [N]
losses = tf.nn.softmax_cross_entropy_with_logits(logits, labels)
# L.shape      = []
L = tf.reduce_mean(losses)
# W.shape      = [M]  ※実際のネットワークだとここはもっと複雑
W = tf.trainable_variales()
# dW.shape     = [M]  ※実際のネットワークだとここはもっと複雑
dW = tf.gradients(L, W)

この勾配$dW$を用いたパラメーター更新がよく行われています。AdamOptimizerなどを使っていると明示的にgradientsを呼び出すことはありませんが、内部的にはtf.gradientsが利用されています。

ここで損失関数$L$をスカラーではなくベクトルにすると何が起きるでしょうか？
L = tf.reduce_mean(losses)をやめて、L = lossesとするケースです。

# L.shape      = [N]
L = losses
# W.shape      = [M]
W = tf.trainable_variables()
# dW.shape     = [?]  
dW = tf.gradients(L, W)

数学的な感覚でいえば、dWがヤコビアンになって欲しいところです。
つまり、dW.shape=[N, M] となり、$dW$の$i$行目$j$列目の要素$dW_{ij}$は、損失$L$の$i$番目の要素を$W$の$j$番目の要素で微分した値になって欲しいところです。

ところが、実際にはdWのshapeはMになります。これはTensorflowの仕様としてtf.gradientsの第一引数に非スカラー値が渡された場合、その第一引数の総和のスカラー値を微分した値を返すためです。つまり上記の例は以下と等価になります。

L = tf.reduce_sum(losses)
W = tf.trainable_variables()
dW = tf.gradients(L, W)

このように動作するため、tf.gradientsではヤコビアンを得ることができません。

Lop (ベクトルとヤコビアンの積)

ヤコビアンは簡単に求まらないのですが、ベクトルとヤコビアンの積は簡単に求めることができます。
今、入力を$M$次元の$x$, 出力を$N$次元の$y$とする関数$f$について考えます。

 y = f(x)

このときヤコビアン$J$は次のように書けます。

  J = \left(
    \begin{array}{cccc}
      \frac{\partial y_1}{\partial x_1} & \frac{\partial y_1}{\partial x_2} & \ldots & \frac{\partial y_{1}}{\partial x_{M}} \\
      \frac{\partial y_2}{\partial x_1} & \frac{\partial y_2}{\partial x_2} & \ldots & \frac{\partial y_{2}}{\partial x_{M}} \\
      \vdots & \vdots & \ddots & \vdots \\
      \frac{\partial y_N}{\partial x_1} & \frac{\partial y_N}{\partial x_2} & \ldots & \frac{\partial y_{N}}{\partial x_{M}} \\
    \end{array}
  \right)

ベクトルとヤコビアンの積$vjp$とは、$J$に対して左から長さ$N$のベクトル$v$をかけた値を指します。

vjp = v^T J

この演算はtensorflow上では次のように書けます。

# fとxは別で定義されているものとする
y = f(x)
v = tf.placeholder(shape=[N])
vjp = tf.gradients(y, x, grad_ys=v)

このようにgrad_ysが指定されている場合、yはgrad_ysと同じshapeであればスカラーである必要がありません。
なお、この$vjp$はtheanoなどではL-operator (Lop)と呼ばれている演算にあたります。

Lopの内部構造

この節ではLopが内部的にどのように達成されているかを解説します。
Tensorflowの微分計算はchain ruleに基づいて行われます。損失$L$をあるレイヤー$f$の入力$x$で微分する場合、そのレイヤーの出力$y$に関する微分を用いて計算されます。

\frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \frac{\partial y}{\partial x}

よって、$f$の微分関数grad_fは入力として$\frac{\partial L}{\partial y}$ 出力として $\frac{\partial L}{\partial x}$ を取るように構成する必要があります。疑似コードとして書くと以下のようになります。

@ops.RegisterGradient("f")
def grad_f(op, grad_y):
  x = op.inputs[0]
  ....
  grad_x = ...
  return grad_x

このとき、grad_yはback propagationにより、損失$L$に近いレイヤーから順に計算された値が渡されます。
$L$にもっとも近いレイヤーでは、このgrad_yには1が渡されます。例えば、以下の例ではreduce_meanの勾配関数がback propagationの過程では、一番最初に計算されます。

losses = tf.nn.softmax_cross_entropy_with_logits(logits, labels)
L = tf.reduce_mean(losses)
dW = tf.gradients(L, W)

このときreduce_meanの勾配関数にはgrad_yとして1が渡されます。これは一番$L$に近いレイヤーでは, $y=L$となるため、以下のchain ruleの式において、$\frac{\partial L}{\partial y} = 1$ となることからも妥当といえます。

\frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \frac{\partial y}{\partial x}

Lopの計算は、この最初のgrad_yにかけたいベクトル$v$をセットすることで実現されています。
初期のgrad_yにベクトル$v$をセットするということは、最終層の出力$y$と損失$L$の関係が以下であることと等価です。

L = v^T y

この両辺を$x$で微分すると

\frac{\partial L}{\partial x} = v^T \frac{\partial y} {\partial x} = v^T J

となり、$vjp$が計算できていることがわかります。

まとめ

以上のようにTensorflowではLopは、通常の微分と同程度のコストで求めることができます。
しかし、ヤコビアンの左からベクトルをかけるLopとは異なり、ヤコビアンの右からベクトルをかけるRopはもう少し議論が複雑です。
次の機会にはRopの計算方法と周辺議論について書きたいと思います。

Rop関係の参考資料
- https://j-towns.github.io/2017/06/12/A-new-trick.html
- https://github.com/renmengye/tensorflow-forward-ad/issues/2