Skip to content

RAD

steerability.algorithms.output_control.rad

args

control

logger = logging.getLogger(__name__) module-attribute

RAD

Bases: OutputControl

Implementation of RAD (Reward-Augmented Decoding) from Deng and Raffel, 2023.

RAD works in two phases:

  1. Preparation (steer): load an AutoModelForSequenceClassification reward model.
  2. Controlled decoding (get_logits_processors): at each decode step, the top-top_k candidate tokens are scored by the reward model, their reward is clamped to [0, 1], and the candidate logits are shifted by beta * reward while non-candidate logits are masked to -inf.

The score read from the reward model is column score_index of its output, after score_transform ("none" reads a raw logit, "sigmoid" and "softmax" map it into [0, 1]). With invert=True the shift uses 1 - reward, steering away from the scored attribute. The "clamp" normalization is absolute rather than relative to the candidate set, so the shift spread across candidates is beta * (max_reward - min_reward) and benign steps preserve the base distribution. When score_transform="none" the reward is an unbounded logit, so invert=True gives 1 - clamp(v), which saturates to 0 for any logit at or above 1; pass score_transform="sigmoid" to invert a logit head meaningfully.

RAD is a step-level control. steer() loads the reward model into a candidate value, and get_logits_processors() returns a fresh ValueGuidedProcessor per call. Candidates are the top-top_k of the scores this processor receives, so its position in a composed output stack matters. Caller-supplied sampling kwargs (temperature, top_p, repetition penalty) apply around the shift; in particular temperature rescales the effective beta, so a protocol-faithful run passes do_sample=True and nothing else. A value_trace list passed via runtime_kwargs receives one ValueStepRecord per scored step (the candidate ids, their pre-shift scores, and the per-candidate reward).

