Hyperparameter Optimization through Neural Network Partitioning with Christos Louizos - #627

0:08All right, everyone, welcome to another episode of the TwiML AI podcast. I'm your host, Sam Charrington, and today I'm joined by Christos Luisos. Christos is an ML researcher 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. Christos, welcome to the podcast. Thank you very much. It's great to be here. I'm looking forward to jumping into our conversation. This episode is part of our iClear series, ICLR, and we'll be talking about one of your papers accepted at the conference on hyperparameter optimization through neural network partitioning, as well as some of the other accepted papers from your team and colleagues at Qualcomm.

0:50But before we do that, I'd love to hear a little bit about your background and how you came into the field of machine learning and AI. Of course. Originally, I come from Greece. That's where I did my bachelor on computer science and telecommunications. And, you know, during the later years of my bachelor, that was the time when I'd gone first into touch with machine learning. And that was basically the point where I decided this is what I want to do. And after that, I followed a master program here in the Netherlands, in Amsterdam, where I basically met my future at that time, PhD advisor, Professor Max Walling.

1:26And after the master came the PhD here. And now, of course, here at Qualcomm AI Research in Amsterdam. Got it. We've, of course, had Max on the show a couple of times now. What was your dissertation on? So my dissertation was on sort of two things. So are you talking about the master dissertation or the PhD? The PhD. Okay, yeah. The PhD was, I mean, the title was called Probabilistic Reasoning for Uncertainty and Compression in Deep Learning. And that is a very long title of basically saying I work quite a lot with Bayesian neural networks. And we tried to find several sort of applications of them.

1:58The first ones that we were looking into was how we can get reliable uncertainty estimates on our predictions. Besides just getting some kind of answer of this kind of image needs to be a two or three, then how confident we are about that answer. The second part was on sort of compression. Like how can we use the fact that two parameters are uncertain to sort of better quantize them, better prune them and so forth for efficiency purposes. And is that area or those areas, do those continue to be your focus at Qualcomm? Yeah, it is one of my focuses here on Qualcomm. I mean, Qualcomm does quite a lot with efficiency.

2:33And that was like one of the main reasons also why I came here at Qualcomm AI Research in Netherlands, just because of my research interests and expertise were very much aligned to those of Qualcomm. And maybe take a step back and tell us broadly about the areas that you focus on from a research perspective. Yeah. So from a research perspective, I basically work on federated learning or distributed learning. So it has sort of two aspects. The first one is how can we do efficient learning at scale while having access to data that are local on edge devices? So the main premise on why you want to do that is that the data are private and therefore cannot leave the device.

3:15So what happens is that then the computation happens on device. You have some kind of server that might be in the cloud, and that server communicates a specific model to the edge devices to training on their private data and then communicate model updates back to the server. So that's sort of the general pipeline of federated learning. And then in terms of research, there's quite a lot of angles that you can pick up just because it's more of a system kind of approach. So we have done quite a lot of work on how, for example, you can efficiently communicate between the server and the clients via means of various compression schemes.

3:49How can you learn sparse models in this setting? The main benefit being that you can reduce both communication costs and also training costs on the edge device, especially since the edge devices are sort of hardware constrained. A little bit into differential privacy and how can that sort of combine with compression. recently, like our now iClear paper on how we can do efficient hyperparameter optimization in this setting. Yeah, talk a little bit about how hyperparameter optimization is connected to that broader range of topics that you just mentioned. Yeah, of course. The main premise of why we want to do hyperparameter optimization is, of course, you know, neural networks have quite a lot of them.

4:29And if you don't sort of set them right, sometimes your performance is not up for the task. Now, the main problem is that in federated learning, doing multiple training runs is not always desirable. First of all, you have to spend more communication costs and more training costs at the edge, which is usually not very desirable. And besides that, every sort of new training run also spends more privacy, if you look at it from that perspective. So meaning traditionally, when you're training a neural network, you're optimizing hyperparameters by just doing the training process over and over and over again.

