Skip to content

ValueGuidance

steerability.algorithms.output_control.value_guidance

args

control

ValueGuidance

Bases: OutputControl

Value-guided decoding as configuration: score candidate tokens with a value function and shift their logits.

ValueGuidance is the generic over the step shape. It exposes the common value slot through flat Args: a candidate policy selects a small set of next tokens, a per-candidate value scores them, the values are normalized per row, and the selected candidates' logits are shifted by beta * value (optionally masking non-candidates to -inf). A method from the literature is an assignment of a config, not a subclass:

- FUDGE: `value={"kind": "classifier", ...}, policy="top_k", beta=1.0, normalize="none"`.
- ARGS: `value={"kind": "reward_model", ...}, policy="top_k", normalize="none"`.
- RAD-equivalent: `value={"kind": "reward_model", ...}, policy="top_k", k=20,
  normalize="clamp", invert=True, mask_non_candidates=True`.
- SASA-equivalent: `value={"kind": "subspace_margin", ...}, policy="surviving",
  normalize="softmax", mask_non_candidates=False, include_in_scoring=False`.

ValueGuidance is a step-level control: it composes a ValueGuidedProcessor into the decoding stack, so it works alongside other output controls and any decoding driver. steer() resolves the value spec (loading a reward model or classifier, or fitting a probe); a fresh processor is returned per call.

Reference:

Source code in steerability/algorithms/output_control/value_guidance/control.py 
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
class ValueGuidance(OutputControl):
    """Value-guided decoding as configuration: score candidate tokens with a value function and shift their logits.

    `ValueGuidance` is the generic over the step shape. It exposes the `common` value slot through
    flat `Args`: a candidate policy selects a small set of next tokens, a per-candidate value scores
    them, the values are normalized per row, and the selected candidates' logits are shifted by
    `beta * value` (optionally masking non-candidates to `-inf`). A method from the literature is an
    assignment of a config, not a subclass:

        - FUDGE: `value={"kind": "classifier", ...}, policy="top_k", beta=1.0, normalize="none"`.
        - ARGS: `value={"kind": "reward_model", ...}, policy="top_k", normalize="none"`.
        - RAD-equivalent: `value={"kind": "reward_model", ...}, policy="top_k", k=20,
          normalize="clamp", invert=True, mask_non_candidates=True`.
        - SASA-equivalent: `value={"kind": "subspace_margin", ...}, policy="surviving",
          normalize="softmax", mask_non_candidates=False, include_in_scoring=False`.

    `ValueGuidance` is a step-level control: it composes a `ValueGuidedProcessor` into the decoding
    stack, so it works alongside other output controls and any decoding driver. `steer()` resolves
    the value spec (loading a reward model or classifier, or fitting a probe); a fresh processor is
    returned per call.

    Reference:

    - "FUDGE: Controlled Text Generation With Future Discriminators"
      Kevin Yang, Dan Klein
      [https://arxiv.org/abs/2104.05218](https://arxiv.org/abs/2104.05218)

    - "ARGS: Alignment as Reward-Guided Search"
      Maxim Khanov, Jirayu Burapacheep, Yixuan Li
      [https://arxiv.org/abs/2402.01694](https://arxiv.org/abs/2402.01694)
    """

    Args = ValueGuidanceArgs

    # placeholders (filled by steer)
    model: PreTrainedModel | None = None
    tokenizer: PreTrainedTokenizerBase | None = None
    _value = None

    def steer_access(self) -> ModelAccess:
        """`ModelAccess.MODULE`; the value spec resolves against the live model (probe fits,
        placement), which is retained past steer (the generate phase is in-process)."""
        return ModelAccess.MODULE

    def steer(
        self,
        model: PreTrainedModel,
        tokenizer: PreTrainedTokenizerBase | None = None,
        **__,
    ) -> PreTrainedModel:
        """Resolve the value spec, then derive batching / scoring posture from the resolved value."""
        self.model = model
        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        device = next(model.parameters()).device
        self._value = resolve_value(self.value, model=model, tokenizer=self.tokenizer, device=device)

        self.supports_batching = bool(self._value.supports_batching) and self.policy == "top_k"
        self.same_model_forwards = getattr(self._value, "same_model_forwards", False)
        if self._value.scoring_cost == "model_forward" and self.include_in_scoring:
            warnings.warn(
                "ValueGuidance: this value forwards the pipeline model per reference position during "
                "compute_logprobs; set include_in_scoring=False (SASA's default posture) unless scored "
                "logprobs under the shifted distribution are required.",
                UserWarning,
            )
        return model

    def get_logits_processors(self, input_ids, runtime_kwargs, attention_mask=None, **kwargs) -> list:
        """Return a fresh `ValueGuidedProcessor` for this generation (the freshness contract)."""
        if getattr(self, "_value", None) is None:
            raise RuntimeError("ValueGuidance.steer() must run before generation.")
        return [
            ValueGuidedProcessor(
                self._value,
                policy=self.policy,
                k=self.k,
                p=self.p,
                beta=self.beta,
                normalize=self.normalize,
                invert=self.invert,
                mask_non_candidates=self.mask_non_candidates,
                max_candidates=self.max_candidates,
                lm_tokenizer=self.tokenizer,
                model=self.model,
                attention_mask=attention_mask,
            )
        ]

    def cleanup(self) -> None:
        """Release the resolved value and model references."""
        value = getattr(self, "_value", None)
        if value is not None:
            value.cleanup()
        self._value = None
        self.model = None
        self.tokenizer = None
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
RUNTIME_KWARGS_SCHEMA = [] 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
model = None class-attribute instance-attribute
same_model_forwards = False class-attribute instance-attribute
supports_batching = False class-attribute instance-attribute
tokenizer = None class-attribute instance-attribute
cleanup()

