Autores: Daniel Braithwaite, Arissa Yoshida, Rafael Celente y Aman Gupta

En publicaciones anteriores [1,2,3], presentamos el enfoque de Nubank para utilizar modelos fundamentales basados en datos de transacciones con el fin de resolver problemas predictivos [4]. Estos textos describieron cómo estructuramos nuestros datos de transacciones para modelos fundamentales [2], cómo preentrenamos dichos modelos y, finalmente, cómo los ajustamos (joint fusion) para tareas específicas [3]. Es importante destacar que observamos grandes avances en tareas críticas para Nubank. El resultado más significativo fue que estas mejoras no provinieron del uso de fuentes adicionales de datos, sino del aprendizaje de representaciones óptimas de transacciones, en lugar de utilizar features diseñadas manualmente.

Aunque poderosos, estos modelos fundamentales son costosos de entrenar desde el punto de vista computacional. En Nubank, buscamos constantemente maneras de mejorar la eficiencia de datos, tanto para reducir costos como para construir modelos con mejor desempeño. En esta publicación, exploramos cómo un nuevo optimizador, Muon [5], nos está ayudando a alcanzar estos objetivos. Muon ha recibido gran atención en la comunidad de investigación en LLMs, especialmente por ser más eficiente en el uso de muestras para alcanzar un nivel de calidad fijo en el preentrenamiento, en comparación con AdamW (que hasta ahora ha sido la opción estándar en la mayoría de los entrenamientos de gran escala).

La calidad de nuestros modelos fundamentales aumenta en función de la cantidad de datos utilizados, llegando a más de 203 millones de filas. Por ejemplo, en la Figura 1 mostramos cómo el AUC del conjunto de prueba para uno de nuestros modelos más pequeños (24 millones de parámetros) escala en función del número de puntos de datos usados en joint fusion. Incluso mejoras pequeñas, como un incremento de 0,05% en el AUC, son altamente valiosas, ya que podrían significar ahorros de millones de dólares para Nubank. Sin embargo, a medida que el AUC mejora, también lo hace el costo de entrenamiento. El joint fusion [3] con 5 millones de filas tarda alrededor de 12 horas en 8 GPUs NVIDIA A100, mientras que con 40 millones de filas el tiempo se eleva a aproximadamente 95 horas, usando las mismas 8 A100s.

Figura 1 – La calidad del modelo mejora en función del tamaño del dataset

El costo computacional de entrenar estos modelos muestra la importancia de utilizar métodos que aumenten la eficiencia de datos. Al mismo tiempo, también significa que podemos lograr mejor desempeño con el mismo número de pasos de entrenamiento. Existen diversos métodos para mejorar la eficiencia de datos, pero en esta publicación exploramos el uso del optimizador Muon [5] para hacer que el preentrenamiento de nuestros modelos fundamentales sea más eficiente. Como consecuencia, estos modelos mejorados generan ahorro de costos y mejor rendimiento de producto para los clientes de Nubank.

El optimizador Muon [5] representa un cambio significativo frente a los enfoques heurísticos dominantes como AdamW, al introducir un método de optimización de segundo orden, simple y derivado de principios fundamentales. Diseñado específicamente para capas lineales densas de redes neuronales, el mecanismo central de Muon puede describirse como un descenso más pronunciado estructurado en matrices, con regularización mediante norma espectral. Su operación fundamental consiste en “ortogonalizar” la matriz de gradientes de cada capa de pesos, empujando los valores singulares para que estén cercanos a 1. Este proceso preserva la información direccional del gradiente mientras normaliza su magnitud en todas las direcciones, evitando que la optimización sea dominada por unos pocos componentes ruidosos o poco útiles. Este concepto, aunque elegante en teoría, se vuelve práctico gracias al uso de la iteración eficiente de Newton-Schulz [6], que aproxima la ortogonalización sin el costo computacional prohibitivo de una descomposición SVD completa.

