Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Hessian of Perplexity for Large Language Models by PyTorch autograd (Open Source)

Ivan Ilin

TL;DR

This work addresses the infeasibility of computing the full Hessian for large language models by providing a practical workflow to obtain exact Hessians on small parameter subsets using PyTorch autograd and to estimate the Hessian diagonal via Hessian-vector products and Hutchinson's trick. It introduces a formal parameterization of transformer-like blocks, defines the perplexity-based objective $\phi(\mathbf w)=\mathrm{PPL}(\mathcal{R}^{-1}(\mathbf w))$, and demonstrates subset-focused Hessian computations across single layers, whole blocks, and across all layers. The study shows that the Hessian is approximately diagonal in empirical observations and details both the methodology and computational considerations (including memory constraints and additive batch properties) needed to perform Hessian analysis on large models. The results provide a concrete, open-source toolkit for curvature analysis, enabling researchers to study second-order structure, verify diagonal-approximation assumptions, and support quantization and optimization research in transformer-based LLMs.

Abstract

Computing the full Hessian matrix -- the matrix of second-order derivatives for an entire Large Language Model (LLM) is infeasible due to its sheer size. In this technical report, we aim to provide a comprehensive guide on how to accurately compute at least a small portion of the Hessian for LLMs using PyTorch autograd library. We also demonstrate how to compute the full diagonal of the Hessian matrix using multiple samples of vector-Hessian Products (HVPs). We hope that both this guide and the accompanying GitHub code will be valuable resources for practitioners and researchers interested in better understanding the behavior and structure of the Hessian in LLMs.

Hessian of Perplexity for Large Language Models by PyTorch autograd (Open Source)

TL;DR

This work addresses the infeasibility of computing the full Hessian for large language models by providing a practical workflow to obtain exact Hessians on small parameter subsets using PyTorch autograd and to estimate the Hessian diagonal via Hessian-vector products and Hutchinson's trick. It introduces a formal parameterization of transformer-like blocks, defines the perplexity-based objective , and demonstrates subset-focused Hessian computations across single layers, whole blocks, and across all layers. The study shows that the Hessian is approximately diagonal in empirical observations and details both the methodology and computational considerations (including memory constraints and additive batch properties) needed to perform Hessian analysis on large models. The results provide a concrete, open-source toolkit for curvature analysis, enabling researchers to study second-order structure, verify diagonal-approximation assumptions, and support quantization and optimization research in transformer-based LLMs.

Abstract

Computing the full Hessian matrix -- the matrix of second-order derivatives for an entire Large Language Model (LLM) is infeasible due to its sheer size. In this technical report, we aim to provide a comprehensive guide on how to accurately compute at least a small portion of the Hessian for LLMs using PyTorch autograd library. We also demonstrate how to compute the full diagonal of the Hessian matrix using multiple samples of vector-Hessian Products (HVPs). We hope that both this guide and the accompanying GitHub code will be valuable resources for practitioners and researchers interested in better understanding the behavior and structure of the Hessian in LLMs.

Paper Structure

This paper contains 23 sections, 2 theorems, 21 equations, 3 figures, 1 table.

Table of Contents

  1. Introduction
  2. Motivation
  3. Goals of this technical report
  4. Notation
  5. Background and Related Work: Computing the Hessian Matrix
  6. Finite Difference Approximation (Simple, but not robust)
  7. Automatic Differentiation
  8. Symbolic Differentiation (SymPy)
  9. Estimating the Hessian Diagonal Using Hessian-Vector Products and Randomized Probing
  10. Prerequisites and definition of PPL function
  11. Notation for different subsets of parameters
  12. Consider a single linear layer from one single block
  13. All layers from a single block
  14. Consider a single linear layer from all blocks
  15. Consider all linear layers from LLM
  16. ...and 8 more sections

Key Result

Theorem 7.1

$PPL'(\Bar{c}')$ defined in (eq:new_def_of_ppl) has the additive property. In other words, the perplexity computed for the full batch $b$ will be equal to the sum of perplexities, computed for $b$ subsequent samples.

Figures (3)

  • Figure 1: Visual representation of different parts of $\nabla^2 \phi({\bf w})$ for different subsets of parameters. For clarity, we have plotted the absolute values of the entries in all cases to better visualize the magnitude of the elements.
  • Figure 2: Experiments with estimation of the Hessian diagonal elements for q_proj linear layer from the first block of OPT-125M.
  • Figure 3: Experiments with estimation of the Hessian diagonal elements for q_proj linear layer from the first block of OPT-125M.

Theorems & Definitions (3)

  • Theorem 7.1
  • proof
  • Corollary 7.2