Reproduce “Beyond Uniform Credit: Causal Credit Assignment for Policy Optimization” on OpenResearch

About this paper

Mykola Khandoga, Rui Yuan, Vinay Kumar Sankarapu

Lexsi Labs introduces Counterfactual Importance Weighting, a method that enhances policy gradient training for language models by assigning credit to generated tokens based on their causal contribution to a correct answer, rather than uniform weighting. This approach yields consistent accuracy improvements of 0.8 to 1.1 percentage points on the GSM8K mathematical reasoning dataset and significantly accelerates training convergence across different model sizes.

Key insights

• Only a small minority (10.9%) of reasoning spans are found to be highly critical for generating correct answers, while a small percentage (3.5%) act as 'distractors' that, if removed, improve the probability of correctness.
• Counterfactual weighting effectively concentrates gradient mass: high-importance tokens receive 1.6x more gradient mass compared to uniform weighting, while low-importance tokens are diluted by 0.6x, demonstrating a 2.7x relative shift in allocation.
• The method is most effective when task success relies on a concentrated verification signal, such as a single numeric answer, and requires domain-specific span detection, as evidenced by its limited benefit in code generation.

Summary by alphaXiv.