A Mean Field Ansatz for Zero-Shot Weight Transfer

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Chen, Xingyuan, Kuang, Wenwei, Deng, Lei, Han, Wei, Bai, Bo, Reis, Goncalo dos
Format: Preprint
Publié: 2024
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866909289425141760
author Chen, Xingyuan
Kuang, Wenwei
Deng, Lei
Han, Wei
Bai, Bo
Reis, Goncalo dos
author_facet Chen, Xingyuan
Kuang, Wenwei
Deng, Lei
Han, Wei
Bai, Bo
Reis, Goncalo dos
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2408_08681
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle A Mean Field Ansatz for Zero-Shot Weight Transfer
Chen, Xingyuan
Kuang, Wenwei
Deng, Lei
Han, Wei
Bai, Bo
Reis, Goncalo dos
Machine Learning
Numerical Analysis
Probability
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.
title A Mean Field Ansatz for Zero-Shot Weight Transfer
topic Machine Learning
Numerical Analysis
Probability
url https://arxiv.org/abs/2408.08681