Scaling Probabilistic Transformer via Efficient Cross-Scale Hyperparameter Transfer

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Kuang, Penghao, Wu, Haoyi, Tu, Kewei
Natura: Preprint
Pubblicazione: 2026
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866918471165542400
author Kuang, Penghao
Wu, Haoyi
Tu, Kewei
author_facet Kuang, Penghao
Wu, Haoyi
Tu, Kewei
contents Probabilistic Transformer (PT), a white-box probabilistic model for contextual word representation, has demonstrated substantial similarity to standard Transformers in both computational structure and downstream task performance on small models and small to medium sized datasets. However, PT is less robust to hyperparameter choices than standard Transformers, making it harder to scale efficiently. In this work, we follow Maximal Update Parametrization (muP) to rescale PT's parameters, so that hyperparameters optimized on small models can be transferred to larger models without additional tuning. With this approach, we successfully scale PT to models with up to 0.4B parameters. Experiments show that PT consistently outperforms standard transformer under the same parameter budget on Masked Language Modeling (MLM) tasks. We hope this work will contribute to the practical deployment of probabilistic models at substantially larger scales in the future.
format Preprint
id arxiv_https___arxiv_org_abs_2604_25409
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Scaling Probabilistic Transformer via Efficient Cross-Scale Hyperparameter Transfer
Kuang, Penghao
Wu, Haoyi
Tu, Kewei
Computation and Language
Probabilistic Transformer (PT), a white-box probabilistic model for contextual word representation, has demonstrated substantial similarity to standard Transformers in both computational structure and downstream task performance on small models and small to medium sized datasets. However, PT is less robust to hyperparameter choices than standard Transformers, making it harder to scale efficiently. In this work, we follow Maximal Update Parametrization (muP) to rescale PT's parameters, so that hyperparameters optimized on small models can be transferred to larger models without additional tuning. With this approach, we successfully scale PT to models with up to 0.4B parameters. Experiments show that PT consistently outperforms standard transformer under the same parameter budget on Masked Language Modeling (MLM) tasks. We hope this work will contribute to the practical deployment of probabilistic models at substantially larger scales in the future.
title Scaling Probabilistic Transformer via Efficient Cross-Scale Hyperparameter Transfer
topic Computation and Language
url https://arxiv.org/abs/2604.25409