💬 NLP

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

This paper demonstrates that tree-structured feed-forward layers with hard hierarchical routing can serve as scalable, sparse alternatives to dense MLPs in large transformers, achieving performance comparable to dense baselines while activating fewer than 5% of units per token and exhibiting an emergent auto-pruning effect that converts dynamic routing into static structural sparsity.

Original authors: Reza Sedghi, Robin Schiewer, Anand Subramoney, David Kappel

Published 2026-04-13

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Reza Sedghi, Robin Schiewer, Anand Subramoney, David Kappel

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are running a massive, high-end restaurant (the AI model) that serves millions of customers (data) every day.

In a traditional restaurant (a standard AI model), every single chef in the kitchen works on every single order, no matter what it is. If you order a salad, the pastry chef, the grill master, and the soup specialist all start chopping, grilling, and boiling. It's incredibly wasteful. They are all working, but most of them are doing things you don't need for that specific dish. This makes the kitchen slow and expensive to run.

To fix this, researchers usually try a "Mixture of Experts" (MoE) approach. They hire a smart manager (a router) who looks at the order and says, "Okay, for this salad, only the salad chef and the dressing chef work. Everyone else, go take a break." This saves energy, but it requires a complex management system, extra training to keep the chefs balanced, and special equipment to handle the switching.

The New Idea: The "Tree of Choices" Kitchen

This paper introduces a clever new way to organize the kitchen called Fast FeedForward (FFF) layers. Instead of hiring a manager to pick chefs, they restructure the kitchen itself into a giant decision tree.

Here is how it works, using simple analogies:

1. The Decision Tree (The Menu Path)

Imagine the kitchen isn't a room full of chefs, but a series of fork-in-the-road signs.

When an order comes in, it enters the kitchen at the top.
It hits a sign: "Is this a hot dish or a cold dish?"
- If Hot, it goes down the left path.
- If Cold, it goes down the right path.
At the next level, another sign asks: "Does it need a sauce?"
- Yes? Go left. No? Go right.

By the time the order reaches the bottom, it has only passed through one single path of the kitchen. Maybe it only touched 3 or 4 specific stations out of the 100 available. The other 96 stations never even saw the order. They didn't burn any electricity. They didn't use any water.

The Magic: The kitchen is designed so that for any order, only a tiny fraction of the staff (less than 5%) actually does the work. The rest of the "staff" (the parameters) are just there as potential options, but they stay dormant.

2. The "Self-Pruning" Effect (The Garden)

The most surprising discovery in this paper is what happens over time.

Imagine you plant a forest of trees (the kitchen paths). At first, the trees are all growing. But as the restaurant gets busier, the kitchen naturally figures out which paths are useless.

If a specific path (e.g., "Left-Right-Left") is never the right choice for any order, the "signal" to that path gets weaker and weaker.
Eventually, that path dies. It stops working entirely. It becomes a permanent dead branch.
The kitchen effectively "prunes" itself. It turns from a dynamic, changing system into a static, efficient structure where the useless parts are permanently turned off.

The researchers found they didn't need a gardener to cut the dead branches; the system did it automatically because of how the "signs" (math) interacted with the "growth" (learning).

3. Why This Matters (The Result)

The team tested this on huge AI models (some with over 1 billion "brain cells").

Speed: Because they only activate a tiny fraction of the brain for each word, the model runs 8.7 times faster than the old, dense models.
Smarts: Despite using so little energy, the model is just as smart as the heavy, slow version. It can write stories, answer questions, and solve logic puzzles just as well.
Simplicity: They didn't need a complex manager (router) or extra rules to balance the load. The tree structure did the work for them.

The Catch

The paper admits that while this is theoretically amazing, our current computer chips (GPUs) are built to do everything at once (parallel processing). This new "tree" method is a bit like a relay race (sequential), where one step must finish before the next starts. Currently, our hardware isn't perfectly optimized for this specific type of relay race, so the real-world speedup might be slightly less than the math predicts. However, the potential for saving massive amounts of energy and money is huge.

Summary

Think of this paper as inventing a smart, self-organizing maze for AI. Instead of forcing the AI to think with its whole brain for every single thought, it builds a maze where the thought only travels down the one correct path. Over time, the maze automatically removes the dead ends, becoming a lean, mean, super-fast machine that uses a fraction of the energy but keeps all the smarts.

1. Problem Statement

In modern Transformer architectures, particularly Large Language Models (LLMs), the Feed-Forward (FF) or MLP blocks account for the majority of parameters and computational cost. As models scale to billions of parameters, the cost of dense FF blocks becomes a primary bottleneck.

Current Solutions: Mixture-of-Experts (MoE) models have emerged as a popular sparse alternative, activating only a subset of parameters per token. However, MoE systems typically require learned router networks, auxiliary load-balancing losses, and specialized infrastructure, adding complexity and training instability.
The Gap: While "Fast Feedforward" (FFF) layers—using hard binary routing in tree structures—were proposed for small models, their scalability, training dynamics, and performance in large-scale autoregressive language models (LLMs) remained unexplored.

