Improving Small Language Model Reasoning With A* Search

1. 0:00 – 0:07

   Intro

   1. AGAnkit Gupta0:00

      [upbeat music]

2. 0:07 – 0:45

   Why small language models need better reasoning at test time

   1. AGAnkit Gupta0:07

      We're here with Alex at our inaugural YCML at Startup School. Alex, tell us what you're working on.

   2. ABAlexander Braverman0:11

      So today I'm gonna be talking about my recently submitted paper to, um, AI Stats, and it was accepted to a NeurIPS workshop. It's test-time scaling for LLM reasoning via our method of A* Search. So the main problem we focus on is LLM reasoning, however, specifically for small language models. In a lot of cases, small language models are actually preferred to LLMs because of lower cost, lower latency, and also lower compute re-required. However, small language models actually lack the complex reasoning skills that larger l-language models have. This begs our, like, main research question: How can we provide high-quality reasoning to c-resource-constrained models?

   3. AGAnkit Gupta0:44

      Makes sense.

3. 0:45 – 1:30

   Limitations of prior approaches: distillation and MCTS-style search

   1. ABAlexander Braverman0:45

      So previously, there was two avenues for solutions. There was distillation and then Monte Carlo tree search or other tree search methods. So distillation, essentially, you transfer knowledge from a larger teacher model, say, like, a few hundred billion parameters, to a smaller, like, learner model or student model, right? However, this has some issues. For example, you have to have access to a larger teacher model and also run the model. And then also, this training process also takes compute, resources, time, et cetera. So it isn't the most convenient for a lot of applications of small language models. Additionally, there's tree search methods such as MCTS, where you have an LLM sample multiple paths, evaluate them, and then choose, for example, uh, the most correct or most, like, highly rated, uh, answer. However, a lot of times, evaluation actually requires external reward models, which are difficult to train and then also host because they're external, of course.

   2. AGAnkit Gupta1:29

      You might not have the-

4. 1:30 – 1:56

   Test-time A* search: reasoning as a tree search problem

   1. ABAlexander Braverman1:30

      Yes, of course. So at a high level, I will talk about what our method does, test-time A*. So we still like to view reasoning as a search, or a tree search in this case, and we wanna find the correct answer, or in this case, the destination node. For example, in this picture, say sixty is the, uh, correct answer. We wanna traverse the tree and get to sixty. However, instead of using a external reward model, we use the model's own self-critique as a heuristic. And in this case, it allows us to traverse the tree without any external reward models.

5. 1:56 – 2:25

   What “self-critique” means and how it becomes a heuristic

   1. AGAnkit Gupta1:56

      Can you define self-critique?

   2. ABAlexander Braverman1:58

      So for example, right, say I ask an LLM to answer five plus five, right? And then it says ten. And I ask the model, the same model, "Okay, this is the question. This is the answer. What would you give it on a scale of zero to one hundred?" This is a self-critique, for example, right?

   3. AGAnkit Gupta2:09

      Cool.

   4. ABAlexander Braverman2:10

      If it was like, say, if it said ten plus ten is five, it would be wrong. And then we use this, uh, we search the tree using an A*-inspired search cost function that in-in, uh, includes a score from this, like, self-critique.

   5. AGAnkit Gupta2:22

      Okay, cool. Makes sense. So the self-critique is part of the value function.

   6. ABAlexander Braverman2:24

      Of course.

   7. AGAnkit Gupta2:24

      Cool.

6. 2:25 – 2:52

   Algorithm loop: expand nodes, compute cost, pick lowest-cost frontier

   1. ABAlexander Braverman2:25

      Yes. This is a diagram. So essentially, you start off with a question, and then you, uh, generate the root node of the tree, and then you have this iterative cycle where you generate children to a selected node, and then you keep exploring based on the currently high-- or sorry, currently lowest cost, uh, node, which is generated using this formula, which I don't think we have time to go into depth. And then you essentially traverse the tree till you reach like a, like a solution, like the correct solution node, and then you obviously output the final answer.

7. 2:52 – 2:55

   Intuition: self-critique helps prune bad branches early

   1. AGAnkit Gupta2:52

      So what's the intuition for why the self-critique improves this here?