Release the resolved value and model references.

Source code in steerability/algorithms/output_control/value_guidance/control.py 
104
105
106
107
108
109
110
111
112
113
114
def cleanup(self) -> None:
    """Release the resolved value and model references."""
    value = getattr(self, "_value", None)
    if value is not None:
        value.cleanup()
    self._value = None
    self.model = None
    self.tokenizer = None
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
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 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
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
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 
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
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 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
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 
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
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 
117
118
119
120
121
122
123
124
125
126
127
128
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 
161
162
163
164
165
166
167
168
169
170
171
172
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 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
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, attention_mask=None, **kwargs)

Return a fresh ValueGuidedProcessor for this generation (the freshness contract).

Source code in steerability/algorithms/output_control/value_guidance/control.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def get_logits_processors(self, input_ids, runtime_kwargs, attention_mask=None, **kwargs) -> list:
    """Return a fresh `ValueGuidedProcessor` for this generation (the freshness contract)."""
    if getattr(self, "_value", None) is None:
        raise RuntimeError("ValueGuidance.steer() must run before generation.")
    return [
        ValueGuidedProcessor(
            self._value,
            policy=self.policy,
            k=self.k,
            p=self.p,
            beta=self.beta,
            normalize=self.normalize,
            invert=self.invert,
            mask_non_candidates=self.mask_non_candidates,
            max_candidates=self.max_candidates,
            lm_tokenizer=self.tokenizer,
            model=self.model,
            attention_mask=attention_mask,
        )
    ]
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 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
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 []
requirements()

Backend requirements computed from this instance's configuration, per phase.

The default requires Capability.IN_PROCESS_TORCH at generate and, when include_in_scoring is True, at score as well, since remote prompt-logprob computation applies neither live processors nor engine-registered sampling processors to prefill logits. Setting include_in_scoring=False removes the score-phase requirement.

Returns:

Type Description
Requirements

The control's phase-keyed requirements.
Source code in steerability/algorithms/output_control/base.py 
214
215
216
217
218
219
220
221
222
223
224
225
226
def requirements(self) -> Requirements:
    """Backend requirements computed from this instance's configuration, per phase.

    The default requires `Capability.IN_PROCESS_TORCH` at generate and, when
    `include_in_scoring` is True, at score as well, since remote prompt-logprob computation
    applies neither live processors nor engine-registered sampling processors to prefill
    logits. Setting `include_in_scoring=False` removes the score-phase requirement.

    Returns:
        The control's phase-keyed requirements.
    """
    score = needs(Capability.IN_PROCESS_TORCH) if self.include_in_scoring else ()
    return Requirements(generate=needs(Capability.IN_PROCESS_TORCH), score=score)
steer(model, tokenizer=None, **__)

Resolve the value spec, then derive batching / scoring posture from the resolved value.

Source code in steerability/algorithms/output_control/value_guidance/control.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def steer(
    self,
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase | None = None,
    **__,
) -> PreTrainedModel:
    """Resolve the value spec, then derive batching / scoring posture from the resolved value."""
    self.model = model
    self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
    device = next(model.parameters()).device
    self._value = resolve_value(self.value, model=model, tokenizer=self.tokenizer, device=device)

    self.supports_batching = bool(self._value.supports_batching) and self.policy == "top_k"
    self.same_model_forwards = getattr(self._value, "same_model_forwards", False)
    if self._value.scoring_cost == "model_forward" and self.include_in_scoring:
        warnings.warn(
            "ValueGuidance: this value forwards the pipeline model per reference position during "
            "compute_logprobs; set include_in_scoring=False (SASA's default posture) unless scored "
            "logprobs under the shifted distribution are required.",
            UserWarning,
        )
    return model
steer_access()

ModelAccess.MODULE; the value spec resolves against the live model (probe fits, placement), which is retained past steer (the generate phase is in-process).

Source code in steerability/algorithms/output_control/value_guidance/control.py 
55
56
57
58
def steer_access(self) -> ModelAccess:
    """`ModelAccess.MODULE`; the value spec resolves against the live model (probe fits,
    placement), which is retained past steer (the generate phase is in-process)."""
    return ModelAccess.MODULE
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 
105
106
107
108
109
110
111
112
113
114
115
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 ()