API reference

Legacy (1.0 API)

Will code written against reward-lens 1.0 still run? Yes. The 1.0 classes are preserved as a compatibility layer and resolved lazily, so importing the package costs nothing until you reach for one. They are reachable two ways: straight from the top level, from reward_lens import RewardModel, RewardLens, ComponentAttribution, and under the explicit reward_lens.legacy namespace. A few 1.0 tools were not folded into that namespace and live only in their own submodules: from reward_lens.hacking import HackingDetector, from reward_lens.comparison import ModelComparator, and from reward_lens.sae import TopKSAE, SAETrainer. The path from these to the 2.0 kernel is the migration guide.

The submodule-only tools

HackingDetector runs the 1.0 battery of bias tests, reporting a standardized effect size per axis for length, sycophancy, confidence, and the rest.

# class

HackingDetector

Source

reward_lens.hacking.HackingDetector 

HackingDetector(model: RewardModel)

Automated reward hacking vulnerability scanner.

Generates controlled stimuli and measures whether the reward model exhibits biases along known failure dimensions.

Parameters

Name Type Default Description
model RewardModel required A RewardModel instance.

Methods

scan(
    tests: Optional[list[str]] = None,
    prompt: Optional[str] = None,
    response: Optional[str] = None,
    max_length: int = 2048
) -> HackingReport
src
Run a full vulnerability scan over the built-in probe suite.

Each requested dimension is evaluated with its own set of built-in (neutral, biased) pairs and summarized as an effect size with a permutation p-value. To test a bespoke pair, use test_custom_pair.

test_custom_pair(
    prompt: str,
    neutral: str,
    biased: str,
    dimension: str = 'custom',
    max_length: int = 2048
) -> BiasTestResult
src
Test a custom pair for bias.

ModelComparator is the 1.0 cross-model comparison. It predates the frame machinery that 2.0 uses to make such comparisons gauge-safe, so read its numbers as raw coordinates.

# class

ModelComparator

Source

reward_lens.comparison.ModelComparator 

ModelComparator(models: dict[str, RewardModel])

Compare preference computation across multiple reward models.

Parameters

Name Type Default Description
models dict[str, RewardModel] required Dict mapping model names to RewardModel instances.

Methods

compare(
    prompt: str,
    preferred: str,
    dispreferred: str,
    max_length: int = 2048
) -> ComparisonResult
src
Run reward lens and attribution on all models for the same pair.

TopKSAE is the top-k sparse autoencoder, with SAETrainer alongside it; both need the [sae] extra installed.

# class

TopKSAE

Source

reward_lens.sae.TopKSAE

Bases: nn.Module

TopKSAE(d_model: int, n_features: int, k: int = 32)

TopK enforces exactly k features active per input, giving clean sparsity control without the L1 penalty tuning required by ReLU SAEs.

Parameters

Name Type Default Description
d_model int required Input dimension (model hidden size).
n_features int required Dictionary size (number of SAE features).
k int 32 Number of active features per input.

Methods

