Optimal Scaling Needs Optimal Norm
Despite recent progress in optimal hyperparameter transfer under model and
dataset scaling, no unifying explanatory principle has been established. Using
the Scion optimizer, we discover that joint optimal scaling across model and
dataset sizes is governed by a single invariant: the operator norm of the
output layer. Across models with up to 1.3B parameters trained on up to 138B
tokens, the optimal learning rate/batch size pair (eta^{ast}, B^{ast})
consistently has the same operator norm value - a phenomenon we term norm
transfer. This constant norm condition is necessary but not sufficient: while
for each dataset size, multiple (eta, B) reach the optimal norm, only a
unique (eta^{ast}, B^{ast}) achieves the best loss. As a sufficient
condition, we provide the first measurement of (eta^{ast}, B^{ast})
scaling with dataset size for Scion, and find that the scaling rules are
consistent with those of the Adam optimizer. Tuning per-layer-group learning
rates also improves model performance, with the output layer being the most
sensitive and hidden layers benefiting from lower learning rates. We provide
practical insights on norm-guided optimal scaling and release our Distributed
Scion (Disco) implementation with logs from over two thousand runs to support
research on LLM training dynamics at scale.
68e4d8f5e4e093a7044e4ea9
Oleg Filatov
68e4d8f5e4e093a7044e4eaa
Jiangtao Wang
68e4d8f5e4e093a7044e4eab
Jan Ebert
68e4d8f5e4e093a7044e4eac
Stefan Kesselheim
Optimal Scaling Needs Optimal Norm
Despite recent progress in optimal hyperparameter transfer under model and
dataset scaling, no unifying explanatory principle has been established. Using
the Scion optimizer, we discover that joint optimal scaling across model and
dataset sizes is governed by a single invariant: the operator norm of the
output layer. Across models with up to 1.3B parameters trained on up to 138B
tokens, the optimal learning rate/batch size pair (eta^{ast}, B^{ast})
consistently has the same operator norm value - a phenomenon we term norm
transfer. This constant norm condition is necessary but not sufficient: while
for each dataset size, multiple (eta, B) reach the optimal norm, only a
unique (eta^{ast}, B^{ast}) achieves the best loss. As a sufficient
condition, we provide the first measurement of (eta^{ast}, B^{ast})
scaling with dataset size for Scion, and find that the scaling rules are
consistent with those of the Adam optimizer. Tuning per-layer-group learning
rates also improves model performance, with the output layer being the most
sensitive and hidden layers benefiting from lower learning rates. We provide
practical insights on norm-guided optimal scaling and release our Distributed
Scion (Disco) implementation with logs from over two thousand runs to support
research on LLM training dynamics at scale.
Joint optimal scaling of model and dataset sizes in deep learning is governed by the operator norm of the output layer, a phenomenon termed norm transfer, which provides a necessary condition for optimal learning rate and batch size.
Scion optimizer operator norm output layer norm transfer learning rate batch size Adam optimizer per-layer-group learning rates Distributed Scion (Disco) LLM training dynamics