DeAL

steerability.algorithms.output_control.deal

args

control

DeAL

Bases: SearchDriver

Implementation of DeAL (Decoding-time Alignment) from Huang et al., 2024.

DeAL performs controlled text generation through iterative lookahead search and reward-guided beam selection. Unlike training-time alignment methods, DeAL operates purely at inference time to steer language model outputs toward desired behaviors.

The algorithm works in three phases:

1. Lookahead Generation: Generate multiple candidate continuations using beam search from the current context.

2. Reward-based Scoring: Evaluate each candidate continuation using a provided reward function that measures alignment with the desired objective (e.g., helpfulness, safety).

3. Iterative Refinement: Select the top-k highest-scoring beams and repeat the process until termination conditions are met (EOS token, max length, or max iterations reached).

DeAL is a decoding driver implemented as a preset of the generic SearchDriver, mapping its arguments onto the search fields (scorer, segment_len, num_candidates, keep_k, max_iterations, and propose_mode="beam"). The composed logits processors and stopping criteria apply inside every lookahead rollout, which means that a step-level control such as RAD steers every DeAL rollout. The reward_params runtime kwarg is honored per row (one mapping merged into the scorer's params).

Reference:

• "DeAL: Decoding-time Alignment for Large Language Models" James Y. Huang, Sailik Sengupta, Daniele Bonadiman, Yi-an Lai, Arshit Gupta, Nikolaos Pappas, Saab Mansour, Katrin Kirchhoff, Dan Roth https://arxiv.org/abs/2402.06147

Source code in steerability/algorithms/output_control/deal/control.py

class DeAL(SearchDriver):
    """
    Implementation of DeAL (Decoding-time Alignment) from Huang et al., 2024.

    DeAL performs controlled text generation through iterative lookahead search and reward-guided beam selection. Unlike
    training-time alignment methods, DeAL operates purely at inference time to steer language model outputs toward
    desired behaviors.

    The algorithm works in three phases:

    1. **Lookahead Generation**: Generate multiple candidate continuations using beam search from the current context.

    2. **Reward-based Scoring**: Evaluate each candidate continuation using a provided reward function that measures
    alignment with the desired objective (e.g., helpfulness, safety).

    3. **Iterative Refinement**: Select the top-k highest-scoring beams and repeat the process until termination
    conditions are met (EOS token, max length, or max iterations reached).

    DeAL is a decoding driver implemented as a preset of the generic `SearchDriver`, mapping its arguments onto the
    search fields (`scorer`, `segment_len`, `num_candidates`, `keep_k`, `max_iterations`, and `propose_mode="beam"`).
    The composed logits processors and stopping criteria apply inside every lookahead rollout, which means that a
    step-level control such as RAD steers every DeAL rollout. The `reward_params` runtime kwarg is honored per row (one
    mapping merged into the scorer's params).

    Reference:

    - "DeAL: Decoding-time Alignment for Large Language Models"
    James Y. Huang, Sailik Sengupta, Daniele Bonadiman, Yi-an Lai, Arshit Gupta, Nikolaos Pappas, Saab Mansour,
    Katrin Kirchhoff, Dan Roth
    https://arxiv.org/abs/2402.06147
    """

    Args = DeALArgs

    tokenizer: PreTrainedTokenizerBase | None = None

    def __init__(self, *args, **kwargs):
        # route through OutputControl (validate DeALArgs, mirror fields, then _configure)
        OutputControl.__init__(self, *args, **kwargs)

    def _configure(self) -> None:
        """Map DeAL's mirrored args onto the generic `SearchDriver` fields."""
        self.scorer = self.reward_func
        self.segment_len = self.lookahead
        self.num_candidates = self.init_beams
        self.keep_k = self.topk
        # self.max_iterations is already mirrored from DeALArgs
        self.propose_mode = "beam"

    def steer(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase | None = None, **_) -> PreTrainedModel:
        """Lightweight preparation; attach the tokenizer used to decode continuations."""
        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        return model

RUNTIME_KWARGS_SCHEMA = [{'name': 'reward_params', 'type': 'dict', 'scope': 'row', 'help': "Entries merged into the scorer's params mapping on every scoring call of this generation. The per-row form is one mapping; the driver handles one prompt per call, so a batched delivery is a one-element sequence holding that mapping. A per-sample mapping (for example a reference answer under 'reference') reaches the scorer's row through SampleSequenceScorer."}] class-attribute instance-attribute

args = self.Args.validate(*args, **kwargs) instance-attribute

enabled = True class-attribute instance-attribute

include_in_scoring = True class-attribute instance-attribute

keep_k = keep_k instance-attribute

max_iterations = max_iterations instance-attribute

num_candidates = num_candidates instance-attribute

propose_mode = propose_mode instance-attribute

same_model_forwards = False class-attribute instance-attribute

scorer = scorer instance-attribute

segment_len = segment_len instance-attribute

supports_batching = False class-attribute instance-attribute

tokenizer = None class-attribute instance-attribute

cleanup()

Release resources allocated during steer().

Override in subclasses that allocate GPU memory or other resources during steering. Default no-op.

Source code in steerability/algorithms/core/base_control.py

def cleanup(self) -> None:
    """Release resources allocated during `steer()`.

    Override in subclasses that allocate GPU memory or other resources during steering. Default
    no-op.
    """
    pass

clone_for_call(seed=None)

A configuration-preserving shallow clone for one generation call.

The clone shares steer-time artifacts (memories, steering vectors, attached tokenizers) with the original but has its own attribute namespace, so per-call attribute mutation on the clone never races another call using the original. When seed is given and the control defines reseed(seed), the clone's client-side RNG is re-seeded.

Parameters:

Name	Type	Description	Default
seed	int | None	Optional seed forwarded to the clone's reseed().	None

Returns:

Type	Description
BaseControl	The clone.

Source code in steerability/algorithms/core/base_control.py

def clone_for_call(self, seed: int | None = None) -> "BaseControl":
    """A configuration-preserving shallow clone for one generation call.

    The clone shares steer-time artifacts (memories, steering vectors, attached tokenizers)
    with the original but has its own attribute namespace, so per-call attribute mutation on
    the clone never races another call using the original. When `seed` is given and the
    control defines `reseed(seed)`, the clone's client-side RNG is re-seeded.

    Args:
        seed: Optional seed forwarded to the clone's `reseed()`.

    Returns:
        The clone.
    """
    clone = copy.copy(self)
    if seed is not None:
        reseed = getattr(clone, "reseed", None)
        if callable(reseed):
            reseed(seed)
    return clone

decode(input_ids, attention_mask, model, logits_processors, stopping_criteria, runtime_kwargs, session=None, **gen_kwargs)

Run the decoding procedure; return full sequence ids (prompt + continuation).

Source code in steerability/algorithms/output_control/common/drivers/search.py

def decode(self, input_ids, attention_mask, model: PreTrainedModel | None, logits_processors,
           stopping_criteria, runtime_kwargs, session=None, **gen_kwargs) -> torch.Tensor:
    if input_ids.dim() != 2 or input_ids.size(0) != 1:
        raise NotImplementedError("SearchDriver handles one prompt at a time (batch size 1).")
    if self.tokenizer is None:
        raise RuntimeError("SearchDriver requires a tokenizer; steer() must run first.")

    runtime_kwargs = runtime_kwargs or {}
    base_generate = resolve_generate_callable(model, runtime_kwargs, session=session)

    prompt_text = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)
    input_length = input_ids.size(1)
    global_budget = gen_kwargs.pop("max_new_tokens", None)

    # segment_len=None means "one full-budget segment" (best-of-N as a config); each rollout then
    # spans the whole generation budget in a single iteration
    segment_len = self.segment_len if self.segment_len is not None else global_budget
    if segment_len is None:
        raise ValueError(
            "SearchDriver requires a segment length: set `segment_len` or pass `max_new_tokens` "
            "as the generation budget (segment_len=None uses the budget as the segment length)."
        )

    reward_params = {
        **_resolve_reward_params(runtime_kwargs),
        "segment_len": segment_len,
        "num_candidates": self.num_candidates,
        "keep_k": self.keep_k,
        "max_iterations": self.max_iterations,
    }
    tokenizer_eos = getattr(self.tokenizer, "eos_token_id", None)
    eos_ids = {tokenizer_eos} if tokenizer_eos is not None else set()
    if model is not None:
        configured = getattr(model.generation_config, "eos_token_id", None)
        eos_ids.update([configured] if isinstance(configured, int) else (configured or []))

    proposer = SegmentProposer(mode=self.propose_mode)
    frontier = Frontier(
        keep_k=self.keep_k,
        eos_token_id=sorted(eos_ids) or None,
        input_length=input_length,
        max_new_tokens=global_budget,
        pad_token_id=getattr(self.tokenizer, "pad_token_id", None),
    )

    current_ids = input_ids
    kept = None
    for _ in range(self.max_iterations):
        # safe to deepcopy: the composed stacks travel as explicit decode() parameters, never inside gen_kwargs
        rollout_kwargs = copy.deepcopy(gen_kwargs)
        frontier_mask = infer_attention_mask_from_ids(current_ids, self.tokenizer.pad_token_id)
        beams = proposer.propose(
            current_ids,
            n=self.num_candidates,
            segment_len=segment_len,
            processors=logits_processors,
            criteria=stopping_criteria,
            model=model,
            base_generate=base_generate,
            attention_mask=frontier_mask,
            **rollout_kwargs,
        )
        continuations = self.tokenizer.decode(
            beams[:, input_length:], skip_special_tokens=True
        )
        scores = self.scorer(prompt_text, continuations, reward_params)
        if len(scores) != beams.size(0):
            raise RuntimeError(f"Scorer returned {len(scores)} scores for {beams.size(0)} beams.")

        step = frontier.keep(beams, scores)
        kept = step
        if all(step.finished_flags):
            break

        unfinished = [i for i, f in enumerate(step.finished_flags) if not f]
        if not unfinished:
            break
        current_ids = step.kept_ids[unfinished]

    final = frontier.best_ids if frontier.best_ids is not None else kept.kept_ids[0]
    return final.unsqueeze(0) if final.dim() == 1 else final

