Furong Huang - Reasoning as Control: Adaptive Test Time Compute for Planning Agents

0:02 When people talk about self-improvement, often times we frame it as a model capability. The model can actually improve his reasoning and then, you know, started to actually learn from his own output and hopefully acquire some new skills. But I think model is just one layer. At the end of the day, it's going to be an agentic system that is going to be deployed and there are a lot of the infrastructure around the model, how the model is actually built. I think that's also a very interesting layer for us to understand how we can keep improving these infrastructure. And that could be even more important for very kind of vertical applications that, you know, utilizing language models in your very own environments, in your corporate, in your own application areas.

0:49 So, and now let's talk about what really happens in agentic system. Given a task, oh, okay. The standard slides, technical difficulty. so, given a task, right? So, you will be observing, you know, some kind of in you know, state, you know, transitioning from the environment. And then this kind of observation goes through your thinking process and at the end of the day, you wanted to make a decision, right? Maybe at every time you're going to make a decision, what is the best thing to do right now? And, you know, in order for agentic system work, often times because these tasks are very, very complicated, you wanted to decompose them into different agents, right? And you wanted to assign roles for them and then, you know, probably you wanted to design some kind of workflow. Like, you know, nowadays I think we call it loop engineering, right? So, and then for each of the agent in your workflow, and you know, you wanted to decide what is the action at every time step. And in order to do that, you need to do go through a thinking process. Often times a chain of thought, right, for each action decision, the agent has to generate a reasoning trace. As you can see, there is this kind of hierarchy of decision-making during run time, right?

2:08 So, starting from the very low level of thinking traces and token control, all the way to action control at every time step of your decision-making process. And then finally, in the very high level, workflow control, right? Can you optimize them? And I think today I'm going to say foundation models are actually more and more becoming a run time decision-maker. At the end of the day, the self-improving agentic system is, you know, possible that you can improve the model capability, but you could also improve the infrastructure, improve how the model is deployed during run time, across thinking, across actions, and also across workflows.

2:50 So, let's talk about a little bit about what we did, you know, think in each of these levels. Down there in the thinking process, I think we started with the steering problem, where, you know, you have some kind of reward model, and you're hoping that you can steer your large language model to align to the specific reward. Of course, if you're familiar with the math, you all know that you have a closed-form solution. And it turns out this closed-form solution gives you a very short answer to the very complicated question of alignment, which is, you know, whatever your base model gives you, you just add an additional steering signal to it. And this steering signal is nothing but something related to a trajectory reward under some specific optimal policy. I know this is very complicated math, but at the end of the day, if you look at what it's really doing, if you're thinking about next token sampling from a language, you know, auto regressive perspective, you're just looking at the language model's output, and then you just steer it with some kind of external signal that you can learn from your corporate data, right? so now, you know, the actual discrepancy here, and sorry, I mean, I promise the slides looked way better on my laptop.

4:05 But, you know, in this situation, there's this discrepancy about you really want to next token generation. You want it to steer your process as early as possible, but you can not because your reward model does not really support that. You have a trajectory level reward model, but what you really want is a early steering before you actually determine your thinking process. So, what we end up doing, for example, in a work, you know, transfer Q-star, we started to think, "Okay, if, you know, the reward models are really not designed for, you know, scale up, you know, the every token sampling, let's just, you know, do the auto completion and use the reward model in the right way." Of course, this is right, but it's really, really slow. On on an academic setting with my A6000 GPU, you can see if you have to generate 500 tokens, it will take you 14 hours, which is basically impossible. Right? So, you're doing it right, but you're doing it way too slow. So, in a very recent work, we tried to improve it by, you know, actually come up with like a specific model family of reward, and we essentially designed a token-level reward and significantly reduced the cost by orders of magnitude.

