Add pluggable trajectory aggregation with learned GBDT reference implementation - #216
Open
lambdabaa wants to merge 10 commits into
Open
Add pluggable trajectory aggregation with learned GBDT reference implementation#216lambdabaa wants to merge 10 commits into
lambdabaa wants to merge 10 commits into
Conversation
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>
Contributor
|
@lambdabaa could we please rebase to main/v1 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
AbstractTrajectoryAggregatorinbase.pyso aggregation is pluggable at the algorithm level, replacing hardcodedprod/min/meanreductions inParticleFilteringandBeamSearchHardcodedAggregator(wraps the three existing reductions),LearnedGBDTAggregator(recommended, sklearn), andLearnedMLPAggregator(torch, kept as an extensibility example)MLXProcessRewardModel(Apple Silicon) andTransformersProcessRewardModel(includes RoPE buffer repair for transformers ≥5.0) toits_hub/integration/Motivation
ParticleFilteringandBeamSearchpreviously hardcodedprod/min/meanwith no override path. None of the three is a defensible default (prodpenalises long trajectories,minis brittle to a single weak step,meandiscards 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):
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 fromprod/meanat 408 problems.Interface
Test plan
tests/test_aggregators.py— unit tests for all three aggregator classes, async delegate, and ParticleFiltering integrationtests/test_algorithms.py— BeamSearch aggregator wiringtests/test_mlx_prm.py— MLXProcessRewardModel unit tests