Add KL-based guidance for controlled decoding

[georgosgeorgos]

Summary

Add KL divergence-based guidance as a steering mechanism for inference-time scaling algorithms — allowing users to bias generation toward desired behaviors while staying close to the base model's distribution.

Motivation

Current scoring in its_hub relies on process reward models (PRM) or outcome reward models (ORM) to evaluate generated steps/responses. These provide a score but don't directly shape the generation distribution. KL-based guidance adds a soft constraint: steer the model toward high-reward regions while penalizing deviation from a reference distribution (typically the unguided base model), preventing reward hacking and mode collapse.

This is especially relevant for:

• Tool-call selection: guide the model toward known-useful tools without completely overriding its reasoning
• Agentic workflows: keep the agent "on track" without hard constraints that break fluency
• Particle filtering: use KL as part of the weight computation instead of (or alongside) PRM scores

Proposed Design

Core abstraction

A GuidanceModule that computes a modified score:

guided_score(x) = reward(x) - β * KL(π_guided || π_ref)

Where:

• reward(x) is the existing PRM/ORM score
• π_ref is the reference (base model) log-probability
• β controls the strength of the KL penalty
• The KL term can be estimated from log-probs returned by the LM API

Integration points

• Particle filtering: modify weight computation in _apropagate() to include KL penalty
• Best-of-N: adjust ranking scores with KL regularization
• Self-consistency: optionally filter candidates that diverge too far from reference

Configuration

guidance = KLGuidance(
    beta=0.1,                    # KL penalty strength
    reference_model=None,        # None = use same model without guidance
    estimation="token_level",    # or "sequence_level"
)
algorithm = ParticleFiltering(..., guidance=guidance)

Requirements

• LM backend must return log-probabilities (vLLM and OpenAI support logprobs)
• Reference model can be the same model (self-KL) or a separate endpoint
• β should be tunable per-step (annealing schedule) or fixed