[论文解读] NeMo: a toolkit for building AI applications using Neural Modules

Oleksii Kuchaiev, Jason Li+12 more|arXiv (Cornell University)|Sep 14, 2019

Natural Language Processing Techniques参考文献 11被引用 176

一句话总结

NeMo 引入神经模块作为框架无关的构建块，具备神经类型系统以确保语义正确性，从而实现可重用的 ASR 和 NLP 集合，并支持高效的 GPU 训练。

ABSTRACT

NeMo (Neural Modules) is a Python framework-agnostic toolkit for creating AI applications through re-usability, abstraction, and composition. NeMo is built around neural modules, conceptual blocks of neural networks that take typed inputs and produce typed outputs. Such modules typically represent data layers, encoders, decoders, language models, loss functions, or methods of combining activations. NeMo makes it easy to combine and re-use these building blocks while providing a level of semantic correctness checking via its neural type system. The toolkit comes with extendable collections of pre-built modules for automatic speech recognition and natural language processing. Furthermore, NeMo provides built-in support for distributed training and mixed precision on latest NVIDIA GPUs. NeMo is open-source https://github.com/NVIDIA/NeMo

研究动机与目标

• 通过具备明确接口的神经模块，促进 AI 系统的软件工程风格分解。
• 引入神经类型系统，在模型构建过程中执行静态语义和维度检查。
• 提供用于 ASR 和 NLP 任务的可重用集合（例如 nemo_asr、 nemo_nlp），包含预先构建的模块。
• 在 NVIDIA GPUs 上通过混合精度和分布式能力实现高性能训练。
• 在利用 PyTorch 作为后端并能够轻松扩展到新后端的同时，提供框架无关的抽象。

提出的方法

• 将 Neural Module (NM) 定义为一个带类型的参数化构建块，可组合成有向无环图（DAG）以描述前向/后向数据流。
• 实现一个神经类型系统，将张量轴映射到语义轴类型，并在 DAG 构建期间执行静态兼容性检查。
• 提供一个框架无关的核心（NeMo Core），具有惰性执行模型和一个 NeuralFactory 用于实例化并连接 NM。
• 以 PyTorch 作为当前后端，同时概述其他后端的路径以及将 PyTorch nn.Modules 自动转换为 NM 的方法。
• 提供预构建的 NM 集合（例如 nemo_asr、 nemo_nlp），包含用于 ASR（CTC 和基于注意力）和 NLP 任务（NMT、语言建模等）的模板。
• 利用高性能训练技术，如混合精度、梯度累积，以及多 GPU/多节点训练（APEX）。

实验结果

研究问题

• RQ1 在组合复杂 AI 模型时，神经模块与神经类型系统如何提高可重用性和一致性？
• RQ2 模块化、带类型的组件对调试、配置管理以及跨领域重用（ASR 与 NLP）在 AI 应用中的影响是什么？
• RQ3 框架无关的抽象结合 PyTorch 后端，能否在现代 GPU 上实现实际的性能和可扩展性？
• RQ4 像 nemo_asr 和 nemo_nlp 这样的预构建集合在语音和语言任务的快速开发与实验中提供了哪些好处？

主要发现

• NeMo 提供一个神经模块抽象，能够实现可重用、可组合的模型组件，并具有静态类型检查，在 DAG 构建时捕捉语义和维度错误。
• 神经类型系统支持语义继承、轴级约束和显式端口接口，提升跨模块兼容性并减少无声错误。
• NeMo 在 NVIDIA GPUs 上通过混合精度、动态损失缩放和梯度累积实现实际的高性能训练，适用于大批量或分布式设置。
• nemo_asr 集合支持基于 CTC 的和序列到序列（seq2seq）的 ASR 模型，通过 Jasper 风格的体系结构以及可重用的数据层和编码器进行演示，具有成功的端到端训练工作流。
• nemo_nlp 集合支持如 NMT、语言建模和基于 BERT 的微调等任务，能够使用带有预训练组件的 Transformer 为基础的流水线。
• 经验示例：一个基于 Transformer 的 NMT 设置，在 WMT16 英德对中使用 8 GPUs 约 15 小时后，在 newstest2014 上达到 29.2 BLEU / 28.5 SacreBLEU。