[Paper Review] 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 TechniquesComputer Science11 references176 citations

TL;DR

NeMo introduces Neural Modules as framework-agnostic building blocks with a neural type system to ensure semantic correctness, enabling reusable ASR and NLP collections with efficient GPU training support.

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

Motivation & Objective

• Promote a software-engineering style of decomposition for AI systems through neural modules with well-defined interfaces.
• Introduce a neural type system to perform static semantic and dimensional checks during model construction.
• Provide reusable collections (e.g., nemo_asr, nemo_nlp) of pre-built modules for ASR and NLP tasks.
• Enable high-performance training on NVIDIA GPUs with mixed precision and distributed capabilities.
• Offer framework-agnostic abstractions while leveraging PyTorch as backend and enabling easy extension to new backends.

Proposed method

• Define Neural Module (NM) as a typed, parametric building block that can be composed into DAGs to describe forward/backward data flows.
• Implement a Neural Type system that maps tensor axes to semantic axis types and performs static compatibility checks during DAG construction.
• Provide a framework-agnostic core (NeMo Core) with a lazy execution model and a NeuralFactory to instantiate and connect NMs.
• Support PyTorch as the current backend while outlining pathways for additional backends and automatic conversion of PyTorch nn.Modules to NMs.
• Offer pre-built collections (e.g., nemo_asr, nemo_nlp) of NMs with templates for ASR (CTC and attention-based) and NLP tasks (NMT, language modeling, etc.).
• Leverage high-performance training techniques such as mixed precision, gradient accumulation, and multi-GPU/multi-node training (APEX).

Experimental results

Research questions

• RQ1 How can neural modules and a neural type system improve reusability and integrity when composing complex AI models?
• RQ2 What is the impact of modular, typified components on debugging, configuration management, and cross-domain reuse (ASR and NLP) within AI applications?
• RQ3 Can framework-agnostic abstractions coupled with a PyTorch backend deliver practical performance and scalability on modern GPUs?
• RQ4 What benefits do pre-built collections like nemo_asr and nemo_nlp offer for rapid development and experimentation in speech and language tasks?

Key findings

• NeMo provides a neural module abstraction that enables reusable, composable model components with static type checking to catch semantic and dimensionality errors at DAG construction time.
• The neural type system supports semantic inheritance, axis-level constraints, and explicit port interfaces, improving cross-module compatibility and reducing silent errors.
• NeMo achieves practical high-performance training on NVIDIA GPUs with mixed precision, dynamic loss scaling, and gradient accumulation for large-batch or distributed settings.
• The nemo_asr collection supports CTC-based and sequence-to-sequence ASR models, demonstrated via Jasper-like architectures and reusable data layers and encoders, with successful end-to-end training workflows.
• The nemo_nlp collection supports tasks like NMT, language modeling, and BERT-based fine-tuning, enabling Transformer-based pipelines with pre-trained components.
• Empirical example: a Transformer-based NMT setup achieving 29.2 BLEU / 28.5 SacreBLEU on newstest2014 after ~15 hours on WMT16 English-German with 8 GPUs.