Stanford CS229 Machine Learning Spring 2026 Lecture 13: LLMs and Next-Word Prediction Loss

Representation Learning and Embeddings

Representation learning aims to create a mapping (parameterized by $\theta$) that transforms raw input $x$—such as images, text, audio, or video—into a vector $v_{\theta}(x)$ in an $m$-dimensional Euclidean space. These vectors, commonly referred to as embeddings, are designed so that similar inputs are mapped to similar embeddings (points close to each other in space), while dissimilar inputs are mapped to distant embeddings.

Embeddings enable similarity search, where a system can identify similar items by finding the nearest neighbors of a query embedding in the vector space. While embeddings were previously used as inputs for linear classification (linear probing), modern applications favor end-to-end models or direct similarity search.

Tuning Embeddings: Supervised vs. Unsupervised Learning

Supervised Pre-training

In supervised pre-training, a neural network is trained on a large labeled dataset (e.g., ImageNet with 1,000 labels) to predict discrete classes. The representation is extracted from the penultimate layer (the last layer before the final classification head), and the classification head is discarded for downstream similarity tasks.

  • Limitation: This approach requires expensive, large-scale labeled datasets. If the label set is not diverse enough (e.g., only binary labels), the learned representations may not capture enough complex patterns to be useful for other tasks.

Contrastive Learning (Unsupervised)

Contrastive learning learns representations without labels by encouraging the model to recognize different versions of the same input as similar and different inputs as dissimilar.

  1. Augmentation: A single image $x$ is transformed into two different versions, $\hat{x}$ and $\tilde{x}$, using techniques like random cropping (the most critical), flipping, blurring, or adding Gaussian noise.
  2. Objective: The model is trained to minimize the distance between the embeddings of the same image's augmentations ($f_{\theta}(\hat{x})$ and $f_{\theta}(\tilde{x})$) while maximizing the distance between augmentations of different images ($x$ and $z$).

Addressing False Negatives: A potential issue arises when two different images (e.g., two different cats) are treated as a negative pair and pushed apart. However, in large datasets, the majority of random pairs are truly dissimilar (e.g., a cat and a chair), meaning the benefit of pushing random images apart outweighs the "collateral damage" of pushing similar but non-identical images apart.

The SimCLR Objective Function

To implement contrastive learning efficiently, a batch-based approach like SimCLR is used. The process involves sampling a batch of $B$ images, creating two augmentations for each, and constructing a similarity matrix based on the inner product (or cosine similarity) of the normalized embeddings.

Mathematical Formulation

For a given augmentation $\hat{x}_i$, the loss function treats the task as a multiclass classification problem: identifying which of the other $2B-1$ embeddings in the batch is the corresponding pair $\tilde{x}_i$.

The loss for a single column $i$ is expressed as: $$\mathcal{L}_i = -\log \frac{\exp(\text{sim}(\hat{x}_i, \tilde{x}i) / \tau)}{\sum{j=1}^{2B} \exp(\text{sim}(\hat{x}_i, x_j) / \tau)}$$

  • Numerator: Encourages the similarity of the positive pair (the same image's augmentations) to be large.
  • Denominator: Encourages the similarity of all other pairs (negative pairs) to be small.

Practical Optimizations

  • Hard Negative Mining: Randomly sampled negative pairs are often too easy to distinguish. "Hard negatives"—examples that look similar but are not the same (e.g., a photo of a child playing soccer vs. a text about the FIFA World Cup)—are used to make the loss function more demanding and the resulting representations more robust.
  • Data Source Sampling: Sampling the batch $B$ from the same data source (e.g., the same codebase or programming language) creates naturally harder negatives, preventing the loss from becoming too small too quickly.

Semantic Search and RAG

Implementing Semantic Search

Semantic search utilizes pre-computed embeddings to find the most relevant objects for a query $q$:

  1. Indexing: Embeddings for all documents $D_1, \dots, D_n$ are computed and stored in a vector database.
  2. Querying: The query $q$ is embedded using the same model $f_{\theta}$.
  3. Retrieval: The system finds the document $D_i$ that maximizes the inner product between $f_{\theta}(q)$ and $f_{\theta}(D_i)$.

Retrieval-Augmented Generation (RAG)

RAG is used to provide Large Language Models (LLMs) with proprietary or up-to-date information without the need for expensive fine-tuning.

  • Process: When a query is received, the system retrieves a small subset of relevant documents from a private corpus using semantic search. These documents are then provided as part of the prompt (context) to the LLM, which generates an answer based on both the query and the retrieved text.
  • Advantages over Fine-Tuning:
    • Modularity: The retrieval system is separate from the LLM.
    • Data Governance: Access controls can be applied at the retrieval stage (e.g., preventing certain users from retrieving sensitive memos).
    • Ease of Update: Deleting a document from the corpus immediately removes it from the model's potential answers, whereas "un-learning" information from a fine-tuned model is an open research problem.

Alternative Retrieval Methods

While semantic search is prevalent, some frontier models (e.g., from Anthropic) utilize LLMs to generate complex regular expressions for keyword search. This is particularly effective for structured data like code, where function names and file trees provide strict patterns for matching.

Sources