5:04And that's less desirable on an edge type of device. Both for, I hadn't thought about, I imagine that you'd have the communication overhead and you'd also be in many instances, draining a battery by trying to do training locally. I had not thought about spending your privacy budget with differential privacy types of applications. That's an interesting one. The point is that every time you communicate something about your data, you reveal something about your data. Eventually, if you do it enough times, you basically spend most of your privacy budget. And exactly, as you said, traditionally, like random search, grid search, these things are not very desirable in federated learning.

5:45Okay, interesting. So that's the broad motivation. How did you frame the problem for the paper? Yeah, so basically how we framed it is with a lot of inspiration from marginal likelihood type of approaches. Marginal likelihood is a specific technique that you can do when you have a Bayesian model where you can sort of optimize hyperparameters while treating the likelihood of the data after you have marginalized all of the parameters. The main thing is that this kind of objective can be decomposed into a specific manner, which is, you know, you can think of it as basically the learning speed of a specific model, where the kind of intuition is that you basically measures how well a model trained on a specific subset of data generalizes to sort of other subsets of data.

6:35And that kind of intuition is then used to optimize the hyperparameters. The main challenge is that computing this marginal likelihood is usually intractable because it requires you to sort of integrate over all possible parameters or configurations that you have. If you try to apply this on a neural network setting, that's obviously infeasible. But basically, we can still take inspiration from that kind of objective. Then based on that, we devised like a practical kind of algorithm that simulates that kind of idea. And how it works is that, assume that we have a specific kind of dataset, we can split it into K parts.

7:13And similarly, we have access to a specific model, a neural network. We can again partition that neural network into K parts, into K subnetworks, if you'd like. And then basically how we do training the parameters of those networks is that each subnetwork only sees a subset of the full dataset. It never sees the full dataset. In doing that, the nice thing is that then you can use the parts of the data set that you haven't really seen for a specific subnetwork as validation data or like data where you can measure how well does that subnetwork generalize. And then you can use that as a learning signal for your hyperparameters.

7:52And are these subnetworks, are they architecturally defined like a layer is a subnet or something like that? Or are they kind of randomly defined? or tell us about how they are defined? Of course. So what we do is we do the simplest thing possible. So we randomly select a subset of parameters and we treat that as one specific subnetwork. In order to ensure that this random subset that you select, of course, goes from input up until the output, we sort of randomly select a subset of parameters for each layer and concatenate all of those and treat that as one subnetwork. These kind of subnetworks are chosen, then, well, I guess the parameter partitions are, of course, non-overlapping.

8:32But one of the main things that we do is that, let's say, if you have two subnetworks, your first subnetwork is trained on the first chunk of data. But then the second subnetwork, from the sort of marginal likelihood perspective, it can be trained on both the second chunk of data and also the first chunk of data. And to basically simulate that effect in our algorithm, what we do is we reuse the subnetwork that was trained on the first part of data as part of the subnetwork trained on the second part of data. So essentially that subnetwork can act some kind of information stored about the first chunk of data when you are optimizing for the second chunk of data.

9:13So you don't really start from zero, more or less. Does that mean that if you've got K subnets that the first subnet is a superset of the zeroth and the second is a superset of the first and so on? Exactly. And in the end, like the last subnetwork is the subnetwork is basically is the full network. It basically has all of them. And that subnetwork has seen all of the data. And that is the network we use to make like any predictions or downstream tasks. Okay. And so does that mean that validation is better for the smaller subnets and less good for the larger subnets? So basically, on the higher up the chain, so to speak, that you go, I guess you condition on more data, and therefore the effect of the validation data becomes less.

10:01The validation data at the lower end of the chain are those that govern more. I guess that's how the hyperparameters evolve. That's also an artifact of this kind of learning speed perspective. So from the learning speeds perspective, like the best kind of parameters are those that allow the neural network to generalize quickly after seeing a few pieces of data. You mentioned that you define these subnets by the simple random approach. How do you split the data set? That's a good question. So we have experiments in our work where we both do this thing for centralized training. So let's say that you had all of the data centrally.