export_constraint(runtime_kwargs=None)

The control's declarative constrained-decoding source, or None.

A control whose per-step masking compiles from a declarative source returns a ConstraintSource; on a backend advertising Capability.GUIDED_DECODING the pipeline renders it onto the engine's native structured-output parameters in place of the control's live processor. The default returns None, which keeps the control on the live processor mechanism.

Parameters:

Name	Type	Description	Default
runtime_kwargs	dict | None	Per-call parameters supplied to generate().	None

Returns:

Type	Description
ConstraintSource | None	The constraint source, or None.

Source code in steerability/algorithms/output_control/base.py

def export_constraint(self, runtime_kwargs: dict | None = None) -> ConstraintSource | None:
    """The control's declarative constrained-decoding source, or None.

    A control whose per-step masking compiles from a declarative source returns a
    `ConstraintSource`; on a backend advertising `Capability.GUIDED_DECODING` the pipeline
    renders it onto the engine's native structured-output parameters in place of the
    control's live processor. The default returns None, which keeps the control on the live
    processor mechanism.

    Args:
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        The constraint source, or None.
    """
    return None

export_generation_params(runtime_kwargs=None)

The control's sampling-expressible contribution, or None.

A control whose behavior is expressible as normalized generation parameters returns a mapping over a subset of stop_strings, stop_token_ids, max_new_tokens, and min_new_tokens; the pipeline merges it into the call's GenerationParams (stop rules union with the caller's; token bounds only tighten) and does not additionally collect the control's live processors and criteria for that call, so the control executes on every backend through the session's composed stop rules. The default returns None, which keeps the control on the live processor/criteria mechanism.