5:22 And as a result, you can see this is a phenomenal because you can do like, you know, any run-time steering of the model. You do not have to train the model at all. You can even do weak-to-strong guidance. You can use a very small model to guide a very huge model, and then you could even do multi-objective alignment in the sense that if you have very different objective, you can real-time adapt the objective to the specific user need. So, this is very cool, but you know, what about the action, right? So, we talked about steering your thinking process, but at the end of the day, your thinking process is really for service of your action. So, how do you make the decision? well, it is an agentic system, right? So, you know, we often know that agentic systems are learned using this called imitation learning, which is essentially distill from like expert you know, expert demonstration, but this you know, often times run into this kind of stuck loop problem where your agents will repeatedly try some actions for many times until termination or you run out of your token budget, right? So, that's very sad. So, people started to think, can we do some kind of you know, better smarter ways of think you know, learning through some world model training and the specifically, you know, people thought of like imitation or I mean imitating self-reflection. But, this is not actually self-reflection. When you are trying to self-reflect, you're not actually using yourself to reflect. You are actually mimicking some stronger model which teaches you how to reflect.

6:56 And it turns out it works better than imitation learning, but it's not really giving you very satisfying result. So, we did a very simple thing here. We call it genuine self-reflection. Rather than actually imitating how stronger models eat you know, think and self-reflect, we force ourself to do the reflection in terms of choosing which action is better at this time point. And you can imagine this is actually a world model. You're trying to build a mental model about which action is better under what state, but you're also forcing yourself to understand why.

7:33 And this very simple thing, you know, worked really phenomenally. You can improve imitation learning as well as reinforcement learning. It actually improves significantly out of distribution, you know, out of distribution kind of performance. And also, you know, performs way better than previous state of the art. So, like this genuine self-reflection really, really helps. And also, I think one really cool thing is that it actually generalizes to a general reasoning by just understanding how agents work better in a specific environment. You can see this, you know, kind of critical learning does not relate to any general reasoning capability, but somehow it magically improves the general reasoning capability. So, now you know, control decoding this can steer how the agents think.

8:28 And then act can actually help your agent understand how to, you know, under what state, what is the better action. Now, finally, I think what's really exciting is this kind of like workflow control. Autonomously decide what are your, you know, agent roles and how do you actually optimize the topology of the agents, you know, sort of like the workflow optimization. So, I have a very simple example here, but for the sake of time, you know, I'm just basically saying this is a very fixed workflow for like, you know, agentic safety kind of application. It works, you know, really well, but the question here is can you actually design, you know, such a framework such that you have a meta-designer as the query comes, you could come up with the best agentic framework, the best harnessing that is specifically for this very specific task. Well, can you do that, you know, in an autonomous way, right? So, people tried to do this. For example, they tried this thing called one-for-all and they also tried this thing for one-for-each. One-for-all basically means that you're going to find the best workflow in the world for the specific task. And then one for each is that you have to find the adaptability, you know, so that you can design the best flow for each of the query that comes into, you know, your model.

9:49 So, it it turns out we're thinking about this you know, kind of philosophy of foundation model, right? So, we don't want to do actually do one of these two, but we wanted to say can we actually precompute the workflow and sort of like, you know, the training time, but then reuse them during the deployment time. And can we achieve efficiencies through that kind of process? And it turns out what we did is we built a bank. We build a bank of workflows and then during test time, yes, we're going to just like adapt to the specific workflow real time really efficiently. And as you can see the result also worked really well. Like you can see significant improvement of the performance.

10:37 I can share a better version of the slide that is actually rendering correct. But, you know, if you trust me, so you will see like the performance is really nice. So, just to sum up, you know, we talked about self-improving agentic systems. We talked about how, you know, there are three levels of runtime decision makings that you can be very smart about how to allocate your computer resource with the limited budget, how to steer your thinking process as early as possible. So, you do not have to go through the very lengthy kind of process of thinking, but in the wrong way. How do you control action?

11:13 And finally, how do you do the workflow control? And with that, thank you.