Tensorflow: calcular Hessian con respecto a cada muestra

Tengo un tensor X de tamaño M x D. Podemos interpretar cada fila de X como una muestra de entrenamiento y cada columna como una característica.

X se usa para calcular un tensor u de tamaño M x 1 (en otras palabras, u depende de X en el gráfico computacional). Podemos interpretar esto como un vector de predicciones; uno para cada muestra. En particular, la fila m-th de u se calcula utilizando solo la fila m-th de X

Ahora, si ejecuto tensor.gradients(u, X)[0] , tensor.gradients(u, X)[0] un tensor M x D correspondiente al gradiente “por muestra” de u con respecto a X

¿Cómo puedo calcular de forma similar el tensor de arpillera “por muestra”? (es decir, una cantidad M x D x D )


Anexo : la respuesta de Pedro a continuación es correcta. También encontré un enfoque diferente al astackr y desastackr (usando la notación de Peter):

 hess2 = tf.stack([ tf.gradients( tmp, a )[ 0 ] for tmp in tf.unstack( grad, num=5, axis=1 ) ], axis = 2) 

En el ejemplo de Peter, D = 5 es el número de características. Sospecho (pero no lo he comprobado) que lo anterior es más rápido para M grande, ya que omite las entradas cero mencionadas en la respuesta de Peter.

tf.hessians() está calculando la arpillera para los ys y xs proporcionados sin importar las dimensiones. Dado que tiene un resultado de dimensión M x D y xs de dimensión M x D, el resultado será de dimensión M x D x M x D. Pero como las salidas por ejemplar son independientes entre sí, la mayoría de los Hessianos serán cero, es decir, solo una porción en la tercera dimensión tendrá algún valor en absoluto. Por lo tanto, para obtener el resultado deseado, debe tomar la diagonal en las dos dimensiones M , o mucho más fácil, simplemente debe sumr y eliminar la tercera dimensión de la siguiente manera:

 hess2 = tf.reduce_sum( hess, axis = 2 ) 

Código de ejemplo (probado):

 import tensorflow as tf a = tf.constant( [ [ 1.0, 1, 1, 1, 1 ], [ 2, 2, 2, 2, 2 ], [ 3, 3, 3, 3, 3 ] ] ) b = tf.constant( [ [ 1.0 ], [ 2 ], [ 3 ], [ 4 ], [ 5 ] ] ) c = tf.matmul( a, b ) c_sq = tf.square( c ) grad = tf.gradients( c_sq, a )[ 0 ] hess = tf.hessians( c_sq, a )[ 0 ] hess2 = tf.reduce_sum( hess, axis = 2 ) with tf.Session() as sess: res = sess.run( [ c_sq, grad, hess2 ] ) for v in res: print( v.shape ) print( v ) print( "=======================") 

saldrá:

(3, 1)
[[225.]
[900.]
[2025.]]
=======================
(3, 5)
[[30. 60. 90. 120. 150.]
[60. 120. 180. 240. 300.]
[90. 180. 270. 360. 450.]]
=======================
(3, 5, 5)
[[[2. 4. 6. 8. 10.]]
[4. 8. 12. 16. 20.]
[6. 12. 18. 24. 30.]
[8. 16. 24. 32. 40.]
[10. 20. 30. 40. 50.]]

[[2. 4. 6. 8. 10.]
[4. 8. 12. 16. 20.]
[6. 12. 18. 24. 30.]
[8. 16. 24. 32. 40.]
[10. 20. 30. 40. 50.]]

[[2. 4. 6. 8. 10.]
[4. 8. 12. 16. 20.]
[6. 12. 18. 24. 30.]
[8. 16. 24. 32. 40.]
[10. 20. 30. 40. 50.]]]
=======================