Los pesos estructurados abaratan el entrenamiento con Muon a gran escala
En palabras sencillas
Entrenar un modelo de lenguaje requiere muchos recursos informáticos. Este artículo presenta Monarch–Muon, un optimizador que almacena algunos pesos del modelo en piezas más pequeñas y actualiza esas piezas con pasos que tienen en cuenta la estructura de las matrices.
En pruebas de una sola ejecución con dos tipos de modelos, obtuvo una pérdida final de validación menor que AdamW cuando ambos usaban la misma estructura de piezas. En mediciones aparte, el tiempo y la memoria por paso de entrenamiento fueron menores en el tamaño más grande evaluado, usando una H100.
La investigación
¿Puede el preentrenamiento de modelos de lenguaje usar actualizaciones que consideran la estructura matricial, capaces de mejorar la optimización, sin asumir todos los costes de almacenamiento y cálculo de los pesos densos? El artículo evalúa si representar ciertos pesos mediante factores estructurados puede abaratar esas actualizaciones y conservar su ventaja de optimización.
Monarch–Muon aplica actualizaciones que tienen en cuenta la estructura matricial a factores de pesos estructurados para conservar ventajas de optimización y reducir el almacenamiento y el coste de ortogonalización.
Monarch–Muon actualiza factores de pesos estructurados mediante pasos que consideran la estructura matricial. Los investigadores lo compararon con AdamW usando estructuras de factores equivalentes en preentrenamientos densos y de mezcla de expertos sobre FineWeb-Edu.
La pérdida final de validación fue menor con Monarch–Muon en las cuatro comparaciones de una sola ejecución. Esto sugiere que las actualizaciones que consideran la estructura matricial pueden ayudar con las estructuras evaluadas.
Las comparaciones de calidad son ejecuciones únicas con tamaños de modelo concretos, y las mediciones de sistemas por separado no determinan el equilibrio conjunto entre calidad y hardware en entrenamientos distribuidos más largos.
La prueba de convergencia supone actualizaciones polares exactas, mientras que el optimizador práctico las aproxima con un número finito de iteraciones Newton–Schulz.
Análisis
Detalles
- Autores del paper original
- Alexander Miasnikov, Ruslan Kabirov, Gleb Molodtsov, Dmitrii Vlashchuk, Vladimir Aletov
- Enviado a arXiv
- Licencia
- CC BY 4.0