Учитывая модель в стиле Кераса, я хотел бы изучить гессиан потери относительно весов. Такие матрицы, как правило, слишком велики для хранения в памяти, поэтому я был бы доволен, чтобы иметь возможность вычислять векторные произведения вида Hv. В настоящее время я делаю это с помощью метода конечных разностей следующим образом (некоторые очевидные вспомогательные функции не показаны):
def get_weight_gradient(model: Sequential) -> np.ndarray:
"""Compute gradient of loss wrt weights."""
with tf.GradientTape() as tape:
yhat = model(x_train_tf)
weights = model.trainable_weights
loss = k.categorical_crossentropy(y_train, yhat)
grad = tape.gradient(loss, weights)
return flatten_gradient(grad) # stack gradient into 1-D vector
def hvp(model: Sequential, v: np.ndarray, epsilon=10**-6) -> np.ndarray:
"""Compute Hessian vector product with an array v."""
chunked_v = chunk_vector(v, model) # break vector up into model-sized chunks
modelp = clone_model(model) # clone model and copy weights
modelm = clone_model(model)
for i, (tf_w, chunked_vi) in enumerate(zip(modelp.weights, chunked_v)):
tf_w.assign(tf_w + epsilon * chunked_vi)
for tf_w, chunked_vi in zip(modelm.weights, chunked_v):
tf_w.assign(tf_w - epsilon * chunked_vi)
gradp = get_weight_gradient(modelp)
gradm = get_weight_gradient(modelm)
hv = (gradp - gradm) / (2 * epsilon)
return hv
Этот подход неудовлетворителен. Во-первых, я могу сказать, что я страдаю численной ошибкой из-за конечной разности, даже для очень скромных моделей (например, плотная MLP, обученная на MNIST с | P | = 2 * 10 ^ 4). Например, истинная матрица Гессе должна иметь симметрию c, что, в свою очередь, означает, что у меня должно быть <Hu, v> = <u, Hv> для всех векторов соответствующего размера u, v. Когда я сравниваю эти два внутренних продукта для наборов случайных векторов, это равенство является лишь приблизительным: корреляция между LHS и RHS составляет около 95%. Кроме того, я не могу извлечь доминантную собственную пару с помощью силового метода. Так что конечное различие не работает.
Я бы хотел использовать алгоритм Perlmutter (http://www.bcl.hamilton.ie/~barak/papers/nc-hessian.pdf) для вычисления векторного гессенского произведения, то есть для использования тождества:
Hv = \partial_r \grad L(w + r*v) |_{r=0}
, где L - это потеря в зависимости от веса модели w.
. Это кажется более естественным способом использования AD TF. возможности, но я не могу понять, как express это в TF. Какой самый естественный способ вычислить векторное произведение Гессена в TF 2.0?