Parameters:

Name	Type	Description	Default
runtime_kwargs	dict | None	Per-call parameters supplied to generate().	None

Returns:

Type	Description
Mapping[str, Any] | None	The parameter contribution, or None.

Source code in steerability/algorithms/output_control/base.py

def export_generation_params(self, runtime_kwargs: dict | None = None) -> Mapping[str, Any] | None:
    """The control's sampling-expressible contribution, or None.

    A control whose behavior is expressible as normalized generation parameters returns a
    mapping over a subset of `stop_strings`, `stop_token_ids`, `max_new_tokens`, and
    `min_new_tokens`; the pipeline merges it into the call's `GenerationParams` (stop rules
    union with the caller's; token bounds only tighten) and does not additionally collect
    the control's live processors and criteria for that call, so the control executes on
    every backend through the session's composed stop rules. The default returns None, which
    keeps the control on the live processor/criteria mechanism.

    Args:
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        The parameter contribution, or None.
    """
    return None

export_processor_spec(runtime_kwargs=None)

The control's engine-hosted processor form, or None.

A control whose per-step logit math is expressible in an engine's served processor vocabulary returns a ProcessorSpec; on a backend advertising Capability.PER_STEP_LOGIT_SPECS with the spec's kind, the pipeline submits it as a ProcessorSpecEntry in place of the control's live processor. The default returns None, which keeps the control on the live processor mechanism.

