Methods of improving LLM training stability

• URL: http://arxiv.org/abs/2410.16682v1
• Date: Tue, 22 Oct 2024 04:27:03 GMT
• Title: Methods of improving LLM training stability
• Abstract summary: We observe that with a high learning rate the L2 norm of all linear layer outputs can grow with each training step and the model diverges. We show that with the last two methods we can increase learning rate by 1.5x in comparison to an approach based on QK layer normalization only.
• Score: 3.3790712267402374
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Training stability of large language models(LLMs) is an important research topic. Reproducing training instabilities can be costly, so we use a small language model with 830M parameters and experiment with higher learning rates to force models to diverge. One of the sources of training instability is the growth of logits in attention layers. We extend the focus of the previous work and look not only at the magnitude of the logits but at all outputs of linear layers in the Transformer block. We observe that with a high learning rate the L2 norm of all linear layer outputs can grow with each training step and the model diverges. Specifically we observe that QKV, Proj and FC2 layers have the largest growth of the output magnitude. This prompts us to explore several options: 1) apply layer normalization not only after QK layers but also after Proj and FC2 layers too; 2) apply layer normalization after the QKV layer (and remove pre normalization). 3) apply QK layer normalization together with softmax capping. We show that with the last two methods we can increase learning rate by 1.5x (without model divergence) in comparison to an approach based on QK layer normalization only. Also we observe significant perplexity improvements for all three methods in comparison to the baseline model.

Related papers

• Hidden Decoding at Scale: Latent Computation Scaling for Large Language Models [75.80275843320511]
  We study whether an existing backbone can keep improving by allocating more computation to each token while leaving the Transformer backbone fixed.<n>We propose Hidden Decoding, a sequence-length scaling method applied during continued pretraining.
  arXiv  Detail & Related papers  (2026-07-09T07:37:59Z)
• A Stable Whitening Optimizer for Efficient Neural Network Training [101.89246340672246]
  Building on the Shampoo family of algorithms, we identify and alleviate three key issues, resulting in the proposed SPlus method.<n>First, we find that naive Shampoo is prone to divergence when matrix-inverses are cached for long periods.<n>Second, we adapt a shape-aware scaling to enable learning rate transfer across network width.<n>Third, we find that high learning rates result in large parameter noise, and propose a simple iterate-averaging scheme which unblocks faster learning.
  arXiv  Detail & Related papers  (2025-06-08T18:43:31Z)
• Do Language Models Use Their Depth Efficiently? [53.56816097840505]
  We analyze the residual stream of the Llama 3.1 and Qwen 3 family of models.<n>We find that layers in the second half contribute much less than those in the first half.<n>For multihop tasks, we are unable to find evidence that models are using increased depth to compose subresults.
  arXiv  Detail & Related papers  (2025-05-20T04:00:56Z)
• RSQ: Learning from Important Tokens Leads to Better Quantized LLMs [65.5558181902098]
  Layer-wise quantization is a key technique for efficiently compressing large models without expensive retraining. We propose RSQ (Rotate, Scale, then Quantize), which applies rotations to the model to mitigate outliers. We demonstrate that RSQ consistently outperforms baseline methods across multiple downstream tasks and three model families.
  arXiv  Detail & Related papers  (2025-03-03T18:46:33Z)
• LESA: Learnable LLM Layer Scaling-Up [57.0510934286449]
  Training Large Language Models (LLMs) from scratch requires immense computational resources, making it prohibitively expensive. Model scaling-up offers a promising solution by leveraging the parameters of smaller models to create larger ones. We propose textbfLESA, a novel learnable method for depth scaling-up.
  arXiv  Detail & Related papers  (2025-02-19T14:58:48Z)
• The Curse of Depth in Large Language Models [28.37870372690079]
  In large language models, nearly half of the layers are less effective than expected.<n>LayerNorm Scaling (LNS) scales the variance of output of the layer normalization inversely by the square root of its depth.<n>LNS consistently outperforms previous normalization and scaling techniques in enhancing LLM pre-training performance.
  arXiv  Detail & Related papers  (2025-02-09T07:03:36Z)