encode(x: torch.Tensor) -> torch.Tensor
src
Encode input to sparse feature activations.
decode(f: torch.Tensor) -> torch.Tensor
src
Decode sparse features back to input space.
forward(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
src
Full forward pass: encode then decode.
feature_reward_alignments(reward_direction: torch.Tensor) -> torch.Tensor
src
Compute the reward alignment of each feature.

This is the quantity (w_r^T @ d_i) for each feature i, where d_i is the i-th decoder column (feature direction in activation space).

decompose_reward(
    x: torch.Tensor,
    reward_direction: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
src
Decompose the reward score into per-feature contributions.
save(path: str) -> None
src
Save the SAE to disk.
load(path: str, device: str = 'cpu') -> TopKSAE
src
Load a saved SAE.

The lazy top-level names

These are the 1.0 symbols in the package’s _LAZY map. Each is reachable from the top level (from reward_lens import RewardModel) and under the explicit reward_lens.legacy namespace; both resolve to the defining module named below on first access.

# class

RewardModel

Source

reward_lens.model.RewardModel · also reward_lens.RewardModel 

RewardModel(
    model: nn.Module,
    tokenizer: AutoTokenizer,
    adapter: ModelAdapter,
    device: torch.device
)

Wrapper around a HuggingFace reward model for mechanistic interpretability.

Parameters

Name Type Default Description
model nn.Module required The HuggingFace model.
tokenizer AutoTokenizer required The tokenizer.
adapter ModelAdapter required Model-specific adapter for architecture navigation.
device torch.device required Device to run on.

Properties

  • reward_direction : torch.Tensor : The reward head weight vector — the direction in activation space that defines the reward. Shape: (d_model,).
  • reward_bias : float : The reward head bias term.
  • d_model : int : Hidden dimension of the model.
  • n_layers : int : Number of transformer layers.
  • n_heads : int : Number of attention heads per layer.
  • d_head : int : Dimension per attention head.

Methods

from_pretrained(
    model_name_or_path: str,
    device: Optional[str] = None,
    torch_dtype: torch.dtype = torch.bfloat16,
    trust_remote_code: bool = True,
    attn_implementation: Optional[str] = None,
    **kwargs: Any = {}
) -> RewardModel
src
Load a reward model from HuggingFace.

This is the main entry point. It auto-detects the model architecture and selects the appropriate adapter.

tokenize_conversation(
    prompt: str,
    response: str,
    max_length: int = 2048
) -> dict[str, torch.Tensor]
src
Tokenize a (prompt, response) pair using the model’s chat template.
tokenize_raw(text: str, max_length: int = 2048) -> dict[str, torch.Tensor]
src
Tokenize raw text (already formatted).
score(prompt: str, response: str, max_length: int = 2048) -> float
src
Compute the scalar reward for a (prompt, response) pair.
score_pair(
    prompt: str,
    preferred: str,
    dispreferred: str,
    max_length: int = 2048
) -> tuple[float, float]
src
Score both completions in a preference pair.
forward_with_cache(
    prompt: str,
    response: str,
    cache_full_sequences: bool = False,
    max_length: int = 2048
) -> tuple[float, ActivationCache]
src

Run a forward pass and cache all intermediate activations.

This is the workhorse function for all interpretability analyses. It registers hooks on every transformer layer to capture:

  • Residual stream states after each layer
  • Attention sublayer outputs
  • MLP sublayer outputs
forward_with_cache_from_inputs(
    inputs: dict[str, torch.Tensor],
    cache_full_sequences: bool = False
) -> tuple[float, ActivationCache]
src
Run forward pass with caching from pre-tokenized inputs.
tokenize_conversation_batch(
    pairs: list[tuple[str, str]],
    max_length: int = 2048
) -> dict[str, torch.Tensor]
src
Tokenize a batch of (prompt, response) pairs with left-padding.

Left-padding means the final (response-end) token sits at the same position for every sequence in the batch — namely position T-1 where T is the longest sequence. This makes the final-token gather cheap and uniform: we pull index T-1 for every batch row.

forward_with_cache_batch(
    pairs: list[tuple[str, str]],
    batch_size: int = 32,
    max_length: int = 2048,
    capture_heads: bool = False,
    progress: bool = False,
    length_bucket: bool = False
) -> BatchedActivationCache
src
Run a batched forward-with-cache over many (prompt, response) pairs.

This is the population-scale primitive. The single-pair :meth:forward_with_cache runs one tokenize + one forward per pair; for the v2 experiments (n>=150 pairs/dim/model) that’s an order of magnitude too slow. This batches batch_size pairs per forward and concatenates results.

project_onto_reward(hidden_state: torch.Tensor) -> torch.Tensor
src
Project a hidden state onto the reward direction.

This is the core “reward lens” operation: given a hidden state h, compute w_r^T @ h + b_r.

hooks(hook_fns: dict[str, callable])
src
Context manager for temporarily registering hooks.

The hooks are active inside the with block and removed on exit.

# class

ActivationCache

Source

reward_lens.model.ActivationCache · also reward_lens.ActivationCache

@dataclass

Stores activations captured during a forward pass.

Attributes

Name Type Default Description
residual_streams dict[int, torch.Tensor] field(default_factory=dict) Dict mapping layer index -> residual stream tensor at final token. Shape of each tensor: (batch, d_model).
attn_outputs dict[int, torch.Tensor] field(default_factory=dict) Dict mapping layer index -> attention sublayer output at final token. Shape: (batch, d_model).
mlp_outputs dict[int, torch.Tensor] field(default_factory=dict) Dict mapping layer index -> MLP sublayer output at final token. Shape: (batch, d_model).
final_token_positions Optional[torch.Tensor] None The token positions used for reward computation. Shape: (batch,).
raw_residual_streams dict[int, torch.Tensor] field(default_factory=dict) Full sequence residual streams (optional, for patching). Dict mapping layer index -> (batch, seq_len, d_model).
raw_attn_outputs dict[int, torch.Tensor] field(default_factory=dict) Full sequence attention outputs (optional).
raw_mlp_outputs dict[int, torch.Tensor] field(default_factory=dict) Full sequence MLP outputs (optional).
attn_head_outputs dict[int, torch.Tensor] field(default_factory=dict) Optional per-head attention outputs at final token. Dict mapping layer index -> (batch, n_heads, d_head). Populated only when capture_heads=True. Available for adapters that implement get_attn_o_proj.
# class

BatchedActivationCache

Source

reward_lens.model.BatchedActivationCache · also reward_lens.BatchedActivationCache

@dataclass

Activation cache for a batch of (prompt, response) pairs.

Identical layout to :class:ActivationCache but every tensor carries an explicit batch dimension as its leading axis. The single-pair cache is really a degenerate batched cache with B=1; we keep both classes distinct so downstream code can be explicit about whether it expects a batched tensor (and handle it vectorised) or a single-pair tensor.

Use :meth:slice to project to a single-pair :class:ActivationCache.

Attributes

Name Type Default
residual_streams dict[int, torch.Tensor] field(default_factory=dict)
attn_outputs dict[int, torch.Tensor] field(default_factory=dict)
mlp_outputs dict[int, torch.Tensor] field(default_factory=dict)
attn_head_outputs dict[int, torch.Tensor] field(default_factory=dict)
final_token_positions Optional[torch.Tensor] None
rewards Optional[torch.Tensor] None

Properties

  • batch_size : int

Methods

slice(i: int) -> ActivationCache
src
Return a single-pair ActivationCache view at batch index i.
# class

RewardLens

Source

reward_lens.lens.RewardLens · also reward_lens.RewardLens 

RewardLens(model: RewardModel)

Compute reward lens analysis for preference pairs.

The reward lens projects the residual stream at each layer onto the reward direction to trace how preference forms across the depth of the model.

Parameters

Name Type Default Description
model RewardModel required A RewardModel instance.

Methods

trace(
    prompt: str,
    preferred: str,
    dispreferred: str,
    max_length: int = 2048
) -> RewardLensResult
src
Run reward lens analysis on a preference pair.

Runs forward passes on both completions, caches all intermediate residual stream states, and projects each onto the reward direction.

trace_single(
    prompt: str,
    response: str,
    max_length: int = 2048
) -> tuple[np.ndarray, np.ndarray]
src
Run reward lens on a single completion (not a pair).
# func

reward_lens_plot

Source

reward_lens.lens.reward_lens_plot · also reward_lens.reward_lens_plot 

reward_lens_plot(
    model: RewardModel,
    prompt: str,
    preferred: str,
    dispreferred: str,
    save_path: Optional[str] = None,
    max_length: int = 2048,
    title: Optional[str] = None
) -> RewardLensResult

Convenience function: run reward lens and plot in one call.

Parameters

Name Type Default Description
model RewardModel required A RewardModel instance.
prompt str required The user prompt.
preferred str required The preferred completion.
dispreferred str required The dispreferred completion.
save_path Optional[str] None Optional path to save the plot.
max_length int 2048 Maximum sequence length.
title Optional[str] None Optional custom title.

Returns

Type Description
RewardLensResult The RewardLensResult.
# class

ComponentAttribution

Source

reward_lens.attribution.ComponentAttribution · also reward_lens.ComponentAttribution 

ComponentAttribution(model: RewardModel)

Decompose reward scores into per-component signed contributions.

This is the observational decomposition — it tells you which components contribute most to the reward, but not whether they are causally necessary. For causal analysis, use ActivationPatcher.

Parameters

Name Type Default Description
model RewardModel required A RewardModel instance.

Methods

attribute(
    prompt: str,
    preferred: str,
    dispreferred: str,
    max_length: int = 2048
) -> ComponentResult
src
Compute per-component reward attribution for a preference pair.

For each component c (embedding, each attention sublayer, each MLP sublayer), computes: contribution_c = w_r^T @ output_c

where output_c is the component’s contribution to the residual stream at the final token position.

attribute_single(prompt: str, response: str, max_length: int = 2048) -> ComponentResult
src
Compute per-component attribution for a single completion.
attribute_heads(
    prompt: str,
    preferred: str,
    dispreferred: str,
    max_length: int = 2048
) -> ComponentResult
src
Per-head attention attribution (head granularity).

Decomposes each layer’s attention contribution into per-head terms by capturing the input to o_proj and projecting each head’s slice through its own o_proj weight slice. Yields n_layers * n_heads head components plus the per-layer MLPs.

# class

ActivationPatcher

Source

reward_lens.patching.ActivationPatcher · also reward_lens.ActivationPatcher 

ActivationPatcher(model: RewardModel)

Causal intervention via activation patching on reward models.

For a preference pair, this tool identifies which components are causally necessary for the reward model’s preference, by swapping component activations between the preferred and dispreferred completions.

Parameters

Name Type Default Description
model RewardModel required A RewardModel instance.

Methods

patch_all_components(
    prompt: str,
    preferred: str,
    dispreferred: str,
    mode: Literal['noising', 'denoising', 'zero', 'mean'] = 'noising',
    max_length: int = 2048,
    show_progress: bool = True,
    mean_corpus: Optional[list[tuple[str, str]]] = None
) -> PatchingResult
src
Patch every attention and MLP component and measure the effect.
patch_single_component(
    prompt: str,
    preferred: str,
    dispreferred: str,
    layer_idx: int,
    component_type: Literal['attn', 'mlp'],
    mode: Literal['noising', 'denoising', 'zero', 'mean'] = 'noising',
    max_length: int = 2048,
    mean_corpus: Optional[list[tuple[str, str]]] = None
) -> float
src
Patch a single specific component and return the effect.

Useful for targeted investigation of specific layers.

patch_all_heads(
    prompt: str,
    preferred: str,
    dispreferred: str,
    mode: Literal['noising', 'denoising'] = 'noising',
    max_length: int = 2048,
    show_progress: bool = True
) -> PatchingResult
src
Patch every attention head individually (head-level granularity).

Implementation: for each layer × head, install a forward-pre-hook on o_proj that replaces the slice corresponding to head h with the source-side per-head input (or zero), leaving every other head untouched. This isolates the causal effect of one head at a time.

patch_all_components_mean(
    prompt: str,
    preferred: str,
    dispreferred: str,
    corpus_pairs: list[tuple[str, str]],
    max_length: int = 2048,
    show_progress: bool = True
) -> PatchingResult
src
Mean-ablation patching.

Replaces each component’s output with the mean of its activations over a user-supplied corpus. Differs from zero ablation when the component has a non-trivial DC offset (which is typical for MLPs post-layernorm). When in doubt, prefer mean over zero.

# class

PathPatcher

Source

reward_lens.path_patching.PathPatcher · also reward_lens.PathPatcher 

PathPatcher(model: RewardModel)

Two-hop head-level path patching.

Parameters

Name Type Default Description
model RewardModel required RewardModel.

Methods

patch(
    prompt: str,
    preferred: str,
    dispreferred: str,
    sender: ComponentSpec,
    receiver: ComponentSpec,
    mode: Literal['noising', 'denoising'] = 'noising',
    max_length: int = 2048
) -> PathPatchResult
src
Run a 2-hop path patch.

Algorithm (noising mode):

  1. Run the source forward (dispreferred) and cache the sender’s head output and the receiver’s input/output.
  2. Run the target forward (preferred) twice: a. Clean: get the original differential. b. Patched: install a forward-pre-hook on the receiver that replaces only the contribution flowing from the sender’s head with the source-side value, leaving every other path untouched.

For head-level senders we splice the head’s output into the receiver’s input via the residual stream: the difference (source_head_out - target_head_out) is added/subtracted at the receiver’s input position.

Returns a PathPatchResult.

# class

PathPatchResult

Source

reward_lens.path_patching.PathPatchResult · also reward_lens.PathPatchResult

@dataclass

No docstring. See the source for detail.

Attributes

Name Type Default
sender ComponentSpec
receiver ComponentSpec
mode str
original_differential float
patched_differential float
path_effect float
# class

DistortionAnalyzer

Source

reward_lens.distortion.DistortionAnalyzer · also reward_lens.DistortionAnalyzer 

DistortionAnalyzer(model: RewardModel)

Analyze reward model for predicted hacking vulnerabilities.

This implements the distortion index framework from “Reward Hacking as Equilibrium under Finite Evaluation.” It predicts which quality dimensions are under-covered by current evaluation and thus likely to be hacked.

Parameters

Name Type Default Description
model RewardModel required A RewardModel instance.

Methods

compute_distortion_index(
    quality_dimensions: list[str],
    evaluation_probes: dict[str, list[PreferencePair]],
    cross_dimension_probes: Optional[list[tuple[PreferencePair, list[str]]]] = None,
    coverage_threshold: float = 0.3,
    distortion_threshold: float = 0.5,
    max_length: int = 2048
) -> DistortionReport
src
Compute distortion index predicting hacking severity per dimension.
analyze_agentic_amplification(
    base_dimensions: list[str],
    tool_count: int,
    base_probes: dict[str, list[PreferencePair]],
    max_length: int = 2048
) -> DistortionReport
src
Analyze how distortion amplifies in agentic settings.

As agents gain access to more tools, quality dimensions scale combinatorially while evaluation typically scales linearly. This method estimates the amplified distortion.

compare_evaluation_strategies(
    quality_dimensions: list[str],
    strategies: dict[str, dict[str, list[PreferencePair]]],
    max_length: int = 2048
) -> dict[str, DistortionReport]
src
Compare distortion across different evaluation strategies.

Useful for deciding which evaluation probes to invest in.

# class

DistortionReport

Source

reward_lens.distortion.DistortionReport · also reward_lens.DistortionReport

@dataclass

Result of distortion index analysis.

Attributes

Name Type Default Description
quality_dimensions list[str] Names of the quality dimensions analyzed.
per_dimension_distortion dict[str, float] Distortion index per dimension (higher = more hackable).
under_covered_dimensions list[str] Dimensions with distortion above threshold.
predicted_hacking_severity float Overall predicted severity (0-1).
coverage_matrix np.ndarray Matrix showing which probes cover which dimensions.
effective_coverage dict[str, float] Effective coverage per dimension (0-1).
recommendations list[str] field(default_factory=list) Suggested additional probes to reduce distortion.

Methods

print_summary() -> None
src
Print a formatted distortion analysis summary.
plot(save_path: Optional[str] = None, figsize: tuple[int, int] = (12, 5)) -> None
src
Plot distortion analysis visualization.
# class

DivergenceAwarePatching

Source

reward_lens.divergence_patching.DivergenceAwarePatching · also reward_lens.DivergenceAwarePatching

Bases: ActivationPatcher

DivergenceAwarePatching(
    model: RewardModel,
    distribution_estimator: Optional[DistributionEstimator] = None
)

Activation patcher with divergence detection and constrained patching.

Extends the base ActivationPatcher to:

  1. Detect when interventions create OOD representations
  2. Classify divergences as harmless vs pernicious
  3. Provide reliability scores for causal claims

Parameters

Name Type Default Description
model RewardModel required A RewardModel instance.
distribution_estimator Optional[DistributionEstimator] None Optional pre-computed distribution estimator.

Methods

fit_distribution(
    prompts: list[str],
    responses: list[str],
    max_length: int = 2048,
    show_progress: bool = True
) -> None
src
Fit the activation distribution from clean data.

Call this before patch_with_divergence_check for best results.

patch_with_divergence_check(
    prompt: str,
    preferred: str,
    dispreferred: str,
    mode: Literal['noising', 'denoising', 'zero', 'mean'] = 'noising',
    divergence_threshold: float = 2.0,
    max_length: int = 2048,
    show_progress: bool = True,
    mean_corpus: Optional[list[tuple[str, str]]] = None
) -> DivergenceAwarePatchingResult
src
Patch all components and check for divergent representations.
constrained_patch(
    prompt: str,
    preferred: str,
    dispreferred: str,
    mode: Literal['noising', 'denoising'] = 'noising',
    cl_weight: float = 0.1,
    n_optimization_steps: int = 50,
    max_length: int = 2048
) -> torch.Tensor
src
Find constrained patch that stays close to training distribution.

Uses a Counterfactual Latent (CL) loss to find a patched activation that achieves the intervention goal while minimizing divergence.

This is more expensive than standard patching but provides more reliable causal evidence.

# class

DivergenceAwarePatchingResult

Source

reward_lens.divergence_patching.DivergenceAwarePatchingResult · also reward_lens.DivergenceAwarePatchingResult

@dataclass

Bases: PatchingResult

Extended patching result with divergence analysis.

Carries every field of :class:PatchingResult and adds the divergence fields below.

Attributes

Name Type Default Description
divergence_info list[DivergenceInfo] field(default_factory=list) Per-component divergence analysis.
has_pernicious_divergence bool False Whether any intervention has pernicious divergence.
reliability_score float 1.0 Overall reliability of causal claims (0-1).
divergent_components list[str] field(default_factory=list) List of components with significant divergence.

Methods

print_divergence_summary() -> None
src
Print a summary of divergence analysis.
plot_with_divergence(
    save_path: Optional[str] = None,
    figsize: tuple[int, int] = (14, 8),
    title: Optional[str] = None
) -> None
src
Plot patching results with divergence overlay.
# class

MisalignmentCascadeDetector

Source

reward_lens.cascade.MisalignmentCascadeDetector · also reward_lens.MisalignmentCascadeDetector 

MisalignmentCascadeDetector(model: RewardModel)

Detect correlated misalignment patterns across multiple dimensions.

This goes beyond HackingDetector’s individual bias tests to find systemic vulnerabilities where failures in one dimension predict failures in others.

Parameters

Name Type Default Description
model RewardModel required A RewardModel instance.

Methods

detect_cascade(
    dimensions: Optional[list[str]] = None,
    custom_tests: Optional[dict[str, list[dict]]] = None,
    correlation_threshold: float = 0.5,
    n_bootstrap: int = 100,
    max_length: int = 2048
) -> CascadeReport
src
Run comprehensive cascade detection across misalignment dimensions.
cross_validate_with_hacking(
    hacking_report: HackingReport,
    cascade_report: Optional[CascadeReport] = None,
    max_length: int = 2048
) -> dict[str, float]
src
Correlate hacking biases with misalignment dimensions.

Tests the paper’s key finding: reward hacking onset correlates with emergent misalignment.

# class

CascadeReport

Source

reward_lens.cascade.CascadeReport · also reward_lens.CascadeReport

@dataclass

Result of misalignment cascade analysis.

Attributes

Name Type Default Description
dimensions_tested list[str] List of misalignment dimensions tested.
per_dimension_scores dict[str, float] Mean reward delta for each dimension.
correlation_matrix np.ndarray Pairwise correlation between dimensions.
cascade_risk_score float Overall systemic vulnerability (0-1).
correlated_pairs list[tuple[str, str, float]] Pairs of dimensions with significant correlation.
cascade_clusters list[list[str]] Groups of dimensions that fail together.
primary_failure_mode Optional[str] The dimension most predictive of others.
recommendations list[str] field(default_factory=list) Suggested mitigations.

Methods

print_summary() -> None
src
Print a formatted cascade analysis summary.
plot(save_path: Optional[str] = None, figsize: tuple[int, int] = (12, 10)) -> None
src
Plot correlation heatmap and dimension scores.
# class

RewardConflictAnalyzer

Source

reward_lens.conflict.RewardConflictAnalyzer · also reward_lens.RewardConflictAnalyzer 

RewardConflictAnalyzer(model: RewardModel)

Analyze conflicts between reward terms.

This helps identify when reward structures may cause models to hide reasoning (in-conflict terms) vs when optimization is safe (aligned or orthogonal terms).

Parameters

Name Type Default Description
model RewardModel required A RewardModel instance.

Methods

analyze_conflicts(
    reward_terms: dict[str, torch.Tensor],
    aligned_threshold: float = 0.5,
    orthogonal_threshold: float = 0.2,
    conflict_threshold: float = -0.3
) -> ConflictReport
src
Analyze conflicts between reward term directions.
learn_term_directions(
    term_pairs: dict[str, list[tuple[str, str, str, str]]],
    max_length: int = 2048
) -> dict[str, torch.Tensor]
src
Learn reward term directions from contrastive pairs.

For each term, provide pairs where the preferred response is better on that specific dimension. The direction is learned as the average difference in final-layer activations.

analyze_multi_objective_model(
    objective_names: Optional[list[str]] = None
) -> ConflictReport
src
Analyze conflicts in a multi-objective reward model.

For models like ArmoRM that have multiple output dimensions, this extracts the direction for each objective and analyzes conflicts.

# class

ConflictReport

Source

reward_lens.conflict.ConflictReport · also reward_lens.ConflictReport

@dataclass

Complete reward term conflict analysis.

Attributes

Name Type Default Description
term_names list[str] Names of analyzed reward terms.
term_directions dict[str, torch.Tensor] The direction vectors for each term.
pairwise_analysis list[TermPairAnalysis] Analysis for each pair of terms.
relationship_matrix np.ndarray Matrix of relationship classifications.
similarity_matrix np.ndarray Matrix of cosine similarities.
in_conflict_pairs list[tuple[str, str]] List of term pairs that are in conflict.
overall_conflict_score float 0-1 score of overall conflict severity.
monitorability_risk float Risk that models will hide reasoning.
recommendations list[str] field(default_factory=list) Overall recommendations.

Methods

print_summary() -> None
src
Print a formatted conflict analysis summary.
plot(save_path: Optional[str] = None, figsize: tuple[int, int] = (12, 5)) -> None
src
Plot similarity matrix and relationship diagram.
# func

quick_conflict_check

Source

reward_lens.conflict.quick_conflict_check · also reward_lens.quick_conflict_check 

quick_conflict_check(
    model: RewardModel,
    term_pairs: dict[str, list[tuple[str, str, str]]],
    max_length: int = 2048
) -> ConflictReport

Convenience function for quick conflict analysis.

Parameters

Name Type Default Description
model RewardModel required RewardModel instance.
term_pairs dict[str, list[tuple[str, str, str]]] required Dict mapping term name to list of (prompt, preferred, dispreferred).
max_length int 2048 Maximum sequence length.

Returns

Type Description
ConflictReport ConflictReport.
# class

ConceptExtractor

Source

reward_lens.concepts.ConceptExtractor · also reward_lens.ConceptExtractor 

ConceptExtractor(model: RewardModel)

Extract and analyze concept vectors from reward model activations.

This enables understanding which abstract concepts (confidence, formality, agreement, etc.) influence reward scores, and identifying potential reward hacking vulnerabilities.

Parameters

Name Type Default Description
model RewardModel required A RewardModel instance.

Methods

extract_concepts(
    concept_pairs: dict[str, list[tuple[str, str, str]]],
    layer: Optional[int] = None,
    max_length: int = 2048,
    show_progress: bool = True
) -> dict[str, torch.Tensor]
src
Extract concept direction vectors from contrastive pairs.

For each concept, the direction is computed as the average difference between positive and negative example activations.

analyze_reward_alignment(
    concept_vectors: dict[str, torch.Tensor],
    alignment_threshold: float = 0.3,
    hacking_concepts: Optional[list[str]] = None
) -> ConceptAlignmentReport
src
Analyze how concepts align with the reward direction.

Concepts that strongly align with reward but represent surface properties (like confidence or verbosity) may indicate hackable biases.

intervene_on_concept(
    prompt: str,
    response: str,
    concept_vector: torch.Tensor,
    strength: float = 1.0,
    layer: Optional[int] = None,
    max_length: int = 2048
) -> float
src
Intervene on a concept and measure reward change.

Adds or subtracts the concept direction from activations to see how it affects reward.

# class

ConceptAlignmentReport

Source

reward_lens.concepts.ConceptAlignmentReport · also reward_lens.ConceptAlignmentReport

@dataclass

Report on concept-reward alignment analysis.

Attributes

Name Type Default Description
concepts list[ConceptInfo] List of analyzed concepts.
reward_aligned_concepts list[str] Concepts that significantly align with reward.
anti_reward_concepts list[str] Concepts that significantly oppose reward.
high_risk_concepts list[str] Concepts that may enable reward hacking.
overall_hacking_risk float Combined hacking risk from all concepts.
recommendations list[str] field(default_factory=list) Suggested actions.

Methods

print_summary() -> None
src
Print a formatted summary of concept-reward alignment.
plot(save_path: Optional[str] = None, figsize: tuple[int, int] = (12, 6)) -> None
src
Plot concept alignment visualization.
# func

quick_concept_analysis

Source

reward_lens.concepts.quick_concept_analysis · also reward_lens.quick_concept_analysis 

quick_concept_analysis(
    model: RewardModel,
    concept_pairs: Optional[dict[str, list[tuple[str, str, str]]]] = None,
    max_length: int = 2048
) -> ConceptAlignmentReport

Convenience function for quick concept analysis.

Parameters

Name Type Default Description
model RewardModel required RewardModel instance.
concept_pairs Optional[dict[str, list[tuple[str, str, str]]]] None Optional custom concept pairs. Uses defaults if None.
max_length int 2048 Maximum sequence length.

Returns

Type Description
ConceptAlignmentReport ConceptAlignmentReport.
# class

SAETrainer

Source

reward_lens.sae.SAETrainer 

SAETrainer(
    d_model: int,
    n_features: Optional[int] = None,
    k: int = 32,
    lr: float = 0.0003,
    batch_size: int = 4096,
    device: str = 'cuda'
)

Train a TopK SAE on collected activations.

Parameters

Name Type Default Description
d_model int required Hidden dimension.
n_features Optional[int] None Dictionary size (typically 4x, 8x, or 16x d_model).
k int 32 Number of active features per input.
lr float 0.0003 Learning rate.
batch_size int 4096 Training batch size.
device str 'cuda' Device for training.

Methods

train(
    activations: torch.Tensor,
    n_epochs: int = 5,
    log_every: int = 100,
    show_progress: bool = True
) -> TopKSAE
src
Train the SAE on collected activations.