[NeurIPS Best Paper] 1000 Layer Networks for Self-Supervised RL — Kevin Wang et al, Princeton

So, welcome to Latent Space. We are basically trying to provide the best possible sort of podcast experience of NeurIPS for people who are not here. Congrats on your paper. How does it feel?

Yeah, it was very exciting. We had a poster yesterday, and today we'll have an oral talk. We just got mobbed—there were a lot of people. It was 3 hours straight of waves of people coming by to talk to us.

I've never received a best paper award. Do you just find out on the website, or what?

I just woke up one day and checked my email.

They just tell you?

They were like, “Oh, you've been awarded best paper.”

But maybe you knew from the reviews as well, right?

Yeah, we knew from the reviews that we did well. But there's a difference between doing well in the reviews and getting best paper, so we didn't actually know that part.

Yeah. Okay, so I skipped a little bit. Maybe we can go one by one and introduce who you are and what you did on the team.

I'm Kevin. I was an undergraduate at Princeton, and I just graduated. I guess I led the project—I started it—and I was very happy to collaborate with Ishaan, Michał, and Ben as well.

Right. Were you all in the same research group? What's your social context? And thanks to Ellen for booking you guys.

We're all from Princeton. This project actually started from an IW seminar—an independent work research seminar—that Ben was teaching. This was one of my first experiences in machine learning research, so it was really valuable to get that experience.

Ishaan was also in that seminar and was working on adjacent things, so we collaborated a lot during the seminar. The project turned out to have some pretty cool results, and later on Tomasz, who was working on similar things, also joined the project and became a good collaborator.

I don't know if any of you guys want to chime in on other elements of coming into and deciding on this problem.

Broadly, my lab works on deep reinforcement learning, but historically “deep” meant 2, 3, or 4 layers—not 10,000. When Kevin and Ishaan mentioned that they wanted to try really deep networks, I was skeptical that it was going to work. I've tried this before and it doesn't work; other people have tried this before and it doesn't work. So I was very, very skeptical starting out. I don't know if I made this clear at the time, but that was my prior going in.

Do you view your job as screening—“Hey guys, this probably isn't going to work. You should try a different idea”—or should you be encouraging people even if the idea is dumb?

It's selecting bets.

Yeah. What made you willing to make this bet?

It seemed relatively low-cost. Michał in particular had spent the past year developing infrastructure that made it a lot easier to run some of these experiments. The precedent was that deeper networks could do a whole lot better; that's what the deep learning revolution has been over the last decade.

Yeah, why do we stop making them deeper?

Reinforcement learning was this one anomaly where we continued to use these really shallow networks. That's particularly true in the settings that we were looking at, where you're often starting from scratch.

Any other perspectives you guys want to chime in with?

I guess maybe I should just go over an overview of our project.

Yes. Okay. Sorry, yes.

The way that I view our project is that, if you look at the landscape of deep learning, you have language, vision, and reinforcement learning. In language and vision, we've converged on these paradigms of scaling to massive networks—hundreds of billions of parameters, even trillions of parameters. A lot has been gained from deep learning through that.

But in the third branch of deep learning, deep reinforcement learning, that hasn't yet been the case. I was very surprised coming into Ben's class and seminar when I looked at the networks and saw that they were using a simple 2-layer MLP for frontier, state-of-the-art RL algorithms. I was very curious: Can we design RL algorithms? Can we put together a recipe for RL that allows it to scale in potentially analogous ways to how language and vision might scale?

What we did was start with the knowledge that traditional RL—let's say value-based RL—doesn't really scale. This is pretty clear from the literature. So we tried a different approach to RL called self-supervised RL. Instead of learning a value function, we're learning representations of states, actions, and future states, such that the representations along the same trajectory are pushed together and the representations along different trajectories are pushed apart.

This is a different approach to RL that allows us to learn in a self-supervised manner. We can solve tasks and reach goals without any human-crafted reward signal. We know that self-supervised learning is scalable in these different areas of deep learning, so can self-supervised RL scale in similar ways?

When we first tried it, it actually didn't work. We made the networks deeper, and performance totally degraded. But I separately realized that there was also other work in our literature where people had tried residual connections. There were a few other architectural components that we had to put into the recipe. Then, all of a sudden, one day I ran an experiment, and there was this one environment in which doubling the depth didn't really do anything, but doubling the depth again with these different components suddenly skyrocketed performance.

