TEON: Tensorized Orthonormalization Beyond Layer-Wise Muon for Large Language Model Pre-Training
The Muon optimizer has demonstrated strong empirical performance in pre-training large language models by performing matrix-level gradient (or momentum) orthogonalization in each layer independently. In this work, we propose TEON, a principled generalization of Muon that extends orthogonalization beyond individual layers by modeling the gradients of a neural network as a structured higher-order tensor. We present TEON's improved convergence guarantee over layer-wise Muon, and further develop a practical instantiation of TEON based on the theoretical analysis with corresponding ablation. We evaluate our approach on two widely adopted architectures: GPT-style models, ranging from 130M to 774M parameters, and LLaMA-style models, ranging from 60M to 1B parameters. Experimental results show that TEON consistently improves training and validation perplexity across model scales and exhibits strong robustness under various approximate SVD schemes.
Code (0)
등록된 구현이 없습니다.
Similar Papers 제목 키워드 기반
On Batch Orthogonalization Layers
Batch normalization has become ubiquitous in many state-of-the-art nets. It accelerates training and yields good performance results. However, there are various other alternatives to normalization, e.g. orthonormalizatio…
Efficient Finite Initialization for Tensorized Neural Networks
We present a novel method for initializing layers of tensorized neural networks in a way that avoids the explosion of the parameters of the matrix it emulates. The method is intended for layers with a high number of node…
Tensor NetworksA Note on Coding and Standardization of Categorical Variables in (Sparse) Group Lasso Regression
Categorical regressor variables are usually handled by introducing a set of indicator variables, and imposing a linear constraint to ensure identifiability in the presence of an intercept, or equivalently, using one of v…
regressionVariable SelectionImproving Molecular Graph Neural Network Explainability with Orthonormalization and Induced Sparsity
Rationalizing which parts of a molecule drive the predictions of a molecular graph convolutional neural network (GCNN) can be difficult. To help, we propose two simple regularization techniques to apply during the traini…
Graph Neural NetworkTensor Decomposition based Personalized Federated Learning
Federated learning (FL) is a new distributed machine learning framework that can achieve reliably collaborative training without collecting users' private data. However, due to FL's frequent communication and average agg…
DiversityFederated LearningModel OptimizationPersonalized Federated Learning+1