Two scoring paths back the value. When efficient=True and the reward model is decoder-only and shares the language model's vocabulary, steer() builds a CachedRewardModelValue that memoizes the reward model's prefix activations across steps (the paper's O(km) unidirectional path). When a precondition or a steer-time smoke forward fails, RAD emits one UserWarning naming the failed precondition and falls back to RewardModelValue(shared_vocab=True), which produces the same scores at higher cost. When the vocabularies differ, RAD uses RewardModelValue(shared_vocab=False), which decodes candidates to text and re-encodes with the reward-model tokenizer. Toggling efficient changes speed only, not scores.

Reference:

- "Reward-Augmented Decoding: Efficient Controlled Text Generation With a Unidirectional Reward Model"
  Haikang Deng, Colin Raffel
  [https://arxiv.org/abs/2310.09520](https://arxiv.org/abs/2310.09520)
Source code in steerability/algorithms/output_control/rad/control.py 
 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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
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
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
class RAD(OutputControl):
    """Implementation of RAD (Reward-Augmented Decoding) from Deng and Raffel, 2023.

    RAD works in two phases:

    1. **Preparation (`steer`)**: load an `AutoModelForSequenceClassification` reward model.
    2. **Controlled decoding (`get_logits_processors`)**: at each decode step, the top-`top_k`
       candidate tokens are scored by the reward model, their reward is clamped to `[0, 1]`, and the
       candidate logits are shifted by `beta * reward` while non-candidate logits are masked to
       `-inf`.

    The score read from the reward model is column `score_index` of its output, after `score_transform`
    (`"none"` reads a raw logit, `"sigmoid"` and `"softmax"` map it into `[0, 1]`). With `invert=True`
    the shift uses `1 - reward`, steering away from the scored attribute. The `"clamp"` normalization
    is absolute rather than relative to the candidate set, so the shift spread across candidates is
    `beta * (max_reward - min_reward)` and benign steps preserve the base distribution. When
    `score_transform="none"` the reward is an unbounded logit, so `invert=True` gives `1 - clamp(v)`,
    which saturates to 0 for any logit at or above 1; pass `score_transform="sigmoid"` to invert a
    logit head meaningfully.

    RAD is a step-level control. `steer()` loads the reward model into a candidate value, and
    `get_logits_processors()` returns a fresh `ValueGuidedProcessor` per call. Candidates are the
    top-`top_k` of the scores this processor receives, so its position in a composed output stack
    matters. Caller-supplied sampling kwargs (temperature, `top_p`, repetition penalty) apply around
    the shift; in particular temperature rescales the effective `beta`, so a protocol-faithful run
    passes `do_sample=True` and nothing else. A `value_trace` list passed via `runtime_kwargs`
    receives one `ValueStepRecord` per scored step (the candidate ids, their pre-shift scores, and
    the per-candidate reward).

    Two scoring paths back the value. When `efficient=True` and the reward model is decoder-only and
    shares the language model's vocabulary, `steer()` builds a `CachedRewardModelValue` that memoizes
    the reward model's prefix activations across steps (the paper's O(km) unidirectional path). When
    a precondition or a steer-time smoke forward fails, RAD emits one `UserWarning` naming the failed
    precondition and falls back to `RewardModelValue(shared_vocab=True)`, which produces the same
    scores at higher cost. When the vocabularies differ, RAD uses `RewardModelValue(shared_vocab=False)`,
    which decodes candidates to text and re-encodes with the reward-model tokenizer. Toggling
    `efficient` changes speed only, not scores.

    Reference:

        - "Reward-Augmented Decoding: Efficient Controlled Text Generation With a Unidirectional Reward Model"
          Haikang Deng, Colin Raffel
          [https://arxiv.org/abs/2310.09520](https://arxiv.org/abs/2310.09520)
    """

    Args = RADArgs

    RUNTIME_KWARGS_SCHEMA: list[dict] = [
        {
            "name": "value_trace",
            "type": "list",
            "scope": "call",
            "help": "A caller-owned list that receives one ValueStepRecord per scored step of this call.",
        },
    ]

    tokenizer: PreTrainedTokenizerBase | None = None
    _value = None

    beta: float

    def steer_access(self) -> ModelAccess:
        """`ModelAccess.MODULE`; the reward model's placement follows the live model at steer time
        (the generate phase is in-process)."""
        return ModelAccess.MODULE

    def steer(
        self,
        model: PreTrainedModel,
        tokenizer: PreTrainedTokenizerBase | None = None,
        **__,
    ) -> None:
        """Load the reward model and build the candidate value.

        Loads an `AutoModelForSequenceClassification` reward model and builds a
        `CachedRewardModelValue` when `efficient` is set and the reward model is decoder-only and
        shares the language model's vocabulary (verified by a smoke forward), otherwise a
        `RewardModelValue`. Derives `supports_batching` from the resolved value. Performs no network
        downloads or filesystem writes (model loading may hit the HF cache).

        Args:
            model (PreTrainedModel): The base language model to be steered.
            tokenizer (PreTrainedTokenizerBase | None): Tokenizer for the base model.
        """
        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        device = next(model.parameters()).device

        reward_model, rm_tokenizer = load_sequence_classifier(
            self.reward_model_id,
            device=device,
            hf_model_kwargs=self.reward_model_kwargs,
        )

        shared_vocab = self._vocab_matches(rm_tokenizer, self.tokenizer)
        self._value = self._build_value(reward_model, rm_tokenizer, shared_vocab)
        self.supports_batching = self._value.supports_batching

    def _build_value(self, reward_model, rm_tokenizer, shared_vocab: bool):
        """Select the candidate value: cached (with smoke-test + degrade), shared-vocab, or text."""
        if not shared_vocab:
            return RewardModelValue(
                reward_model, rm_tokenizer,
                score_index=self.score_index, score_transform=self.score_transform,
                shared_vocab=False,
            )

        stateless = RewardModelValue(
            reward_model, rm_tokenizer,
            score_index=self.score_index, score_transform=self.score_transform,
            shared_vocab=True,
        )
        if not self.efficient:
            return stateless

        if not self._is_unidirectional(reward_model):
            warnings.warn(
                "RAD: the reward model is not decoder-only (no past_key_values/cache_position support); "
                "falling back to the stateless reward value.",
                UserWarning,
            )
            return stateless

        cached = CachedRewardModelValue(
            reward_model, rm_tokenizer,
            score_index=self.score_index, score_transform=self.score_transform,
        )
        if not self._cached_smoke_ok(cached):
            warnings.warn(
                "RAD: the cached reward-model forward failed its smoke test; falling back to the "
                "stateless reward value.",
                UserWarning,
            )
            return stateless
        return cached

    @staticmethod
    def _vocab_matches(rm_tokenizer, lm_tokenizer) -> bool:
        """Whether the reward-model and language-model tokenizers share a vocabulary."""
        if lm_tokenizer is None:
            return False
        try:
            return rm_tokenizer.get_vocab() == lm_tokenizer.get_vocab()
        except Exception:
            return False

    @staticmethod
    def _is_unidirectional(reward_model) -> bool:
        """Whether the reward model's forward accepts `past_key_values` (and thus `cache_position`).

        A decoder-only sequence classifier threads `past_key_values` and absorbs `cache_position`
        through a `**kwargs` catch-all; an encoder classifier (BERT/RoBERTa) accepts neither. The
        cached forward's smoke test at steer time is the final gate.
        """
        import inspect

        try:
            params = inspect.signature(reward_model.forward).parameters
        except (TypeError, ValueError):
            return False
        if "past_key_values" not in params:
            return False
        has_var_keyword = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
        return "cache_position" in params or has_var_keyword

    def _cached_smoke_ok(self, cached: CachedRewardModelValue) -> bool:
        """Run one tiny cached forward to confirm the reward model supports the cached path."""
        from steerability.algorithms.output_control.common.values.base import StepContext

        device = cached._device
        prefix = torch.zeros(1, 2, dtype=torch.long, device=device)
        candidates = torch.zeros(1, 1, dtype=torch.long, device=device)
        ctx = StepContext(
            prefix_ids=prefix,
            candidate_ids=candidates,
            lm_tokenizer=self.tokenizer,
            attention_mask=torch.ones(1, 2, dtype=torch.long, device=device),
        )
        try:
            cached.score(ctx)
        except Exception as exc:
            logger.debug("RAD cached smoke forward failed: %s", exc)
            cached._cached_ids = None
            cached._cache = None
            return False
        cached._cached_ids = None
        cached._cache = None
        return True

    @property
    def scoring_path(self) -> str | None:
        """Which scoring path `steer()` resolved, or `None` before `steer()`.

        Returns `"cached"` for a `CachedRewardModelValue` (the paper's O(km) unidirectional path),
        `"stateless"` for a shared-vocabulary `RewardModelValue`, and `"text"` for a
        `RewardModelValue` over a different vocabulary (the decode-and-re-encode path).
        """
        if self._value is None:
            return None
        if isinstance(self._value, CachedRewardModelValue):
            return "cached"
        return "stateless" if self._value.shared_vocab else "text"

    def get_logits_processors(self, input_ids, runtime_kwargs, attention_mask=None, **kwargs) -> list:
        """Return a fresh `ValueGuidedProcessor` implementing RAD's reward-augmented shift.

        Candidates are the top-`top_k` of the scores this processor receives; non-candidate tokens
        are masked to `-inf`; candidate rewards are clamped to `[0, 1]` (inverted when `invert` is
        set) and the candidate logits are shifted by `beta * reward`. A `value_trace` list in
        `runtime_kwargs` receives one `ValueStepRecord` per scored step; `compute_logprobs` replays
        also append when `include_in_scoring` is True.
        """
        if self._value is None:
            raise RuntimeError("RAD.steer() must run before generation (reward model not loaded).")
        return [
            ValueGuidedProcessor(
                self._value,
                policy="top_k",
                k=self.top_k,
                beta=self.beta,
                normalize="clamp",
                invert=self.invert,
                mask_non_candidates=True,
                lm_tokenizer=self.tokenizer,
                attention_mask=attention_mask,
                trace=(runtime_kwargs or {}).get("value_trace"),
            )
        ]

    def cleanup(self) -> None:
        """Release the reward model and tokenizer to free memory."""
        if self._value is not None:
            self._value.cleanup()
        self._value = None
        self.tokenizer = None

        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        logger.debug("RAD cleanup completed")
RUNTIME_KWARGS_SCHEMA = [{'name': 'value_trace', 'type': 'list', 'scope': 'call', 'help': 'A caller-owned list that receives one ValueStepRecord per scored step of this call.'}] class-attribute instance-attribute
args = self.Args.validate(*args, **kwargs) instance-attribute
beta instance-attribute
enabled = True class-attribute instance-attribute
include_in_scoring = True class-attribute instance-attribute
same_model_forwards = False class-attribute instance-attribute
scoring_path property

Which scoring path steer() resolved, or None before steer().

Returns "cached" for a CachedRewardModelValue (the paper's O(km) unidirectional path), "stateless" for a shared-vocabulary RewardModelValue, and "text" for a RewardModelValue over a different vocabulary (the decode-and-re-encode path).

supports_batching = False class-attribute instance-attribute
tokenizer = None class-attribute instance-attribute
cleanup()

Release the reward model and tokenizer to free memory.

Source code in steerability/algorithms/output_control/rad/control.py 
248
249
250
251
252
253
254
255
256
257
258
259
def cleanup(self) -> None:
    """Release the reward model and tokenizer to free memory."""
    if self._value is not None:
        self._value.cleanup()
    self._value = None
    self.tokenizer = None

    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    logger.debug("RAD cleanup completed")
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 implementing RAD's reward-augmented shift.

Candidates are the top-top_k of the scores this processor receives; non-candidate tokens are masked to -inf; candidate rewards are clamped to [0, 1] (inverted when invert is set) and the candidate logits are shifted by beta * reward. A value_trace list in runtime_kwargs receives one ValueStepRecord per scored step; compute_logprobs replays also append when include_in_scoring is True.

Source code in steerability/algorithms/output_control/rad/control.py 
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
def get_logits_processors(self, input_ids, runtime_kwargs, attention_mask=None, **kwargs) -> list:
    """Return a fresh `ValueGuidedProcessor` implementing RAD's reward-augmented shift.

    Candidates are the top-`top_k` of the scores this processor receives; non-candidate tokens
    are masked to `-inf`; candidate rewards are clamped to `[0, 1]` (inverted when `invert` is
    set) and the candidate logits are shifted by `beta * reward`. A `value_trace` list in
    `runtime_kwargs` receives one `ValueStepRecord` per scored step; `compute_logprobs` replays
    also append when `include_in_scoring` is True.
    """
    if self._value is None:
        raise RuntimeError("RAD.steer() must run before generation (reward model not loaded).")
    return [
        ValueGuidedProcessor(
            self._value,
            policy="top_k",
            k=self.top_k,
            beta=self.beta,
            normalize="clamp",
            invert=self.invert,
            mask_non_candidates=True,
            lm_tokenizer=self.tokenizer,
            attention_mask=attention_mask,
            trace=(runtime_kwargs or {}).get("value_trace"),
        )
    ]
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, **__)

Load the reward model and build the candidate value.

Loads an AutoModelForSequenceClassification reward model and builds a CachedRewardModelValue when efficient is set and the reward model is decoder-only and shares the language model's vocabulary (verified by a smoke forward), otherwise a RewardModelValue. Derives supports_batching from the resolved value. Performs no network downloads or filesystem writes (model loading may hit the HF cache).

Parameters:

Name Type Description Default
model PreTrainedModel

The base language model to be steered.

 required
tokenizer PreTrainedTokenizerBase | None

Tokenizer for the base model.

 None
Source code in steerability/algorithms/output_control/rad/control.py 
 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
115
def steer(
    self,
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase | None = None,
    **__,
) -> None:
    """Load the reward model and build the candidate value.

    Loads an `AutoModelForSequenceClassification` reward model and builds a
    `CachedRewardModelValue` when `efficient` is set and the reward model is decoder-only and
    shares the language model's vocabulary (verified by a smoke forward), otherwise a
    `RewardModelValue`. Derives `supports_batching` from the resolved value. Performs no network
    downloads or filesystem writes (model loading may hit the HF cache).

    Args:
        model (PreTrainedModel): The base language model to be steered.
        tokenizer (PreTrainedTokenizerBase | None): Tokenizer for the base model.
    """
    self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
    device = next(model.parameters()).device

    reward_model, rm_tokenizer = load_sequence_classifier(
        self.reward_model_id,
        device=device,
        hf_model_kwargs=self.reward_model_kwargs,
    )

    shared_vocab = self._vocab_matches(rm_tokenizer, self.tokenizer)
    self._value = self._build_value(reward_model, rm_tokenizer, shared_vocab)
    self.supports_batching = self._value.supports_batching
steer_access()

ModelAccess.MODULE; the reward model's placement follows the live model at steer time (the generate phase is in-process).

Source code in steerability/algorithms/output_control/rad/control.py 
81
82
83
84
def steer_access(self) -> ModelAccess:
    """`ModelAccess.MODULE`; the reward model's placement follows the live model at steer time
    (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 ()

utils

PrefixRewardTrainSpec dataclass

Training settings for a prefix-scored reward model.

Attributes:

Name Type Description
max_length int

Maximum token length per training example (right-padded).
batch_size int

Examples per optimizer step.
epochs int

Passes over the data.
learning_rate float

AdamW learning rate.
weight_decay float

AdamW weight decay.
seed int

Seed for shuffling and initialization of the fresh scalar head.
log_every int

Log the running loss every this many optimizer steps.
Source code in steerability/algorithms/output_control/rad/utils/reward_training.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
@dataclass
class PrefixRewardTrainSpec:
    """Training settings for a prefix-scored reward model.

    Attributes:
        max_length: Maximum token length per training example (right-padded).
        batch_size: Examples per optimizer step.
        epochs: Passes over the data.
        learning_rate: AdamW learning rate.
        weight_decay: AdamW weight decay.
        seed: Seed for shuffling and initialization of the fresh scalar head.
        log_every: Log the running loss every this many optimizer steps.
    """

    max_length: int = 64
    batch_size: int = 64
    epochs: int = 1
    learning_rate: float = 1e-5
    weight_decay: float = 0.01
    seed: int = 0
    log_every: int = 50
batch_size = 64 class-attribute instance-attribute
epochs = 1 class-attribute instance-attribute
learning_rate = 1e-05 class-attribute instance-attribute
log_every = 50 class-attribute instance-attribute
max_length = 64 class-attribute instance-attribute
seed = 0 class-attribute instance-attribute
weight_decay = 0.01 class-attribute instance-attribute

prefix_reward_loss(predictions, labels, attention_mask)

The cumulative squared-error loss over prefixes, averaged over the batch.

For one sequence of length l with prediction r_t at every prefix t (1-indexed) and label r, the loss is sum_t t * (r_t - r)^2 / S_l with S_l = l * (l + 1) / 2. Positions with attention_mask == 0 are excluded and l counts the included positions, so a fully masked row contributes nothing. The loss is computed in float32 regardless of the input dtypes.

Parameters:

Name Type Description Default
predictions Tensor

Per-position predictions [B, T].

 required
labels Tensor

Per-sequence labels [B].

 required
attention_mask Tensor

[B, T] mask; 0 positions are excluded.

 required

Returns:

Type Description
Tensor

A float32 scalar tensor, the mean per-sequence prefix loss over the non-empty rows.
Source code in steerability/algorithms/output_control/rad/utils/reward_training.py 
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
def prefix_reward_loss(
    predictions: torch.Tensor,
    labels: torch.Tensor,
    attention_mask: torch.Tensor,
) -> torch.Tensor:
    """The cumulative squared-error loss over prefixes, averaged over the batch.

    For one sequence of length `l` with prediction `r_t` at every prefix `t` (1-indexed) and label
    `r`, the loss is `sum_t t * (r_t - r)^2 / S_l` with `S_l = l * (l + 1) / 2`. Positions with
    `attention_mask == 0` are excluded and `l` counts the included positions, so a fully masked row
    contributes nothing. The loss is computed in float32 regardless of the input dtypes.

    Args:
        predictions: Per-position predictions `[B, T]`.
        labels: Per-sequence labels `[B]`.
        attention_mask: `[B, T]` mask; `0` positions are excluded.

    Returns:
        A float32 scalar tensor, the mean per-sequence prefix loss over the non-empty rows.
    """
    predictions = predictions.float()
    mask = attention_mask.to(torch.float32)  # [B, T]
    positions = torch.cumsum(mask, dim=1) * mask  # 1-indexed prefix length at each kept position
    lengths = mask.sum(dim=1)  # [B]
    normalizer = lengths * (lengths + 1) / 2  # S_l per row

    squared_error = (predictions - labels.unsqueeze(1).float()) ** 2  # [B, T]
    per_row = (positions * squared_error).sum(dim=1) / normalizer.clamp_min(1.0)  # [B]

    nonempty = lengths > 0
    if not torch.any(nonempty):
        return per_row.sum() * 0.0
    return per_row[nonempty].mean()

prefix_rewards(model, input_ids, attention_mask)

Per-position rewards [B, T] from a sequence-classification model with a scalar head.

Runs the backbone (getattr(model, model.base_model_prefix)) once, applies model.score at every position rather than only the pooled last token, and returns the sigmoid. Training and the prefix-tracking diagnostics use this instead of the classifier's own forward, which pools the last token only.

Parameters:

Name Type Description Default
model

An AutoModelForSequenceClassification with a single-logit score head.

 required
input_ids Tensor

Token ids [B, T].

 required
attention_mask Tensor

[B, T] mask forwarded to the backbone.

 required

Returns:

Type Description
Tensor

Per-position rewards [B, T] in [0, 1].
Source code in steerability/algorithms/output_control/rad/utils/reward_training.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def prefix_rewards(
    model,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
) -> torch.Tensor:
    """Per-position rewards `[B, T]` from a sequence-classification model with a scalar head.

    Runs the backbone (`getattr(model, model.base_model_prefix)`) once, applies `model.score` at every
    position rather than only the pooled last token, and returns the sigmoid. Training and the
    prefix-tracking diagnostics use this instead of the classifier's own `forward`, which pools the
    last token only.

    Args:
        model: An `AutoModelForSequenceClassification` with a single-logit `score` head.
        input_ids: Token ids `[B, T]`.
        attention_mask: `[B, T]` mask forwarded to the backbone.

    Returns:
        Per-position rewards `[B, T]` in `[0, 1]`.
    """
    backbone = getattr(model, model.base_model_prefix)
    hidden = backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state  # [B, T, H]
    logits = model.score(hidden).squeeze(-1)  # [B, T]
    return torch.sigmoid(logits)

train_prefix_reward_model(backbone_name_or_path, texts, labels, output_dir, *, spec=None, device=None)

Fine-tune a scalar-head sequence classifier on texts with prefix_reward_loss and save it.

Builds the classifier from the causal-LM backbone with num_labels=1 (the head class is selected by sequence_classifier_class, so a Granite backbone resolves), scales the fresh scalar head down by HEAD_INIT_SCALE so its initial sigmoid output is unsaturated, right-pads with the backbone tokenizer (pad falls back to eos), trains with AdamW under bf16 autocast and a linear decay schedule, and writes the model and tokenizer to output_dir. The saved config.json records num_labels: 1 and the head's architecture name, so load_sequence_classifier(output_dir) reloads it.

Labels are the dataset's native attribute (for civil_comments, toxicity in [0, 1]); the head predicts the attribute, and RAD's invert=True turns the prediction into a reward.

Parameters:

Name Type Description Default
backbone_name_or_path str

HF id or local path of the causal-LM backbone.

 required
texts Sequence[str]

Training texts.

 required
labels Sequence[float]

Per-text labels in [0, 1], aligned with texts.

 required
output_dir str | Path

Directory the trained model and tokenizer are written to.

 required
spec PrefixRewardTrainSpec | None

Training settings; defaults to PrefixRewardTrainSpec().

 None
device str | device | None

Device to train on; defaults to CUDA when available, else CPU.

 None

Returns:

Type Description
Path

output_dir as a Path.
Source code in steerability/algorithms/output_control/rad/utils/reward_training.py 
119
120
121
122
123
124
125
126
127
128
129
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
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
def train_prefix_reward_model(
    backbone_name_or_path: str,
    texts: Sequence[str],
    labels: Sequence[float],
    output_dir: str | Path,
    *,
    spec: PrefixRewardTrainSpec | None = None,
    device: str | torch.device | None = None,
) -> Path:
    """Fine-tune a scalar-head sequence classifier on `texts` with `prefix_reward_loss` and save it.

    Builds the classifier from the causal-LM backbone with `num_labels=1` (the head class is selected
    by `sequence_classifier_class`, so a Granite backbone resolves), scales the fresh scalar head down
    by `HEAD_INIT_SCALE` so its initial sigmoid output is unsaturated, right-pads with the backbone
    tokenizer (pad falls back to eos), trains with AdamW under bf16 autocast and a linear decay
    schedule, and writes the model and tokenizer to `output_dir`. The saved `config.json` records
    `num_labels: 1` and the head's architecture name, so `load_sequence_classifier(output_dir)`
    reloads it.

    Labels are the dataset's native attribute (for `civil_comments`, toxicity in `[0, 1]`); the head
    predicts the attribute, and RAD's `invert=True` turns the prediction into a reward.

    Args:
        backbone_name_or_path: HF id or local path of the causal-LM backbone.
        texts: Training texts.
        labels: Per-text labels in `[0, 1]`, aligned with `texts`.
        output_dir: Directory the trained model and tokenizer are written to.
        spec: Training settings; defaults to `PrefixRewardTrainSpec()`.
        device: Device to train on; defaults to CUDA when available, else CPU.

    Returns:
        `output_dir` as a `Path`.
    """
    spec = spec or PrefixRewardTrainSpec()
    output_dir = Path(output_dir)
    device = torch.device(device) if device is not None else torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.manual_seed(spec.seed)

    tokenizer = AutoTokenizer.from_pretrained(backbone_name_or_path)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    kwargs = {"num_labels": 1, "dtype": torch.float32}
    model = sequence_classifier_class(backbone_name_or_path, **kwargs).from_pretrained(backbone_name_or_path, **kwargs)
    model.config.pad_token_id = tokenizer.pad_token_id
    with torch.no_grad():
        model.score.weight.mul_(HEAD_INIT_SCALE)
    model = model.to(device)
    model.train()

    texts = list(texts)
    labels = torch.tensor(list(labels), dtype=torch.float32)
    num_examples = len(texts)
    steps_per_epoch = (num_examples + spec.batch_size - 1) // spec.batch_size
    total_steps = steps_per_epoch * spec.epochs

    optimizer = AdamW(model.parameters(), lr=spec.learning_rate, weight_decay=spec.weight_decay)
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=max(total_steps, 1))
    use_bf16 = device.type == "cuda"

    generator = torch.Generator().manual_seed(spec.seed)
    step = 0
    for epoch in range(spec.epochs):
        order = torch.randperm(num_examples, generator=generator)
        for start in range(0, num_examples, spec.batch_size):
            batch_index = order[start:start + spec.batch_size]
            batch_texts = [texts[i] for i in batch_index.tolist()]
            batch_labels = labels[batch_index].to(device)
            encoded = tokenizer(
                batch_texts,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=spec.max_length,
            ).to(device)

            optimizer.zero_grad()
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=use_bf16):
                predictions = prefix_rewards(model, encoded["input_ids"], encoded["attention_mask"])
                loss = prefix_reward_loss(predictions, batch_labels, encoded["attention_mask"])
            loss.backward()
            optimizer.step()
            scheduler.step()

            step += 1
            if step % spec.log_every == 0:
                logger.info("prefix-reward training epoch %d step %d/%d loss %.4f",
                            epoch, step, total_steps, loss.item())

    model.eval()
    output_dir.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(str(output_dir))
    tokenizer.save_pretrained(str(output_dir))
    logger.info("saved prefix-reward model to %s", output_dir)
    return output_dir

reward_training

Train a prefix-scored reward model, the recipe RAD's cached path needs (Deng and Raffel, 2023, §2.1).

RAD scores the reward model at every prefix during decoding, so the reward model must be trained to predict the sequence-level attribute from any prefix, not only the complete text. This module fits a scalar-head sequence classifier on a causal-LM backbone with the paper's cumulative prefix loss: the squared error between the per-prefix prediction and the sequence label, weighted by prefix length and averaged so a full-length sequence and a short one contribute comparably.

The label is the dataset's native attribute (for detoxification the civil_comments toxicity score in [0, 1]), so the trained head predicts toxicity; RAD's invert=True turns the prediction into a reward that steers away from it. Training is a plain PyTorch loop under bf16 autocast, with no Trainer or TRL reward trainer (those implement pairwise preference losses, not this regression).

HEAD_INIT_SCALE = 0.01 module-attribute
logger = logging.getLogger(__name__) module-attribute
PrefixRewardTrainSpec dataclass

Training settings for a prefix-scored reward model.

Attributes:

Name Type Description
max_length int

Maximum token length per training example (right-padded).
batch_size int

Examples per optimizer step.
epochs int

Passes over the data.
learning_rate float

AdamW learning rate.
weight_decay float

AdamW weight decay.
seed int

Seed for shuffling and initialization of the fresh scalar head.
log_every int

Log the running loss every this many optimizer steps.
Source code in steerability/algorithms/output_control/rad/utils/reward_training.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
@dataclass
class PrefixRewardTrainSpec:
    """Training settings for a prefix-scored reward model.

    Attributes:
        max_length: Maximum token length per training example (right-padded).
        batch_size: Examples per optimizer step.
        epochs: Passes over the data.
        learning_rate: AdamW learning rate.
        weight_decay: AdamW weight decay.
        seed: Seed for shuffling and initialization of the fresh scalar head.
        log_every: Log the running loss every this many optimizer steps.
    """

    max_length: int = 64
    batch_size: int = 64
    epochs: int = 1
    learning_rate: float = 1e-5
    weight_decay: float = 0.01
    seed: int = 0
    log_every: int = 50
batch_size = 64 class-attribute instance-attribute
epochs = 1 class-attribute instance-attribute
learning_rate = 1e-05 class-attribute instance-attribute
log_every = 50 class-attribute instance-attribute
max_length = 64 class-attribute instance-attribute
seed = 0 class-attribute instance-attribute
weight_decay = 0.01 class-attribute instance-attribute
prefix_reward_loss(predictions, labels, attention_mask)

The cumulative squared-error loss over prefixes, averaged over the batch.

For one sequence of length l with prediction r_t at every prefix t (1-indexed) and label r, the loss is sum_t t * (r_t - r)^2 / S_l with S_l = l * (l + 1) / 2. Positions with attention_mask == 0 are excluded and l counts the included positions, so a fully masked row contributes nothing. The loss is computed in float32 regardless of the input dtypes.

Parameters:

Name Type Description Default
predictions Tensor

Per-position predictions [B, T].

 required
labels Tensor

Per-sequence labels [B].

 required
attention_mask Tensor

[B, T] mask; 0 positions are excluded.

 required

Returns:

Type Description
Tensor

A float32 scalar tensor, the mean per-sequence prefix loss over the non-empty rows.
Source code in steerability/algorithms/output_control/rad/utils/reward_training.py 
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
def prefix_reward_loss(
    predictions: torch.Tensor,
    labels: torch.Tensor,
    attention_mask: torch.Tensor,
) -> torch.Tensor:
    """The cumulative squared-error loss over prefixes, averaged over the batch.

    For one sequence of length `l` with prediction `r_t` at every prefix `t` (1-indexed) and label
    `r`, the loss is `sum_t t * (r_t - r)^2 / S_l` with `S_l = l * (l + 1) / 2`. Positions with
    `attention_mask == 0` are excluded and `l` counts the included positions, so a fully masked row
    contributes nothing. The loss is computed in float32 regardless of the input dtypes.

    Args:
        predictions: Per-position predictions `[B, T]`.
        labels: Per-sequence labels `[B]`.
        attention_mask: `[B, T]` mask; `0` positions are excluded.

    Returns:
        A float32 scalar tensor, the mean per-sequence prefix loss over the non-empty rows.
    """
    predictions = predictions.float()
    mask = attention_mask.to(torch.float32)  # [B, T]
    positions = torch.cumsum(mask, dim=1) * mask  # 1-indexed prefix length at each kept position
    lengths = mask.sum(dim=1)  # [B]
    normalizer = lengths * (lengths + 1) / 2  # S_l per row

    squared_error = (predictions - labels.unsqueeze(1).float()) ** 2  # [B, T]
    per_row = (positions * squared_error).sum(dim=1) / normalizer.clamp_min(1.0)  # [B]

    nonempty = lengths > 0
    if not torch.any(nonempty):
        return per_row.sum() * 0.0
    return per_row[nonempty].mean()
prefix_rewards(model, input_ids, attention_mask)

Per-position rewards [B, T] from a sequence-classification model with a scalar head.

Runs the backbone (getattr(model, model.base_model_prefix)) once, applies model.score at every position rather than only the pooled last token, and returns the sigmoid. Training and the prefix-tracking diagnostics use this instead of the classifier's own forward, which pools the last token only.

Parameters:

Name Type Description Default
model

An AutoModelForSequenceClassification with a single-logit score head.

 required
input_ids Tensor

Token ids [B, T].

 required
attention_mask Tensor

[B, T] mask forwarded to the backbone.

 required

Returns:

Type Description
Tensor

Per-position rewards [B, T] in [0, 1].
Source code in steerability/algorithms/output_control/rad/utils/reward_training.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def prefix_rewards(
    model,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
) -> torch.Tensor:
    """Per-position rewards `[B, T]` from a sequence-classification model with a scalar head.

    Runs the backbone (`getattr(model, model.base_model_prefix)`) once, applies `model.score` at every
    position rather than only the pooled last token, and returns the sigmoid. Training and the
    prefix-tracking diagnostics use this instead of the classifier's own `forward`, which pools the
    last token only.

    Args:
        model: An `AutoModelForSequenceClassification` with a single-logit `score` head.
        input_ids: Token ids `[B, T]`.
        attention_mask: `[B, T]` mask forwarded to the backbone.

    Returns:
        Per-position rewards `[B, T]` in `[0, 1]`.
    """
    backbone = getattr(model, model.base_model_prefix)
    hidden = backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state  # [B, T, H]
    logits = model.score(hidden).squeeze(-1)  # [B, T]
    return torch.sigmoid(logits)
train_prefix_reward_model(backbone_name_or_path, texts, labels, output_dir, *, spec=None, device=None)

Fine-tune a scalar-head sequence classifier on texts with prefix_reward_loss and save it.

Builds the classifier from the causal-LM backbone with num_labels=1 (the head class is selected by sequence_classifier_class, so a Granite backbone resolves), scales the fresh scalar head down by HEAD_INIT_SCALE so its initial sigmoid output is unsaturated, right-pads with the backbone tokenizer (pad falls back to eos), trains with AdamW under bf16 autocast and a linear decay schedule, and writes the model and tokenizer to output_dir. The saved config.json records num_labels: 1 and the head's architecture name, so load_sequence_classifier(output_dir) reloads it.

Labels are the dataset's native attribute (for civil_comments, toxicity in [0, 1]); the head predicts the attribute, and RAD's invert=True turns the prediction into a reward.

Parameters:

Name Type Description Default
backbone_name_or_path str

HF id or local path of the causal-LM backbone.

 required
texts Sequence[str]

Training texts.

 required
labels Sequence[float]

Per-text labels in [0, 1], aligned with texts.

 required
output_dir str | Path

Directory the trained model and tokenizer are written to.

 required
spec PrefixRewardTrainSpec | None

Training settings; defaults to PrefixRewardTrainSpec().

 None
device str | device | None

Device to train on; defaults to CUDA when available, else CPU.

 None

Returns:

Type Description
Path

output_dir as a Path.
Source code in steerability/algorithms/output_control/rad/utils/reward_training.py 
119
120
121
122
123
124
125
126
127
128
129
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
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
def train_prefix_reward_model(
    backbone_name_or_path: str,
    texts: Sequence[str],
    labels: Sequence[float],
    output_dir: str | Path,
    *,
    spec: PrefixRewardTrainSpec | None = None,
    device: str | torch.device | None = None,
) -> Path:
    """Fine-tune a scalar-head sequence classifier on `texts` with `prefix_reward_loss` and save it.

    Builds the classifier from the causal-LM backbone with `num_labels=1` (the head class is selected
    by `sequence_classifier_class`, so a Granite backbone resolves), scales the fresh scalar head down
    by `HEAD_INIT_SCALE` so its initial sigmoid output is unsaturated, right-pads with the backbone
    tokenizer (pad falls back to eos), trains with AdamW under bf16 autocast and a linear decay
    schedule, and writes the model and tokenizer to `output_dir`. The saved `config.json` records
    `num_labels: 1` and the head's architecture name, so `load_sequence_classifier(output_dir)`
    reloads it.

    Labels are the dataset's native attribute (for `civil_comments`, toxicity in `[0, 1]`); the head
    predicts the attribute, and RAD's `invert=True` turns the prediction into a reward.

    Args:
        backbone_name_or_path: HF id or local path of the causal-LM backbone.
        texts: Training texts.
        labels: Per-text labels in `[0, 1]`, aligned with `texts`.
        output_dir: Directory the trained model and tokenizer are written to.
        spec: Training settings; defaults to `PrefixRewardTrainSpec()`.
        device: Device to train on; defaults to CUDA when available, else CPU.

    Returns:
        `output_dir` as a `Path`.
    """
    spec = spec or PrefixRewardTrainSpec()
    output_dir = Path(output_dir)
    device = torch.device(device) if device is not None else torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.manual_seed(spec.seed)

    tokenizer = AutoTokenizer.from_pretrained(backbone_name_or_path)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    kwargs = {"num_labels": 1, "dtype": torch.float32}
    model = sequence_classifier_class(backbone_name_or_path, **kwargs).from_pretrained(backbone_name_or_path, **kwargs)
    model.config.pad_token_id = tokenizer.pad_token_id
    with torch.no_grad():
        model.score.weight.mul_(HEAD_INIT_SCALE)
    model = model.to(device)
    model.train()

    texts = list(texts)
    labels = torch.tensor(list(labels), dtype=torch.float32)
    num_examples = len(texts)
    steps_per_epoch = (num_examples + spec.batch_size - 1) // spec.batch_size
    total_steps = steps_per_epoch * spec.epochs

    optimizer = AdamW(model.parameters(), lr=spec.learning_rate, weight_decay=spec.weight_decay)
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=max(total_steps, 1))
    use_bf16 = device.type == "cuda"

    generator = torch.Generator().manual_seed(spec.seed)
    step = 0
    for epoch in range(spec.epochs):
        order = torch.randperm(num_examples, generator=generator)
        for start in range(0, num_examples, spec.batch_size):
            batch_index = order[start:start + spec.batch_size]
            batch_texts = [texts[i] for i in batch_index.tolist()]
            batch_labels = labels[batch_index].to(device)
            encoded = tokenizer(
                batch_texts,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=spec.max_length,
            ).to(device)

            optimizer.zero_grad()
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=use_bf16):
                predictions = prefix_rewards(model, encoded["input_ids"], encoded["attention_mask"])
                loss = prefix_reward_loss(predictions, batch_labels, encoded["attention_mask"])
            loss.backward()
            optimizer.step()
            scheduler.step()

            step += 1
            if step % spec.log_every == 0:
                logger.info("prefix-reward training epoch %d step %d/%d loss %.4f",
                            epoch, step, total_steps, loss.item())

    model.eval()
    output_dir.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(str(output_dir))
    tokenizer.save_pretrained(str(output_dir))
    logger.info("saved prefix-reward model to %s", output_dir)
    return output_dir