NeuroTrails: Training with Dynamic Sparse Heads as the Key to Effective Ensembling

TL;DR

NeuroTrails introduces a dynamic sparse multi-head architecture that enhances ensemble performance and robustness in deep learning tasks while significantly reducing computational resources.

Contribution

It proposes a novel, model-agnostic training paradigm with dynamic sparsity that improves ensemble effectiveness without multiple full models.

Findings

Achieves higher accuracy on ImageNet with fewer parameters.

Enhances zero-shot robustness in language models.

Effective across convolutional and transformer architectures.

Abstract

Model ensembles have long been a cornerstone for improving generalization and robustness in deep learning. However, their effectiveness often comes at the cost of substantial computational overhead. To address this issue, state-of-the-art methods aim to replicate ensemble-class performance without requiring multiple independently trained networks. Unfortunately, these algorithms often still demand considerable compute at inference. In response to these limitations, we introduce $\textbf{NeuroTrails}$, a sparse multi-head architecture with dynamically evolving topology. This unexplored model-agnostic training paradigm improves ensemble performance while reducing the required resources. We analyze the underlying reason for its effectiveness and observe that the various neural trails induced by dynamic sparsity attain a $\textit{Goldilocks zone}$ of prediction diversity. NeuroTrails displays efficacy with convolutional and transformer-based architectures on computer vision and language tasks. Experiments on ResNet-50/ImageNet, LLaMA-350M/C4, among many others, demonstrate increased accuracy and stronger robustness in zero-shot generalization, while requiring significantly fewer parameters.