Question about the Squaring of the Hessian Diagonal in SparseGPT

[xiejingcheng]

I have a question regarding the code and pseudocode in your paper "SparseGPT: Massive Language Models Can be Accurately Pruned in One-Shot".

Specifically, in line 96 of sparsegpt.py, the formula tmp = W1 ** 2 / (torch.diag(Hinv1).reshape((1, -1)) ** 2) squares the diagonal elements of the Hessian matrix. Similarly, in line 8 of the pseudocode (Algorithm 1) in the paper, the diagonal elements of the Hessian matrix are also squared.

However, in the original formula (e.g., Equation 3), the diagonal elements of the Hessian matrix are not squared. Could you kindly clarify why squaring the diagonal elements is necessary in the implementation and pseudocode?

Thank you for your time, and I appreciate the important work you've done with SparseGPT.

[LoonLi]

I don't know if you are still have this question but I also think about this question recently.

This square calculation is actually derived from the formula for calculating Cholesky decomposition, which is used to pre-calculate the Hessian synchronization. Let A be the inverse of the Hessian and L be the lower triangle matrix. The calculation for L can then be implemented as:

import torch

def cholesky_decomposition_from_scratch(A):
    # Verify that A is square
    n = A.shape[0]
    if A.shape[1] != n:
        raise ValueError("Input matrix must be square")
    
    # Create a tensor to store the result
    L = torch.zeros_like(A)
    
    # Perform the Cholesky decomposition
    for i in range(n):
        # Compute the diagonal element L[i,i]
        s = A[i, i] - torch.sum(L[i, :i] ** 2)
        if s <= 0:
            raise ValueError("Input matrix is not positive-definite")
        L[i, i] = torch.sqrt(s)
        
        # Compute the non-diagonal elements in column i
        for j in range(i+1, n):
            L[j, i] = (A[j, i] - torch.sum(L[j, :i] * L[i, :i])) / L[i, i]
    
    return L

Here we can see that we took a torch.sqrt for calculating the diagonal elements while other elements are calculated without torch.sqrt. Consequently, we need to take the square calculation for the diagonal elements.