8. 2:55 – 3:43

   Walkthrough on a math word problem: branching + critique-guided traversal

   1. ABAlexander Braverman2:55

      Um, so essentially, you see here, this thirty-four is lower than the seventy-one. In this case, they're the same depth, so the only factor differentiating them is the critique, right? So this part of the tree has been essentially pruned because we saw this, uh, this-- the LLM was critiquing, and its own LLM realized this, uh, like, answer was essentially hallucinated almost. Where it couldn't keep exploring here because, of course, it was much, like, higher than this. Yes. So this can be actually seen better in this example with a, like a w-math word problem. So first, you generate the root node, which is just an LLM call to this answer. And then at each step, you expand it into two branches, two children, uh, using just, uh, this in the prompt and then the self-critique in the prompt. Additionally, you just, uh, traverse the tree. In this case, you can see the final answer. The correct one is on the bottom left. Does this visualize it well?

   2. AGAnkit Gupta3:42

      Makes sense.

9. 3:43 – 4:20

   Benchmark results on GSM8K and MATH500: accuracy gains from A* self-critique search

   1. ABAlexander Braverman3:43

      Awesome, yeah. Uh, and then for experiments, we, uh, check on some various mathematical reasoning benchmarks, including GSM8K and MATH 500. Here's GSM8K. In red is our methodology. You can see here it increased the reasoning accuracy compared to these other test-time scaling methods.

   2. AGAnkit Gupta3:58

      Can you briefly explain what GSM8K is?

   3. ABAlexander Braverman4:00

      So it's a mathematical reasoning dataset, uh, that's around high school level. It depends, like, what you define, but it's just, like, basic mathematical reasoning problems.

   4. AGAnkit Gupta4:07

      And so these are basically, like, off-the-shelf small models, and these are ones in which you use your tree search method.

   5. ABAlexander Braverman4:12

      So it all uses the same model.

   6. AGAnkit Gupta4:14

      Okay.

   7. ABAlexander Braverman4:14

      It's Qwen three four billion. However, these are just other algorithms to test time, uh, to do test-time reasoning.

   8. AGAnkit Gupta4:20

      Makes-

10. 4:20 – 4:58

    Efficiency ablation: accuracy vs tokens and wall-clock time

    1. ABAlexander Braverman4:20

       Yes. However, uh, of course, with test-time scaling, there's a big concern of you're using obviously more time and also more resources to increase accuracy. So we did an ablation to see for the amount of tokens and the amount of time, are we being efficient in how much we're actually improving the accuracy of the model? So on the right is the me-- uh, tokens, and on the left is the time. So for the same amount of tokens in purple is us, we actually increased by two to, sorry, three to four percentage points. And on the right, for the same amount of time, we're also increasing by three percentage points.

    2. AGAnkit Gupta4:50

       Makes-

    3. ABAlexander Braverman4:50

       This shows that not only are we just... Sorry, we're not just only increasing the time used or just other, like, kind of cheating factors almost. Yeah.

    4. AGAnkit Gupta4:56

       Makes sense. Yeah.

11. 4:58 – 5:42

    Ongoing direction: causal reinforcement learning for better credit assignment in agents

    1. ABAlexander Braverman4:58

       Additionally, uh, if you wanted to see the raw numbers. Yeah. And that's for that paper. And then this is some ongoing work. I'm still working on, uh, LLM reasoning. However, more, I'm now more focused on RL, uh, for LLM reasoning, and sp-specifically causal RL. So classical RL is mainly based on correlation. However, like, if I'm driving a car, for example, and I turn right and I crash into a tree, I know that that turn of the wheel is what crashed the car, right?

    2. AGAnkit Gupta5:22

       Right.

    3. ABAlexander Braverman5:23

       But in a lot of aspects with LLM agents, there's a lot of, like, moving parts. The LLM can't always, uh, like, assign credit to where credit's due, and it's actually known as the problem in RL as credit assignment. It's a very, like, big open problem.

    4. AGAnkit Gupta5:34

       So instead of MCTS as your iteration method, you're using RL.

    5. ABAlexander Braverman5:37

       Well, this is for an ongoing, this is new work.

    6. AGAnkit Gupta5:39

       For-

    7. ABAlexander Braverman5:39

       Yeah, this is, uh, separate from test-time scaling. Yeah.

    8. AGAnkit Gupta5:41

       Right.

12. 5:42 – 6:11

    Wrap-up: related prior papers and contact info

    1. ABAlexander Braverman5:42

       Uh, and so this is working on counterfactual reasoning to essentially isolate, like, and give credit, do credit assignment, uh, to create, like, agents that go from just, like, like correlation to knowing, like, cause and effect. 'Cause correlation does not equal causation, and to have true intelligence, true reasoning agents, you need, like, this ca-causation. Like, you need this understanding. Yeah. And then in terms of other work, these are two other previous papers. This is a systems, uh, paper, and this is a LLM hallucination paper. My website and LinkedIn are here if this is relevant. Thank you so much.

    2. AGAnkit Gupta6:09

       Thanks so much for coming here. This was awesome. Thanks.

Episode duration: 6:11