[论文解读] Neural Tangents: Fast and Easy Infinite Neural Networks in Python
Neural Tangents 提供基于 JAX 的 API，通过分析 NNGP/NTK 内核或有限宽度近似，来定义、训练和分析无限宽神经网络，支持可扩展的 CPU/GPU/TPU。
Neural Tangents is a library designed to enable research into infinite-width neural networks. It provides a high-level API for specifying complex and hierarchical neural network architectures. These networks can then be trained and evaluated either at finite-width as usual or in their infinite-width limit. Infinite-width networks can be trained analytically using exact Bayesian inference or using gradient descent via the Neural Tangent Kernel. Additionally, Neural Tangents provides tools to study gradient descent training dynamics of wide but finite networks in either function space or weight space. The entire library runs out-of-the-box on CPU, GPU, or TPU. All computations can be automatically distributed over multiple accelerators with near-linear scaling in the number of devices. Neural Tangents is available at www.github.com/google/neural-tangents. We also provide an accompanying interactive Colab notebook.
研究动机与目标
- 介绍 Neural Tangents，一个用于研究无限宽神经网络的开源库。
- 提供一个高级 API，用于指定复杂架构并获取解析的无限宽内核（NNGP/NTK）。
- 在无限宽极限下同时实现精确贝叶斯推断和梯度下降式的训练。
- 为缺乏解析形式的内核支持蒙特卡洛近似。
- 展示在各种架构和数据集上的可扩展性和实际有用性。
提出的方法
- 定义将标准张量运算（密集层、卷积、非线性）映射到与 NNGP 和 NTK 相应的核运算的转换。
- 在无限宽网络中提供解析内核计算，以及在解析形式不可用时的基于 MC 的内核近似。
- 提供一个 API，通过精确内核在无限宽极限下执行贝叶斯推断或梯度下降式训练。
- 通过泰勒级数近似，将连续时间训练动力学和有限宽权重空间的探索结合起来。
- 利用 XLA/JAX 在 CPU/GPU/TPU 上执行，具备自动设备分配和批处理。
实验结果
研究问题
- RQ1如何在无限宽极限下高效地计算并使用 NNGP 和 NTK 内核来覆盖广泛的神经网络体系结构？
- RQ2是否可以通过解析或蒙特卡洛内核估计，利用无限宽对换来准确分析有限宽网络？
- RQ3跨设备和架构的解析内核计算的性能和可扩展性如何？
- RQ4在无限宽推断和基于内核的训练下，不同架构（例如，全连接、CNN、WideResNet）的比较如何？
- RQ5哪些实用工具和工作流可以加速对无限宽神经网络的研究？
主要发现
- Neural Tangents 通过解析的 NNGP/NTK 内核为多种架构的无限宽网络提供精确推断。
- 蒙特卡洛内核估计在解析内核不可用时提供灵活的近似，宽度增加或样本数量增多时趋近解析值。
- 在实验中，无限宽推断可以在合成数据上与有限宽集合训练密切匹配，说明实际保真度。
- 该库在 CIFAR-10 上的评估显示收敛性和在各架构上的性能层次，其中类似 WideResNet 的结构在测试变体中表现最佳。
- 实现与额外的加速器近线性扩展，并支持批量处理、操作融合以及基于 2D 卷积的协方差计算以提高效率。
更好的研究，从现在开始
从论文设计到论文写作，大幅缩短您的研究时间。
无需绑定信用卡
本解读由 AI 生成，并经人工编辑审核。