Scalable Complexity Control Facilitates Reasoning Ability of LLMs

TL;DR

This paper shows that controlling model complexity through initialization and regularization improves the reasoning ability and scaling performance of large language models across different sizes and data amounts.

Contribution

It introduces a simple complexity control method via initialization rate and weight decay that enhances LLM scaling laws and reasoning capabilities.

Findings

Complexity control improves LLM scaling law performance.

Constant initialization rate accelerates scaling law descent.

Pretrained 2.4B models with different hyperparameters show consistent gains.

Abstract

The reasoning ability of large language models (LLMs) has been rapidly advancing in recent years, attracting interest in more fundamental approaches that can reliably enhance their generalizability. This work demonstrates that model complexity control, conveniently implementable by adjusting the initialization rate and weight decay coefficient, improves the scaling law of LLMs consistently over varying model sizes and data sizes. This gain is further illustrated by comparing the benchmark performance of 2.4B models pretrained on 1T tokens with different complexity hyperparameters. Instead of fixing the initialization std, we found that a constant initialization rate (the exponent of std) enables the scaling law to descend faster in both model and data sizes. These results indicate that complexity control is a promising direction for the continual advancement of LLMs.