10:40And then what we do is we basically split that data set into kind of KIID partitions randomly again. The main kind of one of the main motivations for our work was federated learning. And the point being that there, the data are sort of naturally partitioned because its device that participates in training has a specific chunk of data. And then how we construct sort of these partitions is by randomly picking clients or S devices and grouping them together. And then that group is sort of one specific kind of data chunk. Ah, got it. On the dataset side, you're also splitting into K as well, but that K is made up of an aggregate of data from across the different devices.

11:24Exactly. Again, it's non-overlapping split. Yeah. And I'm imagining that maybe right now this is doing some kind of simple aggregation, but maybe in the future you might want to do kind of topology aware aggregation of these edge devices so that you're minimizing latencies and things like that? Yeah, there's a lot of ways on how can one improve. You could argue if what you want to improve is like speed, how quickly did you train? Then yeah, you could organize that it's kind of data clusters, basically devices that have similar latency to the server so that they all communicate quickly. Another thing that you could optimize for is that the subset of devices that you select or that you cluster sort of collectively have an IID partition of the data of the overall data distribution.

12:13Right. So that would optimize for privacy then. Privacy and most importantly, sort of gradient or like how quickly we'll be able to optimize because then getting a sort of gradient from this kind of subset will be similar to a gradient you would have gotten if you randomly sampled your full data set. Right. Interesting. Interesting. So what type of data did you use for developing the system and for benchmarking? For this, we basically followed like prior work, especially most importantly like prior work that uses marginal likelihood type of approaches to like optimize specific hyperparameters and so that being said it's mostly kind of vision type of use cases so we had for example c410 we had dynemasnet and we had some experiments on mnist variants and most importantly like one of the set of hyperparameters that we optimize so now that goes a little bit more into detail there, is parameters of affine augmentations.

13:15And a way to see if the parameters of the affine augmentations that you learned make sense is to somehow transform your dataset in a way that these affine augmentations make sense. So to give you an example, we had two kinds of datasets from prior work, which is like rotated C for 10, and sort of rotated M is 10. What those datasets are is that we apply some kind of rotation to the actual image and then, you know, still maintain the same label. Then if you train while optimizing this affine kind of augmentation parameters, what do you ideally want is that you do recover somehow the rotation that happens.

13:54Because if you do, you sort of undo that effect that you did on your transform data. And so is this optimization, you mentioned the motivation here is to get away from multiple training iterations to do hyperparameter optimization. Is this a one-step process or are there still multiple steps? So how we do is we do all of this throughout training in an alternating manner. We have the subnetworks and then let's say you pick your sort of clients that will participate in training. Then these clients are sort of allowed to optimize a specific subnetwork. So what they do is they get all of the subnetworks.

14:34They optimize the subnetworks, the parameters of the subnetworks that they can optimize on their local data. And then besides sort of these fine-tuned subnetworks, what they also communicate is the gradient with respect to the hyperparameters when they used basically subnetworks that they cannot train on. Collectively gathering both of these things, what happens is that, you know, you do a single gradient step on the parameters of the subnetworks. And based on that, then you do another gradient step with the hyperparameters condition the subnetworks that you have. But iteratively through the training process as opposed to the first sequential.

15:09Yes, exactly. And one of the nice things that happened in the federated setting is that overall, despite the fact that we optimize both parameters and hyperparameters jointly, our communication costs, like cumulative communication costs, actually go down if you didn't optimize any hyperparameters. And that is mostly due to the kind of partitioning that we do for the network. Its client doesn't need to communicate an update on the full network. Its client only communicates an update on the subnetwork that it can train on and the gradient with respect to the hyperparameters, which usually the hyperparameters are just a handful.

15:45Can you talk a little bit about how you benchmark performance, how you benchmark the system overall? The primary kind of performance check we had was accuracy on some kind of test set. One interesting point is that if you do this kind of rotated versions of existing data sets that I talked about, if you do not learn like a proper augmentation policy, and that could be, for example, rotating image according to the rotation that exists in the data set, your accuracy will suffer. So, for example, like the baseline would be, okay, we have some kind of specific network that doesn't do any augmentation, gets performance A, but then we can learn that augmentation with our technique and then we get performance B, where, you know, of course, B is larger than A.

