Plato’s Cave (Jan Pietersz Saenredam after Cornelis Cornelisz van Haarlem 1604)

II. Neural Networks

representation learningrepresentation learning universal approximatorsuniversal approximators double descentdouble descent

Namely, in §3.1.2 From Localist to Distributional Representations, we cover XYZ and in §3.1.3 From Linear Perceptrons to Non-Linear Multi Layer Perceptrons. we increase the expressivity of model architecture from linear perceptrons otherwise known as single-layer networks to non-linear multi-layer perceptronsmulti-layer perceptron or deep neural networksdeep neural network. Moreover, because deep networks ${\hat{f}}_{\theta }$ are no longer linear, in §3.1.4 From Direct to Iterative Optimization we transition from from direct optimizatio methods that solve for the exact solution to iterative ones that solve for approximate solutions.

After Part II. Neural Networks, we will be ready for Part III. Scaling Networks where we finally complete our quest in building our own language model such as ChatGPT by following Nanochat and subsequently, our own deep learning framework capable of running said language model like PyTorch 2, by implementing borscht.compile() and borscht.distributed(). Part three follows the 2020-2025 time period of what is colloquially known as the the age of scalingage of scaling where researchers focused on scaling up the generality of transformer architecture by adding additional phases after pre-trainingpre-training stochastic parrots. That is, the mid-trainingmid-training and post-trainingpost-training phases which respectively produce helpful assistantsassistants with reinforcement learning from human feedback and slow thinking reasonersreasoners with verifiable rewards.

But before we get there we have some work to do.
Let’s go.

II. Table of Contents

3. Self-Supervised Pre-Training with Deep Networks

In which we transition from manual feature engineering implemented by single layer networks to automatic feature learning by rediscovering various inductive biases that are parameterized by neural network architectures including the FFN, CNN, RNN, LSTM, and GPT. Such parameters are optimized with iterative methods using the approximation from the language of differential calculus afforded by borscht.func and borscht.Tensor.backward(). At the end, we will have rediscovered the so-called age of research for ourselves, which spanned form the years 2012 to 2019.

3.1 Learning Static Representations with FFNs

$\hookleftarrow$ Table of Contents

In which we transition from machine learning to deep learning, which focuses on learning distributional representations automatically with deep nets; introduce such deep nets as multi-layer perceptrons, the non-linear extensions of linear models with the ReLU and Embed primitives; transition our optimization techniques to handle non-linear functions with iterative descent methods.

3.1.1 From Feature Engineering to Representation Learning

$\hookleftarrow$ Table of Contents (§3.1 Learning Representations with FFNs)

Deep learningdeep learning builds off from such foundation, namely, programming of software 2.0 with the three step process of selecting a model architecture, scoring it with a loss, and optimizing such loss. However, deep learning focuses on the class of models ${\hat{f}}_{\theta }$ known as deep neural networks or multi-layer perceptrons that introduce many layers of non-linearities that make them well-suited for learning representationsrepresentation learningSee (Bengio et al. 2012) and (LeCun et al. 2015).

localist representationlocalist representation distributional representationdistributional representation

Localist representations such as one hot vectors are used to program the software 2.0 of machine learning, whereas distributional representationsdistributional representation such as embeddingsembeddings are used to program the software 2.0 of deep learning. In this chapter, we will learn about the neural network primitive of Embed.

Stanford CS224N: NLP with Deep Learning | Lecture 1 - Intro and Word Vectors

Stanford CS224N: NLP with Deep Learning | Lecture 2 - Word Vectors and Language Models

Namely, in §3.1.2 From Localist to Distributional Representations, we cover XYZ and in §3.1.3 From Linear Perceptrons to Non-Linear Multi Layer Perceptrons. we increase the expressivity of model architecture from linear perceptrons otherwise known as single-layer networks to non-linear multi-layer perceptronsmulti-layer perceptron or deep neural networksdeep neural network. Moreover, because deep networks ${\hat{f}}_{\theta }$ are no longer linear, in §3.1.4 From Direct to Iterative Optimization we transition from from direct optimizatio methods that solve for the exact solution to iterative ones that solve for approximate solutions.

In §3. Self-Supervised Pre-Training with Deep Networks, we will implement different neural network architectures that all express various inductive biasesinductive biases. Such inductive biases were all explored during the 2012-2019 time period of what is coloquially known as the age of researchage of research. Keeping the goal of next-token prediction constant, we start from the standard feedforward neural networksFFN (FFN), and then transition to more novel architectures such as convolutional neural networksCNN (CNN), recurrent neural networksRNN (RNN), and finally, generative pre-trained transformersGPT (GPT), specifically by reproducing GPT-2Presented in Language Models are Unsupervised Multitask Learners (Radford et al. 2019) following Nanogpt. Because deep neural networks are no longer linear functions of their parameters, we will need to transition from direct methods to iterative ones with descent-based searches. Starting from gradient descentgradient descent, momentummomentum, weight decayweight decay, AdamWAdamW.

