Quantizing Transformers by Helping Attention Heads Do Nothing with Markus Nagel - #663

0:09All right, everyone, welcome to another episode of the Twimble AI podcast. I am your host, Sam Charrington. Today, I'm joined by Marcus Nagel. Marcus is a research scientist at Qualcomm AI Research. Before we get going, be sure to take a moment to hit that subscribe button wherever you're listening to today's show. Marcus, welcome to the podcast. Thanks a lot, Sam. I'm happy to be here. I'm looking forward to digging into our conversation. We'll be talking about some of your recent research that you're presenting at NeurIPS, as well as some of the other Qualcomm research that appeared at NeurIPS.

0:45Before we dig into those papers, though, I'd love to have you share a little bit about your background. Yeah, so kind of my background and how I got into AI started with doing an undergrad in computer science in Germany. And like doing my degree on computer science, I was always interested in data-driven algorithms, especially like topics like computer vision, pattern recognition, and things like that. So for my bachelor's thesis, I actually did a very interesting internship at CERN. There at CERN, that's where the Large Heteron Collider is. And there I did a project on pattern recognition for beam loss monitoring.

1:22So that was kind of my first real ML kind of driven project. And after that, I was still excited about the topics and that brought me to Amsterdam. And in Amsterdam, I have a master's degree in artificial intelligence at the University of Amsterdam. And of course, yeah, we studied a lot of interesting topics there, like machine learning, reinforcement learning, computer vision, natural language processing. And then like I went more and more towards the research side. Also in my master's thesis, I worked on some more computer vision related topic. It was called event fissure vectors. So it was an encoding of diverse visual streams of photos, for example.

2:06And we also submitted that one then later to a conference, first to CVPR, where unfortunately I had my first rejection as a researcher, but that's common. Important milestone. It's an important milestone and very common in the field. And then later we resubmitted it and it got accepted at BMVC. Awesome. And what's the focus of your research nowadays? Now I'm at Qualcomm AI Research. And actually, I had an interesting journey there and to become a real researcher at Qualcomm AI Research. So I was always after my master's between doing a PhD or going to industry. And actually, it brought me to a third option, which is a startup called Cypher.

2:48That's a startup from Professor Max Welling, which you probably know, and also like Taco Cohen and Tim and Blankeford, who have been earlier on your podcast, like a couple of years back, I think they were co-founders of that startup. And that is sort of what in the end, like brought me to Qualcomm. So we did a lot of interesting machine learning problems there, a little bit more on the applied side. But then in 2017, Cypher got acquired by Qualcomm. And that kind of brought me back into doing more academic research and working on less applied projects. Wow, it's hard to believe it's been six years already.

3:26Time flies indeed. And so what are your current research interests? My current research interests focus in general on efficiency, efficiency of machine learning. And specifically, I'm focusing mainly on inference efficiency. So how can we make machine learning models as efficient as possible? And some of the techniques for that is quantization, pruning. There are also knowledge distillation and other techniques in the fields to make inference as efficient as possible. I guess that brings us right to the first couple of papers that we want to talk about today. These are papers that you're a co-author of and just a couple of the papers that Qualcomm authors have at this year's NeurIPS.

4:13What I found interesting about the first one that we want to talk to is that I've had many conversations with Qualcomm researchers over the years about quantization and how we've been improving quantization methods and more broadly in the field, kind of taking over lots of different modalities. And this first paper is applying the quantization research to transformers. The paper specifically is improve transformer, transformers quantizability by getting rid of the outliers. So let's dig into that one. Talk about the problem that you observed with regards to outliers and quantizability for transformers.

4:56Yeah, sounds great. Yeah, it's indeed a very interesting research. And actually sort of in the title of our paper, we have three different components. So the full title is Quantizable Transformers, removing outliers by helping attention heads do nothing. And kind of the first part of the title, Quantizable Transformers, basically is our motivation. So we want to make transformers easy to quantize because historically they're not very easy to quantize, especially the activations. And that's sort of what the second part is about removing outliers. The outliers is really what is the problem in quantization of transformers.

