中文

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