[论文解读] Hash Layers For Large Sparse Models

Stephen Roller, Sainbayar Sukhbaatar+2 more|arXiv (Cornell University)|Jun 8, 2021

Topic Modeling参考文献 36被引用 48

一句话总结

Hash Layers 将 tokens 路由到固定的专家 FFN，通过预计算哈希，在 perplexity 上达到与 Switch Transformers 和 BASE Layers 相当或更好的表现，且不需要路由参数或额外目标。

ABSTRACT

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.

研究动机与目标

• 在 Transformer 中推动可扩展的稀疏路由，而无需可学习的路由参数。
• 提出一种基于哈希的路由机制，将 token 映射到固定的专家模块。
• 在语言建模和对话任务中评估哈希策略及多哈希扩展。
• 将性能和效率与 Switch Transformers 与 BASE Layers 进行比较。
• 分析哈希选择如何影响训练动态和下游微调。

提出的方法

• 在 Transformer 层中用一组类似 MoE 的专家 FFN 替代原有的 FFN，通过固定的基于哈希的路由函数来访问。
• 预先固定哈希函数（无学习的路由参数），并将 token 路由到每个 token 的单一专家。
• 探索多种哈希策略（随机、平衡、双字母、前一 token、基于位置、Oracle/未来、聚类、分散等）及其对性能的影响。
• 引入多哈希层（MultiHash Layers），组合若干哈希函数以选择子参数并连接它们的输出。
• 在 Pushshift Reddit、RoBERTa+cc100en、Wikitext-103 和 BST 对话任务上进行实验，并与密集基线、Switch Transformers 和 BASE Layers 进行比较。
• 给出哈希函数选择、层位置和多哈希效应的消融分析。

实验结果

研究问题

• RQ1 固定的基于哈希的路由方案是否能够在稀疏 Transformer 模型中达到与可学习路由相当的性能？
• RQ2 哪些哈希策略（随机、平衡、基于 token、n-gram、基于聚类、基于 Oracle）在 perplexity 与训练效率之间取得最佳权衡？
• RQ3 把基于哈希的路由放在后部 Transformer 层是否比放在前部层更能提升性能？
• RQ4 在固定计算预算下，多哈希路由对模型性能有何影响？
• RQ5 与 BASE Layers 相比，Hash Layers 在准确性和训练/更新效率方面在规模化下的表现如何？

主要发现

模型	配置	参数	有效 PPL	测试 PPL
Baseline Transformer	layers=11, d=1024, D=4096	222M	24.90	24.96
Wider Transformer	layers=11, d=2048, D=6144	755M	23.32	23.38
Deeper Transformer	layers=22, d=1536, D=4096	755M	22.72	22.78
Switch Transformer	layers=11,modules=1x64, load_bal=0.1	751M	23.65	23.73
Hash Layer	layers=11,modules=1x64	751M	23.16	23.23
Switch Transformer	layers=11,modules=1x128, load_bal=0.1	1.28B	23.52	23.58
Hash Layer	layers=11,modules=1x128	1.28B	22.89	22.95
Switch Transformer	layers=11,modules=5x16, load_bal=0.01	852M	23.19	23.25
Switch Transformer	layers=11,modules=5x16, load_bal=0.1	852M	23.00	22.93
Hash Layer	layers=11,modules=5x16	852M	23.21	23.27

• Hash Layers 在 Reddit 和 RoBERTa+cc100en 上的表现优于或等同于 Switch Transformers，且参数数量相近。
• 在相同计算预算下，Hash Layers 使用 64 或 128 个模块时 perplexity 更低于对应的 Switch 版本。
• 无需学习门控的哈希路由具有鲁棒性、实现简单，并避免负载均衡损失和路由优化难题。
• 将 Hash Layer 放在后层比前层带来更好的 perplexity，与 BASE Layers 的发现一致。
• 多哈希（2–8 哈希）在保持总参数不变的前提下提升了性能，相比单一哈希。
• 与 BASE Layers 相比，Hash Layers 在相同 perplexity 下提供更高的每秒更新量，因为避免了路由中的去相关步骤。