Then, in §4. The Engine of Automatic Differentiation, you will evolve the implementation of borscht from a numerical linear algebra library like Numpy from §2. The Multidimensional Parallelized Tensor into a full blown batteries-included deep learning framework like TORCH which will be able to run the training and inference of the said nets just trained in §3. This means implementing the optimizers for neural network primitives whose evaluations are accelerated with massively parallel processors known as graphics processing unitsGPU (GPUs) and whose gradients are automatically evaluated with an automatic differentationautomatic differentiation engine.

Consult last few chapters of (Jurafsky and Martin 2026)If Jurafsky and Martin is not your cup of tea, take a look at (Eisenstein 2019)., and (Bishop and Bishop 2024)Similarly, you can also take a look at (Goodfellow et al. 2016) for deep learning.

junk

Newton’s classical mechanics must be understood before Einsten’s general relativity. Similarly, so too of Rosenblatt’s machine learning before that of Bengio, Hinton, and LeCun’s deep learning. First, we understand that language models such as Nanogpt, ChatGpt, and Claude are specified as next-token predictors $f:V^{\ast }\to V$ which although decide on a final classifying token $\mathbf{y}\in V$ for the given context $({\mathbf{x}}_1,{\mathbf{x}}_2,\dots ,{\mathbf{x}}_t)\in V^{\ast }$ (todo $V^t$?), have a step of inference by modeling language stochastically with a conditional random variable $Y|X_1,X_2,\dots ,X_N$ endowed with a probability distribution $p_{Y|X_1,X_2,\dots ,X_N}(\mathbf{y}|{\mathbf{x}}_1,{\mathbf{x}}_2,\dots ,{\mathbf{x}}_t)$ where such distribution corresponds to an array $\mathbf{y}\in {\mathbb{R}}^V$ which can be represented by a borscht.Tensor with .shape of (V). This was the first step in our transition from classical software 1.0 to the jazz of software 2.0, namely, by transitioning from variables to random variables in §1.1 From Certain to Uncertain Knowledge. (expand with coterell formalization)

Second, we understanding that such language models probabilistically specified as autoregressive next-token predictors $f$ are implemented by estimating the parameterization of their functions with weights $\mathbf{w}$ and bias $b$ so that $f_{\theta }$ where $\theta :=\{\mathbf{w},b\}$ is a function that takes in a one-hot vector $\mathbf{x}$ as input (expand, context?), evaluates the weighted sum with $\mathbf{w}$ which is the standard inner product of the dot product, and then either uses the dot product directly in the case of linear regression with §1.3 Parameterizing Classification with Linear Regression or further composes the output with a squashing function $\sigma (\cdot )$ in §1.4 Parameterizing Classification with Logistic Regression. In either case, both models $f$ are specified dimensionally, and, in the case of linear models, can be implemented with the direct optimization of linear algebra’s linear systems as well

3.1.2 From Linear Perceptrons to Non-Linear Multi Layer Perceptrons

$\hookleftarrow$ Table of Contents (§3.1 Learning Representations with FFNs)

• wed sep.23 after the 3.1 source of flow
• work on architecture 3.1.3 (merge the various representations) and optimizers 3.1.4
• including synthesis from intermezzo

junk1 With a solid understanding of how statistical learning recovers stochastic maps of data in higher dimensional coordinate spaces, you are now ready to start training your first deep neural networks. The subdiscipline of training deep neural nets follows the same practice of given some dataset $d=\{(x^{(i)},y^{(i)}){\}}_{i=1}^n$, declaring your model’s forward pass $\hat{y}=f(\mathbf{x};\theta )$, it’s scoring function $\ell (y,\hat{y},\theta )$ and resulting backward pass ${\ell }^{\prime }(y,\hat{y},\theta )$, and iteratively optimizing such score with gradient descent, where ${\theta }^{t+1}:={\theta }^t-\eta \nabla \mathcal{L}(d,\theta )$. The only difference as we’ll shortly see soon is that the inductive bias used to model the function $f$ includes non-linearities $\phi$ with many layers of composition. This second part of the book focuses on the unsupervised/self-supervised task of sequence learning building up to the transformer architecture behind GPT2, but let’s train our first nets using the same tasks from part one to ensure we understand the core foundation of deep neural nets.

junk2 In part 2 we modify and increase the expressivity of the function class by including non-linearities $\phi$. The feedforward neural network simply put is a series of linear and non-linear layers of the form $f:{\mathbb{R}}^{d_0}\to {\mathbb{R}}^{d_{l-1}}$ so $$f(\mathbf{x};w):=W^{(l-1)}\circ (\phi \circ W^{(l-2)})\circ \cdots \circ (\phi \circ W^{(0)})\circ \mathbf{x}$$

where $W^{(i)}\in {\mathbb{R}}^{d_{i-1}\times d_i}$, $\phi :{\mathbb{R}}^{d_i}\to {\mathbb{R}}^{d_{i+1}}$ is an elementwise nonlinearity, and $w:=(W^{(0)},\dots ,W^{(l-1)})$. Conceptually, the linear layers are performing linear transformations that rotate, reflect, shear, and scale space, whereas the nonlinear transformations perform transformations that squash and twist space.

