A Mean Field Ansatz for Zero-Shot Weight Transfer
This work addresses the problem of reducing computational costs for AI researchers and practitioners by offering a theoretical foundation for an incremental improvement in model scaling techniques.
The paper tackles the high cost of pre-training large language models by providing a theoretical explanation for zero-shot weight transfer, showing that weights of different-sized networks share a common distribution under certain assumptions, which was empirically validated on MLPs and LLMs like GPT-3 and Llama-3.1.
The pre-training cost of large language models (LLMs) is prohibitive. One cutting-edge approach to reduce the cost is zero-shot weight transfer, also known as model growth for some cases, which magically transfers the weights trained in a small model to a large model. However, there are still some theoretical mysteries behind the weight transfer. In this paper, inspired by prior applications of mean field theory to neural network dynamics, we introduce a mean field ansatz to provide a theoretical explanation for weight transfer. Specifically, we propose the row-column (RC) ansatz under the mean field point of view, which describes the measure structure of the weights in the neural network (NN) and admits a close measure dynamic. Thus, the weights of different sizes NN admit a common distribution under proper assumptions, and weight transfer methods can be viewed as sampling methods. We empirically validate the RC ansatz by exploring simple MLP examples and LLMs such as GPT-3 and Llama-3.1. We show the mean-field point of view is adequate under suitable assumptions which can provide theoretical support for zero-shot weight transfer.