Optimizerの出力にはTensorFlowの自動微分が効かない

とってもどうでもいい知見を得たのでメモ。

うまく微分できる例

v = constant(3.0)
with GradientTape() as v_gt:
    v_gt.watch(v)
    y = v**2
    print(y)
    dfdv = v_gt.gradient(y, v)
    print(dfdv) # 6.0

\frac{\partial f(v)}{h}=\frac{\partial f(v)}{\partial v}\frac{\partial v}{\partial h}

の$\partial v/\partial h$を計算したいのに、Kerasで用意されているOptimizerを挟むとうまく微分できなくなります。

うまく微分できない例

v = Variable(3.0)
h = Variable(0.1)
opt = keras.optimizers.SGD(h)
opt.build([v])
dfdv = Variable(2.0)

with GradientTape() as h_gt:
    h_gt.watch(h)
    opt.update_step(dfdv,v)
    print(v)
    dvdh = h_gt.gradient(v,h)
    print(dvdh) # -2.0となってほしいのに実際はNoneになる

これは、KerasのOptimizerが、変数vの更新を破壊的操作によって行なっているからです。
うまく微分できる例では、yはvの計算によって誕生しているため、両者の関係性が繋がり、微分成分が生じました。しかし、この例では、vはhとは別の系統から誕生しているので、その後vにどのような処理がなされても、両者の関係性を繋げることができません。

実際、KerasのSGDのソースコードはこんな感じになっていて、変数に対して破壊的操作を施していることが分かります。

sgd.py

variable.assign_add(-gradient * lr)

参考↓

TensorFlowでは、破壊的操作を通じて変数を演算しても、微分の成分は伝播されないようです。破壊的操作でも微分の成分が伝播されるように工夫できないんでしょうかね。

というわけで、$\partial v/\partial h=$更新後の変数に対するハイパーパラメータの勾配を計算するためには、「更新後の変数」が「ハイパーパラメータ」によって誕生している必要があるのです。