junk3 Let’s now turn out attention to the function bodies of these $h^{[i]}$’s with a deep neural netork of $1$ hidden layer, carrying out the task of price regression and sentiment classification so that $f$ has the form $f:{\mathbb{R}}^d\to \mathbb{R}$. With the statistical learning foundation from part one, we will simply present the forward pass $f(\mathbf{x})$, the loss function $\mathcal{L}(d,\theta )$, and the backward pass ${\ell }^{\prime }(y,\hat{y},\theta )$.

Forward Pass

The functions of deep neural networks will take the form of $f(\mathbf{x}):=g\circ \varphi (\mathbf{x})$ with $g$ being a linear classifier on feature extractor

$$\begin{array}{rl}\varphi & :{\mathbb{R}}^{d_1}\to {\mathbb{R}}^{d_n}\\ \varphi (\mathbf{x})& :=h^{[l]}\circ h^{[l-1]}\circ \cdots h^{[1]}\circ h^{[0]}(\mathbf{x})\end{array}$$ where $l$ is the number of compositional layers, and each intermediate function $h^{[i]}$ has the form of $h^{[i]}:{\mathbb{R}}^{d_i}\to {\mathbb{R}}^{d_{i+1}}$.

But what is a neural network? | Deep learning chapter 1 (3Blue1Brown 2017)

The next natural question to ask then is whether the logistic regression model is considered a deep neural network? The answer is that technically yes, it can be considered a degenerative deep neural network with $0$ hidden layersIn the same way that a list can be considered a degenerative binary tree or graph. These so-called hidden layers automate the construction of the representation through learning, so that the model not only discovers the mapping from representation to output, but also the representation itself. The functions of deep neural networks will take the form of $f(\mathbf{x}):=g\circ \varphi (\mathbf{x})$ with $g$ being a linear classifier on feature extractor $\varphi :{\mathbb{R}}^{d_1}\to {\mathbb{R}}^{d_n}$, $\varphi (\mathbf{x}):=h^{[l]}\circ h^{[l-1]}\circ \cdots h^{[1]}\circ h^{[0]}(\mathbf{x})$ where $l$ is the number of compositional layers, and each intermediate function $h^{[i]}$ has the form of $h^{[i]}:{\mathbb{R}}^{d_i}\to {\mathbb{R}}^{d_{i+1}}$. Each hidden layer successively and gradually lifts the complexity and abstraction of the data’s representation.

Together, these two aspects of learning non-linear, representations form the essence of deep learning.

"""
model: Neural Language Models (Bengio et al. 2003) URL: https://www.jmlr.org/papers/volume3/bengio03a/bengio03a.pdf

Dimension key:
# windows
B: batch size
T: sequence length

# input/output
V: vocabulary size
E: embedding dimension
D: model dimension
"""
import torch
import matplotlib.pyplot as plt
%matplotlib inline
# from jaxtyping import ...
g = torch.Generator().manual_seed(1337) # for .randn()

B, T = 32, 3
V, E, D = 27, 10, 200

# step: 0/200000, loss 27.63208770751953
# -> expected loss = nll = p(c) = -torch.tensor(1/V=27).log() = 3.2958
# -> self.W = torch.randn() is sampling from N(0, 1)
# -> self.W * [gain/sqrt(D_in)] (torch.init_kaimingnormal())

# residuals + normalization + Adam/RMSprop has made initialization less fragile
# -> b/c initialization is fragile/intractable with *deep* neural networks

class Linear:
    def __init__(self, D_in, D_out, bias=True):
        self.W_DiDo = torch.randn((D_in, D_out), generator=g) * (5/3)/D_in**0.5 # kaiming init (He et al. 2015)
        self.b_Do = torch.zeros(D_out) if bias else None

    def __call__(self, X_Di):
        self.X_Do = X_Di @ self.W_DiDo
        if self.b_Do is not None:
            self.X_Do += self.b_Do
        self.out = self.X_Do
        return self.X_Do

    def parameters(self):
        return [self.W_DiDo] + ([] if self.b_Do is None else [self.b_Do])

class BatchNorm1D:
  
  def __init__(self, dim, eps=1e-5, momentum=0.1):
    self.eps = eps
    self.momentum = momentum
    self.training = True
    # parameters (trained with backprop)
    self.gamma = torch.ones(dim)
    self.beta = torch.zeros(dim)
    # buffers (trained with a running 'momentum update')
    self.running_mean = torch.zeros(dim)
    self.running_var = torch.ones(dim)
  
  def __call__(self, x):
    # calculate the forward pass
    if self.training:
      xmean = x.mean(0, keepdim=True) # batch mean
      xvar = x.var(0, keepdim=True) # batch variance
    else:
      xmean = self.running_mean
      xvar = self.running_var
    xhat = (x - xmean) / torch.sqrt(xvar + self.eps) # normalize to unit variance
    self.out = self.gamma * xhat + self.beta
    # update the buffers
    if self.training:
      with torch.no_grad():
        self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * xmean
        self.running_var = (1 - self.momentum) * self.running_var + self.momentum * xvar
    return self.out
  
  def parameters(self):
    return [self.gamma, self.beta]