Parameters:

Name	Type	Description	Default
runtime_kwargs	dict | None	Per-call parameters supplied to generate().	None

Returns:

Type	Description
ProcessorSpec | None	The processor spec, or None.

Source code in steerability/algorithms/output_control/base.py

def export_processor_spec(self, runtime_kwargs: dict | None = None) -> ProcessorSpec | None:
    """The control's engine-hosted processor form, or None.

    A control whose per-step logit math is expressible in an engine's served processor
    vocabulary returns a `ProcessorSpec`; on a backend advertising
    `Capability.PER_STEP_LOGIT_SPECS` with the spec's kind, the pipeline submits it as a
    `ProcessorSpecEntry` in place of the control's live processor. The default returns
    None, which keeps the control on the live processor mechanism.

    Args:
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        The processor spec, or None.
    """
    return None

export_state()

Steer-time state to persist when freezing, keyed by logical name.

Values are typed artifacts (SteeringVector, Probe, ProbeSet, a Memory implementation, CheckpointArtifact, LoRAArtifact, torch.Tensor, or a Path to an on-disk product). Must be called after steer(). The default returns an empty mapping (nothing to persist).

Returns:

Type	Description
dict[str, Any]	Mapping from logical name to artifact value.

Source code in steerability/algorithms/core/base_control.py

def export_state(self) -> dict[str, Any]:
    """Steer-time state to persist when freezing, keyed by logical name.

    Values are typed artifacts (`SteeringVector`, `Probe`, `ProbeSet`, a `Memory`
    implementation, `CheckpointArtifact`, `LoRAArtifact`, `torch.Tensor`, or a `Path` to an
    on-disk product). Must be called after `steer()`. The default returns an empty mapping
    (nothing to persist).

    Returns:
        Mapping from logical name to artifact value.
    """
    return {}

fit_identity()

The object whose canonical form defines this control's fit digest, or None.

The digest detects recipe edits that invalidate frozen artifacts, so the returned object should cover exactly the inputs a re-fit would consume (training data, fit specs) and exclude inert application parameters (strengths, multipliers). The default returns None (no fits).

Returns:

Type	Description
Any | None	The fit-identity object, or None.

Source code in steerability/algorithms/core/base_control.py

def fit_identity(self) -> Any | None:
    """The object whose canonical form defines this control's fit digest, or None.

    The digest detects recipe edits that invalidate frozen artifacts, so the returned
    object should cover exactly the inputs a re-fit would consume (training data, fit
    specs) and exclude inert application parameters (strengths, multipliers). The default
    returns None (no fits).

    Returns:
        The fit-identity object, or None.
    """
    return None

frozen_form(state)

The (registry method key, constructor kwargs) of this control's frozen form.

state is the output of export_state(); artifact values inside the returned kwargs are replaced by store references at encode time. When steer_fits() is empty and state is empty, the recipe is the frozen form and the control's own key and recipe args are returned unchanged. A control whose freeze expands into several entries may return a list of (method key, kwargs) pairs instead.

Parameters:

Name	Type	Description	Default
state	dict[str, Any]	The output of export_state().	required

Returns:

Type	Description
tuple[str, dict[str, Any]]	The frozen method key and its constructor kwargs.

Raises:

Type	Description
NotFreezableError	If the control produces fits or state but declares no frozen form.

