TLDR: NeMo is a new approach for modularizing deep neural networks (DNNs) during training. Unlike previous methods, NeMo operates at the neuron level, making it applicable to diverse architectures like Transformers and scalable to large models using contrastive learning. It significantly reduces module size and computational overhead while maintaining or improving accuracy, as demonstrated across various vision and text-related tasks. This innovation promises more efficient and reusable DNNs.
Deep Neural Networks, or DNNs, have become a cornerstone of modern software systems. However, the process of building and training these complex models can be incredibly expensive and time-consuming. While reusing existing models helps reduce these costs, simply using an entire pre-trained model can lead to significant inefficiencies, especially during the inference stage where the model is used to make predictions.
This challenge has led to a growing interest in DNN modularization, a concept borrowed from software engineering. The idea is to break down a large DNN model into smaller, reusable modules, each responsible for a specific function. Among the various approaches, the ‘modularizing-while-training’ (MwT) paradigm has shown great promise. Unlike methods that modularize a model after it’s fully trained, MwT integrates the modularization process directly into the training phase, leading to more effective and efficient results.
However, existing MwT techniques have faced limitations. They primarily focus on smaller Convolutional Neural Network (CNN) models and operate at the convolutional kernel level. This makes them less suitable for the diverse and often much larger DNNs used today, particularly Transformer-based models, which are at the forefront of many advanced applications.
Introducing NeMo: A Breakthrough in DNN Modularization
To overcome these limitations, researchers have proposed NeMo, a novel approach that stands for NEuron-level MOdularizing-while-training. NeMo is designed to be both scalable and more broadly applicable across different DNN architectures. Its key innovation lies in operating at the ‘neuron level’ – the fundamental building block common to all DNNs. This neuron-level granularity ensures that NeMo can be seamlessly applied to Transformer models, as well as various other DNN architectures.
Furthermore, NeMo incorporates a sophisticated modular training method based on ‘contrastive learning’. This technique, combined with an effective composite loss function, allows NeMo to scale efficiently to large-scale models. The process involves two main phases: modular training and structured modularization.
During the modular training phase, NeMo starts with a randomly initialized model and integrates a ‘neuron identifier’. This identifier, composed of several mask generators, learns to pinpoint which neurons are responsible for specific tasks or classes. A unique combined loss function, incorporating both classification accuracy and modularity (cohesion and coupling), guides this training. Cohesion measures how well neurons within a module work together for a specific task, while coupling assesses the overlap of neurons between different modules. By optimizing these factors using cross-entropy and contrastive learning, NeMo ensures that modules are distinct and efficient.
Once training is complete, the structured modularization phase begins. NeMo generates ‘masks’ for each function, which are essentially blueprints indicating which neurons are relevant. These masks are then used to physically remove irrelevant neurons from the trained model, resulting in smaller, tailored sub-models. For instance, if a model was trained to classify 10 different types of images, NeMo could extract a module specifically for identifying ‘cats’ and ‘dogs’, removing all the unnecessary neurons related to other classes.
Why NeMo is a Game-Changer
NeMo offers several significant advantages. By operating at the neuron level, it provides a finer granularity for modularization compared to previous methods, leading to more compact and efficient modules. The contrastive learning-based loss function is also more effective for optimization, especially with large models, and simplifies the tuning process by requiring only one adjustable hyperparameter, unlike older methods that needed three.
Also Read:
- Optimizing Data Mixtures for Language Models with Bayesian Approaches
- Enhancing MLLM Accuracy: A New Method for Controlled Image Captioning
Experimental Validation and Real-World Impact
Comprehensive experiments were conducted on a variety of models, including two Transformer-based models (ViT and DeiT) and four CNN models, across widely-used image classification datasets. The results consistently demonstrated NeMo’s superiority. For Transformer-based models, NeMo achieved an average performance gain of 1.72% in module classification accuracy and a remarkable 58.10% reduction in module size (Neuron Retention Rate). This means NeMo can create modules that are significantly smaller while maintaining or even improving accuracy.
When compared to the state-of-the-art MwT method, NeMo consistently showed higher accuracy, lower neuron/kernel retention rates, better cohesion, and lower coupling. For example, in Transformer models, NeMo reduced the number of operations (FLOPs) by up to 66.73% compared to MwT for on-demand reuse tasks, without sacrificing accuracy. Similar improvements were observed with CNN models, where NeMo achieved smaller modules and up to 25.46% fewer FLOPs.
NeMo’s practical benefits were further highlighted through case studies. In a vision-related task, reusing a ‘hand’ module from a ViT model for a Rock-Paper-Scissors classification task resulted in comparable accuracy while retaining only 47.6% of neurons and reducing FLOPs by 62.54%. For text-related tasks, modularizing a CodeBERT model for Java code clone detection allowed for comparable performance with only 43.58% of neurons retained, significantly cutting inference costs.
While NeMo shows great promise and generalizability across various DNN architectures and scales, challenges remain, particularly with generative models and very large-scale datasets like ImageNet. However, this research, detailed in the paper NeMo: A Neuron-Level Modularizing-While-Training Approach for Decomposing DNN Models, marks a significant step towards more efficient, reusable, and manageable deep learning models, paving the way for advanced model sharing platforms and reduced computational overhead in real-world applications.
