Title: Part 2 Test-Time Compute Scaling
Authors/Date: Akanksha Bhardwaj, Azalia Mirhoseini — Stanford CS329A
Video ID: -Ggc37xLj_Y  |  URL: https://www.youtube.com/watch?v=-Ggc37xLj_Y
Playlist: https://www.youtube.com/playlist?list=PLangBM27OtEA
------------------------------------------------------------------------

0:05 So today, we are going to
talk about inference scaling.
0:09 We briefly hinted on this
in the first lecture,
0:13 so there are three stages
for LLM kind of development.
0:17 The first one is pre-training.
0:19 This stage is historically
takes a lot more time.
0:22 It could take several months.
0:24 A whole lot of GPUs that
are used for training these.
0:30 And then there's the
fine tuning step.
0:32 Usually, this is a less,
or historically has been
0:36 a less compute-intense process.
0:39 If you are talking
about trillions
0:41 of data, tokens in
the pre-training side,
0:45 on the fine tuning side, this
is orders of magnitude less data
0:48 that we use for fine tuning.
0:50 And then there is
the inference part,
0:52 where we get to use the model.
0:54 And we are going to
see different ways
0:58 that at inference
time, we can make
1:00 the model be better and become
more useful without changing
1:05 the parameters of the model
and without any fine tuning
1:10 of the models.
1:11 So let's first start with the
Large Language Monkeys paper.
1:16 Again, you are familiar with the
infinite monkey theorem by now
1:20 because we discussed
it on Monday.
1:24 But to recap, the idea of this
paper is that we have an LLM
1:29 and we have a problem as input.
1:31 Instead of generating the
response for this problem
1:35 only once, we are going to
repeatedly ask the same input
1:39 problem again and again.
1:41 This could be a
number of times, it
1:43 could be 10 times,
100 times, and so on,
1:45 and we will see how that
affects the performance.
1:48 And then we are going to have
a verifier that can tell us
1:51 which one of these responses
were the correct response,
1:55 and then that is going to be
the output of this system.
2:01 The reason we like that is that
by doing this repeated sampling,
2:06 we can make models
that are inferior.
2:08 For example, a Llama 3-8b or 70b
that is not as capable as GPT-4o
2:16 with single attempt.
2:17 But by doing this
repeated sampling
2:20 and selecting the correct one
among the generated candidates,
2:25 we can significantly
improve the performance
2:29 of these inferior models
and make them better
2:31 than these larger
and proprietary ones.
2:35 And these are-- here
you can see examples
2:38 across a really hard math coding
and different styles of math
2:44 problem question
answering problems.
2:49 Is this clear to everyone?
2:52 It just seems like the models
and smaller models already
2:56 know the answers to
these hard problems,
2:58 and just by doing this
repeated sampling,
3:01 we are eliciting and
surfacing those answers.
3:06 It's just they just don't tell
us that in the first trial.
3:13 Now, this paradigm
of repeated sampling,
3:19 works well across
pretty much any domain
3:21 that we have tried in so far,
including agentic benchmarks
3:26 such as SWE-bench, which mimics
how a software engineer goes
3:30 about editing code and
creating new patches.
3:36 So here, we are
seeing the coverage,
3:38 meaning the fraction
of problems that
3:43 are solved by, at least,
one of the samples.
3:46 And here on the x-axis, we have
the number of samples going from
3:49 1 to 1,000.
3:51 And what we are seeing here is
that a models such as DeepSeek
3:57 can--
4:00 I believe this was the
DeepSeek-V3 model can outperform
4:05 models such as Claude 3.5 or
o1 preview after 1,000 samples.
4:12 It can solve more coding
and software engineering
4:14 problems than those.
4:16 And the reason
this is interesting
4:17 is that, we are doing--
4:20 in this case for
coding problems,
4:23 if we have unit
tests that we can
4:26 use to select which one of
these samples is correct, then
4:29 we have an end-to-end
automated way
4:32 to create a more capable
model, in this case,
4:37 out of DeepSeek and
open-source model.
4:44 Now, again, we saw pre-training
scaling laws in the last class,
4:49 and that was, we can predictably
reduce the test laws given
4:56 these equations in this paper.
4:58 And these equations might
change given the setup,
5:01 but there is a predictable
path to increase data, amount
5:06 of compute use, and the
number of parameters
5:08 in the model and use that to
predict, or reduce the test
5:14 laws.
5:15 So this is for pre-training,
but it turns out,
5:18 we can define a similar
scaling laws, but this time
5:22 for inference, or
test time compute,
5:25 and that is something that
we showed in this paper.
5:27 Basically, we're showing
that the relationship
5:30 between coverage and
the number of samples
5:34 that are being drawn in
parallel from the model
5:37 follow an exponential power law.
5:39 So assume this says,
coverage, meaning
5:42 how many problems are solved by,
at least, one of the samples,
5:46 and then k is the
number of samples.
5:48 And a and b are these
kind of coefficients
5:51 that we can find through
the curve-fitting parameters
5:55 that we can find through this
scaling behavior of the model.
6:03 So concretely, as we scaled the
number of samples across a range
6:09 of models from
Llama 3 8B, 7, 8B,
6:13 and models like Gemma or Pythia.
6:16 And this is for various
range of parameter sizes,
6:20 going from 70 million
parameters all the way
6:23 to 70 billion parameters.
6:25 We can see that we can
predict this power law
6:29 and have a predicted curve
that for the most part,
6:35 very closely follows the actual
practical coverage that we
6:41 are seeing as we increase
the number of samples.
6:44 So basically, the
way to think about it
6:46 is that we can predict to
achieve a certain coverage,
6:53 how many samples we
are going to need
6:56 and how much resources we
should allocate to achieve that.
7:06 And again, these curves
show that this behavior
7:10 is true for a range of models
and a range of parameter sizes.
7:15 Again, one of the very
interesting properties
7:19 that we saw here is that
for a very small model,
7:21 like a 70 million parameter, we
also see this scaling behavior.
7:26 And, of course, this is true
for much larger models as well,
7:29 and it is true across domains.
7:33 Now, we looked a bit more
into why we see this power law
7:41 behavior in repeated
sampling because it
7:47 might be at first a
little non-intuitive,
7:50 because if the probability of
the pass at one being correct,
7:58 the probability of a correct
response to a question
8:02 is, for example, p,
then we can calculate
8:07 the probability
of, at least, one
8:09 of the answers being correct at
k for an individual problems.
8:13 So for a problem i, if we have
the pass at one for problem i,
8:19 we can calculate directly what
is pass at k for problem i
8:24 with this exponential equation.
8:28 I hope this is
clear for everybody,
8:30 but basically, this 1 minus p is
the probability of the first k
8:35 answers not being
correct, and then
8:37 1 minus that is the probability
of, at least, one of these k
8:40 answers being correct, so this
is for an individual problem.
8:45 But what we are seeing that
across a suite of problems,
8:49 we are seeing the power
law, scaling law, meaning
8:57 across these problems,
we are seeing
9:00 the equations are changing.
9:02 So the question that
we wanted to answer
9:04 is what behavior should we see
across the problem examples
9:11 in a data set that justifies
our power law scaling laws?
9:20 And by looking into that
further, what was observed
9:26 is that in order to have the
scaling laws that we observed,
9:32 the sufficient and
necessary condition for it
9:35 is that we have a long
tail of hard problems.
9:42 Basically, from a
mathematical perspective,
9:45 we are going to need that
because this is the power law
9:48 scaling, this is the per problem
exponential scaling laws,
9:53 and this is like this sort.
9:55 We need to have this curve for
the past at one of problems
10:00 in order to justify
this power law scaling.
10:05 And empirically, when we
looked into the problems,
10:09 we saw that this
condition holds.
10:12 And again, what
this condition is
10:14 is that we have a long tail of
really hard problems on the most
10:21 complex side of things.
10:23 For example, let's take a look
at this model and the count
10:30 of problems that are
solved at pass at one
10:33 and add more and more samples.
10:36 It seems like across
all these problems,
10:38 a large portion of
problems are simpler
10:41 and they're solved
in pass at one.
10:43 And as we increase
the complexity
10:45 of the problems, where the
pass at one probability
10:48 is going to become
lower and lower,
10:50 we have a long tail of problems
that have that property.
10:56 So these curves
that you see here
10:58 from high at simpler problems
and low at harder and harder
11:03 problems, meaning lower
and lower paths at one
11:05 is the type of problems,
is the type of behavior
11:08 that we are seeing
across these tasks,
11:11 and that is why we
can see the scaling
11:15 laws that we have observed.
11:20 Now, what this means for the
science and the engineering
11:30 around LLMs is that while
previously companies
11:34 would spend hundreds
of millions of dollars,
11:37 or billions of dollars
on pre-training alone,
11:42 and then a lot
less on fine tuning
11:44 and then almost like
nothing at each inference
11:48 call because it used
to be like a one time
11:50 back and forth between the user
and the query and the model.
11:53 Now we have a new
paradigm, where
11:56 we can spend a whole lot
more compute on inference,
12:00 and use that to increase
the capability of model.
12:03 And this compute the inference
compute can be done offline.
12:06 We can release our agents
to go solve a problem
12:10 and just keep generating
tokens and keep
12:13 improving the quality of the
answers that they generate.
12:20 Now, when we talk about
this repeated sampling,
12:24 something that's
important to have in mind
12:29 is the need for
automated verification
12:32 because if you have
a bunch of samples,
12:35 we still need to know which one
of them are the correct ones
12:38 and how to go about
selecting that.
12:43 In certain domains, verification
is easier than the others.
12:48 For example, for certain
types of math problems,
12:52 you can use formal proofs.
12:55 So given a proof,
the strategy, we
12:59 can run them by these
formal proof software tools
13:03 to verify whether each
step is correct or not.
13:06 We can write unit
tests for coding.
13:09 And in some cases, writing a
unit test arguably, arguably,
13:15 is a much simpler task than
writing the entire program that
13:18 solves a problem,
so the unit test
13:21 can be written, for
example, by humans,
13:24 and we can use that
as our verifier.
13:29 There are other cases
like AI as a compiler,
13:33 and this is one of the projects
in the lab and the direction
13:37 that I'm very excited
about, and that
13:39 is when we can generate
compiler lower level
13:42 code, like CUDA for a source
code like PyTorch with LLM.
13:48 And the reason
this is verifiable
13:50 is that, for example
here, we can always
13:56 compare the output of CUDA
with the source PyTorch
14:01 and see if they're the
same for any given input.
14:04 So we don't need if we want
to say whether the model has
14:08 generated the correct CUDA code
for our source PyTorch code,
14:13 we can just see whether their
outputs match for any inputs.
14:21 For those of you who
are not familiar CUDA,
14:24 I hope you all are, but
it's a lower level language
14:28 for writing more hardware
aware and optimized
14:31 for hardware code on GPUs.
14:38 For example, here,
we have a data set,
14:43 or benchmark on CUDA
generation called KernelBench.
14:47 And here, again, we are seeing
the same linear improving
14:52 coverage in this case
for CUDA code generation
14:56 as we increase the
number of samples
14:58 that we take from the model.
15:00 And in this case,
we have this, again,
15:02 by default, perfect
verifier that we
15:05 can use for measuring the
correctness of our CUDA
15:10 generated code.
15:13 And there are other
examples of this.
15:18 Basically, any
kind of translation
15:20 between two languages.
15:21 If you want to port from Python,
or C++ to Java, or vice versa,
15:27 again, this is true.
15:28 You can measure the
equivalency much easier.
15:34 Now there are domains that
we don't have verifier for.
15:38 And in those domains,
what we observe
15:40 is that there is a
large gap between best
15:43 of n methods such
as majority vote
15:46 voting and model-based
rankers and what is
15:52 the true coverage of the model.
15:54 Let's take a look
at one example here.
15:58 So here, we compared a
few different methods.
16:01 We are going to learn
about reward models
16:03 or verifiers on Monday
more, but let's take
16:06 a look at these examples.
16:08 So we have a number of
samples, and then we
16:11 are measuring the
success rate here,
16:13 like basically what would be
the output of this system?
16:17 Majority voting,
or the green one
16:20 is we just see which answer
among the generated answers
16:24 has appeared the most,
so based on that,
16:28 we just select that
answer or that response
16:31 as the output of our system.
16:33 Does that make sense?
16:34 It's just majority voting.
16:36 And the blue line
here is coverage.
16:39 This is assuming we had a
perfect selector, or verifier
16:43 in this case.
16:44 And as you can see here,
there is a large difference
16:47 between majority voting, which
plateaus after 10 or 50 samples,
16:53 and what we could
possibly get if we
16:55 had a perfect verifier, so
this is a very, very large gap.
17:01 Especially, and this gap
is even more pronounced
17:05 if you have harder problems.
17:06 So for example, the
capital MATH here, data set
17:09 is a harder, more complex
data set than GSM8K
17:13 And here, we compared
other methods
17:16 like using a reward model
that measures, gives
17:25 a reward per answer, and we
take the best of those, or ways
17:30 to combine these tools.
17:31 We will learn more about
these reward models.
17:34 Basically, they're
LLM-based LLMs
17:37 that are trained to score
the response quality.
17:44 And even with those,
there is a large gap.
17:47 And this is what we call the
generation verification gap.
17:52 So the generation,
it turns out, we
17:54 can generate a lot
of good responses,
17:57 but the verification
is basically this gap.
18:01 This tells us how
much we can actually
18:02 capture those responses
or those correct answers.
18:09 The reason methods such as
majority voting doesn't work is
18:13 that, as we look into
the quality of answers,
18:17 it turns out that for some
problems that are hardest
18:21 to solve, sometimes we see that
they're solved maybe even once,
18:27 or two times, or three times
across the 1,000 or 10,000
18:32 samples that we
take from the model.
18:34 So these correct problems are
actually very rarely solved
18:40 by the model with the samples
that we are taking from it,
18:43 so it makes sense that a
majority voting mechanism cannot
18:48 capture those
because they're rare.
18:51 But for many of the
problems in this case,
18:53 especially for the
GSM8K, which is simpler
18:56 because the answers, or
responses are simpler
19:00 to generate, majority voting
works for many of the problems,
19:04 but even in this case,
for some of the hardest
19:07 problems, the frequency of
correct answers actually
19:12 pretty low.
19:14 So that makes it harder from
a verification perspective
19:16 how do you find them?
19:18 How do you find that
which ones are correct?
19:22 Now, I want you to--
19:24 here are some
discussion questions
19:28 that I want to put it here
for something like two minutes
19:32 or so, and I want you
to think about it.
19:35 And hopefully, we'll
have some volunteers
19:41 and volunteer some of you
to answer some of them.
19:46 So two minutes for this.
19:48 Maybe we talk after.
19:52 [SIDE CONVERSATION]
19:54 OK, let's get back.
19:57 Who wants to answer question 1?
20:01 I'm very interested
in how you would
20:03 build on this, if you were
to do some research here,
20:08 or what takeaways you
got from these papers.
20:17 Anyone?
20:21 Yes.
20:22 One takeaway is that sampling
any answers and issues.
20:31 If you have a good
verifier, it's greatly
20:34 helps with your accuracy.
20:36 And the second takeaway is that
the quality of the verifier
20:40 is very important
in this approach.
20:44 That's right.
20:49 Anything else?
20:51 Yes.
20:52 [INAUDIBLE] it is possible to
write a model to [INAUDIBLE]
20:56 that way make examples.
20:59 For example, like the model
[INAUDIBLE] descriptions
21:04 and the model for the
solution that is a [INAUDIBLE]
21:12 even if you're not sure
the solution is right.
21:16 Are you talking about
revising a solution--
21:19 Yeah.
21:20 --by next round?
21:22 We're going to hear
about it in a minute.
21:25 Sounds good.
21:26 Yes.
21:27 I'm looking at it
two directions.
21:29 One is the [INAUDIBLE]
and accuracy.
21:32 And for that, I would look
at a hybrid approach, where
21:35 I would develop a
knowledge graph along
21:38 with advanced documentation.
21:42 And then any returns
to the rank method.
21:45 When [INAUDIBLE] comes time the
score is lesser on the accuracy
21:49 and more or less
rely on the knowledge
21:53 to give me more accurate
answers, that way you
21:55 can increase the speed in
the answer while [INAUDIBLE].
22:01 So you're suggesting use
RAC to boost or compare?
22:08 With the knowledge first,
and then only when they
22:13 feel that the accuracy is
not coming up to the mark,
22:15 we go through the
right documentation
22:16 because that also lowers
the time, and there is--
22:21 So you're asking that can help
with efficiency in the time
22:25 domain?
22:25 Is trying to reduce the
dependency and damage
22:29 to increase the
accuracy and precision.
22:33 One more.
22:34 Maybe you were waiting.
22:36 Yeah, so something that I've
thought about also is just--
22:41 and even if it's easy to medium
level of difficulty with respect
22:47 to verification or
verifiable domains,
22:51 it would be cool to see how
pass at k changes with respect
22:57 to k, when we employ initial
search techniques, almost
23:01 like a self-study approach.
23:02 If there's some way
to maybe consolidate,
23:04 let's let the model
explore this domain,
23:07 gather insights as
if we're not asking
23:10 it to sample any solutions yet.
23:12 And maybe that can give
way to some output context,
23:15 and we could then funnel
into parallel sampling.
23:17 We might see maybe better
trends with respect--
23:20 Absolutely.
23:21 And we're going to learn
about some of them.
23:22 There are ways to
improve test time scaling
23:25 beyond just repeated sampling,
including self-study, search,
23:29 and tool use.
23:31 Maybe one more.
23:34 Yes.
23:36 Well, one interesting
extension might
23:37 be to [INAUDIBLE] verifiers.
23:41 Without a good verifier,
then the best you can do
23:45 is [INAUDIBLE] what you showed
in the previous slide, which
23:47 usually doesn't work very
well for hard problems.
23:50 But maybe there
are some problems,
23:52 where it's hard to verify the
answer is correct because it's
23:56 too expensive, or it's
just too hard naturally,
24:00 but it's easy to
find an answer wrong,
24:03 so you can project
like bad answers.
24:06 Maybe that would be more
obvious in terms of testing.
24:08 And then you can just
keep [INAUDIBLE] something
24:11 [INAUDIBLE].
24:13 So one interesting
direction would
24:15 be trying to look at
different domains, where
24:18 there is much difficulty.
24:21 And also maybe verification
is not fully accurate,
24:23 or maybe it's a bit
lost, but verify
24:25 it's its own model
because it's a simulation.
24:28 It's like a physics, or a
molecular dynamics simulation,
24:31 or something, and then how
you would deal with that.
24:34 Basically, there are
other ways to filter out
24:37 the incorrect answers.
24:39 Maybe use a simulation,
or some other tool use,
24:45 or another model to do that,
depending on the domain.
24:49 So all of the data set here,
the 10,000 different samples per
24:53 problem is available
on Hugging Face.
24:57 If anybody is
interested to go around
24:59 how to shrink that
generation verification gap,
25:05 that could be potentially a good
research project for the course.
25:10 Yes, last question.
25:12 Yeah.
25:13 Can one try to generate multiple
verifiers like maybe generate
25:17 10,000 verifiers and then do
majority voting on the result
25:22 of the bad verifiers?
25:24 Yes.
25:25 That's also a very
interesting research question.
25:28 So on Monday, we're going to
learn about this weaver, which
25:31 is a work that we've
done on ensembling,
25:34 but a week supervised
ensembling of verifiers.
25:37 We don't have 1,000 of them, but
we can have 10 or 20 of them,
25:41 but that direction is
definitely also interesting.
25:44 It's very costly, though,
from a compute perspective.
25:47 Maybe we should move on to the--
any other about the other two
25:52 questions, anything?
25:54 We discussed some of
it in our discussion,
25:56 but anything that you
want to share last minute?
26:00 Yes.
26:01 What is up with the investigated
whether the coverage
26:04 number is real, like software
engineering tasks, for example?
26:08 Or could you pass
the test cases,
26:09 but it's actually
about [INAUDIBLE]?
26:12 For example, for
the math problems,
26:14 we manually looked into
the answered questions,
26:17 and it was something
above like a 90s, 7%,
26:22 8% by manually checking that,
but that's also always a case
26:27 for failure mode, if your unit
tests don't have true coverage
26:31 of the code, so the quality
of the verifier matters.
26:37 Yes.
26:38 Maybe it's moved up to be
able to generate [INAUDIBLE],
26:43 so you're verifying each node?
26:45 Yeah, let's talk about
that and what's coming next
26:49 and how we can make different
ways of test time scaling.
26:55 So with the monkeys
framework, we
26:58 saw that repeated
sampling is effective,
27:01 but there are other
ways to scale sampling.
27:04 And let's take a
look at this paper.
27:06 Scaling LLM test time
compute optimally
27:09 can be more effective than
scaling model parameters.
27:12 We going to touch on
certain aspects of it that
27:15 is most relevant to test time
scaling, so there are two ways.
27:20 Parallel sampling
is one way for us
27:22 to generate multiple
answers per input question.
27:27 But there is another way to
generate, or generate answers,
27:31 or use test time scaling
to improve the quality,
27:34 and that is when we do
sequential revisions.
27:38 For example, the model comes
up with one initial steps
27:42 of the problem, or
one initial approach
27:44 for solving the problem,
and it keeps revising that,
27:48 or it keeps improving
and adding to that.
27:52 And that is the sequential
approach, so instead
27:55 of asking the model
multiple times,
27:57 we let the model know that it
can keep revising its answers,
28:01 look into that from
a different angles,
28:03 and continue doing so
until it's confident
28:08 that it's ready to
generate an answer
28:10 and then finally output that.
28:12 Yes.
28:12 So is this by our own
prompting or is it the models?
28:17 In this case, let's assume
by our own prompting,
28:20 but when you look at
the reasoning models,
28:23 they are now trained.
28:24 Internally, we see this
kind of behavior in them,
28:27 so they do this revisions
sometimes in their own answers.
28:35 Now, there is another way
to scale test time compute,
28:39 and that is how we
select the answers.
28:43 So for the parallel sampling,
we looked into the outcome
28:47 like there are n
different responses.
28:50 We look into the outcome and
see if they're correct or not.
28:53 We mentioned unit
tests and so on,
28:56 but there's also this other
way of learning based,
28:59 or outcome reward models, where
they look into the outcome
29:03 and say whether
they're correct or not.
29:06 Of course, their
accuracy, or quality
29:08 can be limited, especially
if we go to new domains.
29:12 But outcome-based reward
model is another way
29:16 that they look into
the final answer
29:17 and they generate a score.
29:19 Basically, these models are
trained to give a reward based
29:23 on an output.
29:25 Then there is also the
process reward models.
29:28 So the difference between
process and outcome
29:31 is that process
reward models are
29:33 designed to give a score per
step of the generated solution.
29:38 For example, we have
a math problems,
29:40 there are five steps
that the model takes
29:43 before generating a solution.
29:48 Process reward models, again, we
learn about this more on Monday,
29:51 I believe, are trained to
give a score for the steps
29:56 of the generation.
29:58 So that's the difference
between the two.
30:02 Now, for best of n
sampling with a verifier,
30:05 again, the problem is easy.
30:06 We generate parallel
samples, and then we
30:09 can use an outcome-based
reward model
30:12 to see which one has
the highest score
30:15 and then take that
as our answer.
30:20 Then there are other ways that
we can bring these reward models
30:26 into the processing, into our
test time scaling framework.
30:30 For example, here, let's assume
at each level of answering
30:35 this question, we have a
budget of four samples.
30:40 We take these four
samples, and then here,
30:43 use a Process Reward Model
or PRM to in this case, we
30:49 select the top two with the
highest score from our PRM.
30:53 Then again, we sample
from them and we continue.
31:00 And we can have a threshold
for selecting the top two
31:08 kind of responses in
this case and expanding
31:11 the tree from there.
31:12 So here, we are doing
a beam search approach
31:16 that uses the process
reward model to guide
31:19 how we are sampling our model.
31:21 Is that clear?
31:24 Yes.
31:25 How do you choose the options?
31:27 Based on the PRM score.
31:28 So the process reward model
is trained to generate
31:33 a score imagine
between 0 and 1 based
31:36 on the step given the question,
and/or just the step alone.
31:44 Is the PRM fine tuned for that?
31:47 Yes, PRMs are generally also
fine tuned from language model.
31:51 As you can imagine,
they would work better
31:53 for in-domain types of tasks, so
if you have a certain benchmark
32:01 that you want to optimize
for, of course, training
32:04 a PRM on a subset of
that task helps you
32:07 with test time
scaling on that task
32:09 better than across new task.
32:12 But because they're
language model based,
32:14 they do exhibit some
generalization across tasks
32:17 as well.
32:19 Yes.
32:21 For the process reward,
is this a reward
32:24 that's being
computed in practice?
32:27 I don't know what the research
is per token that is being done,
32:31 or maybe blocks of tokens
and then it's being computed.
32:34 Yes.
32:35 It's not per token.
32:36 It's per step.
32:39 And this is something
we're going to learn about
32:41 in future lectures.
32:42 Imagine, these are
like a math problem,
32:44 and then there are
steps to answer that.
32:48 One way to do it is to
have humans annotate
32:50 the quality of each step like
yes or no, this is good or bad,
32:54 and then training the model
based on each step, which
32:57 could be like you can assume
it's like one intellectually
33:02 meaningful chunk of the
result, or it could be you
33:05 can define it as a
sentence and so on.
33:13 And, of course, as
it's very intuitive,
33:19 if you combine the two of
outcome based reward model
33:23 and parallel based upon
sequential revisions,
33:26 you get better results.
33:27 This field is
still there's a lot
33:29 of unknown questions around
this, but combining the two.
33:36 For example, one simple
way to combine that
33:39 is instead of one sample
and revisions, you do two
33:46 and you get to
two final answers,
33:49 and then you use the best one,
either using your PRM or ORM
33:55 to select how the two are.
33:57 But you're augmenting
each parallel sample
34:00 with these sequential
revisions, and you
34:02 can use a PRM to
guide that search
34:05 for each sequential revision.
34:08 Something that's interesting
is that the PRM, there
34:12 are many off the shelf PRM's
that you can use for your tasks.
34:16 So you can go and say,
take a pre-trained PRM
34:18 and use that to guide the
search for your own task.
34:22 And you can potentially
find a PRM that
34:24 is most relevant to the type
of task that you're doing,
34:28 and you want to do
test time scaling on,
34:30 or you can train
your own PRM and ORM
34:34 based on your own data as well.
34:40 So this paper looked into the
math data set and something like
34:45 12k train set 500 test
questions, used the palm model.
34:51 And what was interesting
here is that they defined
34:54 a notion of difficulty
for measuring
35:00 these different
scaling behaviors,
35:02 and for that, they
looked into the pass
35:04 at one performance of the
model on each problem.
35:08 And by just measuring
that, they define OK, these
35:11 are the problems
that are hard, these
35:13 are the problems based on
how many of the generated
35:17 answers are correct across a
body of a number of samples.
35:20 They had five different bins for
the complexity of the problem,
35:26 and they did that, as we see to
get some conclusion about what
35:30 test time scaling works better
for harder or easier task.
35:36 They also train their
own process reward model,
35:39 and even they fine tune a model
for creating these revisions.
35:43 But right now, every model that
we use pretty much it's capable.
35:48 It's just as long as
it's instruction-tuned,
35:50 is capable of doing revisions,
if you ask the model to do so.
35:57 And there are some
interesting observations here.
36:00 For example, here is
the majority voting
36:04 for these problems.
36:05 The accuracy as the
generation budget
36:08 increases, and then the
purple one is the ORM one.
36:14 And the green one is the PRM
one, and then the blue one,
36:21 if you really try to optimize
different ways of mixing
36:25 and matching things.
36:27 And this area is the one that
still is an open research
36:31 like how do we do that?
36:33 How do we go about mixing
and matching revisions
36:36 and parallel scaling to
optimize the test time
36:41 scaling because at
the end of the day,
36:42 we want for each accuracy, we
want the minimum generation
36:46 budget to hit that accuracy?
36:49 And in a few slides ahead,
we talk about one method
36:53 to optimize this scaling and
budget performance per budget.
37:00 They also looked into these
different bins of difficulty,
37:03 so number 5 is the most
difficult questions
37:06 because the accuracy is
the lowest and number 1.
37:09 And then here on
the right, there's
37:11 these colors is the
sequential to parallel ratio.
37:16 And while they're saying that
OK for the easier problems,
37:20 it seems like this darkest the
darkest purple color achieves
37:26 highest accuracy,
and that's where
37:29 we have a lot more sequential
tokens compared to parallel.
37:34 When the problems are more
difficult, and to some extent,
37:37 those are some of the
more interesting problems,
37:40 this optimality
is harder to say.
37:42 It's harder to say whether it's
sequential or the ratio changes
37:47 even from bucket four
to bucket five on how
37:51 you should allocate that.
37:57 Another observation from
this paper was that--
38:02 and this is still
a work in progress,
38:04 and I believe still the
true observation as of today
38:08 is that for easy and
medium questions,
38:11 it seems like
additional test time
38:14 compute can be more
favorable than scaling
38:18 pre-training of the model.
38:21 So here, we have these two.
38:24 The easy are the
green, medium is blue,
38:27 and orange is the
hard questions.
38:31 It seems if we look into the
ratio of inference tokens
38:34 to pre-training tokens, it is
positive for the easy and medium
38:46 problem, but for
the hardest problem,
38:48 it seems larger models
like a more pre-training,
38:52 or larger models
still do better if you
38:56 are from a tokens
optimized perspective.
39:01 And this is something
that we have
39:03 seen in many of our research
in the lab or around the lab
39:07 as well.
39:09 It seems like models, simpler
open source models, or smaller
39:14 models are becoming
more and more useful
39:17 as we take more test time
compute, apply more test
39:20 time compute on them.
39:22 But still, for the very,
very hard problems,
39:24 still the frontier
models, which presumably
39:27 use more pre-training
and they are larger,
39:30 they do better even if we had a
whole lot like infinite budget
39:35 for test time scaling.
39:36 Of course, we don't
have infinite budget,
39:38 but any reasonably
large test time scale.
39:41 Yes.
39:42 In this case, how do you
inference for this pre-training?
39:46 Because it seems
like for inference,
39:48 you use it much
like [INAUDIBLE],
39:50 whereas for pre-training
tokens, it's you use it once
39:53 and you can use it twice.
39:54 Yes, that is true.
39:56 But in general, from
a-- but this is still--
40:00 so basically, yes,
pre-training is done once,
40:04 but scaling is every
time we run the model.
40:17 But still, this is an
interesting observation
40:20 from the perspective of, are
we done doing pre-training
40:24 because we can just keep
sampling from the models we
40:27 have, or no, pre-training
is still helpful
40:31 and we should do more and more?
40:34 Not everyone can
pre-train a large model.
40:38 So even if
pre-training was better
40:44 than test time
scaling in any scale,
40:46 not everyone can do that because
it's a very expensive process.
40:49 So this is very interesting
that for many problems,
40:52 we can do test
time scaling, even
40:54 if we do it every time for a
new problem, because that's
40:58 feasible.
40:59 But also, on the other hand,
for the hardest problems,
41:02 we still need to
do pre-training,
41:04 and better models
still do better.
41:07 That makes sense.
41:07 I was just curious
precisely in the graph
41:10 when it says like the ratio is 1
for retraining of, for instance,
41:15 retraining.
41:15 Does that mean that it's
[INAUDIBLE] training,
41:19 there was also one training
used for a single task?
41:23 No.
41:23 No, it's not exactly that.
41:25 It's more like some
ratio of pre-training
41:27 to inference compute.
41:29 It's not a one-to-one because it
doesn't make sense in that case.
41:34 My question is, why is it
that those larger models still
41:37 count better?
41:40 Can we take a smaller model
and yet make it an expert
41:43 on hot questions?
41:45 You can always fine tune
a model to become better
41:50 on a certain set of
questions, but we
41:53 are talking about a more general
perspective of let's have
41:57 a whole lot of data
that's general, not
41:59 specific to a specific task and
just compare the common recipes
42:08 that we use for training models.
42:10 Let's just follow that
and then see how much
42:13 they do better at test time.
42:22 Any other questions?
42:25 So maybe we discussed some
of these already, but here
42:30 are some more questions
that I give you another two
42:35 minutes, let me check the time.
42:36 Yes, we do have time,
to think about this.
42:41 This helps you absorb
some of the topics
42:45 that we learned, or heard
about in the previous slides,
42:50 and then we discuss
them in two minutes.
42:56 OK, should we start?
42:59 Anyone wants to add
any comments here,
43:01 or any questions you might have?
43:06 Yes.
43:07 So do you try sequential and try
too-- what about tree search?
43:12 And in that case,
you're doing the--
43:16 You're mixing the both.
43:17 Yeah, you're kind
of mixing, right?
43:19 Yeah.
43:19 There is no universal
answer that can be there.
43:22 But the general sentiment
was that does better,
43:26 like some sort of combining
the two than doing--
43:29 [INAUDIBLE] you do, so reduce
a lot of the compute that
43:33 you're doing or reuse
[INAUDIBLE] that you're doing
43:35 because it's a tree--
43:37 Yes, yes.
43:39 That is when we discussed the
beam search stuff and the PRM,
43:43 like bringing a PRM to
cut the tree at each level
43:49 and explore only the
more promising branches.
43:57 Anything else?
43:58 Yes.
43:59 [INAUDIBLE] as the thinking
that that's the easier problem,
44:06 the more perhaps you
have towards the success
44:08 of solving it.
44:09 Therefore, it could be
more sequential [INAUDIBLE]
44:13 problem because you would--
there are few students solving,
44:16 so you would need
many parallel process
44:18 to find at least
one viable solution.
44:22 And we started thinking
a little bit also
44:24 like to do this in a chain
of thoughts kind of way,
44:26 where you are doing
a chain of thought,
44:28 but verifying through all the
steps into our chain of thoughts
44:31 instead, and then pruning
from these branches
44:34 that are violating anything.
44:37 Yeah.
44:38 So in terms of sequential,
you're saying that for easier
44:41 ones, easier problems--
for harder problems,
44:45 you mean you need more
exploration potentially
44:48 for easier--
44:49 To the correct answer.
44:51 Or for easier problems,
because like any kind of path
44:57 might lead us to
the correct answer,
44:59 we might not need to explore
different parallel solutions.
45:03 Yeah, that is an intuitive kind
of way of thinking about this.
45:10 Any other thoughts?
45:19 In general, this is a
very interesting kind
45:22 of problem to think about.
45:23 How do you-- let's assume that
these tokens that you generate
45:27 is like a knob that you
have to increase problem,
45:30 the quality of the responses,
how do you allocate this knob?
45:34 And how do you elicit these
generations from the model?
45:38 So let's think about
in the following,
45:40 I'm going to talk about
one way to optimize
45:43 these kind of test time scaling
to generate optimized answers.
45:47 And that is the Archon paper,
an architecture search framework
45:51 for inference time scaling.
45:55 So the problem
state-- and your TA
45:57 is one of the
co-authors of this work.
46:00 So the problem in this
paper is that, again, we
46:03 are interested to see
how we can mix and match
46:07 different inference time
scaling methods to build
46:10 the optimize the frontier
of capability or correctness
46:14 versus cost.
46:18 And the other thing that--
46:23 again, all of these questions
basically saying the same thing
46:26 that how we can optimize
assigning inference
46:29 compute two different tasks.
46:32 And basically, design mechanisms
that gets us to the high quality
46:37 answers while not wasting too
many tokens and generations.
46:43 So here is the arc--
like in this case,
46:46 we thought of inference
scaling in this project
46:49 as an inference
architecture design problem.
46:56 So here is how the
Archon framework works.
47:00 On the input side, we had
a set of target benchmarks
47:03 that we wanted to get good
results on these benchmarks.
47:06 So these are the benchmarks
we are optimizing for.
47:09 There is an inference
called budget.
47:11 And then there's also a
set of available LLMs,
47:14 because we can always mix and
match different models also
47:18 to get to a correct answer
for a input question.
47:23 And then we have a set of
inference time techniques, which
47:25 I'm going to describe
in the following slides.
47:29 And then we designed
this optimizer.
47:33 It's called itest.
47:35 It's inference time architecture
search that basically tells us
47:40 how to put together
these different models
47:43 and different inference
time techniques together,
47:45 whether it's parallel sampling,
sequential, and all that,
47:49 in order to get
high-quality results given
47:53 the budget, the inference
called budget that we have-- how
47:56 to optimize the quality
of results based on that.
48:00 And then the output
is this architecture
48:02 that mixes and matches
these different techniques
48:04 and models together.
48:08 So let's take a look at
different inference time
48:11 operations.
48:12 So first, we have the general--
we call it the generation,
48:16 but this is basically just
sampling from the model.
48:24 And this can be part of the
repeated sampling approach.
48:28 So if we generate n--
48:31 we use n generations
at a time, that
48:33 means we have generated n
different responses to our input
48:37 question.
48:38 So that's one type of
inference time scaling.
48:41 Another one is fusion, which was
surprisingly a very effective
48:45 method.
48:46 So fusion is, let's say, we
have K different responses,
48:53 and we want to generate
one output response out
48:56 of these K. Fusion
takes an LLM and asks
49:02 an LLM the same exact
question that I just said.
49:06 Basically, it says,
here's the input question.
49:09 Here are K different responses
to this input question.
49:13 Given all of this, generate or
synthesize one output response
49:20 for the input question.
49:21 Basically, we are showing all of
the K generations to the model,
49:25 such that the LLM can generate
an output answer while being
49:30 aware of all the possible ways
that this question is answered.
49:38 And this is-- it goes under the
category of sequential update.
49:44 We have this parallel
samples, but then we
49:46 are kind of fusing
them into one answer.
49:49 Then there is the
critic, which basically,
49:56 for any given response,
we can ask a model
49:59 to generate or
describe the weaknesses
50:04 or strengths of these response.
50:06 And that could be one kind
of inference scaling action.
50:12 And then we have the ranker.
50:14 This is-- again, these
are all prompting based
50:16 optimization for the model.
50:18 We haven't trained the model
to do this specific test time
50:21 scaling task.
50:22 We just ask the model in
the instruction prompt that,
50:27 can you rank these
generated responses
50:29 from the previous step
based on their quality?
50:38 And then the verifiers
generates the-- returns
50:41 the response while also talk
about the reasoning of why
50:45 it came to this kind of way
to score a certain answer.
50:52 I'll get to the
next two unit tests
50:55 generation and
evaluation in a bit.
50:57 But just to show you how
effective these different
51:00 inference time scaling
techniques are,
51:03 let's take a look at this
win rate for what tasks--
51:13 we don't have the task here.
51:15 But it's one of the task
of one of these benchmarks
51:19 that I don't know if
you remember the task.
51:23 Yeah.
51:26 Some question answering
reasoning benchmark.
51:31 So here, let's take a look
at how these things work.
51:34 So first, we have--
51:36 let's say, we have a number of
repeated samples from the model.
51:41 It goes from 1 to
10 on the left side.
51:44 Let's focus on the
graph on the left side.
51:46 And then the orange one is
the kind of worst possible way
51:51 of approaching it.
51:52 We just randomly
select one and look
51:55 at the quality and the
correctness of this response.
51:59 The green one is when we
use the model to rank them,
52:04 to rank the responses, and
we select the top one ranked
52:07 by the model and report that.
52:13 That already does better
than the orange one.
52:16 The blue one is
Oracle selection,
52:20 meaning let's assume
we are not mixing
52:23 these responses in any way.
52:26 We just select the
best one generated.
52:30 Let's say, we have
an Oracle verifier
52:31 that can tell us which
one of these is correct.
52:34 And that obviously
does better than just
52:37 asking the model to rank them
because it's an Oracle verifier.
52:41 And then we get to
the red one, which is
52:44 where when we use the fusion.
52:46 So basically, for each point
here, let's say, we are at 5,
52:52 the red data point is basically
we have given these five
52:55 responses to the model and
ask the model to synthesize
52:59 one answer, one final answer,
given all this information.
53:03 And it's so interesting, and
this paradigm is so powerful
53:06 that that can on its own
improve the quality of responses
53:11 over Oracle selection.
53:15 And then finally, the
purple one is where we--
53:20 instead of fusing
all the responses,
53:27 we first selected the
top five, and then we
53:33 went about fusing them.
53:34 So it seems like some sort
of a filtering of the two
53:37 does better.
53:41 Now on the right side, we are
doing a different experiment,
53:44 and that is we are--
53:47 on the left side.
53:48 We had only one model
generating everything.
53:50 On the left side,
we had-- we were
53:52 using-- we were adding models.
53:54 So first one is just one
model generating one answer,
53:57 two the baseline is two
different models generating
54:00 one answer, and 10 is
10 different models,
54:03 each generating one answer.
54:05 And we see that, of
course, the random result
54:09 becomes worse because the
model we looked on the left
54:12 was I think it was among the
better models in this ensemble.
54:17 But we still see the same
trend, even if these responses
54:22 are coming from different
models, basically the same trend
54:25 of how these all
follow each other.
54:29 Yes.
54:30 Which model is fusing?
54:32 If you're using 10
different models,
54:34 you're asking somebody
else to rank or fuse your--
54:38 I think it's the
model number one
54:39 in this case, the
same [INAUDIBLE] that
54:41 was in all of those.
54:43 And this was the way we
created the ensemble was first
54:47 with the best model and then
added the second best model
54:49 and so on.
54:56 Now let's take a
look at a couple
54:58 other inference time techniques.
55:01 And that is unit test
generation and evaluation.
55:04 I think one of you was
mentioning something like this
55:09 earlier as well, that
where we can ask the model
55:11 to generate these
unit tests, say
55:14 when it gets to coding task
or math, math problem-solving
55:18 or other reasoning tasks.
55:22 And the unit test
evaluation is even crazier.
55:24 Like instead of
running the unit test,
55:27 you're asking the model to
evaluate a generated answer
55:32 against a unit test.
55:36 For example, here
is a prompt check
55:40 for-- so the coding
problem in this case
55:43 is check for balanceness of
round brackets in an input
55:47 string.
55:47 That's a coding question.
55:49 And then here is
an example input.
55:51 Here is the type of unit test
that the model could generate.
55:56 For example, given a string
with an odd number of brackets,
56:00 the solution should output
no, which is like a good unit
56:03 test in this case.
56:04 Another one is when a closing
bracket is encountered,
56:08 it must match the most
recently open brackets
56:11 that hasn't been matched yet.
56:12 That's another kind of
reasonable unit test.
56:15 And then you can ask the
model to actually generate
56:17 the code to test this as well.
56:20 Now let's take a
look at the types.
56:22 I don't know if you guys
see this well but the type.
56:25 But this is so
big, there's no way
56:26 for me to put it
here with large font.
56:30 But this is the type of
architectures that Archon found.
56:34 So basically now think
about it this way.
56:37 We have different ways
of mixing and matching
56:40 models and inference
time techniques,
56:42 and we can define these
layers of optimization.
56:45 For example, the first layer
in this case is generation.
56:48 And then these generations
are from different models
56:51 in this case.
56:52 And then we have a
critic and a ranker.
56:56 So basically, the critic will
criticize these generations
57:00 and then we rank them and
then be the optimizer found
57:04 that we could have a bunch
of users at this point,
57:07 like each fuses these
responses into one output.
57:13 And then we can go from there,
continue this critique ranker.
57:19 And this was like a rather
complicated architecture.
57:21 And the way that we
optimize this architecture
57:24 is that we did some
kind of pre-processing
57:28 to reduce the space
of optimization.
57:32 For example, the
fact that we found
57:33 this sequence of generation
critic ranker and user
57:38 do really well.
57:39 We did that offline, so we
limited the search space
57:43 just because doing this
search is very expensive.
57:46 We need to run a lot
of inference calls.
57:48 But then the way we optimize it
that we had a held out data set
57:52 that-- we had a training data
set that we could optimize
57:56 against that and create
this accuracy versus
57:59 the number of calls that we are
making to these models, kind
58:03 of as a metric that we
are optimizing against.
58:07 We want to maximize accuracy.
58:09 And for any given budget,
inference budget that we have.
58:16 And this is for some
of the coding problems.
58:19 This was an optimized
architecture.
58:21 Basically, we generate
a ton of samples.
58:24 We generate unit tests
and then we evaluate them.
58:28 Now, an interesting
property here
58:33 is that we found that
stacking more and more layers
58:38 of these inference
time techniques,
58:40 or making it deep
in a way, actually
58:42 does help with our accuracy.
58:49 In this case, for example,
the right most pink graph
58:53 is where we have an
ensemble, meaning we have
58:57 different models in our loop.
58:59 And we have three layers
of critiques and fusers
59:02 and a final fuser layer.
59:06 And this is significantly
better across many
59:08 of these tasks over
both just using
59:12 a simple the best
model only once,
59:16 or the best model eight times
with one layer of fusion.
59:19 So it seems like these
additional layers
59:22 that we are adding
here is actually--
59:23 just like in deep
learning, we are
59:26 adding layers in pre-training
and the model gets better.
59:28 Seems like these careful
kind of additions
59:31 of these inference layers
are helping the model
59:33 become more and more accurate.
59:35 Across many of the
tasks this is true.
59:40 And so now the
optimizer itself, we
59:43 used an Bayesian optimizer again
to simplify the search space.
59:48 We did a few things.
59:50 For example, we
limited the optimizer
59:53 to only use one inference
time technique per layer,
59:57 and the first layer was
always just the generator.
1:00:00 And then we had this cascade
of ranker, critic, verifier.
1:00:06 They could go
anywhere, but we made
1:00:09 certain kind of modifications
to the search space.
1:00:12 For example, a critic should
go before a ranker or a fuser
1:00:16 always because we did some
testing locally and saw
1:00:20 this is a better choice.
1:00:22 And then a unit
test generator must
1:00:24 be followed by an evaluator up.
1:00:27 And again, we use a Bayesian
optimizer to do this.
1:00:32 This is often a very
powerful methodology,
1:00:39 a ton of good open
source software
1:00:41 to do Bayesian
optimization for you.
1:00:43 Basically, you have
this search space
1:00:45 that you want to configure it
with different kind of choices
1:00:50 for the architecture.
1:00:51 And the output side,
you have the accuracy
1:00:53 and you want to
optimize for that.
1:00:55 And this does better than greedy
search or random selection.
1:00:59 We are much more sample
efficient in terms
1:01:01 of number of configurations
that Archon searches
1:01:04 until it reaches a final
kind of architecture,
1:01:09 and we can define different
objectives for it.
1:01:12 For example, available
models, inference time,
1:01:16 inference, core budget is
one of them or some of them.
1:01:22 The thing that I want you to
pay attention to in this graph
1:01:26 is that this interesting
property that Archon
1:01:29 is, at the end, generating
one response at the very end.
1:01:32 So it's like you're optimizing
paths at one at the very end.
1:01:36 And what was interesting
here is that while using just
1:01:42 the open source models,
we could match or exceed
1:01:48 these closed source
models at the time,
1:01:51 the frontier closed
source model at the time,
1:01:53 by a large margin across
many of these tasks in terms
1:01:58 of paths at one.
1:01:59 Here's the average path
at one improvement.
1:02:02 And not only that,
we could train
1:02:04 Archon to be task specific.
1:02:06 So when we are doing the
Bayesian optimization,
1:02:09 we only look into a certain
task and we test on that,
1:02:14 or we could train Archon
to be general purpose.
1:02:17 So we wanted it to do
better good across many--
1:02:21 to do well across many tasks.
1:02:23 And even that the general
purpose one does really,
1:02:26 really well and better
than the frontier
1:02:29 models across these tasks.
1:02:31 And that was a very
interesting observation
1:02:35 that you can design these
inference time architectures
1:02:38 that do well beyond the
task or the limited tasks
1:02:42 that they're trained on.
1:02:44 So on average, in this case,
we were outperforming GPT-4.0
1:02:48 or Claude 3.5 Sonnet in passage
one by an average of 14.1%
1:02:54 across these instruction
following reasoning and math
1:02:58 encoding problems.
1:03:00 Here are some questions, that
I want you to think about.
1:03:05 But since we're out of time.
1:03:07 Hopefully, you think
about it yourself.
1:03:11 And hope you have
a happy weekend.