5:34As there are in all types of transformers, there's many outliers. And these outliers are the problem because you have the typical tradeoff in quantization between the clipping arrow and the rounding arrow. So it's these range versus precision tradeoff. And if we have a lot of outliers, while the main part of the distribution is fairly well behaving, then we can't represent both of them well. And that is what in the end can break quantization, and specifically in transformers, the activation quantization. And the third part basically of the title says attention heads doing nothing. And this is actually related to a very interesting finding we had, and that is actually that the root cause of the outliers comes actually from attention heads wanting to do a no update behavior.

6:24Was the impact of outliers and the disruption that it was causing from a quantization perspective, was this a known problem or was this something that was discovered as part of the research? Now, that was really a known problem. So that is already known since we had in 21 already a EMNLP paper on the topic. And there's also many other works which recognizes them. What is the main difference between our work compared to prior work is that most prior work is really trying to work around these outliers. For example, saying I don't quantize some outliers and keep a few of them in floating point, or it's, for example, saying certain type of activations where the type, like, for example, after the residual addition, there are a lot of outliers, then some work keeps these activations in floating point.

7:21But this is more really like a workaround, but they're not really tackling and analyzing what is the root cause of these outliers. Okay. And so again, in your case, you found that to be the behavior revolving around the attention heads. How did you identify that as the kind of the root cause? Yeah, so we looked like into really the behavior of the outliers. So we did a lot of studies on where do they occur and when do they occur and like some interesting patterns we found. And some of these patterns were also known before is that these outliers, they are only in very specific channels. And if I look over multiple layers, so multiple attention blocks, they're usually like in the similar or same channels across multiple blocks.

8:10And also they're always associated to very specific tokens. So if you're talking about language transformers, then it's often some type of delimiter tokens. So for example, can be the separator token in case we have them in our embeddings, otherwise some punctuations like dots and commas. And if you're talking about vision transformers, we also show these kinds of tokens are usually like background tokens. So, and we found it very interesting. We were wondering like what are encoded in these kind of meaningless tokens almost, because like they don't have much information, a background token or like a comma or dot.

8:47And we were wondering what actually happens there. And that kind of brought us like digging really deep into what's happening inside the attention head. And what we found is that the outliers, they're actually associated that like other tokens put all their attention to these non-meaningful tokens. So basically it means like instead of if I don't want to do anything, I'm not attending randomly to all tokens, but I'm only attending to these special tokens. And what we then see in the value matrix, so in attention, right, we have the attention weights, which gets multiplied then by the values. What we noticed is that the values of the special tokens, they're often zero or extremely close to zero.

9:32So what it actually means is like our outliers correlate to attending to tokens which have in the value zeros. And then like multiplying these two things obviously leads to a zero output and then like to no update on the hidden representation. So it wasn't just the delimiters. it was that the delimiters were creating relationships with other lower information tokens and that was causing the problem? Yeah, exactly. It's more like if a token, so for example, think about we have a sequence and we have like multiple layers, usually like 20, 30 or more blocks in LLM. And then we slowly aggregate information and refine it.

10:12And at some point, especially towards the end of the LLM, it could be that for certain tokens, we just don't want to do an update anymore because probably the representation is good enough or is where it wants to be. Then in attention, there's not any mechanism to easily say, I'm good. I don't want to update it anymore. And that's why we hypothesize it tries to attend to other tokens which have no information because their values are zeros. And then basically, if it attends to them, it can achieve that no update behavior, which Otherwise, there's not a different way to do that. So then what you do in this paper is create an explicit no update behavior, correct?

10:56Exactly, or almost correct. So that's what we're trying to add to the attention head in our method. So for that, we introduce two methods, which are fairly simple. So basically to do a no update, it means I want to be able to achieve exact zeros or exact ones to do a no update without having outliers. Maybe it's good to explain before shortly how the outliers relate to the no updates. So we established now basically that we want to do a no update behavior and that can be done by attending to tokens which have no information. But then in order to get actually only attending to these tokens and not attending to anything else, we basically need to have in the attention rates where we have the softmax.

