Повторное использование GradientTape для нескольких якобианских вычислений

Question

Я пытаюсь вычислить якобиан выходов нейронной сети TensorFlow относительно ее входов. Это легко достигается с помощью метода tf.GradientTape.jacobian. Вот тривиальный пример, приведенный в документации TensorFlow:

with tf.GradientTape() as g:
  x  = tf.constant([1.0, 2.0])
  g.watch(x)
  y = x * x
jacobian = g.jacobian(y, x)

Это хорошо, если я хочу только вычислить якобиан одного экземпляра входного тензора x. Тем не менее, мне нужно многократно оценивать этот якобиан для различных случаев x. Для нетривиального вычисления Якобиана (например, для глубоко сверточной нейронной сети с нелинейными функциями активации) это невероятно дорого, чтобы повторно выполнить расчет GradientTape и оценить метод jacobian. Я знаю из документации TensorFlow , что градиенты (и, следовательно, якобиан) вычисляются с помощью автоматического c дифференцирования. Я должен представить, что существует некоторая внутренняя память аналитического градиента сети (вычисляемая автоматическим дифференцированием c), которая оценивается при заданных входных данных.

Мой вопрос: правильно ли я предположить, что TensorFlow строит и хранит (хотя бы части) аналитических градиентов, необходимых для вычисления якобиана? И если да, то есть ли способ сохранить этот аналитический градиент и переоценить якобиан с новыми входами, не восстанавливая его с помощью метода GradientTape?

«Постоянный» GradientTape, похоже, не решает эту проблему : он допускает только повторную оценку одного экземпляра GradientTape в отношении нескольких внутренних аргументов вычисления.

Answer 1

Правильно ли я полагаю, что TensorFlow строит и хранит (по крайней мере, части) аналитических градиентов, необходимых для вычисления якобиана?

Нет - я думаю, что вы что-то неправильно понимаете автоматы c дифференциация.

В то время как каждая элементарная операция в tf «знает» о производной аналитики c ее вывода по отношению к входу, когда вычисляются фактические значения градиента или якобиана, числовые значения примыкают (производные вывода ) передаются операции на обратном проходе, а затем вычисляются более числовые значения с использованием формулы analyti c для каждой элементарной операции и правила цепочки.

И если да, то есть ли способ сохранить этот аналитический градиент и переоценить якобиан с новыми входами, не восстанавливая его с помощью метода GradientTape?

Нет. Если вы хотите вычислить градиент или якобиан на новом входе, вам нужно будет снова выполнить весь расчет. Для глубоких нейронных сетей это невозможно.

Кстати, если вы берете градиенты функции потерь вашей нейронной сети относительно параметров вашей сети, время для вычисления градиентов будет O (1) стоимость вычисления потерь сам. Это обратное распространение, и оно является частью красоты автоматического дифференцирования в обратном режиме c. Но если ваша сеть имеет N выходов и вы хотите вычислить полный якобиан вашей сети, это будет стоить O (N) времени вычисления выходов вашей сети. Возможно, поэтому вычисление якобиана так дорого.

Например, если вы обучаете сеть на MNIST, и ваша сеть имеет 10 выходов, которые вы объединяете в одну функцию потерь, вычисляя градиенты функция потерь займет O (1) времени, но вычисление якобиана из 10 выходов относительно параметров займет O (10) времени.