Getting this to work was very nontrivial. Usually, when we need to think about hyperparameter optimization, we try changing A to see if it makes things better, then changing B to see whether it makes things better. If we just made the depth bigger, it got worse. If we just added residual connections, it didn't get better. It was really this combination of factors that Kevin and I figured out that made this work.

As a precursor to that, we also tried scaling along different dimensions: scaling the batch size and scaling the width of the network, meaning the hidden layers. Those were pretty similar to just scaling depth naively. Once we started introducing residual connections and these specific architectural choices, that's when we saw these significant jumps in performance—these critical depths at which performance multiplied by a pretty huge factor. That's where we really noticed unlocking significant performance gains, as opposed to scaling just along width, which did yield some performance improvements. When you look at the number of parameters that your network has as you grow width, it's roughly quadratic as opposed to something like growing depth. So it's more parameter-efficient, also more sample-efficient from the experiments that we conducted.

Nice. In some ways, you're sort of replicating stuff that is seen in the wild but on a very small model that you can study. Would you say that's right?

Yeah. To add to what Kevin said earlier, we saw these huge performance improvements in language models and image-generation models by making them larger and deeper, which seems very intuitive.

That's why our work draws from foundational research, like residual networks, which employ residual connections to avoid vanishing gradients. That's something we show in some of our ablations further down in the paper, probably in the appendices. We ran experiments without these residual connections, borrowing concepts that have existed in other fields and applying them to this setting with RL to show that they work.

Before Ben has to go, I'll leave the last word to him. What additional work does this inspire that you want to push on next?

I think there's one thing I'd clarify about the paper, and then I'll directly answer the question. A lot of people reading the title are like, “Wow, big networks—they're great. I'll take big networks and”

“You solved it now. We can just go.”

Yeah, we say, “Take big networks, add them to PPO, add them to SAC, add them to your favorite reinforcement learning algorithm.” But I think that's actually not the main conclusion. The main conclusion is that using big networks not only requires these architectural tricks, but also, as Kevin mentioned before, requires using a different objective.

This objective doesn't actually use rewards in it. So there's another word in the title—“reinforcement learning”—that might also be a little bit of a misnomer, because we aren't directly trying to maximize rewards.

Our code doesn't have a line of code saying, “Maximize rewards here.” So, at the end of the day, is this a reinforcement learning method? I don't know. It looks much more similar to the self-supervised methods in other areas of machine learning.

I think the work really stands at some interesting intersection of reinforcement learning and self-supervised learning research. We had this little figure on the bottom left of the poster, which was a screenshot of a slide from Yann LeCun talking about how to build intelligent systems, and whether that's going to be done by unsupervised learning, supervised learning, or reinforcement learning. I think what the paper really suggests is that the boundary between these things is really blurry, and maybe the keys to building intelligent systems are going to be leveraging insights from all of them.

Yeah, the LeCun take.

Exactly. (Laughter)

Well, thank you for your time. I know you have to go soon for Jon.

Yeah, thank you so much for coming. I think that insight about blurring things is interesting. You were talking about the abstraction layer of representation learning. I don't know if that triggers anything in terms of the mix between self-supervised and reinforcement learning. Is that something fundamental that you've discovered, or that people don't understand when they read the paper?

Yeah, I think the best way that I would explain it is that we know standard RL is not super scalable. Why can this different approach, or different objective for RL, be scalable? I think it's because we're fundamentally shifting the burden of learning from something like Q-learning, or regressing to TD errors, which we know is quite spurious, noisy, and biased, to fundamentally a classification problem.

We're trying to classify whether a future state is along the same trajectory or along a different trajectory. We do this with representation learning, and we know that classification, cross-entropy loss, and representation learning are scalable in the deep learning literature. If we think about language and some of the objectives there, in some sense we're blurring the lines: we're doing reinforcement learning—it's still an actor-critic reinforcement learning algorithm, like a goal-conditioned reinforcement learning algorithm—but the burden of learning, of solving that RL task, shifts to something more similar to the objectives that you might see in language and vision, which we know have scaled so much.