class Tanh:
    def __call__(self, X_BD):
        self.X_BD = torch.tanh(X_BD)
        # plt.hist(self.X_BD.view(-1).tolist(), 50); # distribution of weights
        # plt.imshow(self.X_BD.abs() > 0.99, cmap='gray', interpolation='nearest') # vanishing gradients
        self.out = self.X_BD
        return self.X_BD
    
    def parameters(self):
        return []

model = [
    Linear(T * E, D, bias=False), BatchNorm1D(D), Tanh(),
    Linear(D, D, bias=False), BatchNorm1D(D), Tanh(),
    Linear(D, V, bias=False), BatchNorm1D(V)
]

C = torch.randn((V,E), generator=g)
params = [C] + [p for l in model for p in l.parameters()]
for p in params:
    p.requires_grad = True

To summarize, in order to transition from manually handcrafting features to automatically learning them, we transition our model parameterizations from single layer to multi-layer perceptrons with non-linearities that learn a progressive sequence of semantically meaningful linear transformations and bases upon optimization (which we’ll learn and implement in 3.1.3 From Direct to Iterative Optimization). In fact, because of such non-linearities, multi-layer perceptrons can in theory approximate any function as stated by the universal approximation theoremThis is the software 2.0 parallel to software 1.0’ turing machines which are capable of universal computation. See https://en.wikipedia.org/wiki/Universal_Turing_machine. Put shortly, a multi-layer perceptron is

$$\begin{array}{rl}{f}_{\theta }:V& \stackrel{I}{\to }{\mathbb{R}}^{|V|}\stackrel{W_1}{\to }{\mathbb{R}}^d\stackrel{\phi }{\to }{\mathbb{R}}^d\stackrel{W_2}{\to }{\mathbb{R}}^{|V|}\stackrel{\mathrm{softmax}}{\to }{\Delta }^{|V|}\\ x& \mapsto \mathrm{softmax}\circ W_2\circ \phi \circ W_1\circ I(x)\end{array}$$

where $W_1$ and $W_2$ are supported by ${\mathbb{R}}^{d\times |V|}$ and ${\mathbb{R}}^{|V|\times d}$ respectively, $\phi (\cdot ):=\tanh (\cdot )=\sinh (\cdot )/\cosh (\cdot )=\exp (x)-\exp (-x)/\exp (x)+\exp (-x)$, and of course $\text{softmax}(z{)}_i=\exp (z_i)/{\sum }_{j=1}^k\exp (z_j)$. To score such parameters, we continue to maintain the cross entropy loss because the overall classification from input to output space remains unchanged. That is, the function we’d like to recover is a next-token predictor that categorizes a possible next output token given an input sentence, and so the cross entropy loss of such a model is,

$$\begin{gathered} \begin{array}{rl}\ell :{\Delta }^{|V|}\times {\Delta }^{|V|}& \to \mathbb{R}\\ (p,q)& \mapsto p(x)\log q(x)\end{array} \\ \begin{array}{rl}\mathcal{L}:\Theta & \to \mathbb{R}\\ \theta & \mapsto -\frac{1}{n}\sum _{(x^{(i)},y^{(i)})}\ell [I(y^{(i)}),f_{\theta }(x^{(i)})]\end{array} \end{gathered}$$

where $\theta :=(W_1,W_2)$ and so $\Theta ={\mathbb{R}}^d\times {\mathbb{R}}^d$. Now that we’ve parameterized next-token prediction with a multi-layer perceptron and have scored such parameters with the cross-entropy loss, let us move on to it’s optimization.

3.1.3 From Direct to Iterative Optimization

$\hookleftarrow$ Table of Contents (§3.1 Learning Representations with FFNs)

import borscht
import borscht.nn as nn
import borscht.nn.functional as F

class Linear(nn.Module):
  def __init__(self, Di: int, Do: int):
    self.W_DiDo = borscht.randn((Di,Do))
    self.b_Do = borscht.randn((Do))

  def __call__(self, x_Di: borscht.Tensor):
    out_Do = x_Di @ self.W_DiDo
    if self.b_Do is not None:
      out_Do += self.b_Do

    return out

class Tanh():
  def __call__(self, X_BD: borscht.Tensor):
    return borscht.tanh(X_BD)

V = 2
T = 1
E = 10

model = [Linear(T*E, D), Tanh()]
yhat = model(x)
loss_scalar = F.cross_entropy(yhat,y)

Similar to the linear perceptron used as a probabilistic model, the direct method of gaussian elimination does not suffice. Why so? (insert explanation)

Stanford CS229 Machine Learning | Lecture 8: Neural Networks 2 (Backprop) Stanford CS224N: NLP with Deep Learning | Lecture 3 - Backpropagation, Neural Network

