Skip to content

Add pluggable trajectory aggregation with learned GBDT reference implementation - #216

Open
lambdabaa wants to merge 10 commits into
Red-Hat-AI-Innovation-Team:v0from
lambdabaa:feat/trajectory-aggregator
Open

Add pluggable trajectory aggregation with learned GBDT reference implementation#216
lambdabaa wants to merge 10 commits into
Red-Hat-AI-Innovation-Team:v0from
lambdabaa:feat/trajectory-aggregator

Conversation

@lambdabaa

Copy link
Copy Markdown

Summary

  • Introduces AbstractTrajectoryAggregator in base.py so aggregation is pluggable at the algorithm level, replacing hardcoded prod/min/mean reductions in ParticleFiltering and BeamSearch
  • Ships HardcodedAggregator (wraps the three existing reductions), LearnedGBDTAggregator (recommended, sklearn), and LearnedMLPAggregator (torch, kept as an extensibility example)
  • Adds MLXProcessRewardModel (Apple Silicon) and TransformersProcessRewardModel (includes RoPE buffer repair for transformers ≥5.0) to its_hub/integration/
  • 231 new tests covering aggregators, algorithm integration, and both PRM implementations

Motivation

ParticleFiltering and BeamSearch previously hardcoded prod/min/mean with no override path. None of the three is a defensible default (prod penalises long trajectories, min is brittle to a single weak step, mean discards position). This PR makes the choice explicit and extensible.

Learned aggregator evaluation

Nested 5-fold CV across 408 mixed-difficulty MATH problems (Levels 1–5, N=8 trajectories per problem):

Aggregator Mean acc 95% CI
random 0.380 [0.306, 0.454]
min 0.490 [0.419, 0.562]
prod / mean 0.498 [0.434, 0.577]
learned_mlp 0.495 [0.428, 0.562]
gbdt 0.517 [0.454, 0.581]

GBDT is the only model that consistently clears the fixed baselines across all 5 folds. Its decision surface concentrates ~78% of feature importance on min — effectively learning a continuous threshold on the worst step, which is the dominant signal at this data scale. MLP is statistically indistinguishable from prod/mean at 408 problems.

Interface

from its_hub.aggregators import HardcodedAggregator, LearnedGBDTAggregator

# Drop-in replacement — existing code unchanged (defaults to HardcodedAggregator("prod"))
pf = ParticleFiltering(sg=sg, prm=prm, aggregator=HardcodedAggregator("min"))

# Learned aggregator
agg = LearnedGBDTAggregator("its_hub/aggregators/checkpoints/gbdt_agg.pkl")
pf = ParticleFiltering(sg=sg, prm=prm, aggregator=agg)

# Custom aggregator
class MyAggregator(AbstractTrajectoryAggregator):
    def aggregate(self, step_scores: list[float]) -> float: ...

Test plan

  • tests/test_aggregators.py — unit tests for all three aggregator classes, async delegate, and ParticleFiltering integration
  • tests/test_algorithms.py — BeamSearch aggregator wiring
  • tests/test_mlx_prm.py — MLXProcessRewardModel unit tests

Add step_scores field to Path, plug AbstractTrajectoryAggregator into
BeamSearch.__init__ (default: HardcodedAggregator('prod')), accumulate
step scores per level, and use aggregator for final trajectory selection.

Signed-off-by: lambdabaa <aria@caa.columbia.edu>
ParticleFiltering now accepts and forwards an aggregator parameter to
ParticleGibbs, completing the constructor update from rt-1e8.

Add tests/test_aggregators.py covering HardcodedAggregator (prod/min/mean
math, empty input, unknown reduction), async aaggregate delegation,
LearnedMLPAggregator (dummy checkpoint load, forward pass, sigmoid bounds),
and ParticleFiltering aggregator integration.

Signed-off-by: lambdabaa <aria@caa.columbia.edu>
Implements AbstractProcessRewardModel via mlx-lm with 4-bit quantized
weights (default: Qwen/Qwen2.5-Math-PRM-7B). Scores step-by-step trajectories
by computing P(correct) at step-boundary positions using good/bad token logits.
MLX import is optional and guarded with a clear error message.

Signed-off-by: lambdabaa <aria@caa.columbia.edu>
Signed-off-by: lambdabaa <aria@caa.columbia.edu>
… unit tests

tests/test_algorithms.py: TestBeamSearchAggregatorIntegration — four tests
mirroring ParticleFiltering: accepts aggregator param, defaults to prod,
custom ZeroAggregator produces valid result, prod vs min both return dicts.

tests/test_mlx_prm.py: six mocked tests for MLXProcessRewardModel — interface
conformance, missing-mlx ImportError, score/ascore shape, batch length,
scalar vs list dispatch, and order preservation.

Signed-off-by: lambdabaa <aria@caa.columbia.edu>
Signed-off-by: lambdabaa <aria@caa.columbia.edu>
Signed-off-by: lambdabaa <aria@caa.columbia.edu>
- Export TransformersProcessRewardModel from its_hub.integration so users
  can import it alongside MLXProcessRewardModel and LocalVllmProcessRewardModel
- Update MLXProcessRewardModel docstring to explicitly document that it does
  not work with Qwen2.5-Math-PRM-7B (classifier score head rejected by mlx_lm)
  and to point to TransformersProcessRewardModel as the correct alternative
- Add Gas City beads ignore rules to .gitignore

Signed-off-by: lambdabaa <aria@caa.columbia.edu>
…nsformers 5.x meta-init

Transformers >=5.0 uses meta-tensor initialisation during from_pretrained.
Non-persistent buffers (inv_freq, cos_cached, sin_cached) on Qwen2RotaryEmbedding
are materialised as zeros instead of being computed from the RoPE formula.
This causes all attention Q/K to be NaN, collapsing every PRM score to the
constant 0.50003338 (pure bias term of the score head).

Fix: _repair_rotary_embeddings() iterates all RotaryEmbedding modules after
model load and recomputes inv_freq + cos/sin caches from the stored base/dim.
Also: add trailing <extra_0> separator so each step gets its own position token
(previously N-1 separators for N steps caused the last step's score to be
duplicated from the second-to-last position).

Includes: seed=42 width=16 MLP checkpoint trained on 200 MATH-Hard problems.
Signed-off-by: lambdabaa <aria@caa.columbia.edu>
GBDT is the only learned aggregator that consistently clears the fixed
baselines in nested 5-fold CV (0.517 vs 0.498 for prod/mean), selected
for min-threshold behaviour that matches the dominant signal at this
corpus scale.  MLP is retained as a code example of how to extend the
interface with a torch model.

Changes:
- Add LearnedGBDTAggregator to aggregators/learned.py (requires sklearn,
  handles ImportError gracefully like the MLP's torch dependency)
- Export LearnedGBDTAggregator from aggregators/__init__.py
- Add gbdt_agg.pkl checkpoint (GradientBoostingClassifier, max_depth=2,
  n_estimators=200, trained on 408 mixed-difficulty problems)
- Update mlp_agg.pt to checkpoint retrained on combined L1-5 corpus
- Add TestLearnedGBDTAggregator suite mirroring MLP test coverage
- Gitignore .claude/ directory

Signed-off-by: Ari Aye <ari.aye@gatesfoundation.org>
Signed-off-by: lambdabaa <aria@caa.columbia.edu>
@beatsmonster

Copy link
Copy Markdown
Contributor

@lambdabaa could we please rebase to main/v1

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants