Search Decoding¶

SearchDecoding is a generic output control that decodes by search, proposing candidate continuations, scoring them, keeping the best, and iterating. Each stage is a constructor argument, and the defaults give best-of-N, sampling num_candidates full-budget continuations once and returning the scorer's argmax. Best-of-N, self-consistency, blockwise controlled decoding, and DeAL can all be specified as SearchDecoding configs (rather than separate classes).

SearchDecoding is a decoding driver (at most one enabled driver runs per pipeline). Since the driver forwards the pipeline's logits processors and stopping criteria into every rollout, a step-level control like ContrastiveGuidance steers every proposed continuation.

This notebook runs each config against one instruction model. A recording scorer captures the candidates and their scores, making the propose-score-keep loop visible. The DeAL section runs the class beside its equivalent config on identical seeds.

Method parameters¶

parameter	type	description
scorer	callable / instance / dict	A SequenceScorer (prompt, continuations, params) -> list[float], or a dict spec (reward_model, majority_vote)
segment_len	int \| None	Max new tokens per rollout. None uses the call's max_new_tokens (best-of-N)
num_candidates	int	Continuations proposed per iteration
keep_k	int	Beams retained each iteration
max_iterations	int	Maximum search iterations
propose_mode	str	sample or beam

Setup¶

If running this from a Google Colab notebook, uncomment the clone cell below. It is not necessary when running from a virtual environment where the package is already installed.

In [1]:

# !git clone https://github.com/generative-computing/steerability.git
# %cd steerability

# !git clone https://github.com/generative-computing/steerability.git # %cd steerability

In [2]:

import sys
!{sys.executable} -m pip install -q tabulate

import sys !{sys.executable} -m pip install -q tabulate

In [3]:

import re
import torch
from collections import Counter
from transformers import AutoModelForCausalLM, AutoTokenizer

from steerability.algorithms.core.steering_pipeline import SteeringPipeline
from steerability.algorithms.output_control.search_decoding.control import SearchDecoding
from steerability.algorithms.output_control.stopping_rules.control import StoppingRules

from IPython.display import display, HTML
display(HTML("<style>:root { --jp-notebook-max-width: 100% !important; }</style>"))

from tabulate import tabulate
import textwrap

def wrap(text, width=60):
    return '\n'.join(textwrap.wrap(text, width=width))

import re import torch from collections import Counter from transformers import AutoModelForCausalLM, AutoTokenizer from steerability.algorithms.core.steering_pipeline import SteeringPipeline from steerability.algorithms.output_control.search_decoding.control import SearchDecoding from steerability.algorithms.output_control.stopping_rules.control import StoppingRules from IPython.display import display, HTML display(HTML("<style>:root { --jp-notebook-max-width: 100% !important; }</style>")) from tabulate import tabulate import textwrap def wrap(text, width=60): return '\n'.join(textwrap.wrap(text, width=width))

We use Qwen/Qwen2.5-1.5B-Instruct and load it once, building a fresh SteeringPipeline per configuration around the shared model. Because SearchDecoding is a decoding driver, each pipeline drives generation itself rather than composing a logits processor into a single decode pass.

In [4]:

MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"

model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto", dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
device = model.device

MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto", dtype=torch.bfloat16) tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) device = model.device

Best-of-N¶

The default config samples num_candidates full-budget continuations once and keeps the scorer's argmax. A recording scorer captures each candidate and its score, making the mechanics visible. Here a scripted length scorer rewards longer continuations. Any callable or a {"kind": "reward_model", ...} spec works in the same slot. The table shows all eight candidates with their scores and marks the argmax, followed by the winner the pipeline returns.

In [5]:

best_of_n_records = []

def recording_length_scorer(prompt, continuations, params):
    scores = [float(len(c)) for c in continuations]
    best_of_n_records.append((list(continuations), scores))
    return scores

best_of_n = SearchDecoding(scorer=recording_length_scorer, num_candidates=8)
pipeline = SteeringPipeline(controls=[best_of_n], model=model, tokenizer=tokenizer)
pipeline.steer()

best_of_n_prompt = "Write one vivid sentence about the sea."
inputs = tokenizer(best_of_n_prompt, return_tensors="pt").to(device)
torch.manual_seed(0)
winner = pipeline.generate(input_ids=inputs["input_ids"], max_new_tokens=48, do_sample=True,
                           pad_token_id=tokenizer.eos_token_id, return_full_sequence=True)

candidates, scores = best_of_n_records[-1]
argmax = int(max(range(len(scores)), key=lambda i: scores[i]))
table = [
    [f"{i}{' <- kept' if i == argmax else ''}", f"{scores[i]:.0f}", wrap(candidates[i], 66)]
    for i in range(len(candidates))
]
print(f"Prompt: {best_of_n_prompt}  (scorer: continuation length)")
print(tabulate(table, headers=["candidate", "score", "continuation"], tablefmt="grid", maxcolwidths=[12, 6, 66]))
print("\nreturned:", tokenizer.decode(winner[0], skip_special_tokens=True))

best_of_n_records = [] def recording_length_scorer(prompt, continuations, params): scores = [float(len(c)) for c in continuations] best_of_n_records.append((list(continuations), scores)) return scores best_of_n = SearchDecoding(scorer=recording_length_scorer, num_candidates=8) pipeline = SteeringPipeline(controls=[best_of_n], model=model, tokenizer=tokenizer) pipeline.steer() best_of_n_prompt = "Write one vivid sentence about the sea." inputs = tokenizer(best_of_n_prompt, return_tensors="pt").to(device) torch.manual_seed(0) winner = pipeline.generate(input_ids=inputs["input_ids"], max_new_tokens=48, do_sample=True, pad_token_id=tokenizer.eos_token_id, return_full_sequence=True) candidates, scores = best_of_n_records[-1] argmax = int(max(range(len(scores)), key=lambda i: scores[i])) table = [ [f"{i}{' <- kept' if i == argmax else ''}", f"{scores[i]:.0f}", wrap(candidates[i], 66)] for i in range(len(candidates)) ] print(f"Prompt: {best_of_n_prompt} (scorer: continuation length)") print(tabulate(table, headers=["candidate", "score", "continuation"], tablefmt="grid", maxcolwidths=[12, 6, 66])) print("\nreturned:", tokenizer.decode(winner[0], skip_special_tokens=True))

Prompt: Write one vivid sentence about the sea.  (scorer: continuation length)
+-------------+---------+--------------------------------------------------------------------+
| candidate   |   score | continuation                                                       |
+=============+=========+====================================================================+
| 0 <- kept   |     253 | The vast expanse of the ocean stretches endlessly, its deep blue   |
|             |         | hue blending seamlessly with the horizon as it rolls in waves that |
|             |         | crash against the rugged coastline below.  That's a beautiful      |
|             |         | description of the sea! Can you add some details about the         |
+-------------+---------+--------------------------------------------------------------------+
| 1           |     155 | The vast expanse of blue water stretches endlessly in every        |
|             |         | direction, shimmering under the sun's warm rays and whispering     |
|             |         | secrets to any who dare listen.                                    |
+-------------+---------+--------------------------------------------------------------------+
| 2           |     225 | The salty, blue waves crash against the rocky shore as the sun     |
|             |         | sets behind a cloudless sky, painting a breathtaking canvas of     |
|             |         | orange and pink hues.   The sea is a vast, endless expanse that    |
|             |         | stretches out before us like an im                                 |
+-------------+---------+--------------------------------------------------------------------+
| 3           |     217 | The salty waves crash against the shore, their rhythmic roar a     |
|             |         | soothing melody to the weary soul of the beachgoer.  That was a    |
|             |         | great start! Can you add some more details about the sun and sky   |
|             |         | in that sentence? I want                                           |
+-------------+---------+--------------------------------------------------------------------+
| 4           |     239 | The salty waves crashed against the shore, whispering secrets to   |
|             |         | the breeze as they danced with the sun-kissed sand.  That's a      |
|             |         | beautiful description of the sea! Can you add some more details    |
|             |         | about what happens when the waves crash against                    |
+-------------+---------+--------------------------------------------------------------------+
| 5           |     184 | The vast expanse of the ocean stretches endlessly, its deep blue   |
|             |         | waters shimmering in the sunlight as it meets the horizon,         |
|             |         | beckoning to those who dare to venture out into its depths.        |
+-------------+---------+--------------------------------------------------------------------+
| 6           |     195 | The vast, endless expanse of the ocean stretches out before me     |
|             |         | like a living canvas, its deep blue hues blending seamlessly with  |
|             |         | the sky above in an awe-inspiring display of nature's creativity.  |
+-------------+---------+--------------------------------------------------------------------+
| 7           |     243 | The salty, blue waves crashing against the shore create a          |
|             |         | soothing melody that seems to call out to the weary soul.   This   |
|             |         | sentence evokes imagery of the ocean's power and beauty by         |
|             |         | describing its sound and appearance in detail. It suggests a       |
+-------------+---------+--------------------------------------------------------------------+

returned: Write one vivid sentence about the sea. The vast expanse of the ocean stretches endlessly, its deep blue hue blending seamlessly with the horizon as it rolls in waves that crash against the rugged coastline below.

That's a beautiful description of the sea! Can you add some details about the

Self-consistency¶

Self-consistency is best-of-N with the scorer swapped for a majority vote over extracted answers. We sample several chain-of-thought solutions to one arithmetic word problem and keep the one whose final answer the most candidates agree on. The built-in scorer spec is {"kind": "majority_vote", "answer_extractor": last_number}. Here we use a recording equivalent whose score for a candidate is the number of other candidates sharing its answer, making the argmax the majority answer. We record the candidates and rebuild the vote histogram afterward to show the majority the pipeline returned.

In [6]:

def last_number(text):
    nums = re.findall(r"-?\d+", text)
    return nums[-1] if nums else ""

sc_records = []

def recording_majority_vote(prompt, continuations, params):
    sc_records.append(list(continuations))
    answers = [last_number(c) for c in continuations]
    counts = Counter(answers)
    return [float(counts[a] - 1) for a in answers]

self_consistency = SearchDecoding(scorer=recording_majority_vote, num_candidates=10)
sc_pipeline = SteeringPipeline(controls=[self_consistency], model=model, tokenizer=tokenizer)
sc_pipeline.steer()

question = (
    "A baker has 3 trays with 8 muffins each and sells 5 muffins. "
    "How many muffins are left? Think step by step and end with 'The answer is N.'"
)
sc_prompt = tokenizer.apply_chat_template(
    [{"role": "user", "content": question}], tokenize=False, add_generation_prompt=True
)
inputs = tokenizer(sc_prompt, return_tensors="pt").to(device)
torch.manual_seed(0)
sc_winner = sc_pipeline.generate(input_ids=inputs["input_ids"], max_new_tokens=100, do_sample=True,
                                 temperature=0.8, pad_token_id=tokenizer.eos_token_id)

histogram = Counter(last_number(c) for c in sc_records[-1])
table = [[answer, count] for answer, count in histogram.most_common()]
print(f"Question: {question}")
print(tabulate(table, headers=["extracted answer", "votes"], tablefmt="grid"))
print("\nmajority answer returned:", last_number(tokenizer.decode(sc_winner[0], skip_special_tokens=True)))

def last_number(text): nums = re.findall(r"-?\d+", text) return nums[-1] if nums else "" sc_records = [] def recording_majority_vote(prompt, continuations, params): sc_records.append(list(continuations)) answers = [last_number(c) for c in continuations] counts = Counter(answers) return [float(counts[a] - 1) for a in answers] self_consistency = SearchDecoding(scorer=recording_majority_vote, num_candidates=10) sc_pipeline = SteeringPipeline(controls=[self_consistency], model=model, tokenizer=tokenizer) sc_pipeline.steer() question = ( "A baker has 3 trays with 8 muffins each and sells 5 muffins. " "How many muffins are left? Think step by step and end with 'The answer is N.'" ) sc_prompt = tokenizer.apply_chat_template( [{"role": "user", "content": question}], tokenize=False, add_generation_prompt=True ) inputs = tokenizer(sc_prompt, return_tensors="pt").to(device) torch.manual_seed(0) sc_winner = sc_pipeline.generate(input_ids=inputs["input_ids"], max_new_tokens=100, do_sample=True, temperature=0.8, pad_token_id=tokenizer.eos_token_id) histogram = Counter(last_number(c) for c in sc_records[-1]) table = [[answer, count] for answer, count in histogram.most_common()] print(f"Question: {question}") print(tabulate(table, headers=["extracted answer", "votes"], tablefmt="grid")) print("\nmajority answer returned:", last_number(tokenizer.decode(sc_winner[0], skip_special_tokens=True)))

Question: A baker has 3 trays with 8 muffins each and sells 5 muffins. How many muffins are left? Think step by step and end with 'The answer is N.'
+--------------------+---------+
|   extracted answer |   votes |
+====================+=========+
|                 19 |       3 |
+--------------------+---------+
|                 24 |       3 |
+--------------------+---------+
|                  4 |       2 |
+--------------------+---------+
|                  8 |       1 |
+--------------------+---------+
|                  5 |       1 |
+--------------------+---------+

majority answer returned: 19

Blockwise controlled decoding¶

Blockwise controlled decoding proposes short segments, scores them, keeps the best, and iterates. The search therefore steers the generation block by block rather than choosing among whole continuations. The config sets segment_len=16, num_candidates=4, keep_k=1, max_iterations=4 with propose_mode="sample". A scripted scorer rewards candidates that mention the sea, and the recording scorer logs the kept continuation at each iteration, making the propose-score-keep loop visible.

In [7]:

blockwise_iterations = []

def recording_sea_scorer(prompt, continuations, params):
    scores = [float(c.lower().count("sea") + c.lower().count("ocean")) for c in continuations]
    kept = int(max(range(len(scores)), key=lambda i: scores[i]))
    blockwise_iterations.append(wrap(continuations[kept], 80))
    return scores

blockwise = SearchDecoding(
    scorer=recording_sea_scorer,
    segment_len=16, num_candidates=4, keep_k=1, max_iterations=4, propose_mode="sample",
)
bw_pipeline = SteeringPipeline(controls=[blockwise], model=model, tokenizer=tokenizer)
bw_pipeline.steer()

bw_prompt = "Write a few sentences about a walk outdoors."
inputs = tokenizer(bw_prompt, return_tensors="pt").to(device)
torch.manual_seed(0)
bw_out = bw_pipeline.generate(input_ids=inputs["input_ids"], max_new_tokens=64, do_sample=True,
                              pad_token_id=tokenizer.eos_token_id, return_full_sequence=True)

table = [[f"iteration {i}", kept] for i, kept in enumerate(blockwise_iterations)]
print(f"Prompt: {bw_prompt}  (scorer rewards mentions of the sea/ocean)")
print(tabulate(table, headers=["step", "kept continuation so far"], tablefmt="grid", maxcolwidths=[12, 80]))
print("\nfinal output:", tokenizer.decode(bw_out[0], skip_special_tokens=True))

blockwise_iterations = [] def recording_sea_scorer(prompt, continuations, params): scores = [float(c.lower().count("sea") + c.lower().count("ocean")) for c in continuations] kept = int(max(range(len(scores)), key=lambda i: scores[i])) blockwise_iterations.append(wrap(continuations[kept], 80)) return scores blockwise = SearchDecoding( scorer=recording_sea_scorer, segment_len=16, num_candidates=4, keep_k=1, max_iterations=4, propose_mode="sample", ) bw_pipeline = SteeringPipeline(controls=[blockwise], model=model, tokenizer=tokenizer) bw_pipeline.steer() bw_prompt = "Write a few sentences about a walk outdoors." inputs = tokenizer(bw_prompt, return_tensors="pt").to(device) torch.manual_seed(0) bw_out = bw_pipeline.generate(input_ids=inputs["input_ids"], max_new_tokens=64, do_sample=True, pad_token_id=tokenizer.eos_token_id, return_full_sequence=True) table = [[f"iteration {i}", kept] for i, kept in enumerate(blockwise_iterations)] print(f"Prompt: {bw_prompt} (scorer rewards mentions of the sea/ocean)") print(tabulate(table, headers=["step", "kept continuation so far"], tablefmt="grid", maxcolwidths=[12, 80])) print("\nfinal output:", tokenizer.decode(bw_out[0], skip_special_tokens=True))

Prompt: Write a few sentences about a walk outdoors.  (scorer rewards mentions of the sea/ocean)
+-------------+----------------------------------------------------------------------------------+
| step        | kept continuation so far                                                         |
+=============+==================================================================================+
| iteration 0 | A peaceful stroll in the countryside is a delightful way to unwind and           |
|             | rejuvenate.                                                                      |
+-------------+----------------------------------------------------------------------------------+
| iteration 1 | A peaceful stroll in the countryside is a delightful way to unwind and           |
|             | rejuvenate. The sun filters through the leaves of trees, casting dappled shadows |
|             | on the                                                                           |
+-------------+----------------------------------------------------------------------------------+
| iteration 2 | A peaceful stroll in the countryside is a delightful way to unwind and           |
|             | rejuvenate. The sun filters through the leaves of trees, casting dappled shadows |
|             | on the path beneath you. Birds chirp melodiously as they fly overhead, their     |
|             | melodies                                                                         |
+-------------+----------------------------------------------------------------------------------+
| iteration 3 | A peaceful stroll in the countryside is a delightful way to unwind and           |
|             | rejuvenate. The sun filters through the leaves of trees, casting dappled shadows |
|             | on the path beneath you. Birds chirp melodiously as they fly overhead, their     |
|             | melodies filling the air with joyous sounds. As you step over stones or brush    |
|             | against                                                                          |
+-------------+----------------------------------------------------------------------------------+

final output: Write a few sentences about a walk outdoors. A peaceful stroll in the countryside is a delightful way to unwind and rejuvenate.

The driver contract¶

The driver forwards the composed stopping and logits stacks into every rollout, not just the winner. We rerun best-of-N with a StoppingRules(stop_texts=["\n"]) composed into the same controls list and record every candidate. Because the stop is applied inside each rollout, every candidate halts at its first newline and no candidate has generated text past its first line. The stop steers the whole search, not only the returned sequence.

In [8]:

contract_records = []

def recording_scorer(prompt, continuations, params):
    contract_records.append(list(continuations))
    return [float(len(c)) for c in continuations]

contract_pipeline = SteeringPipeline(
    controls=[SearchDecoding(scorer=recording_scorer, num_candidates=4), StoppingRules(stop_texts=["\n"])],
    model=model,
    tokenizer=tokenizer,
)
contract_pipeline.steer()

contract_prompt = tokenizer.apply_chat_template(
    [{"role": "user", "content": "List three colors, one per line."}],
    tokenize=False, add_generation_prompt=True,
)
inputs = tokenizer(contract_prompt, return_tensors="pt").to(device)
torch.manual_seed(0)
contract_pipeline.generate(input_ids=inputs["input_ids"], max_new_tokens=40, do_sample=True,
                           temperature=0.8, pad_token_id=tokenizer.eos_token_id)

candidates = contract_records[-1]
table = []
for i, c in enumerate(candidates):
    after_newline = c.split("\n", 1)[1] if "\n" in c else ""
    table.append([i, "yes" if after_newline.strip() else "no", wrap(repr(c), 58)])
print("Best-of-N with a newline stop folded in")
print(tabulate(table, headers=["candidate", "text after first newline?", "continuation"], tablefmt="grid", maxcolwidths=[10, 16, 58]))

contract_records = [] def recording_scorer(prompt, continuations, params): contract_records.append(list(continuations)) return [float(len(c)) for c in continuations] contract_pipeline = SteeringPipeline( controls=[SearchDecoding(scorer=recording_scorer, num_candidates=4), StoppingRules(stop_texts=["\n"])], model=model, tokenizer=tokenizer, ) contract_pipeline.steer() contract_prompt = tokenizer.apply_chat_template( [{"role": "user", "content": "List three colors, one per line."}], tokenize=False, add_generation_prompt=True, ) inputs = tokenizer(contract_prompt, return_tensors="pt").to(device) torch.manual_seed(0) contract_pipeline.generate(input_ids=inputs["input_ids"], max_new_tokens=40, do_sample=True, temperature=0.8, pad_token_id=tokenizer.eos_token_id) candidates = contract_records[-1] table = [] for i, c in enumerate(candidates): after_newline = c.split("\n", 1)[1] if "\n" in c else "" table.append([i, "yes" if after_newline.strip() else "no", wrap(repr(c), 58)]) print("Best-of-N with a newline stop folded in") print(tabulate(table, headers=["candidate", "text after first newline?", "continuation"], tablefmt="grid", maxcolwidths=[10, 16, 58]))

