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.
Figures & tables
Figure 1: MLP layers and sparse replacements. a) The standard MLP feed-forward block. b) A top-k MoE block, featuring a router that recruits experts (linear-activation-linear sub-networks) on a per-token level. Here, k=2 experts process each input token and their outputs are integrated via a weighted sum. c) An FFF block that contains P=4 parallel trees of depth D=2 , shown in yellow. Tree output is integrated via a concatenation of the individual sparse tree outputs and a final linear projection. A single path through the tree (green) is activated through hard routing decisions to produce the layer output.
Figure 2: Training loss curves for 125M-parameter GPT-style models comparing dense feed-forward (FF), Fast Feedforward (FFF), and sparsity-matched Mixture-of-Experts (MoE) layers.
Figure 3: Utilization in a d=5 tree. Path utilization as the normalized fraction of samples visiting each path for FFF (left) and FFF post (right). The x-axis denotes the leaf node index.
Figure 4: Validation accuracy under statistical path pruning. We evaluate different fixed pruning ratios by statistically pruning paths and measuring validation accuracy. Each curve corresponds to a different sparsity (static, permanent) ratio.
Model
SlimPajama
WikiText
OPT-1.3B (300B tokens)
13.11
14.3
OPT-1.3B (26B tokens)
15.83
17.2
OPT-1.3B FFF(d=4, 26B tokens)
16.61
18.10
OPT-1.3B FFF(d=4, 26B tokens, FT)
13.37
14.55
OPT-1.3B FFF(d=6, 26B tokens)
17.57
20.01
OPT-1.3B FFF(d=6, 26B tokens, FT)
14.11
15.02
Table 1: Perplexity (PPL) comparison of OPT-1.3B variants across two datasets and different training regimes and FFF configurations. Lower is better.
Model
HellaSwag
Obqa
WinoGrande
ARC-c
ARC-e
boolq
piqa
Avg
OPT-1.3B (300B tokens) †
53.65
33.40
59.59
29.44
50.80
60.83
72.36
51.44
Pythia-1.0B (1T tokens) †
47.16
31.40
53.43
27.05
48.99
57.83
69.21
48.30
Pythia-1.4B (1T tokens) †
52.01
33.20
57.38
28.50
54.00
63.27
70.95
51.33
TinyLlama v1.0 (2T tokens) †
59.20
36.00
59.12
30.12
55.25
57.83
73.29
52.99
TinyLlama v1.1 (2T tokens) †
61.47
36.80
59.43
32.68
55.47
55.99
73.56
53.63
OPT-1.3B FF
33.7
27.2
49.6
24.4
41.7
59.9
65.2
43.1
Table 2: Zero-shot performance on commonsense reasoning tasks. † Results as reported in the original papers and not reproduced in this work, i.e. OPT-1.3 (300B tokens): ( Zhang et al., 2022 ) ; Pythia: ( Biderman et al., 2023 ) and TinyLLaMA: ( Zhang et al., 2024 ) .
Model
MMLU
DROP
OPT-1.3B (300B tokens)
24.90
14.32
Pythia-1.0B (1T tokens) †
25.70
4.25
Pythia-1.4B (1T tokens) †
25.41
12.27
TinyLlama v1.0 (2T tokens) †
25.34
15.34
TinyLlama v1.1 (2T tokens) †
26.58
15.31
OPT-1.3B FF (retrained)
25.87
15.27
Table 3: Few-shot performance on reasoning and comprehension tasks. Results are accuracy (%) for MMLU (5-shot) and F1 (%) for DROP (3-shot). † indicates results reported from the original papers and not reproduced in this work.
Model
Overall Sparsity
MLP block Sparsity
Relative FLOPs
Relative FLOPs/S
Dense FF (GPT)
0%
0%
1.0 ×
1.0 ×
FFF (d=13)
61%
99%
0.0016 ×
0.0084 ×
FFF (d=6)
58%
94%
0.0952 ×
0.45 ×
FFF (d=4)
52%
83%
0.2666 ×
0.78 ×
Table 4: Efficiency metrics for three selected model configurations, including MLP-layer sparsity, overall model sparsity, relative FLOPs, and relative FLOPs/S compared to a dense GPT baseline.
Figure 5: Relative speedup and FLOPS/s for all configurations, ranging from depth 0 (dense FF) to depth 11. Results are averaged over five runs. d indicate depth of tree.
Model
layer execution (ms)
Speedup
Dense FF (GPT)
3.03
1.0 ×
FFF (d=13)
0.57
5.3 ×
FFF (d=6)
0.35
8.7 ×
FFF (d=4)
1.08
2.8 ×
Table 5: Speedup relative to the dense GPT baseline for three selected model on single A100 configurations.
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Variant
Tree depth / #experts
PPL (best)
FF
–
23.7786±0.0493 (23.7409)
FFF
3
25.9104±0.0690 (25.8447)
5
27.6623±0.3203 (27.2926)
7
29.3686±0.2763 (29.0518)
FFF post
3
29.6824±0.0203 (29.6640)
5
29.8080±0.5093 (29.0890)
Appendix
Table 6: Per-run scalar PPL aggregated by strategy and subgroup. Entries are mean ± std, averaged over multiple runs with different random seeds (best). Lower is better.
Figure 6: training curves 1.3 billion parameter
Figure 7: Training loss curves for OPT-style models with 2.7B and 6.7B parameters trained from scratch on a reduced SlimPajama subset. For each model size, we compare a dense feed-forward baseline with a Forest Router configuration of depth 4 (approximately 84% sparsity). Due to compute constraints, these experiments are conducted on a limited dataset and are intended to evaluate scalability and training stability rather than final convergence.
Figure 8: Extension of Fast Feedforward (FFF) layers to object detection with DETR on the CPPE-5 dataset. Right: evaluation performance measured by mAP@50 , defined as mean Average Precision at an IoU threshold of 0.5. Left: training loss curves. Results are averaged over multiple runs with different random seeds. FFF-based models achieve performance comparable to dense feed-forward baselines under a fixed training budget.
Figure 9: Empirical analysis of auto-pruning dynamics in a toy setting. Left: The gradient signal value received by a randomly selected tree node during backpropagation during training, showing persistent negative values after an initial training phase. Right: Histogram of the corresponding node output activations, which follow an approximately zero-mean normal distribution. Both measurements are computed using batch size of 10,000 to reduce stochastic variability.
Figure 10: Exploiting static sparsity in OPT-125M model.
Hyperparameter
Value
Model
Base model
facebook/opt-125m
Initialization
From scratch
Sequence length
1024
Tokenizer
GPT-style tokenizer (offline tokenization)
Dataset
Appendix
Table 7: Hyperparameters and training configuration for OPT-125M.
Figure 11: Binary trees with different branch weights and corresponding induced path distributions.
Figure 12: Visualization of FFF-induced space partitioning on a 2D checkered-field classification task. Left: routing-induced decision boundaries mapped back to the input space; each split corresponds to a node-level routing decision, and the combined boundaries show how the forest partitions the plane into routed regions. Right: PCA projection of representations grouped by routing outcomes; colors indicate distinct routed regions (sets of active paths) that contribute differently to the output.
Transformer feed-forward networks (FFNs) are the within-token MLP blocks whose activations, structure, efficiency, and role in model computation are studied.