TorchNTK:一个用于计算 PyTorch 模型神经正切核的库
机器学习
2022-05-26 v1
摘要
我们介绍 torchNTK,一个用于在 PyTorch 框架中计算神经网络模型经验神经正切核(NTK)的 Python 库。我们提供了一种计算多层感知机 NTK 的高效方法。我们将显式微分实现与自动微分实现进行比较,后者具有将库实用性扩展到 PyTorch 支持的任何架构(如卷积网络)的优势。该库的一个特性是我们向用户暴露逐层 NTK 分量,并表明在某些情形下逐层计算更具内存效率。我们进行了初步实验以展示该软件的用例并探查 NTK。
引用
@article{arxiv.2205.12372,
title = {TorchNTK: A Library for Calculation of Neural Tangent Kernels of PyTorch Models},
author = {Andrew Engel and Zhichao Wang and Anand D. Sarwate and Sutanay Choudhury and Tony Chiang},
journal= {arXiv preprint arXiv:2205.12372},
year = {2022}
}
备注
19 pages, 5 figures