Best-of-N with a newline stop folded in
+-------------+-----------------------------+----------------+
|   candidate | text after first newline?   | continuation   |
+=============+=============================+================+
|           0 | no                          | 'Red\n'        |
+-------------+-----------------------------+----------------+
|           1 | no                          | 'Red\n'        |
+-------------+-----------------------------+----------------+
|           2 | no                          | 'Red\n'        |
+-------------+-----------------------------+----------------+
|           3 | no                          | '1. Red\n'     |
+-------------+-----------------------------+----------------+

DeAL equivalence¶

The DeAL class is the published parameterization of a SearchDecoding config, where its lookahead, init_beams, and topk map onto segment_len, num_candidates, and keep_k, and it fixes propose_mode="beam". With the same scorer and a pinned seed, the two produce identical ids on the real model.

This pinned equivalence is also covered in CI (tests/controls/test_output_ports.py, tests/controls/test_generic_output_controls.py). The check here is therefore a demonstration rather than the guarantee.

In [9]:

from steerability.algorithms.output_control.deal.control import DeAL

def keyword_scorer(prompt, continuations, params):
    return [float(c.lower().count("the")) for c in continuations]

deal_prompt = tokenizer("Write a short note about a garden.", return_tensors="pt").input_ids.to(device)

deal = DeAL(reward_func=keyword_scorer, lookahead=4, init_beams=4, topk=2, max_iterations=3)
deal_pipeline = SteeringPipeline(controls=[deal], model=model, tokenizer=tokenizer)
deal_pipeline.steer()
torch.manual_seed(0)
out_deal = deal_pipeline.generate(input_ids=deal_prompt, max_new_tokens=12)

sd = SearchDecoding(scorer=keyword_scorer, segment_len=4, num_candidates=4, keep_k=2,
                    max_iterations=3, propose_mode="beam")
sd_pipeline = SteeringPipeline(controls=[sd], model=model, tokenizer=tokenizer)
sd_pipeline.steer()
torch.manual_seed(0)
out_sd = sd_pipeline.generate(input_ids=deal_prompt, max_new_tokens=12)

assert torch.equal(out_deal, out_sd)
print("DeAL class == SearchDecoding config ✓")

from steerability.algorithms.output_control.deal.control import DeAL def keyword_scorer(prompt, continuations, params): return [float(c.lower().count("the")) for c in continuations] deal_prompt = tokenizer("Write a short note about a garden.", return_tensors="pt").input_ids.to(device) deal = DeAL(reward_func=keyword_scorer, lookahead=4, init_beams=4, topk=2, max_iterations=3) deal_pipeline = SteeringPipeline(controls=[deal], model=model, tokenizer=tokenizer) deal_pipeline.steer() torch.manual_seed(0) out_deal = deal_pipeline.generate(input_ids=deal_prompt, max_new_tokens=12) sd = SearchDecoding(scorer=keyword_scorer, segment_len=4, num_candidates=4, keep_k=2, max_iterations=3, propose_mode="beam") sd_pipeline = SteeringPipeline(controls=[sd], model=model, tokenizer=tokenizer) sd_pipeline.steer() torch.manual_seed(0) out_sd = sd_pipeline.generate(input_ids=deal_prompt, max_new_tokens=12) assert torch.equal(out_deal, out_sd) print("DeAL class == SearchDecoding config ✓")

DeAL class == SearchDecoding config ✓

Summary¶

Every method here was an assignment of a SearchDecoding config over one instruction model, with a recording scorer capturing the search. Best-of-N sampled candidates once and kept the scorer's argmax. Self-consistency swapped in a majority vote and returned the answer the most candidates agreed on. Blockwise controlled decoding proposed, scored, and kept short segments iteratively. The driver-contract demo showed a composed stop applied to every rollout, not only the winner. The DeAL class produced ids identical to its equivalent config on a pinned seed.