Enhancing Fast Feed Forward Networks with Load Balancing and a Master Leaf Node

TL;DR

This paper introduces load balancing and Master Leaf techniques into fast feedforward networks, significantly improving their accuracy and consistency by leveraging Mixture of Experts concepts.

Contribution

It presents a novel integration of MoE-inspired load balancing and Master Leaf methods into FFFs, enhancing performance and simplifying training.

Findings

Up to 16.3% accuracy increase in training

Up to 3% accuracy increase in testing

Reduced variance in results

Abstract

Fast feedforward networks (FFFs) are a class of neural networks that exploit the observation that different regions of the input space activate distinct subsets of neurons in wide networks. FFFs partition the input space into separate sections using a differentiable binary tree of neurons and during inference descend the binary tree in order to improve computational efficiency. Inspired by Mixture of Experts (MoE) research, we propose the incorporation of load balancing and Master Leaf techniques into the FFF architecture to improve performance and simplify the training process. We reproduce experiments found in literature and present results on FFF models enhanced using these techniques. The proposed architecture and training recipe achieves up to 16.3% and 3% absolute classification accuracy increase in training and test accuracy, respectively, compared to the original FFF architecture. Additionally, we observe a smaller variance in the results compared to those reported in prior research. These findings demonstrate the potential of integrating MoE-inspired techniques into FFFs for developing more accurate and efficient models.