• TrAct: Making First-layer Pre-Activations Trainable [65.40281259525578]
  We consider the training of the first layer of vision models and notice the clear relationship between pixel values and update magnitudes. An image with low contrast has a smaller impact on learning than an image with higher contrast. A very bright or very dark image has a stronger impact on the weights than an image with moderate brightness.
  arXiv  Detail & Related papers  (2024-10-31T14:25:55Z)
• LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding [13.747101397628887]
  We present an end-to-end solution to speed-up inference of large language models (LLMs) We apply layer dropout, with low dropout rates for earlier layers and higher dropout rates for later layers, and an early exit loss where all transformer layers share the same exit. We show that this training recipe increases the accuracy of early exit at earlier layers, without adding any auxiliary layers or modules to the model.
  arXiv  Detail & Related papers  (2024-04-25T16:20:23Z)
• Layer-wise Linear Mode Connectivity [52.6945036534469]
  Averaging neural network parameters is an intuitive method for the knowledge of two independent models. It is most prominently used in federated learning. We analyse the performance of the models that result from averaging single, or groups.
  arXiv  Detail & Related papers  (2023-07-13T09:39:10Z)
• Layer Grafted Pre-training: Bridging Contrastive Learning And Masked Image Modeling For Label-Efficient Representations [130.05189514598996]
  Mask Image Modeling (MIM) and Contrastive Learning (CL) demonstrate that self-supervision is powerful to learn good representations. In this paper, we make the empirical observation that a naive joint optimization of CL and MIM losses leads to conflicting gradient directions. Inspired by experimental observations, we find that MIM and CL are suitable to lower and higher layers, respectively. We propose a surprisingly simple, "sequential cascade" fashion: early layers are first trained under one MIM loss, on top of which latter layers continue to be trained under another CL loss.
  arXiv  Detail & Related papers  (2023-02-27T20:52:10Z)
• Improved Convergence Guarantees for Shallow Neural Networks [91.3755431537592]
  We prove convergence of depth 2 neural networks, trained via gradient descent, to a global minimum. Our model has the following features: regression with quadratic loss function, fully connected feedforward architecture, RelU activations, Gaussian data instances, adversarial labels. They strongly suggest that, at least in our model, the convergence phenomenon extends well beyond the NTK regime''
  arXiv  Detail & Related papers  (2022-12-05T14:47:52Z)
• NormFormer: Improved Transformer Pretraining with Extra Normalization [31.250781670447257]
  NormFormer architecture adds three normalization operations to each layer. Extra operations incur negligible compute cost. For masked language modeling, NormFormer improves fine-tuned GLUE performance by 1.9%.
  arXiv  Detail & Related papers  (2021-10-18T16:47:45Z)
• LV-BERT: Exploiting Layer Variety for BERT [85.27287501885807]
  We introduce convolution into the layer type set, which is experimentally found beneficial to pre-trained models. We then adopt an evolutionary algorithm guided by pre-training accuracy to find the optimal architecture. LV-BERT model obtained by our method outperforms BERT and its variants on various downstream tasks.
  arXiv  Detail & Related papers  (2021-06-22T13:20:14Z)
• Continual Learning in Deep Networks: an Analysis of the Last Layer [7.405588182540833]
  Three factors affect catastrophic forgetting in the output layer: (1) weights modifications, (2) interferences, and (3) projection drift. We show that the best-performing output layer type depends on the data distribution drifts or the amount of data available.
  arXiv  Detail & Related papers  (2021-06-03T13:41:29Z)
• Training with Multi-Layer Embeddings for Model Reduction [0.9046327456472286]
  We introduce a multi-layer embedding training architecture that trains embeddings via a sequence of linear layers. We show that it allows reducing d by 4-8X, with a corresponding improvement in memory footprint, at given model accuracy.
  arXiv  Detail & Related papers  (2020-06-10T02:47:40Z)
• Lipreading using Temporal Convolutional Networks [57.41253104365274]
  Current model for recognition of isolated words in-the-wild consists of a residual network and Bi-directional Gated Recurrent Unit layers. We address the limitations of this model and we propose changes which further improve its performance. Our proposed model results in an absolute improvement of 1.2% and 3.2%, respectively, in these datasets.
  arXiv  Detail & Related papers  (2020-01-23T17:49:35Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.