11:44we need to achieve exact zeros and potentially exact ones if we only want to attend to a single no-information token. But in order that the output of the softmax is zero or very close to zero, we need to have a very big dynamic range in the input because they are relatively related to each other. And in order to have actually at the attention high dynamic range, there's one thing or two things which amplify that to get that is normally there's always layer normalization between either on the residual or right before attention. And layer normalization is dampening significantly the dynamic range.

12:23But then to still have a high dynamic range, it means that the outliers in the layer before need to be even bigger in order to achieve these exact zeros. And also another kind of complicating fact is that softmax will never give exact zeros, which also means during training, we will always have a gradient signal, which then if you want to achieve that exact zero, basically always tells the layer before, make these outliers bigger and bigger and bigger. And that's basically exactly the behavior we want to counteract. And we want to slightly change attention to easily represent these exact zeros.

13:01We introduced two methods for doing so. One we call clipped softmax. So that's actually a fairly simple method. What we are doing is we take a vanilla softmax, then we slightly stretch the output of the softmax so that it goes from, let's say, minus 0.01 to plus 1.01, and then we are clipping it back to be between 0 and 1. So what that actually means is then with a finite dynamic range in the input side, this changed version of softmax can actually achieve exact zeros and exact ones during inference and also during training. And like that, we can then like say we don't want to attend to any of the tokens or only to some tokens with no information.

13:45So the other methods we call gated attention. So that's more explicitly modeling this behavior of representing exact zeros, because some of our early studies also showed that actually having the exact zeros is the main important thing and not the exact ones. So we model that behavior explicitly by like adding one gate to the attention mechanism. And that basically can modulate the attention weights by just like one row in the attention matrix. So that's like the same dimension where we take the softmax over. We can modulate with a sigmoid gate. So it basically means we can with a single parameter or with a single output neuron, we can say, oh yeah, here we don't want to have attention and put that completely to zero.

14:30And did those two methods perform equally or did one turn out to be better than the other in some situations? In most cases perform equally, but not in all cases. I think first, the most interesting fact is that both of the methods, like if we train a floating point model, they basically lead to the same performance in floating point domain as the original model, sometimes even slightly better, but not much. But then what we really saw is the resulting model has significantly less outliers. So in the paper, we showed that by measuring the maximum infinity norm and the average cortosis. And of course, as a result, if we don't have these outliers, then out-of-the-box quantization is significantly better and works well.

15:13And comparing these two methods, in some cases, they were almost the same. or Clipped Softmax was, for example, for the BERT model, we evaluated slightly better. But we had one case and that was on the OPT model. So that's like a language decoder only model. We had for some reason that the Clipped Softmax didn't achieve the desired outcome. So it kind of failed. But for all the other models like BERT and with Clipped Softmax worked well. So in that sense, yeah, gated attention seemed to work better in all three models we tested it on. And how did you ultimately measure the performance of this approach relative to other kind of traditional approaches to quantization?

15:59And how much performance does this kind of get you back relative to the original model that you're using in floating point? Yeah, that's a very good question. So in general, we did not necessarily compare it directly to other quantization methods because actually we placed this paper more not really as a quantization paper. So in a sense, this is not really a quantization method, but this is rather analyzing, like getting rid of the root cause of what makes quantization difficult, at least activation quantization. So in that sense, this work is completely orthogonal to any quantization method.

16:38So for example, we just do 8-bit post-training quantization afterwards to show that the model is easily quantizable. But if you want to get the best performance out, you could use any modern PTQ or QAT algorithm to apply on top of that. For example, GPTQ or other algorithms out here. Right, because all the quantization methods suffered with the same problems with outliers and this is addressing that core root problem. Exactly. Like all quantization methods that want to do activation quantization. Right. But you still want to try to characterize the benefit that this approach offers. How did you do that?