2. Methodology: Fast Feedforward (FFF) Layers

The authors propose replacing standard dense MLP blocks with tree-structured feed-forward layers that enable conditional computation without a separate router network.

Architecture: An FFF layer consists of $P$ parallel perfect binary trees of depth $D$ .
Mechanism:
- Hard Routing: For each input token, the model traverses a single root-to-leaf path in each tree. At each node, a linear projection generates a logit. A hard decision (based on the sign of the logit) determines whether to go left or right.
- Computation: Only the nodes along the active path are evaluated. The outputs of these visited nodes are aggregated (summed) to produce the layer output.
- Sparsity: This induces extreme sparsity. For a tree of depth $D$ , only $D$ nodes are active out of $2^{D+1}-1$ total nodes, allowing sparsity levels exceeding 99%.
Implementation: To ensure compatibility with existing GPU kernels, the authors use a "masked" implementation where all node logits are computed in parallel, but a binary mask (generated via a binary tree search operator) selects only the active path for the final output.

3. Key Contributions & Findings

A. Scalability to Billion-Parameter Models

The paper provides the first systematic demonstration that FFF layers can scale to autoregressive language models with over 1 billion parameters (specifically 1.3B and up to 6.7B in appendices).

Performance: FFF models match the performance of dense feed-forward baselines and sparsity-matched MoE models across language modeling (perplexity) and downstream tasks (zero-shot and few-shot reasoning).
Training Stability: The models exhibit stable training dynamics, converging similarly to dense baselines under both "from scratch" and fine-tuning regimes.

B. Emergent Auto-Pruning Mechanism

A novel discovery is an auto-pruning effect that occurs naturally during training due to the interaction between hard routing and the asymmetric GELU nonlinearity.

The Phenomenon: The combination of hard binary decisions and GELU creates an asymmetric gradient flow. Gradients for nodes with negative logits are suppressed, causing the routing logits to drift.
Result: Over time, entire tree paths become permanently inactive (deactivated). This converts dynamic, input-dependent routing into static structural sparsity.
Benefit: This yields additional computational savings without explicit pruning criteria or auxiliary losses. The model effectively "learns" to prune redundant paths.

C. Routing Imbalance and Architectural Control

The authors observed severe routing imbalance (heavy-tailed distribution), where the leftmost path is selected >50% of the time, mimicking expert collapse in MoE.

Analysis: This imbalance stems from the specific placement of the nonlinearity (GELU) before the sparse output accumulation.
Solution: By moving the nonlinearity to after the accumulation (a variant called FFFpost), the gradient suppression is removed, leading to balanced tree utilization.
Insight: Interestingly, the authors found that the imbalanced (auto-pruning) version often performed slightly better or equal to the balanced version, suggesting that resource competition and dynamic pruning are beneficial for performance.

D. Efficiency and Runtime

Speedup: Despite the sequential nature of tree traversal, the massive reduction in active parameters leads to significant speedups.
Results: A tree depth of 6 (approx. 94% sparsity) achieved an 8.7x speedup in wall-clock time compared to dense feed-forward layers on A100 GPUs, while maintaining comparable perplexity.
Trade-off: Deeper trees (e.g., depth 13) offer higher sparsity but face diminishing returns due to traversal overhead. Depth 6 is identified as the optimal trade-off.

4. Experimental Results

Language Modeling: On OPT-1.3B models trained on ~26B tokens, FFF models (depth 4 and 6) achieved perplexity within 10% of dense baselines.
Downstream Tasks: In zero-shot and few-shot benchmarks (HellaSwag, ARC, MMLU, DROP), FFF models consistently matched or exceeded dense baselines. Notably, on MMLU, sparser FFF models outperformed dense counterparts.
Comparison to MoE: FFF layers achieved performance comparable to MoE models with matched sparsity but without the need for complex router training or load-balancing losses.
Generalization: The approach was successfully extended to vision tasks (DETR object detection), showing similar performance to dense baselines.

5. Significance

This work establishes tree-structured feed-forward layers as a practical, scalable, and controllable alternative to dense MLPs and MoE architectures.

Simplicity: It eliminates the need for separate router networks and auxiliary balancing losses.
Efficiency: It offers a mechanism to achieve extreme sparsity (>90%) with predictable, structured activation patterns, leading to significant runtime speedups.
Emergent Behavior: The discovery of auto-pruning suggests that hard-routed trees can self-organize into efficient, statically sparse structures, bridging the gap between dynamic conditional computation and static model compression.

In conclusion, the paper demonstrates that replacing dense feed-forward blocks with hard-routed tree structures is a viable path toward more efficient, scalable, and interpretable large language models.