Библиотека PyTorch Hessian Eigenthings 1.0.0 вычисляет разложения матрицы Гессе при линейной памяти

Библиотека PyTorch Hessian Eigenthings вышла в альфа-версии 1.0.0 и предлагает эффективный способ вычисления собственных разложений матрицы Гессе для моделей на PyTorch. Раньше полный расчет требовал памяти квадратично зависящей от числа параметров, что делало анализ невозможным для реальных нейросетей. Новый модуль обходит эту проблему, используя итеративные методы вроде Lanczos или стохастической итерации мощностей, которым нужны только произведение матрицы на вектор. Это снижает затраты на память до линейной зависимости. Связку из произведения матрицы на вектор и этих алгоритмов применяют к реальным моделям от HuggingFace и TransformerLens. Библиотека вычисляет не только собственные значения и векторы, но и следы через Hutch++ плюс плотность спектра посредством стохастического квадратичного Lanczos. Если классификационная задача подразумевает использование обобщенной матрицы Гаусса-Ньютона, следует создать GGNOperator. Также доступна EmpiricalFisherOperator для внешней матрицы градиентов по выборкам. Все эти операторы работают с единым интерфейсом, поэтому алгоритмы применимы к любому из них. При невозможности двойного обратного прохода можно включить метод конечных разностей через метод "finite_difference". Анализ по отдельным блокам модели возможен с помощью параметра param_filter, где подстраиваются имена слоев. Для больших языковых моделей автоподбор ядра происходит через hf_lm_loss_of_output, выбирая Triton на CUDA для ускорения в 3.4 раза и снижения пиковой памяти вдвое, или torch.compile, дающего 2.6 кратного прироста. Принудительный выбор неотлаженного режима работает при параметре fused="eager". Примеры кода доступны для MLP, tiny-GPT2 и модели TransformerLens. Проект использует uv для сборки и тестирования, исходники хранятся на GitHub. Автором разработки является Ной Голмант совместно с Жэвэем Яо, Амиром Голхами, Майклом Махони и Джозефом Гонзалес из лаборатории RISELab на Калифорнийском университете в Беркли. Код итерации мощностей основан на работах группы Yao и Голхами, опубликованных в NeurIPS 2018. Ускоренная стохастическая итерация взятá из статьи группы De Sa на PMLR 2017. Версия 1.0 заимствует идеи у PyHessian, curvlinops и HessFormer.