17:20So we applied like post-training quantization with 8-bit activations to both the original model with the vanilla attention, so softmax-based attention. And then we compared the quantization performance of that one with exactly the same algorithms used to the performance we get when we train the transformer using the clip softmax or gated attention. And what did you find there? We found that both of our methods were significantly better. So for example, for BERT, even like the performance is close to random in 8-bit activations, if you quantize really every layer and all activations, and with our methods, we have close to floating point accuracy.

18:03Awesome. The second paper that you worked on that we wanted to dig into looks at pruning versus quantization. Talk a little bit about the context there. So the context here is that for model efficiency, so maybe first I have to one thing to mention and to distinguish it like in the mental switch from the other work is like the other work really was tackling activation quantization issues. And now we are looking in more general model efficiency and specific into the problem of weight compression. And for weight compression to make the model smaller, there are two common techniques, either pruning or sparsity or quantization.

18:44And we really wanted to look at which of the two is better because like both of them are fairly well-studied fields. also academically as well as from a hardware perspective, but it's very rarely that they are on a very rigorous manner get compared. And what we wanted to do in this work is to compare them as fairly as possible to draw conclusions, what is better from an accuracy perspective. Okay. So tell us about how you approached this problem and what you found. Yeah. So in order to make a fair comparison, we made a couple of assumptions, some base assumptions. We compare in this study like uniform symmetric quantization compared to unstructured magnitude pruning, and we assume that the base data type is floating point 16.

19:33So it means we assume an equal compression rate if we have like 50 % sparsity versus int 8 or 75 % sparsity versus int 4. And that's sort of the base assumption of all the studies. And it's also a very common assumption that people assume FP16 is working without loss inaccuracy. And then basically from there, we compared like pruning versus quantization on various levels. So we started from comparing it on distributions to characterize how the error is. Then we looked at real weight tensors. And then we went to what happens if we add, for example, post-training pruning or post-training quantization algorithms on a layer by layer basis.

20:18And finally, we had a set of comparisons with end-to-end model fine-tuning. And what did these comparisons show you? So basically, these comparisons show that overall quantization is significantly better in most cases than pruning. So what we really can see is that in many of the cases, it is better. And there we only found a few cases in which pruning can be better in quantization. So for example, one interesting study we did is we took all the weights from PyTorch model zoo, as well as like the weights from several LLMs from hugging phase, and we quantized them to a lot of different bit widths and the corresponding sparsity ratios.

21:02And then we looked at them, where is the SNR better for pruning or for quantization. And we noted that only in very specific cases, that's mainly low bit widths, and then also cases where we have many outliers in the weights. In these cases, pruning can perform better. So we can see that when we look at the cortosis. So we kind of transform the data to blot it as a function of the cortosis. And then we could really see if the cortosis is high, which indicates there are many outliers in the distribution. Then it can be that in the two, three-bit scenario, which corresponds to 85 to 90 % sparsity, that in these cases, pruning can perform better.

21:49So how consumable is this result for folks that are trying to create smaller models? Is this saying just go straight to quantization and don't worry about pruning unless you're doing model A, B, and C, or do you still need to kind of look at both for your particular use case because it's very use case dependent? It's not very use case dependent, but one may ask, okay, what happens if we add algorithm to it? Like, for example, if we optimize it with modern post-training quantization algorithms or pruning algorithms. And actually, that was a big part of our work to answer that question. And in order to answer that question, we were deriving upper bounds on the signal-to-noise ratio for pruning and for quantization.

22:37And then actually we read it on part of the data, the same study, but with this upper bound, which tells us basically, it doesn't matter what algorithm will be invented in the future, the SNR will never be better than that. And then we basically compared on many of the weight tensor the data for that. And sort of the findings were fairly similar to what I mentioned earlier on the real weight tensor, but this is a really good insight, which kind of shows that at least for post-training quantization or pruning algorithms that we can like have much stronger guarantees than just looking at some random data or one specific algorithm.