Este diseño basado en principios se traduce directamente en ganancias sustanciales de eficiencia, tanto en datos como en cómputo. Las actualizaciones de momento ortogonalizado de Muon permiten pasos más estables y directos hacia el mínimo de la función de pérdida, y posibilitan que el modelo aprenda más de cada token procesado. Los beneficios son significativos también desde la perspectiva computacional: experimentos de scaling law muestran consistentemente que Muon puede alcanzar una calidad de modelo comparable a la de modelos entrenados con AdamW, consumiendo solo alrededor de la mitad (~52%) de los FLOPs de entrenamiento, lo que equivale a una mejora de aproximadamente 2x en eficiencia computacional [7,8].

Para probar nuestra hipótesis de que Muon puede generar mejores modelos fundamentales para Nubank, preentrenamos varios modelos de 330 millones de parámetros en un dataset de 20 millones de muestras. Comparamos el desempeño del optimizador Muon contra el ampliamente utilizado AdamW en cuatro tasas de aprendizaje diferentes: 1e-4, 2e-4, 1e-3 y 2e-3. La figura a continuación muestra estos resultados. Es importante notar que Muon converge significativamente más rápido que AdamW y alcanza soluciones con pérdidas de validación menores en todas las tasas de aprendizaje probadas.

Al consolidar todas las ejecuciones en una sola comparación, observamos que los tres mejores modelos fueron: Muon 1e-3, Muon 2e-3 y AdamW 1e-3. Cabe resaltar que los entrenamientos con Muon convergieron más rápido que el mejor resultado obtenido con AdamW. Estos hallazgos confirman nuestra hipótesis de que el uso de Muon puede entrenar mejores modelos fundamentales. Un punto adicional: las pérdidas de predicción del próximo token son inusualmente bajas para language modeling, ya que los tokens especializados utilizados en nuestros modelos fundamentales poseen un vocabulario potencial restringido.

En esta publicación mostramos las ventajas de integrar el optimizador Muon en el pipeline de preentrenamiento de los modelos fundamentales de Nubank. Al adoptar Muon, logramos una convergencia más rápida y una calidad de modelos superior en comparación con AdamW, desbloqueando mejoras en la eficiencia de datos y cómputo. Estos avances se traducen directamente en beneficios tangibles para Nubank: reducción de costos de entrenamiento y mejor desempeño de producto, lo que, en última instancia, brinda una mejor experiencia a nuestros clientes. Nuestros resultados confirman que técnicas de optimización sofisticadas como Muon son cruciales para expandir los límites de lo posible con modelos fundamentales a gran escala, garantizando que sigamos innovando de manera eficiente y efectiva.

Referencias

[1] Braithwaite, D., & Udagawa, H. (2025, March 24). Understanding our customers’ finances through foundation models. Building Nubank. https://building.nubank.com/understanding-our-customers-finances-through-foundation-models/

[2] Braithwaite, D., & Udagawa, H. (2025, April 22). Defining an interface between transaction data and foundation models. Building Nubank. https://building.nubank.com/defining-an-interface-between-transaction-data-and-foundation-models/

[3] Braithwaite, D., Cavalcanti, M., & Udagawa, H. (2025, May 14). Fine-tuning transaction user models. Building Nubank. https://building.nubank.com/fine-tuning-transaction-user-models/

[4] Braithwaite, D. T., Cavalcanti, M., McEver, R. A., et al (2025). Your Spending Needs Attention: Modeling Financial Habits with Transformers. arXiv preprint arXiv:2507.23267.

[5] Jordan, K., Jin, Y., Boza, V., You, J., Cesista, F., Newhouse, L., & Bernstein, J. (2024). Muon: An optimizer for hidden layers in neural networks. https://kellerjordan.github.io/posts/muon/

[6] Bernstein, J., & Newhouse, L. (2024). Old optimizer, new norm: An anthology. arXiv preprint arXiv:2409.20325.

[7] Shah, I., Polloreno, A. M., Stratos, K., Monk, P., Chaluvaraju, A., Hojel, A., … & Vaswani, A. (2025). Practical efficiency of muon for pretraining. arXiv preprint arXiv:2505.02222.

[8] Liu, J., Su, J., Yao, X., Jiang, Z., Lai, G., Du, Y., … & Yang, Z. (2025). Muon is scalable for LLM training. arXiv preprint arXiv:2502.16982.

Descubre las oportunidades