Rather than find the optimal set of parameters $\theta$ in closed form solution directly, we need to transition our optimization method from an iterativeiterative optimization one that incrementally finds increasingly better approximations of the optimum with a sequence $[{\theta }^{(1)},{\theta }^{(2)},\dots ,{\theta }^{(t)}]$ that takes $t$ steps. Take a look at the 3Blue1Brown visual lecture below for an intuition pump on where we’re going this chapter.

Similarly, before we iteratively optimize the original cross entropy loss function of our multi-layer perceptron which produces a real-valued score for many parameters that have support in d-dimensional Euclidean space, let us consider loss functions of the one and two dimensional case first. That is, consider some arbitrary non-linear loss function that scores a single parameter

import borscht
def my_1dloss(theta: borscht.Tensor) -> borscht.Tensor:
    loss_scalar = borscht.Tensor(
                    ((theta+2)**2 *
                    (theta-1)**2 *
                    (theta-3)**2) + 0.3
                  )
    return loss_scalar

theta_scalar = borscht.randn()
my_1dloss_scalar = my_1dloss(theta_scalar())

as well as it’s corresponding plot of a curvecurve

import matplot
plt.plot(my_1dloss)

Instead of solving for the $\theta$ that minimizes $\mathcal{L}$ in closed-form directly, we can solve so iteratively. Recalling the derivative’sderivative geometric perspective from high school calculus, the derivative of a function at a certain point is the line tangenttangent to the output, which is the instantaneous rate of change between two points of a secantsecant line. This secant line corresponds to the analytical perspective of the derivativederivative defined by the limitlimit.

Geometrically speaking, starting from some randomly initialized parameter $\theta$ with borscht.randn(todo), we can take iterative steps $[{\theta }^{(1)},{\theta }^{(2)},\dots ,{\theta }^{(t)}]$ towards the minimum by __.

Now, consider another arbitrary non-linear loss function that scores not a single parameter, but rather two parameters

import borscht
def my_1dloss(theta: borscht.Tensor) -> borscht.Tensor:
    loss_scalar = borscht.Tensor(
                    ((theta+2)**2 *
                    (theta-1)**2 *
                    (theta-3)**2) + 0.3
                  )
    return loss_scalar

theta_scalar = borscht.randn()
my_1dloss_scalar = my_1dloss(theta_scalar())

as well as it’s corresponding plot of a surfacesurface

gradientgradient with borscht.grad()

derivativederivative defined by the limitlimit as a linear operator

import matplot
plt.plot(my_1dloss)

3.1_ffn_torch.ipynb

# 1. dataloader
import torch.nn.functional as F
import random

words = open('./data/names.txt', 'r').read().splitlines()
v = sorted(list(set(''.join(words))))
encode = { c:i+1 for i,c in enumerate(v) }
encode['.'] = 0
decode = { i:c for c,i in encode.items() }

def gen_dataset(words):
    X, Y = [], []
    for w in words:
        context = [0] * T;
        for c in w + '.':
            X.append(context)
            Y.append(encode[c])
            # print(''.join(decode[i] for i in context), '-->', decode[encode[c]])
            context = context[1:] + [encode[c]]
    X, Y = torch.tensor(X), torch.tensor(Y) # X:(N,C) Y:(N)
    return X, Y

random.seed(42)
random.shuffle(words)
n1, n2 = int(0.8*len(words)), int(0.9*len(words))
Xtr, Ytr = gen_dataset(words[:n1])
Xdev, Ydev = gen_dataset(words[n1:n2])
Xte, Yte = gen_dataset(words[n2:])

# 2. training loop
N = Xtr.shape[0]
losses, steps = [], []
for step in range(200000):
    # 1. forward
    indices_B = torch.randint(0, N, (B,))
    X_B, Y_B = Xtr[indices_B], Ytr[indices_B]

    X_BD = C[X_B].view(-1, T * E)
    for layer in model:
        X_BD = layer(X_BD)
    loss = F.cross_entropy(X_BD, Y_B)

    # 2. backward
    for layer in model:
        layer.out.retain_grad()
    for p in params:
        p.grad = None
    loss.backward()

    # 3. update
    for p in params:
        p.data += -0.01 * p.grad

    steps.append(step)
    losses.append(loss.log10().item())
    if step % 10000 == 0:
        print(f"step: {step}/{200000}, loss {loss.item()}")

plt.plot(steps, losses)

step: 0/200000, loss 3.811568260192871
step: 10000/200000, loss 2.625439405441284
step: 20000/200000, loss 2.426504135131836
step: 30000/200000, loss 1.8564293384552002
step: 40000/200000, loss 2.563091993331909
step: 50000/200000, loss 1.8718260526657104
step: 60000/200000, loss 2.189483642578125
step: 70000/200000, loss 1.9642211198806763
step: 80000/200000, loss 2.261333465576172
step: 90000/200000, loss 2.2367687225341797
step: 100000/200000, loss 1.7542214393615723
step: 110000/200000, loss 1.9069358110427856
step: 120000/200000, loss 2.2940409183502197
step: 130000/200000, loss 2.030090808868408
step: 140000/200000, loss 2.556252956390381
step: 150000/200000, loss 2.0693819522857666
step: 160000/200000, loss 2.2053565979003906
step: 170000/200000, loss 2.3926899433135986
step: 180000/200000, loss 2.4394421577453613
step: 190000/200000, loss 2.0519824028015137

