Author: Justin Y. Chen

This blog post covers material for the second lecture of Harvard CS 2881r: AI Safety taught by Boaz Barak. This is an outline of the lecture, see the video and slides for more details.

Links

Author Intro

Hi! I’m Justin, a PhD student at MIT in the Theory of Computation Group. I work on research problems at the intersection of algorithm design and machine learning. Recently, I have been crafting algorithms which interact with powerful but untrustworthy AI models (learning-augmented algorithms and LLM evaluation) or are stable to small changes in their inputs (differential privacy and replicability).

Lecture 2: Modern LLM Training

The methods which have succeeded in training modern LLMs have been those which can scale with more data and more compute. The successful approaches are not necessarily designed to make optimal use of available resources but can easily be extended to more and more resources without plateauing in performance. In brief, the predominant techniques are those which are simple, not stupid, and scale. This blog will explore modern machine learning architecture and training for text generation with a focus on why high-level design choices have been made and without delving deeply into technical details.

Model Architecture

The key task which LLMs are based on is next-token prediction (NTP): given a sequence of tokens in a corpus of text, predict the most likely next word to occur. This paradigm for text generation is amenable to scaling. The data available for next-token prediction is ubiquitous, any existing piece of text can be turned into training data simply by masking out tokens and asking the model to predict the missing word. Furthermore, next-token prediction across a myriad of environments is challenging, so models have room to improve and get better with more data and model capacity.

The key to modern machine learning computation is parallelism across multiple axes and scales. Modern LLMs are powered by GPUs, massively parallelizable machines which can perform staggering amounts of computation per second. Computation on GPUs is generally bottlenecked by memory requirements and data movement. Given the scale of LLMs, clusters of up to hundreds of thousands of GPUs are used in tandem during training. Parallelism is leveraged across tokens, training batches, weight layers, and e...

Posts

Wikitag Contributions

Comments