I think that's one of the fundamental insights that we've seen: by approaching RL in this different way, we were able to get so much more out of it. We were able to scale our networks significantly beyond what is standardly used in RL.

Can I jump in? I'll just give a bit more context about the architecture. We use another objective, InfoNCE, so the contrastive loss. However, the architecture is quite similar to previous works and papers, like DrQ, or SimBa and SimBa-2.

We also tweaked the architecture a bit. However, it's not that we invented the wheel for the first time. It's the merging of the architecture and the objective that makes the scale really go up, and performance follows the scale.

I think that's something that we should probably mine deeper. What domains, what industries—if you've applied it to multiple different types of networks or datasets—is there a particular affinity that you think is low-hanging fruit?

Actually, if you look at a lot of our tasks, they're particularly robotics tasks. Personally, I'd be very curious about how work like this could impact the robotics field. My understanding of robotics is that a lot of robotics right now has a few different approaches.

One approach is that we want to train robots using imitation learning, so we try to collect an insane amount of data. We have a ton of human supervision, and we try to scale up this data; we're learning with imitation learning. But on the other hand, perhaps there's another approach, such as goal-conditioned reinforcement learning, where we can actually train robotic agents—and train RL agents—to solve meaningful tasks with absolutely no human supervision or demonstrations.

It's much more scalable, yeah.

So this could serve as an alternate approach. Perhaps instead of scaling data—scaling manual human supervision, which is not super scalable—if there are ways to make goal-conditioned reinforcement learning scalable, we can just scale the architecture, or we can scale—

Because you're focused on your objectives. Yeah.

Right. With certain different objectives, I think that could be very exciting, to see how that can affect a field like robotics, for example.

Yeah. Double-click on just one thing about the efficiency you were talking about. I would expect the deeper it is, the worse it should be, quadratically. I'm not familiar with the pre-existing literature; I'm just working out intuitions.

Basically, what are the trade-offs that you've found that you think you might want to warn people about?

Because you were the guy who mentioned efficiency.

Sure. I was referring to one of the figures on our poster, also in our paper, where we compare the number of parameters that models have as we scale along the axis of depth and as we scale along the axis of width.

Yeah.

From our baseline architecture, the most basic one would be a width of 256: the hidden layers have 256 neurons, and then the depth is 4 hidden layers. The point I was making there is that when you scale along depth, the number of parameters that your model has is going to grow roughly linearly.

Whereas with width, you're making your network outputs wider, and then the input to the next network is also growing as well. The number of parameters your network is then going to have grows approximately quadratically. One of the experiments we did was examining, as we grow the number of parameters in our model by scaling along these 2 different choices, which one, for the same approximate number of parameters, yields better performance.

The depth curve goes like this. It jumps up pretty fast. That's present throughout our paper. With width, it grows a little more slowly. The takeaway is that if you're a bit more resource-constrained, scaling along depth might be better because there are fewer parameters with a smaller model, with a smaller number of learnable parameters.

Width is expensive.

Which is expensive. Exactly. In general, of course, more parameters are also going to be more expensive. So that's another consideration to think about when using these networks, I suppose.

Any other rules of thumb like that that I can extract? This is just the most basic one that I could think of.

I don't know if there are any others.

Yeah, I guess, to your original question about the trade-offs, one of the limitations that we mention is that obviously, if you make the networks bigger, they will take longer to run. If you double the depth, at some level of depth it might take twice as much to make a forward pass through the network.

However, within our paper, for most environments, we're able to saturate—get to almost perfect performance—without needing to get to 1,000 layers. Maybe 64 layers, for example, is sufficient. In this regime, the latency of the network is not necessarily a significant bottleneck.

You can imagine there are a lot of tasks, especially in RL, in which collecting data might be the bottleneck, and making forward passes through our network may not be the bottleneck. In our research, we specifically used the JAX-GCRL environment, which is a JAX-based, GPU-accelerated environment. We can collect thousands of environment trajectories in parallel at the same time, so that we're able to—

Right? This is built in.

This is built in so that we can collect 1,000 trajectories at the same time across all these environments. That makes sure that we have enough data to saturate the learning.

Wow.

That's like more data collection.

Okay, and I don't know if you want to expand on JAX, maybe.

Most people are familiar with Python and less familiar with JAX.

