Scaling Laws for Linear Complexity Language Models (2024)

Abstract: The interest in linear complexity models for large language models is on the rise, although their scaling capacity remains uncertain. In this study, we present the scaling laws for linear complexity language models to establish a foundation for their scalability. Specifically, we examine the scaling behaviors of three efficient linear architectures. These include TNL, a linear attention model with data-independent decay; HGRN2, a linear RNN with data-dependent decay; and cosFormer2, a linear attention model without decay. We also include LLaMA as a baseline architecture for comparison with softmax attention. These models were trained with six variants, ranging from 70M to 7B parameters on a 300B-token corpus, and evaluated with a total of 1,376 intermediate checkpoints on various downstream tasks. These tasks include validation loss, commonsense reasoning, and information retrieval and generation. The study reveals that existing linear complexity language models exhibit similar scaling capabilities as conventional transformer-based models while also demonstrating superior linguistic proficiency and knowledge retention.

## introduction the prosperity of large language models (llms) has necessitated the development of scaling laws @xcite to optimize the tradeoff between increasing model size and expanding training data within finite computational resources. scaling laws empirically study the correlation between model performance and factors including the number of parameters, training tokens, and flops. previous works @xcite @xcite have established power laws to describe ⋆ equal contribution. the corresponding author (email: zhongyiran@gmail.com). linear complexity language models @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite @xcite , have emerged as a promising alternative to traditional transformers in causal language modeling. however, the scalability of these models remains uncertain, which limits their applicability to large language models. to address this concern, in this paper we have developed pre-training scaling laws for efficient large language models. following the approach outlined in @xcite , we have used the training loss as a regression target to establish power law equations against flops and infer the optimal model size and dataset size under constant computation budgets for linear complexity models. our study focuses on investigating three efficient architectures, as detailed in section 2: tnl @xcite , hgrn2 @xcite , and cosformer2 @xcite @xcite . additionally, llama @xcite , in the master row, we present predicted training curves for various architectures, with each subsequent row representing a different architecture. on the left, the training curves for models ranging from 70m to 7b parameters are displayed. from these curves, we extract the envelope of minimum loss per flop, using these data points to estimate the optimal model size (center) for a specified compute budget, and the optimal number of training tokens (right). our contributions are summarized as follows: ## scaling laws the concept of scaling laws involves four key factors: loss l, specifically the cross-entropy loss in a causal language modeling setting; model size n , which is determined by the number of model parameters; dataset size d, calculated as the number of training tokens; and computation budget c, represented by the total flops used for training. n opt and d opt are the optimal model size and dataset size given a fixed computation budget. initially, we establish power law equations between l and c. in this analysis, we adopt the approach introduced by @xcite , treating the training loss as an unbiased estimate of the test loss. subsequently, based on the fitted curve, we ascertain the optimal loss for specific flops, enabling us to obtain coefficients for n opt ∝ c a and d opt ∝ c b . when modeling the scaling trend of loss against factors such as n , d, and c, the original scaling laws (kaplan et al., 2020) utilize the power function l(x) = (x 0 /x) α x , where x represents the factor of interest. subsequent studies @xcite @xcite employ a more general power-law plus constant form, l(x) = ϵ + (x 0 /x) α x , to achieve improved fitting. here, the constant ϵ is interpreted as irreducible loss or the entropy of natural text @xcite . in our case, we have simplified all forms of the power law and unified them into l(x) = β x x α x , which allows for a more intuitive comparison of the scaling capabilities of different models based on coefficients α x and β x . ## related work scaling laws in large language models aim for an ideal balance between increasing the number of parameters and enlarging the training corpus, given limited computation resources @xcite @xcite @xcite . the initial scaling laws @xcite use the test-time cross-entropy loss as a regression target to investigate its power-law correlations with model size, dataset size and training computation budget. hoffmann et al. (2022) use three approaches to find the optimal model size and dataset size given a fixed computation budget. by 1) freezing model size and varying number of training tokens, 2) fixing flops and changing model sizes and dataset sizes and 3) directly solving a constrained optimization equation, they conclude that models and the training corpus should be scaled equally when enlarging computing resources. they use the revised scaling law to train a compute-optimal model, chinchilla, that stands out across various benchmarks. other works extend scaling laws to multiple modalities @xcite , mixture of expert models @xcite and reinforcement learning @xcite . @xcite studied the influence of additional factors such as learning rate, context length, and batch size on the scaling-law coefficients. @xcite studies scaling laws of downstream task performance in a transfer learning setting for the machine translation task. ## conclusion our comprehensive study has demonstrated that linear complexity language models, including tnl, hgrn2, and cosformer2, exhibit competitive scaling capabilities akin to transformer-based models while also showcasing enhanced linguistic proficiency and knowledge retention. with rigorous training across a vast parameter range and extensive evaluation of diverse tasks, our findings validate these models as promising contenders for future large-scale language model development.