AI Safety (CS 2881) Fall 2026 Lecture 3: Modern LLM Training and Inference

Overview

• Speaker: Boaz Barak
• Channel: Boaz Barak
• Main topic: Modern LLM Training, Inference, Scaling Laws, RLHF, and Alignment
• Purpose: Provide computer science students with a rigorous technical foundation in modern LLM training pipelines, resource scaling trade-offs, and alignment techniques relevant to AI safety. This lecture covers the fundamentals of Large Language Model training and inference, exploring deep learning scaling laws, next-token prediction mechanics, transformer architecture and parallelism (tensor, layer, context, and expert sharding), reinforcement learning from human/AI feedback (RLHF/RLAIF), and deliberative alignment for reasoning models.

Topic Map

Scaling in Computer Science vs AI

• Explanation: Comparing the scale of general computer science with AI training runs using Dijkstra's perspective on conceptual hierarchies.
• Key claims:
  • Computing science deals with scale ratios around 10^9 (from bits to megabytes).
  • Large-scale AI training runs can involve 10^25 floating-point operations.
  • Programmers and AI researchers must manage exceptionally deep hierarchies of abstraction.
• Examples:
  • Comparing a supersonic jet plane with a crawling baby (~10^3 speed ratio)
  • Comparing an atom to a billiard ball (~10^9 ratio)
• Terminology:
  • conceptual hierarchies
  • baffling ratio
  • floating-point operations
  • FLOPs
• Why it matters: Highlights why AI models operate at scales vastly exceeding traditional software engineering, requiring novel methods of thinking about abstraction and resource allocation.

Deep Learning and Scaling Laws

• Explanation: Analyzing how intelligence gained per training resource scales in deep learning compared to classical machine learning.
• Key claims:
  • Classical learning methods show diminishing returns and saturate quickly when given infinite resources.
  • Deep learning methods scale effectively over massive increases in compute and data.
  • Chinchilla scaling laws define the optimal allocation between parameter count and token count.
• Examples:
  • Small-scale clever investor vs. multi-billion dollar index fund analogy
• Terminology:
  • Chinchilla scaling law
  • FLOP budget
  • parameter-data trade-off
  • area under the curve
• Why it matters: Explains why modern frontier models require massive capital and compute investments to achieve continuous gains in capability.

Next Token Prediction Mechanics

• Explanation: The core mathematical formulation of autoregressive language models predicting the next token from context.
• Key claims:
  • LLMs take input token sequences from a finite vocabulary and output a probability distribution over the next token.
  • Next token prediction serves as a powerful self-supervised learning signal over vast datasets.
  • Transformers embed tokens into high-dimensional vectors and map them via embedding and unembedding matrices.
• Examples:
  • Predicting x_{n+1} given tokens x_1, x_2, ..., x_n
• Terminology:
  • autoregressive
  • embedding matrix
  • unembedding matrix
  • temperature parameter
  • cross-entropy loss
  • stochastic parrot
• Why it matters: Provides the foundational model architecture and training objective underlying all modern generative language models.

Transformer Architecture and Hardware Parallelism

• Explanation: How transformers utilize GPU hardware through arithmetic intensity and various sharding strategies.
• Key claims:
  • Transformer operations are dominated by matrix multiplication, favoring high arithmetic intensity inside GPU cores.
  • Tensor sharding splits parameter matrices across GPUs.
  • Layer sharding distributes sequential model layers across multiple devices.
• Examples:
  • GPU 1 and GPU 2 communicating via all-reduce operations
• Terminology:
  • tensor sharding
  • layer sharding
  • context sharding
  • expert sharding
  • arithmetic intensity
  • KV cache
• Why it matters: Crucial for understanding how trillion-parameter models are distributed and trained efficiently across massive GPU clusters.

Training Stages: Pretraining and Post-Training

• Explanation: The multi-stage training pipeline from unsupervised token modeling to assistant alignment.
• Key claims:
  • Pretraining models the general token distribution using self-supervised next-token prediction on massive internet text.
  • Post-training teaches the model to act as a helpful and safe assistant.
  • Supervised Fine-Tuning (SFT) and Reinforcement Learning (RLHF/RLAIF) shape model behavior.
• Examples:
  • Transitioning from raw text prediction to dialogue assistant interactions
• Terminology:
  • pretraining
  • post-training
  • supervised fine-tuning
  • SFT
  • RLHF
  • RLAIF
• Why it matters: Clarifies how raw statistical predictors are transformed into conversational, instruction-following AI assistants.

Reinforcement Learning and Policy Gradients

• Explanation: Mathematical formulation of reinforcement learning, policy gradients, and the REINFORCE algorithm.
• Key claims:
  • Policy gradient methods maximize expected rewards for generated rollouts.
  • On-policy training samples responses from the current model policy.
  • Group Relative Policy Optimization (GRPO) normalizes rewards and incorporates KL penalties.
• Examples:
  • Using reward models to score model rollouts during alignment
• Terminology:
  • policy gradient
  • REINFORCE
  • GRPO
  • KL divergence
  • reward model
  • on-policy