I think JAX is getting traction, especially in the RL field, because for online reinforcement learning, getting as much data as you can is the most important. There has to be a PyTorch equivalent, but anyway—

Do you have any tips for other people also exploring this kind of rollout?

Yeah. I can also recommend JAXRL for goal-conditioned RL, but there are also multi-agent JAX implementations and others.

So, going back to our paper, if you look at the plots, we only see this huge performance increase when we cross 50 million transitions. So I think the data is crucial here.

Yeah, I guess, even to build on that, I like drawing analogies to successes in other areas of deep learning. For example, in large language models, the reason why we're able to scale to such large networks is that we found a paradigm in which we can leverage the entire internet-scale of data to learn, right?

Data in RL traditionally has been hard to come by, but now, with these GPU-accelerated environments, we can collect hundreds of millions of time steps of data within just a few hours. I think that this serves as a really good test bed for us to find ways to scale up network capacity and get similar kinds of gains.

I think that has to go. Are you saying that you would do pre-training differently in LLMs? What's the difference in the objective now?

Yeah, very simply, the paradigm that you're referencing is next-word or next-token prediction, right?

It's very robust. [Laughter.]

I mean, how do you change that? Oh, I'm not saying I want to leverage insights from that to apply to RL.

I feel like you should go the other way.

You think you should go the other way?

Maybe. I mean, that would be a very interesting research direction, too. But actually, even on that point, one of the things I was thinking about is that the way our objective works is, in some sense, not exactly next-word prediction, but kind of like next-state prediction, right?

You imagine you're at some current state and taking some current action, and we want to predict whether or not this future state—this certain state—is a future state along the same trajectory or a different trajectory. So, in some sense, we are actually doing some sort of implicit world model—

Implicit world model.

Implicit—like, you know, in—

I don't know if that's a bad word.

Or, like, in language, you do a cross-entropy loss to classify the next token, right? And here we're just doing a binary classification of whether or not a state is a future state.

Yeah, yeah, it's a classification.

And so I do see that there are some parallels here that perhaps we should dig into more deeply and see what is at the core of what enables deep learning to scale. Then we can leverage those insights, distill them, and apply them across all different fields, whether it's language or reinforcement learning.

Did you get my meaning about the world-model stuff?

Yeah. Actually, I think I might have heard Professor Eysenbach talking about this yesterday at a poster. He was explaining to a couple of people that, because this is doing representation learning and trying to learn meaningful representations for a given state and action, but for a given goal, in some sense you can think of it almost like learning a model of the environment, learning a model of the world, but without having to do any sort of next-frame prediction or anything like that, which is a little bit more high-dimensional and complex.

Yeah.

I would say the angle that I'm trying to think about and push is that, instead of learning the next world, you're basically generating a number of candidate possible worlds and classifying them. To your point, that's exactly how I do things. Let's say I'm playing poker and I'm trying to classify what hands you have. There's a range of hands based on what you're doing. The more information I get, the more I resolve to, “Oh, I know exactly what hand you have,” based on what you're showing—or whether you're bluffing.

But that's a different thing. You know what I mean? I feel like that is the ultimate sort of end goal of representation, which is a world. But I don't know if that is too vague compared to the more concrete types of world models that, let's say, the video-generation people are doing.

And then, I guess, one other thing I'm also exploring—you mentioned that deep models are slower or more expensive. That is a trend in the inference world: making models shallower, right? I wonder if this short catchphrase I was thinking about, “deep teacher, shallow student,” would be a good deployment paradigm.

You push the frontier of capabilities with the deep teacher, and then you distill it back.

Yeah.

Actually, this is a good point. If you go to our website, this is one of the future directions that we list at the very bottom.

Okay.

Yeah, we would love to see if we could get similar performance. We achieve state-of-the-art performance on goal-conditioned RL in JAXRL by a significant amount, so it was very exciting to see the frontier of the ability to train RL agents pushed.

If we can do that in a way that is also just as efficient as standard networks, that would be very cool.

Yeah, because training doesn't have to be the same thing that you deploy for inference, right? You know what I mean?

Yeah. So if there are ways to distill down to a smaller model or prune the model and still retain performance, that's a very interesting research direction.

What else are your personal passions?

