Road to Generative AI - Part 3: A deep dive into model initialization

import matplotlib.pyplot as plt
import random
import torch
import torch.nn.functional as F
from unidecode import unidecode

DATASET_PATH = "./datasets/birds/birds.csv"

birds = open(DATASET_PATH, "r").read().splitlines()

print("First 10 birds in the dataset:\n")
print(", ".join(birds[:10]))
print(f"There are {len(birds):,d} birds in the dataset.")

min_length = map(len, birds)
max_length = map(len, birds)
print(f"\nThe shortest character name has {min(min_length)} characters.")
print(f"The longest character name has {max(max_length)} characters.")

from unidecode import unidecode

def clean_name(name):
    """
    Clean the bird name by:
    - Removing leading and trailing whitespaces
    - Converting to lowercase
    - Removing accents
    - Removing special characters
    - Replacing spaces with underscores
    """

    name = name.strip().lower()
    # replace special characters with a space
    name = ''.join(char if char.isalnum() or char.isspace() else ' ' for char in name)
    name = name.replace("`", "_")  # Remove apostrophes
    name = name.replace(" ", "_")
    name = unidecode(name)
    return name

# clean all names in the dataset
birds = list(map(clean_name, birds))

# create a mapping from tokens to indices
unique_tokens = set([c for w in birds for c in w])
SPECIAL_TOKEN = "."
index_to_token = {i: t for i, t in enumerate(unique_tokens, start=1)}
token_to_index = {v: k for k, v in index_to_token.items()}
index_to_token[0] = SPECIAL_TOKEN
token_to_index[SPECIAL_TOKEN] = 0

# log information about the tokenization
print(f"Number of unique tokens: {len(unique_tokens)}")
print(", ".join(sorted(unique_tokens)))
print(f"\nToken mapping: {index_to_token}")

# Model parameters
CONTEXT_SIZE = 3
N_EMBEDDINGS = 10
N_HIDDEN = 200

# Training parameters
TRAINING_SET_PORTION = 0.8
DEVELOPMENT_SET_PORTION = 0.1
TEST_SET_PORTION = 1 - (TRAINING_SET_PORTION + DEVELOPMENT_SET_PORTION)