• Why it matters: Forms the mathematical backbone of model alignment and reasoning enhancement via reinforcement learning.

AI Safety and Deliberative Alignment

• Explanation: Addressing safety vulnerabilities, refusals, and advanced alignment strategies for reasoning models.
• Key claims:
  • Classical safety training relies heavily on refusal-oriented responses and simple SFT/RLHF.
  • Reasoning models benefit from deliberative alignment where chain-of-thought reasoning is incorporated into safety checks.
  • Reinforcing the thinking process allows models to evaluate safety specifications internally before generating answers.
• Examples:
  • Decoding ROT13-encoded adversarial prompts while checking safety policies in the chain of thought
• Terminology:
  • deliberative alignment
  • harmlessness
  • jailbreak performance
  • chain of thought
  • specification compliance
• Why it matters: Outlines cutting-edge techniques for ensuring advanced reasoning models adhere to safety guidelines even under adversarial conditions.

Key Points

Deep learning scales efficiently with compute and data

• Explanation: Unlike classical algorithms that saturate quickly, deep learning intelligence scales predictably with increased FLOPs and dataset size.
• Evidence: Log-log scaling laws in transformer performance across model sizes.
• Practical implication: Training larger models on more data reliably yields better capabilities.

Matrix multiplication drives transformer hardware efficiency

• Explanation: GPUs excel at compute-heavy matrix multiplications, yielding high arithmetic intensity when communication overhead is minimized.
• Evidence: Transformer parameter sharding strategies across GPU clusters.
• Practical implication: Model architecture choices must align with hardware memory bandwidth and core execution limits.

Pretraining vs. Post-training separation

• Explanation: Models first learn statistical language distributions via pretraining, then learn agent behavior and persona via SFT and RL.
• Evidence: Standard industry training pipelines for modern LLMs.
• Practical implication: Capabilities are acquired during pretraining; alignment and usability are instilled during post-training.

Deliberative alignment enhances reasoning safety

• Explanation: Allowing models to reason through safety specifications in their chain of thought significantly improves robustness against sophisticated jailbreaks.
• Evidence: Evaluations showing o1-preview and reasoning models outperforming standard models on jailbreak resilience.
• Practical implication: Future safety mechanisms must evaluate internal reasoning pathways rather than just final outputs.

Frameworks, Models & Processes

Chinchilla Scaling Law

• How it works: Balances model parameter count (N) and training token count (D) against a fixed compute budget to minimize loss.
• Components:
  • Loss function L(D, F)
  • FLOP budget F
  • Optimal parameter-to-token ratio
• When to use: When planning large-scale pretraining runs to ensure optimal resource allocation between model size and data volume.

Policy Gradient (REINFORCE / GRPO)

• How it works: Maximizes expected reward by adjusting policy parameters in proportion to the gradient of log probabilities scaled by normalized rewards.
• Components:
  • Prompt distribution
  • Rollout generation policy
  • Reward function
  • KL divergence penalty
• When to use: During reinforcement learning post-training to align language models with human preferences or reasoning correctness.

Examples & Case Studies

Evaluating reasoning models on ROT13-encoded harmful prompts

• Illustrates: How reasoning models use internal chain-of-thought deliberation to evaluate safety policies before responding.
• Lesson: Models with internal monologue can detect and correctly refuse subtle adversarial bypass attempts.

Actionable Takeaways

• Immediate:
  • Understand the mechanics of next-token prediction and autoregressive sampling.
  • Recognize the trade-offs between tensor, layer, and context sharding in distributed GPU training.
• Strategic:
  • Scaling laws dictate that capability growth requires proportional scaling of compute, data, and parameters.
  • Post-training and alignment methods are essential for shaping raw predictors into safe, instruction-following assistants.
• Questions to investigate:
  • How do emergent reasoning capabilities scale relative to compute during reinforcement learning?
  • What are the limits of deliberative alignment against adaptive adversarial attacks?

Claims Worth Verifying

• Transformer arithmetic intensity is defined as computation divided by memory communication. (technical definition)
• Chinchilla scaling dictates that parameters and tokens should scale equally with compute. (empirical scaling law)

Notable Quotes

"The programmer is in the unique position that this is the only discipline and profession in which such a gigantic ratio, which totally baffles our imagination, has to be bridged by a single technology." (at 2:22) "You are a next token predictor, all of us because we live in time, we all of our actions, our speech, we generate it one second at the time." (at 216:48)

Compressed Summary

• LLMs are trained via autoregressive next-token prediction over massive token corpora.
• Scaling laws govern the optimal trade-off between model parameters and training data size.
• Hardware parallelism relies on tensor, layer, and expert sharding across GPU clusters.
• Alignment and safety are achieved through post-training techniques like SFT, RLHF, and deliberative reasoning.
• Keywords: scaling laws, transformers, token prediction, reinforcement learning, alignment, sharding
• Core insight: Modern LLMs scale predictably through autoregressive next-token prediction on massive compute, while advanced alignment techniques harness internal reasoning to ensure safety and capability.