Source code in steerability/algorithms/core/base_control.py

def frozen_form(self, state: dict[str, Any]) -> tuple[str, dict[str, Any]]:
    """The `(registry method key, constructor kwargs)` of this control's frozen form.

    `state` is the output of `export_state()`; artifact values inside the returned kwargs
    are replaced by store references at encode time. When `steer_fits()` is empty and
    `state` is empty, the recipe is the frozen form and the control's own key and recipe
    args are returned unchanged. A control whose freeze expands into several entries may
    return a list of `(method key, kwargs)` pairs instead.

    Args:
        state: The output of `export_state()`.

    Returns:
        The frozen method key and its constructor kwargs.

    Raises:
        NotFreezableError: If the control produces fits or state but declares no frozen
            form.
    """
    if not state and not self.steer_fits():
        from steerability.algorithms.core.registry import method_key_for

        args = getattr(self, "args", None)
        kwargs = {field.name: getattr(args, field.name) for field in fields(args) if field.init} \
            if args is not None else {}
        return method_key_for(type(self)), kwargs
    raise NotFreezableError(
        f"{type(self).__name__} produces steer-time state but declares no frozen form; "
        "override frozen_form() (and export_state()) to support freezing."
    )

get_logits_processors(input_ids, runtime_kwargs, **kwargs)

The control's logits processors for the current generation.

Called once per generate() / compute_logprobs() call, after input and state controls have prepared the prompt (mirrors StateControl.get_hooks). **kwargs carries attention_mask and the caller's generation kwargs. Returned objects follow the HF LogitsProcessor convention; in-list order is preserved by the composition.

A processor must behave as a function of (prefix_ids, scores). Internal state is permitted only as memoization keyed on the prefix and must re-derive on a prefix mismatch, since drivers may restart, rewind, or reorder sequences, and scoring replays prefixes teacher-forced (subclass common.processors.base.PrefixKeyedProcessor to satisfy this mechanically). Return fresh processor instances from this hook; it is invoked once per call precisely so that per-generation state is isolated.

Parameters:

Name	Type	Description	Default
input_ids	Tensor	The steered prompt token ids [batch, seq_len].	required
runtime_kwargs	dict | None	Per-call parameters supplied to generate().	required

Returns:

Type	Description
list	A list of HF LogitsProcessor-style objects.

Source code in steerability/algorithms/output_control/base.py

def get_logits_processors(self, input_ids: torch.Tensor, runtime_kwargs: dict | None, **kwargs) -> list:
    """The control's logits processors for the current generation.

    Called once per `generate()` / `compute_logprobs()` call, after input and state
    controls have prepared the prompt (mirrors `StateControl.get_hooks`). `**kwargs`
    carries `attention_mask` and the caller's generation kwargs. Returned objects
    follow the HF `LogitsProcessor` convention; in-list order is preserved by the composition.

    A processor must behave as a function of `(prefix_ids, scores)`. Internal state is
    permitted only as memoization keyed on the prefix and must re-derive on a prefix mismatch,
    since drivers may restart, rewind, or reorder sequences, and scoring replays prefixes
    teacher-forced (subclass `common.processors.base.PrefixKeyedProcessor` to satisfy this
    mechanically). Return fresh processor instances from this hook; it is invoked once per call
    precisely so that per-generation state is isolated.

    Args:
        input_ids: The steered prompt token ids `[batch, seq_len]`.
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        A list of HF `LogitsProcessor`-style objects.
    """
    return []

get_stopping_criteria(input_ids, runtime_kwargs, **kwargs)

The control's stopping criteria.

Not applied during scoring (there is no loop to stop). Same call convention as get_logits_processors.

Parameters:

Name	Type	Description	Default
input_ids	Tensor	The steered prompt token ids [batch, seq_len].	required
runtime_kwargs	dict | None	Per-call parameters supplied to generate().	required

Returns:

Type	Description
list	A list of HF StoppingCriteria-style objects.

Source code in steerability/algorithms/output_control/base.py

