mann 0.2.4

Package containing utilities for implementing RSN2/MANN

MANN

MANN, which stands for Multitask Artificial Neural Networks, is a Python package which enables creating sparse multitask models compatible with TensorFlow. This package contains custom layers and utilities to facilitate the training and optimization of models using the Reduction of Sub-Network Neuroplasticity (RSN2) training procedure developed by AI Squared, Inc.

Installation

This package is available through PyPi and can be installed via the following command:

pip install mann

To install the current version directly from GitHub without cloning, run the following command:

pip install git+https://github.com/AISquaredInc/mann.git

Alternatively, you can install the package by cloning the repository from GitHub using the following commands:

# clone the repository and cd into it
git clone https://github.com/AISquaredInc/mann
cd mann

# install the package
pip install .

Mac M1 Users

For those with a Mac with the M1 processor, this package can be installed, but the standard version of TensorFlow is not compatible with the M1 SOC. In order to install a compatible version of TensorFlow, please install the Miniforge conda environment, which utilizes the conda-forge channel only. Once you are using Miniforge, using conda to install TensorFlow in that environment should install the correct version. After installing TensorFlow, the command pip install mann will install the MANN package.

Capabilities

The MANN package includes two subpackages, the mann.utils package and the mann.layers package. As the name implies, the mann.utils package includes utilities which assist in model training. The mann.layers package includes custom Keras-compatible layers which can be used to train sparse multitask models.

Utils

The mann.utils subpackage contains helper functions for performing training and conversion of models using masking layers.

In addition to the functions just mentioned, there is also an ActiveSparsification callback object which enables active sparsification during training rather than solely one-shot sparsification. Note that this callback currently only supports simultaneous training. We are working to support iterative training with this callback as well.

1. mask_model
   • The mask_model function is central to the RSN2 training procedure and enables masking/pruning a model so a large percentage of the weights are inactive.
   • Inputs to the mask_model function are a TensorFlow model, a percentile in integer form, a method - either one of 'gradients' or 'magnitude', input data, and target data.
2. get_custom_objects
   • The get_custom_objects function takes no parameters and returns a dictionary of all custom objects required to load a model trained using this package.
3. remove_layer_masks
   • The remove_layer_masks function takes a trained model with masked layers and converts it to a model without masking layers.
4. add_layer_masks
   • The add_layer_masks function takes an existing model that has non-MANN layers and converts it so that all layers which have an analog in the MANN package. This enables pretrained models to be converted and sparsified.
5. quantize_model
   • The quantize_model function takes in a model and a datatype to quantize the model to.
6. build_transformer_block
   • The build_transformer_block function can be used to build a block in a transformer architecture.
7. build_token_position_embedding
   • The build_token_position_embedding function can be used to build a token and position embedding block for use in a transformer architecture model.
8. The get_task_masking_gradients function retrieves the gradients of masking weights within a model.
9. The mask_task_weights function masks specific task weights within a model.
10. The train_model_iteratively function iteratively trains a model utilizing early stopping and active sparsification on a per-task basis. NOTE that this function only works on models without MultiMaskedConv2D layers.

Layers

The mann.layers subpackage contains custom Keras-compatible layers which can be used to train sparse multitask models. The layers contained in this package are as follows:

1. MaskedDense
   • This layer is nearly identical to the Keras Dense layer, but it supports masking and pruning to reduce the number of active weights.
2. MaskedConv2D
   • This layer is nearly identical to the Keras Conv2D layer, but it supports masking and pruning to reduce the number of active weights.
3. MultiMaskedDense
   • This layer supports isolating pathways within the network and dedicating them for individual tasks and performing fully-connected operations on the input data.
4. MultiMaskedConv2D
   • This layer supports isolating pathways within the network and dedicating them for individual tasks and performing convolutional operations on the input data.
5. MultiDense
   • This layer supports multitask inference using a fully-connected architecture and is not designed for training. Once a model is trained with the MultiMaskedDense layer, that layer can be converted into this layer for inference by using the mann.utils.remove_layer_masks function.
6. MultiConv2D
   • This layer supports multitask inference using a convolutional architecture and is not designed for training. Once a model is trained with the MultiMaskedConv2D layer, that layer can be converted to this layer for inference by using the mann.utils.remove_layer_masks function.
7. SelectorLayer
   • This layer selects which of the multiple inputs fed into it is returned as a result. This layer is designed to be used specifically with multitask layers.
8. SumLayer
   • This layer returns the element-wise sum of all of the inputs.
9. FilterLayer
   • This layer can be turned on or off, and indicates whether the single input passed to it should be output or if all zeros should be returned.
10. MultiMaxPool2D
    • This layer implements Max Pool operations on multitask inputs.

Additional Documentation and Training Materials

Additional documentation and training materials will be added to the BeyondML Documentation Website as we continue to develop this project and its capabilities.

Feature Roadmap

• PyTorch Support
  • We are currently working on building support for PyTorch models and layers into this package
• Fixing issues with iterative training and MultiMaskedConv2D layers
  • As mentioned above, there are issues with finding pert-task gradients with models utilizing MultiMaskedConv2D layers. Future iterations of the technology will address bugs with these kinds of models

Changes

Below are a list of additional features, bug fixes, and other changes made for each version.

Version 0.2.2

• Small documentation changes
• Added quantize_model function
• Added build_transformer_block and build_token_position_embedding_block functions for transformer functionality
• Removed unnecessary imports breaking imports in minimal environments

Version 0.2.3

• Per-task pruning
  • Functionality for this feature is implemented, but usage is expected to be incomplete. Note that task gradients have to be passed retrieved and passed to the function directly (helper function available), and that the model has to initially be compiled using a compatible loss function (recommended 'mse') to identify gradients.
  • It has been found that this functionality is currently only supported for models with the following layers:
    • MaskedConv2D
    • MaskedDense
    • MultiMaskedDense
  • Note also that this functionality does not support cases where layers of an individual model are other TensorFlow models, but supporting this functionality is on the roadmap.
• Iterative training using per-task pruning
  • Functionality for this feature is implemented, but there are known bugs when trying to apply this methodology to models with the MultiMaskedConv2D layer present