Source-linked AI summary
Hash Layers For Large Sparse Models
Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, Jason Weston
TL;DR
Large Transformer MoE models need efficient routing, but learned gating can be difficult and costly. This paper fixes token-to-expert assignments with hash-based sparse routing, finding competitive performance with learned MoE routers while using no routing parameters or assignment optimization. The strongest results favor random or balanced hashes based on local token features, although the authors position learning-to-route as future work.
Problem
MoE routing must remain efficient while avoiding poor or overly specialized expert assignments.
Method
The paper replaces learned MoE routing with fixed hashing from input tokens to expert modules, including multi-hash variants for separate FFN segments.
Results
Hash-based routing is competitive with or outperforms Switch Transformers and BASE Layers across the studied tasks, with local random and balanced hashes performing best among the tested choices.
Takeaways & Limitations
Hash Layers are a strong candidate for very large-model training runs because they are robust, fast, and simpler than learned routing alternatives.
Takeaways & Limitations
The experiments scale only to 4.5B parameters, and the authors regard learning-to-route as future work.
Abstract
from arXiv · showhide
We investigate the training of sparse layers that use different parameters for different inputs based on hashing in large Transformer models. Specifically, we modify the feedforward layer to hash to different sets of weights depending on the current token, over all tokens in the sequence. We show that this procedure either outperforms or is competitive with learning-to-route mixture-of-expert methods such as Switch Transformers and BASE Layers, while requiring no routing parameters or extra terms in the objective function such as a load balancing loss, and no sophisticated assignment algorithm. We study the performance of different hashing techniques, hash sizes and input features, and show that balanced and random hashes focused on the most local features work best, compared to either learning clusters or using longer-range context. We show our approach works well both on large language modeling and dialogue tasks, and on downstream fine-tuning tasks.
1 Introduction
Scaling Transformer models improves performance but makes training and evaluation slower and more difficult to engineer. Hash Layers address MoE routing challenges with a fixed, sparse token-to-expert strategy that is competitive with learned routers.
- Larger Transformer models improve with scale but become slower to train and evaluate, especially at extreme scale.
- MoE models reduce computation by routing each input through a sparse subset of expert weights.
- Poor routing can prevent expert specialization, cause over-specialization, or make routing inefficient.
- Hash Layers use a simple sparse routing strategy that compares favorably with Switch Transformers and BASE Layers without routing parameters, objective changes, or assignment algorithms.
2 Background
Transformer MoE layers replace feedforward networks with multiple expert FFNs and use token-dependent gating to select sparse computation. This increases parameter capacity while keeping per-token computation relatively constant, but makes gating optimization difficult.
- A Transformer computes parallel token representations through layers containing self-attention and feedforward sublayers.
- Replacing an FFN with K expert FFNs lets a gating function mix expert outputs for each token.
- Because gating depends on each token’s hidden state, different tokens can receive different mixtures of experts.
- Sparse MoE routing activates only a few experts, increasing parameters with K while keeping per-token computation constant for single-expert routing.
- Optimizing sparse gating values can be difficult despite the computational advantages of sparse MoE models.
3 Method
The method fixes token-to-expert assignments by hashing input tokens rather than learning a routing network. It evaluates local, contextual, clustered, dispersed, and multi-hash designs while keeping expert computation sparse.
- Hash Layer: Hash routing activates one expert and uses no learned routing parameters.
- Hash Layer: Tokens are assigned to experts through a fixed hash mapping into buckets, while the FFN still processes each token’s hidden state.
- Hash functions: The fixed hash function maps original input tokens to experts in advance, requiring no routing training or adjustable parameters.
- Hash functions: Random Hash assigns tokens randomly, while Balanced assignment addresses frequency-driven imbalance across experts.
- Hash functions: Bigram, Previous Token, Position, Oracle Future, and Predicted Future hashes test how local context, position, and future-token information affect routing.
- Hash functions: Clustered Hash assigns similar tokens to shared centroids, whereas Dispersed Hash distributes tokens within clusters across different buckets.
- MultiHash Layers: MultiHash Layers select parameters independently for multiple FFN segments and concatenate the resulting representations.
4 Related Work
Related MoE methods learn sparse routing through backpropagation, regularization, or assignment optimization. Hash Layers instead use fixed assignments informed by natural-language inductive biases and connect to feature hashing and hash-based neural compression.
- Prior sparse MoE approaches learn gating with backpropagation, sometimes using load-balancing regularization, while other methods solve a linear assignment problem.
- Hash Layers differ by using no learned assignments and exploiting the language-modeling value of n-grams.
- The method is related to feature hashing, where word or n-gram features are mapped into a lower-dimensional feature space.
- Hash-based neural-network compression is a related use of hashing, but differs because Hash Layers increase model parameters rather than compressing them.
5 Experiments
Across experiments, Hash Layers generally outperform or match learned sparse-routing baselines, with advantages depending on module count, layer placement, hashing features, and dictionary size. They also retain strong fine-tuning behavior and show computational benefits relative to BASE Layers.
- Hash vs. Switch routing: Hash Layers outperform Switch by about 0.4-0.5 perplexity on both pushshift.io Reddit and RoBERTa+cc100en with 64 expert FFNs.Both methods use 64 expert FFNs and 751M total parameters.
- Dense vs. Sparse Models: Both sparse models outperform their dense Baseline and Wider Transformer counterparts, while the Deeper Transformer performs better but uses more inference computation.The Deeper Transformer has similar parameter count to the sparse models but uses dense rather than conditional computation.
- Hash layer module size: Increasing from 64 to 128 modules enlarges Hash’s advantage over Switch to about 0.6 perplexity, while smaller module counts produce similar performance.The 128-module models have 1.28B total parameters.
- Hash layer position: Later decoder layers work best for a 64-module Hash Layer, although placing it in layer 1 remains competitive with Switch in later layers.The layer-position comparison uses the pushshift.io Reddit task.
- Multi-layer routing: With routing every other layer and 16 modules per sparse layer, Switch slightly outperforms Hash at load-balancing parameter 0.1, while both reach 23.19 at 0.01.The reported values are 23.00 versus 23.21 at 0.1 and 23.19 for both methods at 0.01.
- Downstream fine-tuning: Fine-tuning results generally preserve the pre-training method order, while multiple Hash Layers yield a small but consistent improvement at constant total parameters.The multi-layer result uses three Hash Layers with random hashes in a 4.5B-parameter model.
- Hash function analysis: Balanced and random hashes perform similarly at 23.16 and 23.22 validation perplexity, whereas clustered hashes perform worse at 23.90 versus 23.22.Balanced assignment is more even across experts, and dispersed hashing restores clustered-hash performance to approximately random-hash levels.
- Hash function analysis: Content features localized to the current token work best: bigrams score 24.19 versus 23.16 for current-token hashing, and previous-token hashing scores 24.16 versus 23.16.Position-only hashing matches the dense baseline, while oracle future-token hashing reaches 1.9 perplexity and predicted future-token hashing reaches 25.02.
6 Conclusion
Hash layers provide a simple, efficient sparse-model approach that is highly competitive with Switch Transformers and BASE Layers while avoiding learned routing machinery. The experiments reach 4.5B parameters, but larger industrial-scale models remain untested.
- Hash layers are highly competitive with Switch Transformers and BASE Layers across varied datasets and settings.
- The routing mechanism requires no extra learning parameters, assignment algorithm, or objective-function changes.
- The authors position hash layers as a strong candidate for training very large models because of their simplicity and robustness.
- Experiments scale to 4.5B parameters but do not reach the scales of large industrial works.
- The authors propose hash layers as a strong baseline for future research on learning-to-route methods.
A Additional Results
The additional results cover multi-hashing, fine-tuning configurations, and Switch load balancing. Together, they examine alternative sparse-model configurations and the sensitivity of Switch performance to balancing choices.
- Multi-hashing keeps the parameter count fixed while splitting and indexing FFN weights with multiple hashes before concatenating them.
- Fine-tuning experiments compare dense and sparse models across multiple configurations on the BST tasks.
- Switch Transformers with 64 modules perform worse without load balancing, making the balancing parameter important.
B.1 Comparisons to Switch
The Switch comparisons use matched standard hyperparameters for Switch and Hash layers, including 64 experts, one MoE layer, and the same Transformer dimensions and optimization schedule.
- Both configurations have approximately 751M total parameters, with 751,224,896 for Switch and 751,159,296 for Hash Layer.
- Both Switch and Hash configurations use 64 expert modules in one MoE layer.
- The compared models use embedding size 1024, FFN size 4096, 16 attention heads, 2 encoder layers, and 11 decoder layers.
- The Switch and Hash runs share maximum learning rate 0.002, 10,000 warmup steps, inverse-square-root scheduling, 100,000 maximum steps, ADAM, and gradient clipping at 1.0.
B.2 Comparisons to Base
The BASE comparisons use matched optimization settings and BASE-specific gradient clipping across Hash and BASE configurations. The experiments run on substantially larger GPU allocations than the standard runs.
- BASE, Hash Layer, and 3x Hash Layer use the same maximum learning rate, warmup, polynomial-decay schedule, maximum steps, optimizer, and gradient clipping.
- The comparisons use maximum learning rate 7.5e-4, 2,000 warmup steps, 62,500 maximum steps, ADAM, and gradient clipping at 0.1.
- BASE comparisons apply BASE’s gradient clipping method, computing gradient norms only from shared parameters to avoid additional device communication.
- Standard experiments used 8 32GB V100 GPUs for roughly 20 hours, while BASE comparisons used 32 V100s for approximately 2 days.
D Societal Impact
The paper notes that Hash Layer may enable much larger models at lower cost than parameter-equivalent dense models, while potentially increasing environmental impact. It also flags potentially offensive Internet-derived text in the datasets as an issue outside the paper’s scope.
- Hash Layer may enable much larger models at a fraction of the cost of parameter-equivalent dense models, but those models may have greater environmental impact.
- The datasets contain varied and potentially offensive text procured from the Internet by third parties.
- Mitigating negative effects from offensive dataset content is identified as important but outside the scope of this paper.