Large language models have been shown to exhibit surprisingly general-purpose language capabilities. Critical to their performance is the sheer scale at which they are trained, with continuous improvements coming from increasingly larger models. However, training such models from scratch is extremely resource-intensive. This proposal aims to investigate methods for more efficiently training large language models by leveraging smaller language models that have already been pretrained. In particular, we propose to exploit the linear algebraic structure of smaller, pretrained networks to better initialize and train the parameters of larger networks.

This work is a collaboration between MIT faculty Yoon Kim and Amazon researchers Rakesh Chada and Sriram Venkatapathy. Funded by Amazon Alexa.

With this project, we propose an algorithm-system co-design framework to make on-device training possible with only 256KB of memory. On-device training faces two unique challenges: (1) the quantized graphs of neural networks are hard to optimize due to low bit-precision and the lack of normalization; (2) the limited hardware resource (memory and computation) does not allow full back- propagation. To cope with the optimization difficulty, we propose Quantization-Aware Scaling to calibrate the gradient scales and stabilize 8-bit quantized training. To reduce the memory footprint, we propose Sparse Update to skip the gradient computation of less important layers and sub-tensors.

This project is led by MIT faculty Song Han. Partly funded by Amazon Alexa.