def get_stopping_criteria(self, input_ids: torch.Tensor, runtime_kwargs: dict | None, **kwargs) -> list:
    """The control's stopping criteria.

    Not applied during scoring (there is no loop to stop). Same call convention as
    `get_logits_processors`.

    Args:
        input_ids: The steered prompt token ids `[batch, seq_len]`.
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        A list of HF `StoppingCriteria`-style objects.
    """
    return []

max_rollouts_per_query()

num_candidates * (1 + (max_iterations - 1) * keep_k).

The first iteration proposes num_candidates continuations from the single-row prompt; each later iteration proposes num_candidates from each of up to keep_k retained beams.

Source code in steerability/algorithms/output_control/common/drivers/search.py

def max_rollouts_per_query(self) -> int:
    """`num_candidates * (1 + (max_iterations - 1) * keep_k)`.

    The first iteration proposes `num_candidates` continuations from the single-row prompt;
    each later iteration proposes `num_candidates` from each of up to `keep_k` retained
    beams.
    """
    return self.num_candidates * (1 + (self.max_iterations - 1) * self.keep_k)

requirements()

Rollouts run through the session, so sampled proposals require nothing beyond the session contract; beam proposals require Capability.BEAM_PROPOSALS.

Source code in steerability/algorithms/output_control/common/drivers/search.py

def requirements(self) -> Requirements:
    """Rollouts run through the session, so sampled proposals require nothing beyond the
    session contract; beam proposals require `Capability.BEAM_PROPOSALS`."""
    if getattr(self, "propose_mode", "sample") == "beam":
        return Requirements(generate=needs(
            Capability.BEAM_PROPOSALS,
            hint="use propose_mode='sample' or run this pipeline on the huggingface backend",
        ))
    return Requirements()

steer(model, tokenizer=None, **_)

Lightweight preparation; attach the tokenizer used to decode continuations.

Source code in steerability/algorithms/output_control/deal/control.py

def steer(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase | None = None, **_) -> PreTrainedModel:
    """Lightweight preparation; attach the tokenizer used to decode continuations."""
    self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
    return model

steer_access()

The model access this instance's steer step requires, on the ModelAccess ladder.

The default is ModelAccess.FACTS (layout and tokenizer only). A control whose steer step generates or scores through the session declares ROLLOUTS, one that captures hidden states declares CAPTURE, and one that touches the model as a live torch.nn.Module declares MODULE. Configuration determines the result. The pipeline hands steer() a session scoped to the declared rung, and the live model only at MODULE. A control may retain the pipeline model beyond steer() only if its generate phase requires Capability.IN_PROCESS_TORCH.

Returns:

Type	Description
ModelAccess	The declared access rung.

Source code in steerability/algorithms/core/base_control.py

def steer_access(self) -> ModelAccess:
    """The model access this instance's steer step requires, on the `ModelAccess` ladder.

    The default is `ModelAccess.FACTS` (layout and tokenizer only). A control whose steer
    step generates or scores through the session declares `ROLLOUTS`, one that captures
    hidden states declares `CAPTURE`, and one that touches the model as a live
    `torch.nn.Module` declares `MODULE`. Configuration determines the result. The pipeline
    hands `steer()` a session scoped to the declared rung, and the live model only at
    `MODULE`. A control may retain the pipeline model beyond `steer()` only if its
    generate phase requires `Capability.IN_PROCESS_TORCH`.

    Returns:
        The declared access rung.
    """
    return ModelAccess.FACTS

steer_fits()

The fit artifacts this instance's steer step will produce, for the steer plan.

Each entry is (artifact, artifact_class), where artifact is the fit source or recipe class name and artifact_class is "direction" or "calibrated". The default is an empty tuple (no fits).

Returns:

Type	Description
tuple[tuple[str, str], ...]	The declared fit artifacts, in declaration order.

Source code in steerability/algorithms/core/base_control.py

def steer_fits(self) -> tuple[tuple[str, str], ...]:
    """The fit artifacts this instance's steer step will produce, for the steer plan.

    Each entry is `(artifact, artifact_class)`, where `artifact` is the fit source or
    recipe class name and `artifact_class` is `"direction"` or `"calibrated"`. The default
    is an empty tuple (no fits).

    Returns:
        The declared fit artifacts, in declaration order.
    """
    return ()