AI LLM

Understanding The Machinery That Trains Deep Learning Models

Deep learning has reshaped the landscape of machine learning. It achieved overwhelming advantages across diverse domains, from image classification, video and text generation, to enabling multi-agent development (agents use foundation autoregressive models trained with SGD as their brains). In order to reduce the cognitive burden on practitioners, the inner

5 min read
Understanding The Machinery That Trains Deep Learning Models

Deep learning has reshaped the landscape of machine learning. It achieved overwhelming advantages across diverse domains, from image classification, video and text generation, to enabling multi-agent development (agents use foundation autoregressive models trained with SGD as their brains).

In order to reduce the cognitive burden on practitioners, the inner workings of deep learning are hidden behind sophisticated deep learning frameworks like PyTorch (which makes deep learning models unfortunately being regarded as black boxes). This includes the forward pass and automatic differentiation.

While PyTorch lowers the bar for people and students in general who want to implement a deep learning model, developed an interest in understanding deep neural networks or seek to reproduce the results of a given paper, understanding how deep neural networks work enables other opportunities, like developing new ideas or deep learning models, and especially knowing what can and cannot be implemented.

Basically, a deep learning network is comprised of a cascade of layers. This includes linear layers, which attempt to learn projection matrices, and convolutional neural layers, which aim at learning kernels contracted with image patches.

Other layers exist that do not have anything to learn, like activation functions. Activation functions introduce non-linearity into the network (which also makes the overall DNN non-convex), thereby enabling the learning of new representations of the inputs in order to get closer to a given target; for example, learning a new representation where the data is linearly separable. Without activation functions, the output must live in the hyperplane spanned by the inputs which is often not enough.

source: https://arxiv.org/abs/1804.00057

Scalar loss functions define how bad the model is at predicting and are used to update the model parameters. Specifically, the gradient with respect to each parameter tells us its rate of change and points in the direction of the highest increase, so we move in the opposite direction when updating the parameters. There is a whole dance required to update the model parameters using the chain rule, and this is automated using autograd, where basically during the forward pass, i.e., when the input enters the deep neural network, a computation graph is created. Then, when reaching the tail, i.e., the loss, the gradient is backpropagated according to the chain rule.

Regarding the loss, in deep learning and also machine learning, the data of interest is often assumed to be sampled from an unknown distribution, and the goal is to learn more about this distribution, i.e., approximate it using a family of probability distributions \( \{ p_{\theta} \}\) (cross-entropy is often used as a measure of divergence). The deep neural network would parametrize this distribution, and the goal is to use the empirical distribution and maximize the likelihood of the data (or the log-likelihood); i.e., we want to learn parameters that are more likely to have generated the data (sometimes we maximize a lower bound instead as computing the likelihood is not always tractable). This is why we need a lot of data in order to converge and hopefully learn a distribution that is closer to the actual unknown distribution \(p_{\text{data}}\).

During the forward pass, the input and the output of a differentiable operator are tensors; even the loss is a tensor. When the input traverses the deep neural network, PyTorch and other deep learning frameworks will try to cast the problem as much as possible in terms of GEMMs (e.g., convolutional neural networks) to enjoy locality and also vectorization, reducing fetch and decode overhead and thereby increasing performance.

During the backward pass, a graph task is created, and DFS is used to execute the graph starting from the root node, which is usually the loss tensor. As said earlier, during the forward pass, the computation graph is created on the fly; PyTorch, for example, uses probing through its dispatching mechanism.

GNNs, LSTMs, RNNs, LLMs, and many deep generative models (used to generate synthetic samples resembling those in the dataset e.g celebrity faces or natural language) are trained using PyTorch or other deep learning frameworks because otherwise you would need to manually craft the backpropagation formulas each time you manipulate the architecture. Thanks to automatic differentiation and the chain rule, you get to focus on the architecture and enjoy PyTorch's flexibility.

A typical piece of code for training a deep neural network is very simple and is listed below. First, we load the data as tensors. We then define the model architecture (i.e layers that the input will cascade through), the loss function (that will drive the optimization), and an optimizer (that will update the parameters after a backward pass). Finally, we execute the training loop. At each step, a batch of data is loaded stochastically and passed through the model during the forward pass. The batch is used to approximate the gradient of the loss, which is then computed during the backward pass. The optimizer subsequently traverses the model parameters and updates them according to gradient descent, or one of its many variants

Basically, almost every programmer can write and train a deep neural network today, from students finishing their bachelor's or master's degrees to software engineers who have only recently developed a profound interest in deep learning. PyTorch, hides a tremendous amount of complexity behind this simple interface, including automatic differentiation, tensor allocation, and the execution of forward and backward operations. If you are interested into learning more about deep learning, you can refer to my notes which attempt to bridge the gap between conceptual understanding and practical implementation by revealing the machinery hidden behind modern deep learning frameworks: https://books.deep-kondah.com/deep-learning-systems/; notes will be updated soon to cover more topics likes multi agentic systems, manifold learning and gpu programming: https://books.deep-kondah.com/deep-learning-systems/:

https://books.deep-kondah.com/deep-learning-systems/

In future posts, we will delve deeper into GPU/CUDA kernels at the microarchitecture level, autoregressive model training (particularly large language models, or LLMs), MoE, inference engines and optimization (and of course security defects).

Share This Post

Check out these related posts

Dissecting the A2A Protocol: Foundations for Interoperable Multi-Agent Systems (MAS)

Neural Information Retrieval & Acceleration of The Nearest-Neighbor Search (NNS)

Dancing With Agents: A Deep Dive into Multi-Agent Systems