The AI Generalist Program : What is the Mathematics Behind the Transformer Model?

We know a concept called vanishing gradients in Neural Networks. In the previous lesson, we talked about gradient descent, backpropagation and the mathematics of selecting the right learning parameter. Gradient descent and backpropagation need differentiation (dy/dx – that we learnt in our school mathematics to find the minima or maxima). Mathematically, for gradient descent, we need to calculate the gradient to optimise the parameters as follows:

Parameters = Parameters – learning rate * gradient of loss function for every change in parameter … (1)

In the case of a Neural Network, the Chain Rule from Calculus is applied as follows to find the gradient by a derivative-based approach:

Gradient of the loss function with respect to a parameter = Derivative of the loss function with respect to hidden layer N * Derivative of hidden layer N with respect to hidden layer N-1 * … so on … (2)

Suppose every derivative is much less than 1 (it can be less than 0.5 as well, based on the activation function); the final gradient, as per the above formula, is almost equal to 0 if there are many layers (Deep Neural Networks) because of multiplication. This is called the Vanishing Gradient Problem.

Have you ever wondered how transformers – the basic building blocks of Language Models – avoid vanishing gradients?

Transformer model architecture from "Attention Is All You Need"
Transformer model architecture from “Attention Is All You Need”

In the case of a transformer, as shown in the following image from the original paper (Attention Is All You Need), the Residual Layer forms the following equation:

y = x + f(x) … (3)

Now, using the optimisation method, from equation (3), we can derive that …

The gradient of the loss function is a function of (I + J)

where I is the Identity Matrix, and J is the Jacobian matrix.

Due to the Vanishing Gradient problem, even if all the elements of matrix J are 0, the Identity Matrix remains. Hence, the gradient of the loss function is never 0. Thus, this layer helps transformers avoid the Vanishing Gradient Problem during training. This is how the transformer model’s architecture is designed to prevent gradients from vanishing, and there is a proper mathematical proof confirming that it will not let vanishing gradients happen while training the transformer. Hence, a transformer model can hold trillions of parameters and be trained with huge datasets.

In the next chapter, I will talk about Explainable AI algorithms, a crucial concept for ensuring that an AI system is trustworthy.

Subscribe to our newsletter to read our articles on the latest trends in AI.

khasnabish Avatar

Posted by

Leave a Reply