Hello Everyone,
I am currently working on a research project in which I need to calculate the Hessian Vector Product (HVP) w.r.t. my model’s parameters. I am able to successfully calculate the HVP, however the calculation slows significantly and chews up a significant amount of memory as the size of the training set I am using increases.
Let me make things more explicit:
	def GetHVPFunc_alt(self, datax, datay, testx, testy):
		params = [ p for p in self.model.parameters() if p.requires_grad ]
		predtest = self.model(testx)
		losstest = calc_loss(predtest, testy)
		grads_test = self.To_List(autograd.grad(losstest, params))

		def HVP(v):
			s = time.time()
			predtrain = self.model(datax)
			losstrain = calc_loss(predtrain, datay)

			v_t = self.Reverse_To_List(v, params)
			out = hvp_alt(losstrain, params, v_t)
			print("Time HVP: ", time.time() - s)
			return out
		return HVP, grads_test

I want to calculate H*v where H is the hessian of the loss function w.r.t. the model parameters. As I increase the number of points in datax and datay (the training set images/labels respectively) the computation time of calculating the model output and the model loss increases as expected, however the calculation time along with the memory usage of hvp_alt grows significantly faster. This is likely because the function calls autograd.grad twice, which I believe creates large graphs that take a lot of memory and that scale with the number of inputs to the model. See the code below:
def hvp(y, w, v):
    if len(w) != len(v):
        raise(ValueError("w and v must have the same length."))
    # First backprop
    first_grads = grad(y, w, retain_graph=True, create_graph=True)
    # Elementwise products
    elemwise_products = 0
    for grad_elem, v_elem in zip(first_grads, v):
        elemwise_products += torch.sum(grad_elem * v_elem)
    # Second backprop
    return_grads = grad(elemwise_products, w, create_graph=False)
    return return_grads

I should note that the code above is not my own and is taken from the following github which comes from the paper I am basing my research on: GitHub - nimarb/pytorch_influence_functions: This is a PyTorch reimplementation of Influence Functions from the ICML2017 best paper: Understanding Black-box Predictions via Influence Functions by Pang Wei Koh and Percy Liang. 1 . I have been using 100 data points from CIFAR-10 to approximate the HVP causing each iteration of my overall code (I am approximating the inverse HVP i.e. H^-1 * v) to finish in ~6 seconds, however when I increase the training set size to 1000 the run time increases to ~60 seconds.
In short, is there a way for me to get an accurate, efficient approximation of the HVP that doesn’t depend so heavily on training set size? Is my current solution of using a subset of the data sufficient or will this result in an inaccurate HVP calculation?
Feel free to ask for clarification if anything was confusing.
Thank you so much,
Ethan