問題設定

勾配を用いる学習では、ミニバッチごとに損失関数$L$のパラメーター勾配 $\nabla_{\theta} L$ を求めます。
また、損失関数$L$は各データ点$z_i$における損失$L_i$の平均で表されます。

L(z) = \frac{1}{N}\sum_{i}^{N} L_i(z_i)

パラメーターを最適化するだけであれば$\nabla_{\theta} L$ を求めるだけで十分なのですが、こちらの記事のような手法を用いる際には、データ点ごとの勾配$\nabla_{\theta} L_i$を計算する必要が生じます。このデータ点ごとの勾配はPer-Example GradientやUnaggregated Gradientと呼ばれることもあるようです。

本記事ではTensorflowにおいてデータ点ごとに勾配を求める際の手法や課題をまとめています。CNNを想定して記載していますが、CNN以外にも転用可能です。

課題と論点

効率を無視すれば$\nabla_{\theta} L_i$ を求めるのは容易です。バッチサイズを1にして勾配を求めれば得られます。
しかし、バッチサイズを1にすると諸々のオーバーヘッドが発生しますし、GPUへ供給されるデータ量が減少するため、GPUの利用効率が減少してしまいます。実際、後述の実験ではバッチサイズ64の勾配計算と比較して、バッチサイズ1の勾配計算はデータ１点あたり5倍の計算時間を要しています。

しかし、issueにもあるようにTensorflowでは$\nabla_{\theta} L_i$を求める簡便なサポートはありませんし、将来的なサポートの提供も難しいでしょう。Tensorflowではbackpropagation時に情報として、$l$番目のレイヤーの出力を$f^l_{N,H,W,C}$とすると、損失関数をその出力の各要素で偏微分したものを伝搬しています。

\delta^l = \Bigl( \frac{\partial L}{\partial f^l_{1,1,1,1}},\frac{\partial L}{\partial f^l_{1,1,1,2}} \dots \frac{\partial L}{\partial f^l_{N,H,W,C}} \Bigr)

よって、各$L_i$に対する微分値は捨てられているため、通常のbackpropagationの結果からは収集できません。
そこで色々な手法でこの$L_i$が捨てられないような処置がとられています。

実装法1　バッチサイズを1にした際の実装

トリビアルな実装ですが、バッチサイズを1にするのがもっとも簡単な方法です。擬似コードは以下のようになります。

logits = inference(images)          # forward計算
losses = loss(logits, labels)
L = tf.reduce_mean(losses)
grads = tf.gradients(L, parameters) # backward計算

$images$のバッチサイズを1にすれば$grad$に格納される勾配が個別のデータ点に対応したものになります。

実装法2 ループを用いた実装

バッチサイズを1にした実装では、N個のデータ点に対してforward計算をN回、backward計算をN回行う必要があります。
しかし、backward計算はまとめられませんが、forward計算は1回で複数のデータ点に対して行えます。

logits = inference(images)          # forward計算
losses = loss(logits, labels)

grads = []
for loss in losses:
    grad = tf.gradients(loss, parameters)
    grads.append(grad)

上記の例ではpython上でループを行っていますが、tensorflowのwhileを用いてループさせることができます。
こちらのissueにwhileを用いた場合のコードが記載されています。下記の比較実験ではtensorflowのwhileを用いていますが、pythonループの方が速い可能性があります。

実装法3 パラメーターコピーを用いる手法

実装法2では、gradientsを何度も呼び出す必要がありました。しかし、計算グラフの構築を工夫することでgradientsの呼び出しを１本化することができます。こちらのissueで議論されています。
個々のデータごとに、異なるパラメーターのコピー$W_{copy}$を用いて推論を行うことで、$W_{copy}$で微分を行うと個々のデータに対する勾配を得ることができます。

list_logits = []
Ws = []
for image in images:
    W_copy = tf.identity(W)
    logit = network.inference(image, W_copy)
    
    list_logits.append(logits)
    Ws.append(W_copy)

logits = tf.concat(list_logits, axis=0)
loss = calc_loss(logits, labels)

grads = tf.gradients(loss, Ws)

この計算によって生成される計算グラフと伝搬される勾配情報は下図のようになります。$\nabla_{W_{copy1}} L_i$は$\nabla_{W} L_i$に等しいため、データ点ごとの勾配を取り出すことが可能になります。

この実装がおそらく最速なのですが実装コストが高いです。ネットワークがコピーしたパラメーターを用いるように改修するのが手間です。

実装法4 Conv2Dのgradientを上書きする

標準のConv2Dの勾配を出す演算Conv2DBackpropFilterは、全データ点に関する勾配を計算します。これを個々のデータ点に関する勾配を計算する関数に差し替えます。少し複雑な手法なので詳細は別の記事に記載します。

パフォーマンス比較

比較に使った実装は、githubに置いてあります。
実験環境は次のとおりです。

Title	Contents
GPU	GTX 1070
CPU	Core(TM) i5-4440
Cuda	Cuda 8.0 cuDNN 5.1.5
Tensorflow	1.2.1

用いたネットワークは入力は64x64の画像、convolutionを7回、fc-layerを3回挟んだものを使っています。
その上で、データ点１個あたりの勾配を求める時間を計測しました。

Implementation	Time per data(msec)
(参考タイム)バッチサイズ64での計算	1.138
バッチサイズ1で計算	5.657
ループを用いた手法	169.55
パラメーターコピーを用いた手法	4.964
Conv2Dのgradientを上書き	19.65

一番上のバッチサイズ64での計算では個々のデータ点の勾配ではなく、64点の勾配平均を求めるのにかかった時間です。性能上限の参考タイムとして記載しています。

バッチサイズ１で計算は遅いのですが、シンプルゆえか他の手法では中々超えることができません。パラメーターコピーを用いた手法は、性能は改善するのですが、実装コストを考えると利用は難しいです。

まとめ

残念ながら劇的に性能を改善する手法は発見できていません。そもそもConv2Dのgradientを上書きする手法以外はConv2DBackpropFilterをN回呼び出す必要がある点は変わらないため、性能改善幅に限界があります。

高速実装には失敗していますが、どのあたりが性能限界を定めているかをConv2Dのgradientを上書きする手法の紹介を通じて別記事に記載する予定です。