Out[2]:

[<matplotlib.lines.Line2D at 0x11099cce0>]

3.2 Iterative Optimization

$\hookleftarrow$ Table of Contents

In which we delve deeper into foundations of gradient-based iterative optimization, building ourselves up to the AdamW optimizer.

3.2.1 Taylor Series and their Derivatives, Gradients, and Hessians

3.2.2 First Order Optimization

3.2.3 Second Order Optimization

3.2.4 Momentum

3.2.5 Weight Decay

3.3 Learning Representations with CNNs

$\hookleftarrow$ Table of Contents

Lorem ipsum dolor sit amet, consectetur adipiscing elit. Proin finibus interdum arcu, ut venenatis dolor rutrum eget. Nulla nec urna tellus. Donec sollicitudin hendrerit sollicitudin. Proin sagittis, velit commodo congue rhoncus, metus turpis hendrerit nisi, a accumsan quam ante ultricies sem. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Ut aliquet ut justo non posuere. Quisque in rutrum nunc. Phasellus vehicula tortor mi, vehicula placerat diam ultrices eget. Aliquam et consequat eros, quis fermentum tortor. Sed eu nisl viverra, auctor velit vel, elementum ante. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Proin aliquet metus sit amet nulla cursus maximus in feugiat mi. Praesent vehicula magna sed scelerisque elementum. Nulla euismod rutrum turpis sed hendrerit.

Lorem ipsum dolor sit amet, consectetur adipiscing elit. Proin finibus interdum arcu, ut venenatis dolor rutrum eget. Nulla nec urna tellus. Donec sollicitudin hendrerit sollicitudin. Proin sagittis, velit commodo congue rhoncus, metus turpis hendrerit nisi, a accumsan quam ante ultricies sem. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Ut aliquet ut justo non posuere. Quisque in rutrum nunc. Phasellus vehicula tortor mi, vehicula placerat diam ultrices eget. Aliquam et consequat eros, quis fermentum tortor. Sed eu nisl viverra, auctor velit vel, elementum ante. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Proin aliquet metus sit amet nulla cursus maximus in feugiat mi. Praesent vehicula magna sed scelerisque elementum. Nulla euismod rutrum turpis sed hendrerit.

3.4 Learning Representations with RNNs

$\hookleftarrow$ Table of Contents

Lorem ipsum dolor sit amet, consectetur adipiscing elit. Proin finibus interdum arcu, ut venenatis dolor rutrum eget. Nulla nec urna tellus. Donec sollicitudin hendrerit sollicitudin. Proin sagittis, velit commodo congue rhoncus, metus turpis hendrerit nisi, a accumsan quam ante ultricies sem. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Ut aliquet ut justo non posuere. Quisque in rutrum nunc. Phasellus vehicula tortor mi, vehicula placerat diam ultrices eget. Aliquam et consequat eros, quis fermentum tortor. Sed eu nisl viverra, auctor velit vel, elementum ante. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Proin aliquet metus sit amet nulla cursus maximus in feugiat mi. Praesent vehicula magna sed scelerisque elementum. Nulla euismod rutrum turpis sed hendrerit.

class RNNCell(nn.Module):
  """
  the job of a 'Cell' is to:
  take input at current time step x_{t} and the hidden state at the
  previous time step h_{t-1} and return the resulting hidden state
  h_{t} at the current timestep
  """
  def __init__(self, config):
    super().__init__()
    self.xh_to_h = nn.Linear(config.n_embd + config.n_embd2, config.n_embd2)

  def forward(self, xt, hprev):
    xh = torch.cat([xt, hprev], dim=1)
    ht = F.tanh(self.xh_to_h(xh))
    return ht

class RNN(nn.Module):
  def __init__(self, config, cell_type):
    super().__init__()
    self.block_size = config.block_size
    self.vocab_size = config.vocab_size
    self.start = nn.Parameter(torch.zeros(1, config.n_embd2)) # the starting hidden state
    self.wte = nn.Embedding(config.vocab_size, config.n_embd) # token embeddings table
    if cell_type == 'rnn':
        self.cell = RNNCell(config)
    elif cell_type == 'gru':
        self.cell = GRUCell(config)
    self.lm_head = nn.Linear(config.n_embd2, self.vocab_size)

  def get_block_size(self):
    return self.block_size

  def forward(self, idx, targets=None):
    device = idx.device
    b, t = idx.size()

    # embed all the integers up front and all at once for efficiency
    emb = self.wte(idx) # (b, t, n_embd)

    # sequentially iterate over the inputs and update the RNN state each tick
    hprev = self.start.expand((b, -1)) # expand out the batch dimension
    hiddens = []
    for i in range(t):
      xt = emb[:, i, :] # (b, n_embd)
      ht = self.cell(xt, hprev) # (b, n_embd2)
      hprev = ht
      hiddens.append(ht)

    # decode the outputs
    hidden = torch.stack(hiddens, 1) # (b, t, n_embd2)
    logits = self.lm_head(hidden)

    # if we are given some desired targets also calculate the loss
    loss = None
    if targets is not None:
      loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)

    return logits, loss

