Hyperparameter Transfer with Mixture-of-Expert Layers

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Jiang, Tianze, Bordelon, Blake, Pehlevan, Cengiz, Hanin, Boris
Format: Preprint
Publié: 2026
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866913151595839488
author Jiang, Tianze
Bordelon, Blake
Pehlevan, Cengiz
Hanin, Boris
author_facet Jiang, Tianze
Bordelon, Blake
Pehlevan, Cengiz
Hanin, Boris
contents Mixture-of-Experts (MoE) layers have emerged as an important tool in scaling up modern neural networks by decoupling total trainable parameters from activated parameters in the forward pass for each token. However, sparse MoEs add complexity to training due to (i) new trainable parameters (router weights) that, like all other parameter groups, require hyperparameter (HP) tuning; (ii) new architecture scale dimensions (number of and size of experts) that must be chosen and potentially taken large. To make HP selection cheap and reliable, we propose a new parameterization for transformer models with MoE layers when scaling model width, depth, number of experts, and expert (hidden) size. Our parameterization is justified by a novel dynamical mean-field theory (DMFT) analysis. When varying different model dimensions trained at a fixed token budget, we find empirically that our parameterization enables reliable HP transfer across models from 51M to over 2B total parameters. We further take HPs identified from sweeping small models on a short token horizon to train larger models on longer horizons and report performant model behaviors.
format Preprint
id arxiv_https___arxiv_org_abs_2601_20205
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Hyperparameter Transfer with Mixture-of-Expert Layers
Jiang, Tianze
Bordelon, Blake
Pehlevan, Cengiz
Hanin, Boris
Machine Learning
Mixture-of-Experts (MoE) layers have emerged as an important tool in scaling up modern neural networks by decoupling total trainable parameters from activated parameters in the forward pass for each token. However, sparse MoEs add complexity to training due to (i) new trainable parameters (router weights) that, like all other parameter groups, require hyperparameter (HP) tuning; (ii) new architecture scale dimensions (number of and size of experts) that must be chosen and potentially taken large. To make HP selection cheap and reliable, we propose a new parameterization for transformer models with MoE layers when scaling model width, depth, number of experts, and expert (hidden) size. Our parameterization is justified by a novel dynamical mean-field theory (DMFT) analysis. When varying different model dimensions trained at a fixed token budget, we find empirically that our parameterization enables reliable HP transfer across models from 51M to over 2B total parameters. We further take HPs identified from sweeping small models on a short token horizon to train larger models on longer horizons and report performant model behaviors.
title Hyperparameter Transfer with Mixture-of-Expert Layers
topic Machine Learning
url https://arxiv.org/abs/2601.20205