重復使用 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各種實例，我需要多次多次評估這個雅可比行列式。 對於非平凡的 Jacobian 計算（例如，對於具有非線性激活函數的深度卷積神經網絡），重復運行 GradientTape 計算並評估jacobian方法的成本非常高。 我從TensorFlow 文檔中知道梯度（以及雅可比）是通過自動微分計算的。 我必須想象有一些網絡分析梯度的內部存儲（通過自動微分計算），在給定的輸入下進行評估。

我的問題：我假設 TensorFlow 構建和存儲（至少部分）計算雅可比行列式所需的分析梯度是否正確？ 如果是這樣，有沒有辦法保存這個分析梯度並用新輸入重新評估雅可比矩陣，而不必通過 GradientTape 方法重建它？

“持久” GradientTape 似乎並沒有解決這個問題：它只允許針對計算的多個內部參數重復評估單個 GradientTape 實例。

Answer 1

也許你覺得這很有幫助：

我需要多次計算任意函數的雅可比。 我的問題是我不恰當地使用GradientTape ，但我發布的代碼可能會幫助你或給你一些見解。 我發布了一個使用基於會話的tf.gradient()函數和現代GriadientTape方法計算雅可比的自包含示例。 在幫助下，我讓它們在彼此相同的數量級內運行。

• 如果您的問題集中在嘗試重用調用之間的中間計算以提高速度，那么我認為尼克的回答更適用。
• 如果您的問題集中在嘗試使 GradientTape 與靜態圖一樣快，那么請確保將其包裝在@tf.function因為它就是這樣做的。

請參閱我的問題： 與 tf.gradients() 相比，用於計算雅可比的 Abysmal tf.GradientTape 性能

Answer 2

假設 TensorFlow 構建和存儲（至少部分）計算雅可比行列式所需的分析梯度，我是否正確？

不 - 我認為你一定對自動微分有誤解。

雖然 tf 中的每個基本運算都“知道”其輸出相對於輸入的解析導數，但在計算實際梯度或雅可比值時，伴隨項（輸出的導數）的數值將傳遞給反向運算通過，然后使用每個基本運算的解析公式和鏈式法則計算更多數值。

如果是這樣，有沒有辦法保存這個分析梯度並使用新輸入重新評估雅可比矩陣，而不必通過 GradientTape 方法重建它？

不。 如果要在新輸入上計算梯度或雅可比，則需要再次執行整個計算。 對於深度神經網絡，沒有辦法解決這個問題。

順便說一句，如果您根據網絡參數對神經網絡的損失函數進行梯度計算，則計算梯度的時間將為 O(1) 計算損失本身的成本。 這是反向傳播，並且是反向模式自動微分之美的一部分。 但是，如果您的網絡有 N 個輸出，並且您想計算網絡的完整雅可比，則計算網絡輸出的時間將花費 O(N)。 這可能就是計算雅可比行列式如此昂貴的原因。

例如，如果您在 MNIST 上訓練網絡，並且您的網絡有 10 個輸出，您將它們組合成一個損失函數，計算損失函數的梯度將花費 O(1) 時間，但計算 10 個輸出的雅可比與參數相關的輸出將花費 O(10) 時間。