16:30So that's for the centralized setting. And in that case, you're considering the augmentation to be a hyperparameter of the network. Exactly. So then the augmentation is sort of trained such that a subnetwork trained on the first chunk of data generalizes well to other chunks of data and so forth. In this kind of rotated case, one way to do that is to actually learn to rotate the image itself while you augment. And in this augmentation case, is there a single hyperparameter, some single rotation value that's applied to all of the images or are there, is it a vector? Are there multiple directions?

17:06Yeah. So basically it's a vector. What we optimize that determines how much do we see or how much do we translate, how much do we rotate the image in both coordinates. And that vector is basically global. So it's randomly augmented according to that vector. Of course, the augmentation will not be exactly the same because you draw a different randomness for each point, but the augmentation strategy overall is here across all of the data points. Awesome. And then in terms of initialization, are you doing anything special there? That turns out to be like an important point. So the thing is, how do you define these subnetworks, right?

17:42So you could argue that you get sort of a subnetwork and you start from, you know, zero initialization. So what we did is that, let's say, basically zero initialization for the parameters that do not participate in training. But it turns out that, at least empirically, that was suboptimal. And what we did is we randomly initialized a network. And then a subnetwork is basically a network where you optimize some parameters and the other set of parameters that you not optimize are set to their initial random values. And that turns out empirically to work much better compared to, let's say, setting zero or the parameters that you do not optimize, like getting a sparse kind of subnetwork.

18:22Did you expect it to be the other way around? Did you expect that sparsity would have some characteristic that the network could take advantage of? Well, ideally, like sparsity would have been nice more from the perspective of efficiency, because you do not really need to store those parameters in memory or depending how your sparsity is defined. Maybe you can also avoid doing some of the computation. Whereas if you sort of don't have that sparsity and you just have them on the initial values, you still have to do everything essentially in terms of computation. So from that perspective, it would have been nice.

18:54But empirically, it seems that we cannot get away, at least for now. who knows, maybe future research shows that we can get away with it. And what about the K parameter that you use, which is itself a hyperparameter? How do you choose that? The thing is that we don't really have a good answer for that, unfortunately. Good thing is that we did an ablation study on that. And of course, there is some effect, but it doesn't seem to be significant. So, I mean, if you really want to squeeze out the best possible performance, then that is something that you need to optimize. But if you go with something reasonable, let's say two, three, or four kind of subnetworks and data partitions, and you sort of assign parameter counts and data points appropriately, then that kind of strategy generalizes reasonably well across all of the different settings that we considered.

19:42Having said that, ideally, I mean, this is a paper about hyperparameter optimization, so we would want eventually to automize that hyperparameter as well. Yeah. Do you envision a way to do it kind of within the same process or is it by definition have to be a separate optimization step? But yeah, it's kind of hard to tell. The good thing is that this kind of hyperparameter optimization method is orthogonal to other kind of hyperparameter optimization methods you can do. For example, in the centralized setting, you could combine this kind of training method with another hyperparameter optimization method that optimizes over K.

20:21Having said that, that is not intuitively the cleanest solution you can have. Ideally, you want to do everything under the same kind of framework, but at least at first sight, it needs more consideration on how this can be done. Cool. So that was the main paper that you worked on for ICLR. There are some others that your team colleagues have worked on that we wanted to talk through. Any particular one you want to talk about first? Sure. I can talk about one that I have thought a little bit about as well. So the first one was TTN, Domain Sift Aware, Pass Normalization in Test Time Adaptation. Basically, this comes from the perspective that at test time or when you deploy an actual model, sort of distribution shifts are kind of natural.

21:03If you consider continuous stream of data, you train on a specific set of data, but in practice, you might not have guarantees that test data that you will apply come from the same distribution. You could think, for example, this ties nicely with a federated setting where you could have like new clients joining in and maybe they have like a new device, a new mobile phone, and the camera sensor is a bit different. So therefore the images look a bit different and not exactly like the images that you use for training. So basically, this paper was sort of a relatively clean idea in that, you know, you have this kind of module usually in your neural network, batch formalization, which estimates some kind of mean and variant statistics over your input data set.