Lorem ipsum dolor sit amet, consectetur adipiscing elit. Proin finibus interdum arcu, ut venenatis dolor rutrum eget. Nulla nec urna tellus. Donec sollicitudin hendrerit sollicitudin. Proin sagittis, velit commodo congue rhoncus, metus turpis hendrerit nisi, a accumsan quam ante ultricies sem. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Ut aliquet ut justo non posuere. Quisque in rutrum nunc. Phasellus vehicula tortor mi, vehicula placerat diam ultrices eget. Aliquam et consequat eros, quis fermentum tortor. Sed eu nisl viverra, auctor velit vel, elementum ante. Orci varius natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Proin aliquet metus sit amet nulla cursus maximus in feugiat mi. Praesent vehicula magna sed scelerisque elementum. Nulla euismod rutrum turpis sed hendrerit.

3.5 Learning Dynamic Representations with GPTs

$\hookleftarrow$ Table of Contents

In which we transition from long-range depencies of recurrent neural nets to the ___ of transformer-based neural nets.

3.5.1 Dynamic Embeddings with Attention

$\hookleftarrow$ Table of Contents (§3.5 Learning Representations with GPTs)

fluffy blue creature roamed the verdant forest. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.

Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.

microgpt: https://gist.github.com/karpathy/8627fe009c40f57531cb18360106ce95

class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        D, H = config.n_embd, config.n_head
        self.ln_1 = nn.LayerNorm(D)
        self.attn = MHA(config)
        self.mlp = FFN(config) # .mlp for HF
        self.ln_2 = nn.LayerNorm(D)

    def forward(self, X_BTD):
        # residuals:
        # - (He et al. 2015 https://arxiv.org/abs/1512.03385)
        # - (Elhage et al. 2021 https://transformer-circuits.pub/2021/framework/index.html)
        X_BTD = X_BTD + self.attn(self.ln_1(X_BTD))
        X_BTD = X_BTD + self.mlp(self.ln_2(X_BTD))
        return X_BTD