23:18Again, from a user perspective, if you're only looking at upper bounds and not also providing lower bounds, I still kind of need to look at both approaches in order to know what's going to be best for my specific data and algorithm, right? Or no? Not necessarily specific. So we have a lower, I mean, the lower bound is the algorithm we have. So that's a very standard algorithm, like for example, Ardorand for quantization and a similar data optimized algorithm for pruning. So that's our lower bound. And we show in the paper, we make a nice plot with some squares, basically where one weight tensor is represented by one square.

23:56So the lower left corner indicates there, what's the SNR for the two achieved algorithms. And then basically we span a box, which kind of goes in the X and Y directions, what is the bound. And basically we know if that bound will intersect the diagonal of that SNR plot, then we know it may depend on the instance of the algorithm. And in the cases it is above or below the diagonal, it shows, says that it's independent of the algorithm, either pruning or quantization would be better. Depends where the box lies. Great. So we also wanted to talk through some of the other papers written by your colleagues there at Qualcomm.

24:37Why don't you walk us through, I think we've got four, walk us through the ones that you think are most interesting there. Yeah. So the third paper, or the first of the other papers I want to talk about is also a paper from the model efficiency team where I'm working in. It's called Scalarization for Multitask and multi-domain at scale. And that one really looks in more the training side. So what happens if we do training with multi-domain and multi-task? Because this kind of training with multi-domain learning and multi-task learning that can be also seen as a way of compression information into one backbone.

25:18And there can also be knowledge transfer between the different tasks or different domains. And that can also improve efficiency. So that's kind of why we're looking into that. Is the idea here that we're able to take advantage of unique properties of multitask and multi-domain learning to make training more efficient? Yeah, it can both make training more efficient, but it can also make inference more efficient. Because if we train a model with multiple tasks, it means a model can do these multiple tasks at the same time instead of running like three independent models for inference at the same time.

25:57So it tackles both efficiencies. And the specific work looks into how to do this multitask and multi-domain learning in the best manner. So a very common baseline for it is scalarization. So it basically means we have these multiple losses for the either multiple tasks or multiple domains. And then they're like summed together for the final loss with scalars. And prior work either uses just uniform weighting of the different losses or sometimes statically tune them. And also there are sometimes some automatic adaptation of the scaled loss based on, for example, gradient information. And we're kind of like, or like my colleagues are taking a slightly different approach, which basically can adapt the scalar weights of the different losses during training.

26:50So what they introduce is population-based training. That basically means that with population-based training, we tune these scalarization weights of the different losses. And how you can think about it is it's a kind of typical evolutionary algorithm which tune these weights and we have a population of models. Then at every step, we train n models at the same time, but n doesn't need to be the amount of tasks. That can be also significantly less than the amount of tasks. And then at the synchronization step, we evaluate how are these n models with different weights per loss doing. And then we have a typical evolution step where we say the top contender, we replicate.

27:37The middle ones from performance wise are surviving. and bottom, the worst models, they are pruned away. And so like that, we basically consistently update the weighting of some of the models then lead to a really good weighting at the end. And what are typical values of N? Like how far are we scaling multitask learning nowadays? What do you tend to see? That really depends on the tasks or the application areas. It can definitely in some cases be, let's say, in the number of tasks or domains like easily like 10 or more, but it's very dependent. In some cases, it's only a few. And then, of course, it's an easier problem.

28:23What's the next paper that you wanted to talk about? I think the next paper to talk about is called EDGIE. That stands for Equivariant Diffusion for Planning with Embodied Agents. And that one basically builds on a very well-known paper called diffusers. And that basically uses diffusion models for a sequence of state and actions. From that paper, the key idea was that you basically learn one model on states and actions. And then at test time, you can use kind of any reward functions to kind of perform your actions. And is the diffusion models have become very popular recently in the computer vision domain with image generation, like stable diffusion and the like?