Yeah. Currently, I'm pursuing a direction of stitching in reinforcement learning. We're trying to generalize reinforcement learning from shorter sub-behaviors so that they are stitched or merged during test time. I think this is one of the last papers that I will tackle during my PhD.

Personally, I'm very curious: can we push the frontier as much as possible? If you look at our paper, we focus on scaling depth, but we notice that scaling width also improves performance. We also find that, by scaling depth, we unlock the ability to scale along batch size as well.

Cool, like—yeah.

Right. For context, in traditional value-based RL, scaling batch size is not super effective. But there is also other work in different areas of deep learning showing that scaling batch size is only most effective when there is a large enough network capacity to take advantage of it.

One hypothesis might be that the reason scaling batch size isn't as effective in traditional RL is because we've been using these tiny networks that haven't been able to capture enough. One of our experiments is that, because we enabled successful training of deep networks, we were able to use this as a great test bed for testing this hypothesis.

We find that, indeed, as we scale network capacity, we also unlock this different dimension of scaling by batch size. All that is to say that I'm very curious for someone with enough compute to take some of these environments, scale up the batch size, scale up depth to the maximum capability, and also scale width.

In the same way that, in language, we're scaling along so many different axes, can we unlock different dimensions of scaling as well? What capabilities can we get, and how far can we push the frontier of training these RL agents by doing that?

Before we pass to Ishaan, when you say “enough compute,” what kind of compute budget did you have? How does it— I just want to see what you guys got.

Good question. We wanted to make sure that this was quite accessible. The nice thing is that all of our experiments, even the 1,000-layer networks, can be run on one single 80-gigabyte H100 GPU.

So that's dollars?

Yeah.

Right, right, right.

Everything can be run on one GPU. But in theory, if we had a distributed training setup and could just blast compute through this and really wanted to push the frontier, it would be very interesting to see how things go.

Yep. Cool.

I've actively been trying to learn as much as I can about vision-language-action models.

Vision-language-action models?

Vision-language.

Vision-language, yeah.

Yeah. I'm curious about applications of representations for these.

Yeah, exactly—for robotics.

Yeah, we just released our episode with General Intuition, where, if you know a bit about their history, they started as a game-playing company. They basically have a vision-language-action model.

I saw a preview. It was very impressive. I'm not sure exactly how transferable it is to embodied use cases, but it doesn't have to be embodied. A screen is fine, you know. I don't know if you have any takes on that.

Yeah, it's definitely an exciting research direction.

Yeah, I think the concept of actions as something that you are outputting is actually not that popular in industry, right? Only because text has completely dominated the last 3 years, and tool calling—which is just another form of structured text—

And I feel like the action research is—I don't know what needs to happen in order to unlock the next phase in that. I don't know if you have anything interesting out here. Shout it out.

Yeah, there's a lot of cool work on leveraging pre-trained VLMs: you freeze them, and then you apply some sort of experts on top of that to output actions. There are also systems for doing hierarchical planning, maybe outputting some higher-level plan. This is a larger network that takes a little longer to do inference, and so it outputs its plans with less frequency, in some sort of chunk. From there, there's some sort of second system that operates a bit faster. I think there's quite a bit of interesting research in that direction. So that's sort of what I'm looking forward to.

Cool. Final question: hardest question you were asked at the poster session, or just favorite encounter—anyone famous that you met?

So I actually haven't gotten a chance to go to the conference that much. I'm actually working full-time now. So far, I literally just got my badge a few moments before the session, so I guess I wouldn't be the best person to answer that question.

No, no, no. People ask you stuff, right? I might pose it that way, because people are asking you or meeting you, and you can just give a vibe of what people are saying.

Yeah, I think people thought it was a very eye-opening paper because the objective is quite simple and elegant. We were able to—I don't want to say overturn, but sort of challenge—the conventional wisdom that RL is not super scalable, push it to such limits as 1,000 layers deep, and see continued improvement in performance. I think the general impression I've gotten is that this could be a really cool direction. If we can build along this direction, we can really scale along all these different dimensions and push the frontier of RL's abilities. I'm very curious to see how that goes.

All right. Well, thank you so much for dropping by. Congrats on the paper again, and good luck in your future work.

Thank you. Thanks for having us.