Dynamic sparsity in tree-structured feed-forward layers at scale

TL;DR

This paper introduces tree-structured sparse feed-forward layers in transformers, enabling scalable, conditional computation that activates less than 5% of units per token while maintaining performance.

Contribution

It demonstrates for the first time that tree-structured sparsity can be effectively applied to autoregressive language models beyond 1B parameters, matching dense baselines.

Findings

Models activate less than 5% of units per token.

Tree-structured sparsity matches dense model performance.

Emergent auto-pruning reduces unused paths over training.

Abstract

At typical context lengths, the feed-forward MLP block accounts for a large share of a transformer's compute budget, motivating sparse alternatives to dense MLP blocks. We study sparse, tree-structured feed-forward layers as drop-in replacements for MLP blocks in deep transformer architectures, enabling conditional computation via hard hierarchical routing without a separate router network. We demonstrate for the first time that this form of tree-structured conditional sparsity can be applied for autoregressive language modeling and downstream question answering, including zero- and few-shot settings, and its scalability beyond 1B parameters. Despite activating fewer than 5% of the feed-forward block's units per token, our models match dense baselines under controlled training and fine-tuning protocols. We further analyze training dynamics and identify an emergent auto-pruning effect: the interaction of hard routing with asymmetric nonlinearities progressively deactivates unused paths, yielding partial conversion of dynamic routing into static structural sparsity. We show that simple architectural choices can modulate this behavior and recover balanced trees without auxiliary losses. Overall, our work demonstrates that tree-structured feed-forward layers provide a scalable and controllable mechanism for sparsifying large transformer models.