The Residual Stream: How Transformers Actually Compute

imagine you're writing a book you start with a draft and want to turn it into a market Ready book you hire a team of an editor marketing specialist and a designer editor fixes grammar marketing adds an appeal for potential readers and designer creates a cover together their Collective effort turns the draft into Market ready products in order to ensure that the essence of the draft is not lost it's added to the output of the team work this example illustrates how residual learning works here draft is an inputed a stack of layers and it's defined as X Market Ready book is a desired output h ofx and f ofx is an incremental improvements that team made to X meaning that instead of directly learning H ofx the network learns residual function f ofx which is basically h ofx Min - x that simplifies learning process and allows to train very deep networks the idea comes from 2015 paper called Deep residual learning for image recognition and the residual learning concept is explained in section three where Al authors say that rather than expect stch layers to approximate H ofx we explicitly let these layers approximate a residual function f ofx which is H of x minus X the original function thus becomes F ofx + x as Illustrated uh by this building block figure uh where X is an input to stack layers and Y is an output and it's an element wise addition of residual function f and input X and uh residual function here uh is uh configured by two layers uh defined by W1 and W2 and r activation function is applied after first layer and second nonlinearity also R is applied after addition f ofx is usually implemented using multiple layers with two and three being most popular choices and when we add more layers it allows model to represent more complex Transformations however with just one single layer uh it's just a linear transformation of X so authors also claim that they haven't seen any uh benefits of this configuration main problem that authors are addressing using residual learning is when the network gets deeper the performance degrades significantly and they also share an example of uh training two networks one is 20 layer and another one is much deeper 56 layer deep network uh they train it on a C r10 challenge as you can see the deeper network has training has higher training error and does test error authors also claim that this problem is not caused by overfitting it's also not caused by vanish gradients problem because they applied uh wch normalization to all networks uh which is one of the most popular solutions to this problem let's talk about resid learning in Transformers Transformer is a residual Network there is a great paper called a mathematical framework for Transformer circuits written by team in anthropic in 2021 that provides great insights into how Transformers work with one of them uh being about residual uh learning so they focus on decoder only model and I took this picture from um tutorial written by J LR that uh clearly explains the architecture and Transformer decoder model uh predicts next token in a sequence given previous tokens uh and consist of a stack of decoder blocks each each one of them is a pair of mask self attention and fit forward fit forward neural network MLP uh and mask self attention consist of multiple heads running in parallel so so coming back to this paper they also have a very good explanation uh about the architecture of this model that start uh with embedding layer of the token followed by what they call a residual block and followed by unembedded so a residual Block in this definition is a pair of attention layer and MLP layer uh and the paper says that each uh attention layer and MLP layer uh reads from a residual stream they call it a residual stream by doing a linear projection and then they add linear projection back to this residual stream so you can see here that for example XI is an input to this attention layer that after doing some calculations add uh the linear projection back to the residual stream and it becomes x i + 1 which is then uh an input to MLP so residual stream is uh just simply a sum of all the outputs of the previous uh layers plus the token embedding and they say that residual stream it's better to think of it as a communication Channel because it doesn't do any processing by itself uh but all the layers communicate through it both attention and MLP layers read and R to residual stream by doing linear projection the residual stream Dimension is bch time sequence length time embedding Dimension attention layer then read it through multiplication by weight matrices w k WQ and WV the result is key query and value matrices that are then passed to calculate attention for each of the head stack of attention matrices are then concatenated and then multiplied by wo that produces ATT tension output Matrix of the same Dimension as of residual string this is the right [Music] operation MLP layer consists of two lanar Transformations it reads from the residual stream by applying a weight Matrix W1 to the input after applying nonlinearity like jello it then performs another matrix multiplication by W2 to get the output of shape batch time sequence lens times embedding Dimension that's written back to the residual string another interesting Insight um shared in this paper is that residual stream has a high dimensionality and uh therefore can be divided into different subspaces and each uh Subspace is dedicated for different kind of information can be dedicated for example this residual stream um shown on the picture is divided into three subspaces and for example example first Subspace might contain information about like a previous token or second Subspace can uh store something like the subject of the sentence etc etc and if layers uh read and write to the same Subspace then they communicate to each other however if they do it to different this joint subspaces then they like don't interact at all uh and another claim that this paper made was that uh basis of uh residual uh space are not anyhow special therefore the basis is non-privileged uh and then this paper published in 2023 in anthropic called privileged bases in the Transformer residual stream uh show that in practice uh it's a wrong claim authors first introduce us uh to a context and talk about another paper called 8bit matrix multiplication for Transformers at scale written in 20022 researchers in this paper examined hidden layers of a transformer and found out lers specific neurons that produce unusually high activation values compared to others up to 20x larger they first appear in 25% of all Transformer layer in small models and then affect all layers when model reaches size of 6.7 billions this outliers are highly systematic at a 6.7 billion scale 150,000 outliers occur per sequence but they are concentrated in only six feature Dimensions across the entire Transformer to understand the importance of this features uh researchers created an experiment where they first um created an alternative version of the model where they set out lers to zero in a head in a hidden state that goes to the attention layer they also say we do this for all layers independently meaning we forward the regular soft Mark probabilities values to avoid cascading errors and isolate the effects due to the outlier features they also created another control version of the model where they set neurons to zero but this time randomly they compare the performance of these three models uh using top one softmax probability and validation perplexity and results are interesting so they see that in a mod in a model where they set outliers to zero they reduced uh softmax probability from from 40% to 20% and increased validation perplexity by 600 to 1,000% like the this is huge and in a model with uh random uh where they set neurons randomly to zero they reduced soft Max probability by only between 0.02 to 0.3% and increased validation perplexity by 0.1% though so this results uh completely indicate the importance like the critical importance of these outliers to the model performance [Music] and a natural question that arises from interpretability Viewpoint is what are these features what do they represent or what purpose do they serve and why are these features bases aligned if you are not sure what I'm talking about when I say that features are basis aligned check out my other video when I talk about fundamental concepts for mechanistic interpretability including privileged bases to answer this research questions team in anthropic uh conducted experiments using their own 200 million parameter model and first they suggested to use a metric to detect uh outliers in activation vector and the reasoning was following so if we think of uh activation vectors as random Direction then the property of this is that every element of this activation Vector is drawn from isotropic gaussian followed by normalization to a unit Vector so if we have uh for a separate token if we have a activation Vector as following A1 A2 A3 Etc then we can think that each of these elements is drawn from gaussian then if we scale the activation Vector to unit Norm uh then we can calculate the curosis which is the measure of tailin is right so if a curosis is three this means like it's a normal distribution with no heavy Tails otherwise if it's much larger than three it uh signals that there are some uh extreme values and we can repeat these uh calculations for several tokens and take the average and uh te anthropic uh shared with us uh the plot they created where they measure the is for each of the layer in their own 200 million parameter parameter model and you can see that the values are very large uh signaling the outliers to test that metric works correctly they also um uh did the rotation of the bases uh to see if um Vector values will be distributed more equally across um basis Direction so uh to illustrate this imagine so this is the um uh basis in which Transformer works it's like 2D with uh bases corresponding to some meaningful features like gender and royalty with the men as an activation Vector so you can see that it has very high value for gender and zero value for royalty uh so it's basis Alliance but if we rotate the basis change it you can see that the values are equally distributed along basis directions and if we plot after doing the rotation to Transformer activations uh we calculate the curosis again and plot its values you can see uh that the values are again three uh so the signals that um activation vectors are gum distributed to understand what causes these large values and activ ations uh Team run few experiment and they had three hypothesis and the first hypothesis is to suspect that layer Norm has something to do with it because uh it has a mean subtraction and per Channel learned weight that both might introduce basis dependency so just to make it more illustrative we have uh value uh activation Vector X then we subtract the mean from it and then we calculate Cal ated by these learned parameters gamma and at beta uh and the they might potentially introduce basis dependency I took this picture from Arena tutorial uh and then researchers what they did to as an experiment they uh introduced modified layer Norm where they uh change the subtraction of the mean and just uh did the scaling and they also used a single learnable scale parameter and they claimed that after this model should be rotation invariant uh when they uh calculated the CES using this uh updated layer Norm what they saw are again these large outliers um that led uh to conclusion that um layer Norm is not causing any outliers hypothesis 2 uh was to think that maybe limited numerical Precision causes model to favor certain directions because this way uh model can prevent large Val values from overshadowing smaller ones so what I mean by this is that let's see just a simple uh large and small numbers summation in Python represented uh using num uh flow 32 uh type and one number is uh 10 to the point of8 a small number is 10 to the uh sorry to 10 the power of minus 2 when we add them both we expect to see this uh value at the bottom however the result is this so we lose the second digit and the setup of this experiment was uh to rotate every parameter Matrix between read and write operations from residual stream on a forward pass and then we uh they check the performance of model using change in loss and the result was that change in loss was absolutely tiny suggesting the model is generally rotation invariant even when accounted for floating Point error so again just to make it more illustrative let's go back to this example so if uh the model is Basis uh aign so we will should see a large values here right and if we rotate uh uh the bases then the values uh values of this Vector should be equally distributed across bases directions finally the last hypothesis uh they tested was that maybe there is something in a combination of limited numerical precision and gradients and the experiment was very similar to the hypothesis 2 test uh but this time they rotated not parameter Matrix but activations again between read and right to residual uh stream but and during training and not forward pass as before and they again check the performance of a model using change in loss they had two setups for this experiment so if this is a residual stream and we have uh we have two layers MLP and attention and many more then in the first uh setup uh these layers they share rotation Matrix so it's the same for each of the layer for every layer in the uh Transformer model in the second set up though every layer has its own rotation Matrix so every layer operates in its own basis and the result of this experiment was that they didn't see any difference between Baseline model between shared rotation and independent rotation so setup one and setup two uh in a on a training loss and the conclusion was that Transformers do not rely on a privileged basis to train and function properly even when floating Point Precision is taken into account and when they also looked at the activation after they after the model reads from residual stream and rotates to new basis and before it passes it to MLP to protest they looked at the curosis of this activation and they again saw very large values and the conclusion here was that basis alignment is an artifact of the computation inside of the Transformer layers more so than from the representation in the residual string