Gandharv Patil

Sravanti Addepalli

Arun Suggala

R. Aravamudhan

Soumya Sharma

Anirban Laha

Aravindan Raghuveer

Karthikeyan Shanmugam

Reward models (RMs) are fundamental to aligning Large Language Models (LLMs) via human feedback, yet they often suffer from reward hacking. … (see more)They tend to latch on to superficial or spurious attributes, such as response length or formatting, mistaking these cues learned from correlations in training data for the true causal drivers of quality (e.g., factuality, relevance). This occurs because standard training objectives struggle to disentangle these factors, leading to brittle RMs and misaligned policies. We introduce Crome (Causally Robust Reward Modeling), a novel framework grounded in an explicit causal model designed to mitigate reward hacking. Crome employs the following synthetic targeted augmentations during training: (1) Causal Augmentations, which are pairs that differ along specific causal attributes, to enforce sensitivity along each causal attribute individually, and (2) Neutral Augmentations, which are tie-label pairs varying primarily in spurious attributes, to enforce invariance along spurious attributes. Notably, our augmentations are produced without any knowledge of spurious factors, via answer interventions only along causal rubrics, that are identified by querying an oracle LLM. Empirically, Crome significantly outperforms standard baselines on RewardBench, improving average accuracy by up to 5.4% and achieving gains of up to 13.2% and 7.2% in specific categories. The robustness of Crome is further testified by the consistent gains obtained in a Best-of-N inference setting across increasing N, across various benchmarks, including the popular RewardBench (covering chat, chat-hard, safety, and reasoning tasks), the safety-focused WildGuardTest, and the reasoning-specific GSM8k.

2025-06-19

ArXiv (preprint)

arxiv.org

Robust Reward Modeling via Causal Rubrics

Pragya Srivastava

Harman Singh

Rahul Madhavan

Sravanti Addepalli

Arun Suggala

Rengarajan Aravamudhan

Soumya Sharma

Anirban Laha

Aravindan Raghuveer

Karthikeyan Shanmugam

Reward models (RMs) for LLM alignment often exhibit reward hacking, mistaking spurious correlates (e.g., length, format) for causal quality … (see more)drivers (e.g., factuality, relevance), leading to brittle RMs. We introduce CROME (Causally Robust Reward Modeling), a causally-grounded framework using targeted augmentations to mitigate this. CROME employs: (1) Causal Augmentations, pairs isolating specific causal attribute changes, to enforce sensitivity, and (2) Neutral Augmentations, tie-labeled pairs varying spurious attributes while preserving causal content, to enforce invariance. Crucially, augmentations target LLM-identified causal rubrics, requiring no prior knowledge of spurious factors. CROME significantly outperforms baselines on RewardBench (Avg +5.4\%, Safety +13.2\%, Reasoning +7.2\%) and demonstrates enhanced robustness via improved Best-of-N performance across RewardBench, WildGuardTest, and GSM8k.

2025-06-10

ICML.cc/2025/Workshop/MoFA (poster)

openreview.net

On learning history-based policies for controlling Markov decision processes

Aditya Mahajan

2024-04-18

Proceedings of The 27th International Conference on Artificial Intelligence and Statistics (published)

openreview.net

Finite time analysis of temporal difference learning with linear function approximation: Tail averaging and regularisation

Prashanth L.A.

Dheeraj M. Nagaraj

We study the finite-time behaviour of the popular temporal difference (TD) learning algorithm, when combined with tail-averaging. We derive … (see more)finite time bounds on the parameter error of the tail-averaged TD iterate under a step-size choice that does not require information about the eigenvalues of the matrix underlying the projected TD fixed point. Our analysis shows that tail-averaged TD converges at the optimal O (1/t) rate, both in expectation and with high probability. In addition, our bounds exhibit a sharper rate of decay for the initial error (bias), which is an improvement over averaging all iterates. We also propose and analyse a variant of TD that incorporates regularisation, and show that this variant fares favourably in problems with ill-conditioned features.

2023-04-11

Proceedings of The 26th International Conference on Artificial Intelligence and Statistics (published)

arxiv.org

Variance Penalized On-Policy and Off-Policy Actor-Critic

Ayush Jain

2021-05-18

Proceedings of the AAAI Conference on Artificial Intelligence (published)

arxiv.org

Finite time analysis of temporal difference learning with linear function approximation: the tail averaged case

Prashanth L.A.