21:45And then the idea is how can we basically do this in a way that is more robust to distribution shifts that can happen. And the idea was that you can sort of simulate after training some kind of distribution shift by applying some kind of augmentation on your existing images. And then you can learn to basically interpolate appropriately between the statistics you have on your clean data and the statistics that you have on your sort of corrupted data. And empirically, this turns out to have beneficial kind of effects in terms of performance when you have some kind of distribution shift at test time.

22:24So let me see if I have the idea here. So you want to apply batch norm, but you're experiencing some kind of domain shift. And so it sounds like you are applying batch norm to the shifted data and identifying how those parameters are shifting and then applying that same shift to the data that you know. The idea is that Batstorm has this scale and shift to sort of move, like, standardize the data. But of course, if the data comes from a different distribution, then this scale of shift doesn't really standardize the data anymore. But if you simulate kind of this effect a little bit during training, you can find like a more robust kind of mean and variance to standardize your data against.

23:06Are you assuming that you know how the domain is going to shift at training time? Well, I mean, this work at training time sort of simulated sort of domain shifts. But from what it seems is that you can still get good performance, even on domain shifts that you haven't really simulated at test time. Okay. I think I was trying to wrap my head around this, the domain shifts. Sorry, can you repeat the last thing that you said we were talking about, whether it happens at training time or test time, right? Yeah. The authors have what they call this kind of post-training phase. So first you train normally, you know, batch normalization as usual.

23:43And then there is the post-training phase where all of the parameters are frozen. And then you just learn a single coefficient per kind of channel in batch normalization that determines whether how much of the original kind of statistics you should use and how much of the sort of statistics at test time you should use. And how this kind of interpolated coefficient is trained by simulating distribution shifts by doing specific augmentations on your data set. When they evaluated this method, you don't necessarily need to use the exact same augmentations. Of course, if you do, then you have more guarantees that the performance will be good.

24:23But empirically, this works so that you can still get better performance when you use just the statistics that were estimated based on just clean data. Interesting. So there were also a couple of papers related to scheduling. One was robust scheduling with Gflow Nets and the other was neural DAG scheduling. Can you talk about those? Yeah, of course. Basically, this is more from a sort of combinatorial kind of optimization perspective. A way to think of this is, let's say, some kind of compiler where you have a specific set of instructions or a computation graph that you want to evaluate. And then the point is, okay, how do I schedule the operations that lead from input to specific output?

25:05The idea here is, ideally, you want to learn this kind of schedule that happens for these operations in a manner that sort of adheres to the specific kind of computation graph that you have, maybe the specific hardware that you have, so that you can sort of minimize latency. There is these two works. The sort of first one, the neural scheduling via one-shot priority sampling. Basically, the authors use like a previous sort of in-house method proposed by Qualcomm, the topoformer, to learn sort of representations over this kind of directed acyclic graph, this DAG over computations. And then they learn some kind of parametrized policy on top to provide priorities for all of these nodes in the graph.

25:46And then the second one, the robust scheduling with GFLOW Nets, basically still attacks sort of scheduling, but from a slightly different perspective in that sort of evaluating how good the schedule is can usually be very costly because you have to run that specific schedule. And if that is a schedule over, for example, training like a trillion parameter model, that is not something that you can very easily evaluate. So there is this whole set of heuristics people try to use. And the idea of this work is that you can basically pick a reasonable heuristic and then you can sample schedules proportional to the reward that you get from that heuristic.

26:26And if you can sample diverse kind of schedules from that, then you can try and see, okay, like some of these might correlate well with the actual task at hand and, you know, might perform well actually in practice compared to just picking, let's say, the top one from the according to the heuristic, which depending how off the heuristic is with respect to the actual task at hand might not do as well. Okay. And then you've got one composite slice transformer. What's that one trying to do? This is also very good work dealing with how you can basically speed up the attention mechanism in transformers.