def build_datasets(words: list[str]) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Build datasets from a list of words by creating input and target tensors.
    
    Args:
        words (list[str]): List of words to build the datasets from.
        
    Returns:
        tuple[torch.Tensor, torch.Tensor]: Tuple containing the input tensor X and target tensor Y.
    """
    # Create a mapping from tokens to indices
    X, Y = [], []
    
    # Create the context for each character in the words
    for w in words:
        context = [0] * CONTEXT_SIZE
        for ch in w + SPECIAL_TOKEN:  # Add special token at the end
            ix = token_to_index[ch]
            X.append(context)
            Y.append(ix)
            # Update the context by shifting it and adding the new index 
            context = context[1:] + [ix]

    # Convert lists to tensors
    X = torch.tensor(X, dtype=torch.int64)
    Y = torch.tensor(Y, dtype=torch.int64)

    return X, Y
    
# Shuffle the words
random.seed(1234)
random.shuffle(birds)

# Split the dataset into training, development, and test sets
train_size = int(TRAINING_SET_PORTION * len(birds))
dev_size = int(DEVELOPMENT_SET_PORTION * len(birds))

X_train, Y_train = build_datasets(birds[:train_size])
X_dev, Y_dev = build_datasets(birds[train_size:train_size + dev_size])
X_test, Y_test = build_datasets(birds[train_size + dev_size:])

# print tensor shapes
print("Training set shape:", X_train.shape, Y_train.shape)
print("Development set shape:", X_dev.shape, Y_dev.shape)
print("Test set shape:", X_test.shape, Y_test.shape)

def create_model(embedding_dim, context_size, layer_size, n_token,
                W1_scale: float = 1.0, b1_scale: float = 1.0,
                W2_scale: float = 1.0, b2_scale: float = 1.0):
    """
    Create a simple feedforward neural network model.
    This model consists of an embedding layer, a hidden layer, and an output layer.

    Args:
        -- Model architecture --
        embedding_dim (int): Dimension of the embedding layer.
        context_size (int): Size of the context window.
        layer_size (int): Size of the hidden layer.
        n_token (int): Number of unique tokens in the vocabulary.

        -- Normalization parameters --
        W1_scale (float): Scale factor for the weights of the first layer. Default is 1.0.
        b1_scale (float): Scale factor for the biases of the first layer. Default is 1.0.
        W2_scale (float): Scale factor for the weights of the second layer. Default is 1.0.
        b2_scale (float): Scale factor for the biases of the second layer. Default is 1.0.

    Returns:
        list[torch.Tensor]: List of model parameters including embeddings, weights, and biases.
            - C (torch.Tensor): Embedding matrix of shape (n_token, embedding_dim).
            - W1 (torch.Tensor): Weights for the first layer of shape (embedding_dim * context_size, layer_size).
            - b1 (torch.Tensor): Bias for the first layer of shape (layer_size,).
            - W2 (torch.Tensor): Weights for the second layer of shape (layer_size, n_token).
            - b2 (torch.Tensor): Bias for the second layer of shape (n_token,).
            - bngain (torch.Tensor): Gain for batch normalization of shape (1, layer_size).
            - bnbias (torch.Tensor): Bias for batch normalization of shape (1, layer_size).
    """

    g = torch.Generator().manual_seed(1234)

    # Embedding matrix
    C = torch.randn((n_token, embedding_dim), dtype=torch.float32, generator=g) # shape (28, 10)

    # Hidden layer
    W1 = torch.randn((embedding_dim * context_size, layer_size), dtype=torch.float32, generator=g) * W1_scale # shape (30, 128)
    b1 = torch.randn((layer_size,), dtype=torch.float32, generator=g) * b1_scale # shape (128,)

    # Output layer
    W2 = torch.randn((layer_size, n_token), dtype=torch.float32, generator=g) * W2_scale # shape (128, 28)
    b2 = torch.randn((n_token,), dtype=torch.float32, generator=g) * b2_scale # shape (28,)

    # Batch normalization parameters
    bngain = torch.ones((1, layer_size), dtype=torch.float32)  # Gain for batch normalization
    bnbias = torch.zeros((1, layer_size), dtype=torch.float32)

    bnmean_running = torch.zeros((1, layer_size), dtype=torch.float32)  # Running mean for batch normalization 
    bnstd_running = torch.ones((1, layer_size), dtype=torch.float32)  # Running std for batch normalization

    params = [C, W1, b1, W2, b2, bngain, bnbias]

    print(f"The model has {sum(p.nelement() for p in params):,d} parameters.")

    for p in params:
        p.requires_grad = True  # Set requires_grad to True to enable backpropagation

    # Add batch normalization parameters
    params += [bnmean_running, bnstd_running]

    return params

# Training parameters
MAX_STEPS = 200_000
BATCH_SIZE = 32
INITIAL_LEARNING_RATE = 0.1
LEARNING_RATE_DECAY = 0.1

def forward(params, X, use_batch_norm: bool = False):
    """
    Forward pass through the
    neural network model.
    Args:
        params (list[torch.Tensor]): List of model parameters including embeddings, weights, and biases.
        X (torch.Tensor): Input tensor of shape (BATCH_SIZE, CONTEXT_SIZE).
        use_batch_norm (bool): If True, applies batch normalization to the hidden layer activations.
    Returns:
        torch.Tensor: Output tensor of shape (BATCH_SIZE, n_token) representing the logits.
    """

    # Unpack parameters for easier access
    C, W1, b1, W2, b2, bngain, bnbias, bnmean_running, bnstd_running = params

    # Forward pass
    with torch.no_grad():
        emb = C[X]  # shape (BATCH_SIZE, CONTEXT_SIZE, N_EMBEDDINGS)
        embcat = emb.view(emb.shape[0], -1)  # concatenate the vectors
        hpreact = embcat @ W1 + b1  # shape (BATCH_SIZE, LAYER_SIZE)
        if use_batch_norm:
            hpreact = (hpreact - bnmean_running) / bnstd_running  # Batch normalization
            # Apply gain and bias
            hpreact = bngain * hpreact + bnbias  # shape (BATCH_SIZE, LAYER_SIZE)
        h = torch.tanh(hpreact)  # shape (BATCH_SIZE, LAYER_SIZE)
        logits = h @ W2 + b2  # shape (BATCH_SIZE, n_token)

    return logits

def train_model(params, X_train, Y_train, X_dev=None, Y_dev=None, break_after_1_step=False, use_batch_norm: bool = False):
    """
    Train the model using the provided parameters and training data.
    
    Args:
        params (list[torch.Tensor]): List of model parameters including embeddings, weights, and biases.
        X_train (torch.Tensor): Input tensor for training.
        Y_train (torch.Tensor): Target tensor for training.
        X_dev (torch.Tensor, optional): Input tensor for development set. Defaults to None.
        Y_dev (torch.Tensor, optional): Target tensor for development set. Defaults to None.
        break_after_1_step (bool, optional): If True, breaks after one training step
        use_batch_norm (bool, optional): If True, applies batch normalization. Defaults to False.
        
    Returns:
        lossi (list[float]): List of loss values recorded during training.
        h (Tensor): Hidden layer activations after the last forward pass.
    """

    # Unpack parameters for easier access
    C, W1, b1, W2, b2, bngain, bnbias, bnmean_running, bnstd_running = params

    # Create a generator for reproducibility
    g = torch.Generator().manual_seed(1234)

    lossi = []

    n_steps = MAX_STEPS

    for i in range(n_steps):

        # Minibatch index
        ix = torch.randint(0, X_train.shape[0], (BATCH_SIZE,), generator=g)  # Randomly sample BATCH_SIZE indices, shape (BATCH_SIZE,)
        Xb, Yb = X_train[ix], Y_train[ix]  # Get the corresponding input and target tensors

        # Forward pass
        emb = C[Xb]  # shape (BATCH_SIZE, CONTEXT_SIZE, N_EMBEDDINGS)
        embcat = emb.view(emb.shape[0], -1)  # concatenate the vectors
        
        # Linear layer
        hpreact = embcat @ W1 + b1  # shape (BATCH_SIZE, LAYER_SIZE)

        # Batch normalization
        if use_batch_norm:
            bnmean_i = hpreact.mean(dim=0, keepdim=True).mean()  # Mean for batch normalization
            bnstd_i = hpreact.std(dim=0, keepdim=True)  # Std for batch normalization
            hpreact = (hpreact - bnmean_i) / bnstd_i  # Batch normalization
            # Apply gain and bias
            hpreact = bngain * hpreact + bnbias  # shape (BATCH_SIZE, LAYER_SIZE)

            # Update running statistics for batch normalization
            with torch.no_grad():
                bnmean_running = bnmean_running * 0.999 + bnmean_i * 0.001
                bnstd_running = bnstd_running * 0.999 + bnstd_i * 0.001

        # Activation function
        h = torch.tanh(hpreact)  # shape (BATCH_SIZE, LAYER_SIZE)

        # Output layer
        logits = h @ W2 + b2  # shape (BATCH_SIZE, n_token)
        loss = F.cross_entropy(logits, Yb)  # compute the cross-entropy loss

        # Backward pass
        for p in params:
            p.grad = None  # Reset gradients to zero before backpropagation
        loss.backward()  # Compute gradients

        # Update parameters using gradient descent
        lr = INITIAL_LEARNING_RATE if i < 100_000 else INITIAL_LEARNING_RATE * LEARNING_RATE_DECAY
        for p in params:
            if p.grad is not None:
                p.data -= lr * p.grad

        # Logging
        if i % 10_000 == 0:
            print(f"{i:8,d}/{n_steps:8,d} | Loss: {loss.item():.4f} | Learning Rate: {lr:.6f}")

        lossi.append(loss.log10().item())

        if break_after_1_step and i == 0:
            print("Breaking after one step as requested.")
            break

    if X_dev is not None and Y_dev is not None:
        # Evaluate on the development set
        with torch.no_grad():
            dev_logits = forward(params, X_dev)
            dev_loss = F.cross_entropy(dev_logits, Y_dev)
            print(f"Development Loss: {dev_loss.item():.4f}")

    return lossi, h

# Create the model without any normalization
params = create_model(N_EMBEDDINGS, CONTEXT_SIZE, N_HIDDEN, len(token_to_index))
# Train the model
lossi, h = train_model(params, X_train, Y_train, X_dev=None, Y_dev=None, break_after_1_step=True)

expected_probability = 1 / len(token_to_index)
expected_loss = -torch.log(torch.tensor(expected_probability))
print(f"Expected probability for a random guess: {expected_probability:.4f}")
print(f"Expected loss for a random guess: {expected_loss.item():.4f}")

# Create the model and reduce the amplitude of the output layer
params = create_model(N_EMBEDDINGS, CONTEXT_SIZE, N_HIDDEN, len(token_to_index), W2_scale=0.1, b2_scale=0.0)
# Train the model
lossi, h = train_model(params, X_train, Y_train, X_dev=None, Y_dev=None, break_after_1_step=True)

fig, ax = plt.subplots(figsize=(8, 4))
plt.hist(h.view(-1).tolist(), 50)
plt.title("Histogram of hidden layer activations")
plt.xlabel("Activation value")
plt.ylabel("Frequency")
plt.show();

plt.figure(figsize=(12, 4))
plt.imshow(h.abs()>0.99, aspect='auto', cmap='gray')
plt.title("Hidden layer activations (absolute values > 0.99)")
plt.xlabel("Neuron index")
plt.ylabel("Batch index")
plt.colorbar(label="Activation > 0.99")
plt.show();

neuron_values = h.sum(dim=0).abs()
print(f"Minimum absolute neuron activation value: {neuron_values.min().item():.4f}")
print(f"Maximum absolute neuron activation value: {neuron_values.max().item():.4f}")

# Create the model and reduce the amplitude of the output layer and the hidden layer
params = create_model(N_EMBEDDINGS, CONTEXT_SIZE, N_HIDDEN, len(token_to_index), W2_scale=0.1, b2_scale=0.0, W1_scale=0.1, b1_scale=0.0)
# Train the model
lossi, h = train_model(params, X_train, Y_train, X_dev=None, Y_dev=None, break_after_1_step=True)

g = torch.Generator().manual_seed(56789)

class Linear:

    def __init__(self, fan_in, fan_out, bias=True, generator=g):
        self.weight = torch.randn(fan_in, fan_out, generator=generator) / fan_in ** 0.5
        self.bias = torch.zeros(fan_out) if bias else None

    def __call__(self, x):
        self.out = x @ self.weight
        if self.bias is not None:
            self.out += self.bias
        return self.out

    def parameters(self):
        return [self.weight] + ([] if self.bias is None else [self.bias])

class BatchNorm1d:

    def __init__(self, dim, eps=1e-5, momentum=0.1):
        self.momentum = momentum
        self.eps = eps
        self.training = True

        # parameters trained with backprop
        self.gamma = torch.ones(dim)
        self.beta = torch.zeros(dim)

        # buffers trained with momentum
        self.running_mean = torch.zeros(dim)
        self.running_var = torch.ones(dim)

    def __call__(self, x):
        if self.training:
            xmean = x.mean(0, keepdim=True) # shape (1, dim)
            xvar = x.var(0, keepdim=True, unbiased=True) # shape (1, dim)
        else:
            xmean = self.running_mean
            xvar = self.running_var

        xhat = (x - xmean) / torch.sqrt(xvar + self.eps)
        self.out = self.gamma * xhat + self.beta

        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):
        self.out = torch.tanh(x)
        return self.out

    def parameters(self):
        return []

n_emb = 10
n_hidden = 100
vocab_size = len(token_to_index)


# model initialization
C = torch.randn((vocab_size, n_emb), generator=g)
layers = [
    Linear(n_emb * CONTEXT_SIZE, n_hidden), BatchNorm1d(n_hidden), Tanh(),
    Linear(n_hidden,             n_hidden), BatchNorm1d(n_hidden), Tanh(),
    Linear(n_hidden,             n_hidden), BatchNorm1d(n_hidden), Tanh(),
    Linear(n_hidden,             n_hidden), BatchNorm1d(n_hidden), Tanh(),
    Linear(n_hidden,             n_hidden), BatchNorm1d(n_hidden), Tanh(),
    Linear(n_hidden,             vocab_size)
]


with torch.no_grad():
    # shrink confidence of last layer
    layers[-1].weight *= 0.1
    # apply tanh gain to other layers
    for layer in layers[:-1]:
        if isinstance(layer, Linear):
            layer.weight *= 5/3


parameters = [C] + [p for layer in layers for p in layer.parameters()]

# display number of total parameters
print(f"Total number of parameters: {sum(p.nelement() for p in parameters):,}")
for p in parameters:
    p.requires_grad = True

fig, ax = plt.subplots(figsize=(8, 4))
plt.hist(h.view(-1).tolist(), 50)
plt.title("Histogram of hidden layer activations")
plt.xlabel("Activation value")
plt.ylabel("Frequency")
plt.show();

neuron_values = h.sum(dim=0).abs()
print(f"Minimum absolute neuron activation value: {neuron_values.min().item():.4f}")
print(f"Maximum absolute neuron activation value: {neuron_values.max().item():.4f}")

x = torch.randn(1000, 10)
w = torch.randn(10, 200)
y = x @ w  # shape (1000, 200)
print(f"x mean = {x.mean().item():.4f}, std = {x.std().item():.4f}")
print(f"y mean = {y.mean().item():.4f}, std = {y.std().item():.4f}")

plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.hist(x.view(-1).tolist(), 50)
plt.title("Histogram of input activations")
plt.xlabel("Activation value")
plt.ylabel("Frequency") 
plt.xlim(-15, 15)
plt.subplot(1, 2, 2)
plt.hist(y.view(-1).tolist(), 50)
plt.title("Histogram of output activations")
plt.xlabel("Activation value")
plt.ylabel("Frequency")
# use same scale for the x-axis
plt.xlim(-15, 15)
plt.show();

x = torch.randn(1000, 10)
SCALING_FACTOR = (1 / x.shape[-1])**0.5
w = torch.randn(10, 200) * SCALING_FACTOR
y = x @ w  # shape (1000, 200)
print(f"x mean = {x.mean().item():.4f}, std = {x.std().item():.4f}")
print(f"y mean = {y.mean().item():.4f}, std = {y.std().item():.4f}")

plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.hist(x.view(-1).tolist(), 50)
plt.title("Histogram of input activations")
plt.xlabel("Activation value")
plt.ylabel("Frequency") 
plt.xlim(-15, 15)
plt.subplot(1, 2, 2)
plt.hist(y.view(-1).tolist(), 50)
plt.title("Histogram of output activations")
plt.xlabel("Activation value")
plt.ylabel("Frequency")
# use same scale for the x-axis
plt.xlim(-15, 15)
plt.show();

W1_scale = (5 / 3) / (N_EMBEDDINGS * CONTEXT_SIZE)**0.5

params = create_model(N_EMBEDDINGS, CONTEXT_SIZE, N_HIDDEN, len(token_to_index), W2_scale=0.1, b2_scale=0.0, W1_scale=0.1, b1_scale=0.0)
# Train the model
lossi, h = train_model(params, X_train, Y_train, X_dev=None, Y_dev=None, use_batch_norm=True)

def plot_loss(lossi):
    import matplotlib.pyplot as plt

    plt.figure(figsize=(10, 4))
    plt.plot(lossi, label='Loss (log scale)')
    plt.xlabel('Training Steps')
    plt.ylabel('Loss (log scale)')
    plt.title('Training Loss Over Time')
    plt.legend()
    plt.grid()
    plt.show()

plot_loss(lossi)

@torch.no_grad()
def split_loss(split, params, use_batch_norm=False):
    """
    Compute the loss for a given split of the data.

    Args:
        split (str): The split to compute the loss for.
        params (list): The model parameters.
        use_batch_norm (bool): Whether to use batch normalization.
    """
    
    x, y = {
        "train": (X_train, Y_train),
        "dev": (X_dev, Y_dev),
        "test": (X_test, Y_test)
    }[split]

    logits = forward(params, x, use_batch_norm)

    loss = F.cross_entropy(logits, y)  # compute the cross-entropy loss

    print(f"{split.capitalize()} Loss: {loss.item():.4f}")

split_loss("train", params)
split_loss("dev", params)

import re

def clean_name(name: str) -> str:
    # remove trailing "."
    name = name.rstrip(".")

    # replace "_s(_|$)" with "'s"
    name = re.sub(r"_s$", "'s", name)
    name = re.sub(r"_s_", "'s ", name)

    # replace "_" with " "
    name = name.replace("_", " ")

    # title each word
    name = " ".join(word.title() if len(word) > 1 else word for word in name.split())

    # replace "'S" with "'s"
    name = name.replace("'S", "'s")

    return name

def generate_names(params, n_names=10, use_batch_norm=False):
    """
    Generate new bird names using the trained model.
    
    Args:
        params (list[torch.Tensor]): List of model parameters including embeddings, weights, and biases.
        n_names (int): Number of names to generate. Default is 10.
        
    Returns:
        None
    """
    
    # sample new bird names from the model
    g = torch.Generator().manual_seed(56789)
    print("Generating new bird names:")
    for _ in range(n_names):
        context = [0] * CONTEXT_SIZE
        out = []
        while True:
            # Convert context to tensor and get embeddings
            context_tensor = torch.tensor(context, dtype=torch.int64).unsqueeze(0)  # shape (1, CONTEXT_SIZE)
            logits = forward(params, context_tensor, use_batch_norm)
            probs = F.softmax(logits, dim=1)  # shape (1, n_token)

            # Sample from the distribution
            ix = torch.multinomial(probs, num_samples=1, generator=g).item()  # Get the index of the sampled token
            out.append(index_to_token[ix])  # Append the token to the output

            if ix == token_to_index[SPECIAL_TOKEN]:  # Stop if we hit the special token
                break

            # Update the context by shifting it and adding the new index 
            context = context[1:] + [ix]

        print("   "+"".join(clean_name("".join(out))))  # Print the generated bird name

generate_names(params, n_names=10)

g = torch.Generator().manual_seed(56789)

class Linear:

    def __init__(self, fan_in, fan_out, bias=True, generator=g):
        self.weight = torch.randn(fan_in, fan_out, generator=generator) / fan_in ** 0.5
        self.bias = torch.zeros(fan_out) if bias else None

    def __call__(self, x):
        self.out = x @ self.weight
        if self.bias is not None:
            self.out += self.bias
        return self.out

    def parameters(self):
        return [self.weight] + ([] if self.bias is None else [self.bias])

class BatchNorm1d:

    def __init__(self, dim, eps=1e-5, momentum=0.1):
        self.momentum = momentum
        self.eps = eps
        self.training = True

        # parameters trained with backprop
        self.gamma = torch.ones(dim)
        self.beta = torch.zeros(dim)

        # buffers trained with momentum
        self.running_mean = torch.zeros(dim)
        self.running_var = torch.ones(dim)

    def __call__(self, x):
        if self.training:
            xmean = x.mean(0, keepdim=True) # shape (1, dim)
            xvar = x.var(0, keepdim=True, unbiased=True) # shape (1, dim)
        else:
            xmean = self.running_mean
            xvar = self.running_var

        xhat = (x - xmean) / torch.sqrt(xvar + self.eps)
        self.out = self.gamma * xhat + self.beta

        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):
        self.out = torch.tanh(x)
        return self.out

    def parameters(self):
        return []

n_emb = 10
n_hidden = 100
vocab_size = len(token_to_index)


# model initialization
C = torch.randn((vocab_size, n_emb), generator=g)
layers = [
    Linear(n_emb * CONTEXT_SIZE, n_hidden), BatchNorm1d(n_hidden), Tanh(),
    Linear(n_hidden,             n_hidden), BatchNorm1d(n_hidden), Tanh(),
    Linear(n_hidden,             n_hidden), BatchNorm1d(n_hidden), Tanh(),
    Linear(n_hidden,             n_hidden), BatchNorm1d(n_hidden), Tanh(),
    Linear(n_hidden,             n_hidden), BatchNorm1d(n_hidden), Tanh(),
    Linear(n_hidden,             vocab_size)
]


with torch.no_grad():
    # shrink confidence of last layer
    layers[-1].weight *= 0.1
    # apply tanh gain to other layers
    for layer in layers[:-1]:
        if isinstance(layer, Linear):
            layer.weight *= 5/3


parameters = [C] + [p for layer in layers for p in layer.parameters()]

# display number of total parameters
print(f"Total number of parameters: {sum(p.nelement() for p in parameters):,}")
for p in parameters:
    p.requires_grad = True

# training parameters
max_steps = 200_000
BATCH_SIZE = 32
loss_i = []
ud = []

# training block
for i in range(max_steps):

    # construct the minibatch
    idx = torch.randint(low=0, high=X_train.shape[0], size=(BATCH_SIZE,), generator=g)
    Xb, Yb = X_train[idx], Y_train[idx]

    # forward pass
    emb = C[Xb] # shape (BATCH_SIZE, CONTEXT_SIZE, N_EMBEDDINGS)
    x = emb.view(emb.shape[0], -1) # shape (BATCH_SIZE, CONTEXT_SIZE * N_EMBEDDINGS)
    for layer in layers:
        x = layer(x)

    # compute the loss
    loss = F.cross_entropy(x, Yb)

    # backward pass
    for layer in layers:
        layer.out.retain_grad()
    for p in parameters:
        p.grad = None
    loss.backward()

    # update
    lr = 0.075 if i < 100_000 else 0.01
    for p in parameters:
        p.data -= p.grad * lr

    # log tracking
    if i % 10_000 == 0:
        print(f"{i:8,d}/{max_steps:8,d} | Loss: {loss.item():.4f} | Learning Rate: {lr:.6f}")
    lossi.append(loss.log10().item())

    with torch.no_grad():
        ud.append(
            [(lr * p.grad.std() / p.data.std()).log10().item() for p in parameters]
        )

    if i == 1_000:
        break

# visualize the content of the hidden layers
plt.figure(figsize=(20, 4))
legends = []
for i, layer in enumerate(layers[:-1]):
    if isinstance(layer, Tanh):
        t = layer.out

        print("layer %d (%10s): mean %+.2f, std %.2f, saturated: %2.2f%%" % (i, layer.__class__.__name__, t.mean(), t.std(), (t > 0.97).float().mean() * 100.))
        hy, hx = torch.histogram(t, density=True)
        plt.plot(hx[:-1].detach().numpy(), hy.detach().numpy())
        legends.append(f"layer {i} ({layer.__class__.__name__})")

plt.legend(legends)
plt.title("Distribution of activations")
plt.show()

# visualize the content of the hidden layer gradients
plt.figure(figsize=(20, 4))
legends = []
for i, layer in enumerate(layers[:-1]):
    if isinstance(layer, Tanh):
        t = layer.out.grad

        print("layer %d (%10s): mean %+f, std %e" % (i, layer.__class__.__name__, t.mean(), t.std()))
        hy, hx = torch.histogram(t, density=True)
        plt.plot(hx[:-1].detach().numpy(), hy.detach().numpy())
        legends.append(f"layer {i} ({layer.__class__.__name__})")

plt.legend(legends)
plt.title("Distribution of gradients")
plt.show()

# visualize the content of the spread of the data over the spread of the gradient
plt.figure(figsize=(20, 4))
legends = []
for i, p in enumerate(parameters):
    t = p.grad
    if p.ndim == 2:
        print('weight %10s | mean %+f | std %e | grad:data ratio %e' % (tuple(p.shape), t.mean(), t.std(), t.abs().mean() / p.abs().mean()))
        hy, hx = torch.histogram(t, density=True)
        plt.plot(hx[:-1].detach().numpy(), hy.detach().numpy())
        legends.append(f"weight {i} ({tuple(p.shape)})")

plt.legend(legends)
plt.title("Weight gradient distribution")
plt.show()

plt.figure(figsize=(20, 4))
legends = []
for i, p in enumerate(parameters):
    if p.ndim == 2:
        plt.plot([ud[j][i] for j in range(min(1000, len(ud)))])
        legends.append('param %d' % i)
plt.plot([0, len(ud)], [-3, -3], 'k') # these ratios should be ~1e-3, indicate on plot
plt.legend(legends)
plt.xlabel("step")
plt.ylabel("log10(gradient std / data std)")
plt.title("Ratio of gradient std to data std over time");