class MHA(nn.Module):
    def __init__(self, config):
        super().__init__()
        T, D, H = config.block_size, config.n_embd, config.n_head
        assert D % H == 0

        self.H = H

        self.c_attn = nn.Linear(D, 3 * D)
        self.c_proj = nn.Linear(D, D)
        self.c_proj.GPT2_SCALE_INIT = 1
        self.register_buffer('bias', torch.tril(torch.ones(T, T)).view(1, 1, T, T)) # tril -> bias for HF

    def forward(self, X_BTD):
        B,T,D = X_BTD.shape
        H = self.H
        # 1. project to learned QKV subspaces Q=WqX, K=WkX, V=WvX
        Wq_DK, Wk_DK, Wv_DK = self.c_attn(X_BTD).split(D, dim=2)
        Q_BHTK, K_BHTK, V_BHTK = Wq_DK.view(B, T, H, D // H).transpose(1, 2), Wk_DK.view(B, T, H, D // H).transpose(1, 2), Wv_DK.view(B, T, H, D // H).transpose(1, 2)

        # 2. evaluate scores A(QKV) = softmax(QK^T/sqrt(d_k))V
        # A_BHTT = Q_BHTK @ K_BHTK.transpose(-2, -1) * (1.0 / math.sqrt(K_BHTK.size(-1)))
        # A_BHTT = A_BHTT.masked_fill(self.bias[:, :, :T, :T]==0, float('-inf'))
        # A_BHTT = F.softmax(A_BHTT, dim=-1) # todo, when dim=-1?
        # S_BHTD = A_BHTT @ V_BHTK
        S_BHTD = F.scaled_dot_product_attention(Q_BHTK, K_BHTK, V_BHTK, is_causal=True)

        # 3. contextualize the embeddings
        S_BTD = S_BHTD.transpose(1, 2).contiguous().view(B, T, D) # performs cat
        S_BTD = self.c_proj(S_BTD)

        return S_BTD

3.5.2 Storing Facts with FFN

$\hookleftarrow$ Table of Contents (§3.5 Learning Representations with GPTs)

class FFN(nn.Module):
    def __init__(self, config):
        super().__init__()
        D = config.n_embd
        self.c_fc = nn.Linear(D, 4*D) # projecting up to extract features from context embeddings
        self.gelu = nn.GELU(approximate='tanh') # (Hendrycks and Gimpel 2016 https://arxiv.org/abs/1606.08415)
        self.c_proj = nn.Linear(4*D, D) # projecting back down to residual pathway
        self.c_proj.GPT2_SCALE_INIT = 1

    def forward(self, X_BTD):
        X_BT4D = self.c_fc(X_BTD)
        X_BT4D = self.gelu(X_BT4D)
        X_BTD = self.c_proj(X_BT4D)
        return X_BTD

3.5.2 Embedding $W_E$ and Unembedding $W_U$

$\hookleftarrow$ Table of Contents (§3.5 Learning Representations with GPTs)

3.5.3 Optimizing Loss with AdamW on Cross Entropy

$\hookleftarrow$ Table of Contents (§3.5 Learning Representations with GPTs)

Intermezzo Three: The Language of Optimization and Approximation

The learning process continues to be a three-step one where we specify a parameterization, a scoring of such parameters, and the optimization of such scoring function

$$\begin{array}{rl}& \mathcal{L}(\theta ):=\stackrel{\text{step two}}{\overbrace{\ell (\theta )}}\circ \stackrel{\text{step one}}{\overbrace{{\hat{f}}_{\theta }(x)}}\\ \stackrel{\text{step three}}{\overbrace{\underset{\theta }{\mathrm{argmin}}}}& \mathcal{L}(\theta )\end{array}$$

While automatically learning features for next-token prediction maintains the cross entropy loss for the scoring of parameters, the parameterization changes so that optimizing such loss of it’s parameters if a process such that a progressive sequence of semantically meaningful linear transformations that squash, shear, and rotate n-dimensional Euclidean space is learned

import borscht
import borscht.nn as nn
import borscht.nn.functional as F

class MLP(nn.Module):
    def __init__(self):
      raise NotImplementedError("")

    def __call__(self, x_E: borscht.Tensor):
      raise NotImplementedError("")

model = MLP(Din, Dout)

# random forward pass
x = borscht.Tensor()
y_logits = model(x)
y_distribution = F.cross_entropy_loss(y_logits)

and now that the model class is non-linear, said optimization changes to an iterative process of gradient descent.

model = MLP(Din, Dout)

for t in range(100):
  # 1. forward
  x = borscht.Tensor()
  y_logits = model(x)
  y_distribution = F.cross_entropy_loss(y_logits, ___)

  # 2. backward
  y_distribution.backward()
  optim.step()

For this iterative optimization process, we earlier appealed to the intuitive geometric perspective of the derivative, where in the one and two-dimensional cases, the derivative of the loss function is viewed as the line and plane that is tangent to the function’s output and points in the direction up the curve and surface’s hill. Although we were unable to visualize the gradient of the n-dimensional case of the actual loss function with n being the size of the parameter space, we could generalize the gradient by appending a partial derivative for each additional parameter that the loss function takes in as an argument. That is,

print(loss.backward())

Now that we intuitively understand how the notions of descent and sensitivity from the language of optimization and approximation are used in the training deep neural networks from the geometric and computational perspective, let us analytically formalize such notions into precise concepts.

I3.1 Sequences

I3.2 Limits of Sequences

I3.3 Limits of Functions

I3.3 Derivatives as Linear Operators

4. The Engine of Automatic Differentiation

$\hookleftarrow$ Table of Contents

automatic differentiationautomatic differentiation optimizersoptimizers network primitivesnetwork primitives parallel programmingparallel programming massively parallel processorsmassively parallel processors graphics processing unit (GPU)GPU CUDA RustCUDA Rust PTXPTX SASSSASS

Finish strong young hacker.

4.1 Automatic Differentation

$\hookleftarrow$ Table of Contents

In which we …

4.1.1 Algorithmic Differentiation

$\hookleftarrow$ Table of Contents (§4.1 Automatic Differentiation)

4.1.2 Implementing Backward for FFN

$\hookleftarrow$ Table of Contents (§4.1 Automatic Differentiation)

4.1.5 Implementing Backward for GPT

$\hookleftarrow$ Table of Contents (§4.1 Automatic Differentiation)

4.2 Gradient-Based Optimization

$\hookleftarrow$ Table of Contents

Tensor.backward() function in order to materialize the gradient of the loss on our parameter weights $w=(W^{(0)},\dots ,W^{(l-1)})$ which gradient descent uses in it’s update rule ${\theta }^{t+1}:={\theta }^t-\alpha \nabla \mathcal{L}(\theta )$. We will now dive deeper into how .backward() is implemented. –>

4.1.4 Stochastic Gradient Descent

$\hookleftarrow$ Table of Contents

4.1.5 Adam

$\hookleftarrow$ Table of Contents (§4.1 Automatic Differentiation)

4.2.3 Implementing Forward Pass for CNN

4.2.4 Implementing Forward Pass for RNN

4.2.5 Implementing Forward Pass for GPT

missing notebook: /opt/build/repo/sitp/notebooks/4.1_makemore.ipynb

missing notebook: /opt/build/repo/sitp/notebooks/4.1_makemore.ipynb

4.3 From Strides to Layouts

$\hookleftarrow$ Table of Contents

4.4 Accelerating Attention

$\hookleftarrow$ Table of Contents