HARGO: Heterogeneity-Aware Reward-Guided Optimization for RL Post-Training of LLMs on HPC Tasks

Tiangang Li Xiangbo Tian

Jul 2026 · arXiv.org · Vol abs/2607.28301 · 0 citations · 30 references

Computer Science

TL;DR

Heterogeneity-Aware Reward-Guided Optimization is proposed, which introduces per-response importance weighting via confidence-modulated advantage: computing a discrimination signal from group-level reward contrast and a confidence signal from reference model log-probabilities, then modulating the advantage before computing per-response weights, without requiring task-type labels.

Abstract

Supervised fine-tuning (SFT) can equip large language models (LLMs) with domain knowledge for high-performance computing (HPC) tasks such as data race detection and benchmark question answering. However, knowledge alone does not guarantee task-appropriate behavior: the same SFT model that correctly classifies 88.65\% of C/C++ data race samples produces verbose, imprecise answers to factual queries, with 65.9\% of MLPerf responses exceeding 40 characters. Reinforcement learning (RL) post-training addresses this gap by optimizing for task-specific rewards rather than token-level imitation. Yet HPC tasks exhibit extreme heterogeneity, with binary classification, factual QA, and semantic generation differing by 58x in answer length, spanning three distinct reward distributions, and showing widely varying SFT accuracy. This makes uniform-weight RL methods such as GRPO suboptimal. We propose HARGO, Heterogeneity-Aware Reward-Guided Optimization, which introduces per-response importance weighting via confidence-modulated advantage: computing a discrimination signal from group-level reward contrast and a confidence signal from reference model log-probabilities, then modulating the advantage before computing per-response weights, without requiring task-type labels. Across four HPC tasks and nine methods, HARGO achieves the best performance on all three primary metrics: WinRate 54.62\%, Data Race F1 91.30\%, and PLP Similarity 0.8558. Ablation confirms complementary contributions from both signals. HARGO establishes the best overall alignment quality among compared methods for heterogeneous HPC tasks.

DOI arXiv Semantic Scholar

View source

Similar papers

Jul 2026

Multi-turn RL with Structural and Performance Aware Rewards for CUDA Kernel Generation

Empirical findings suggest that CudaPerf significantly outperforms strong baselines, including Qwen-3-32B and CUDA Agent by achieving up to 5X and 3.32X improvements in speedup, and 17%&7% improvements in correctness, respectively.

Q. I. Mahmud, Nesreen K. Ahmed, Ali Jannesari · arXiv.org · 0 citations

Jul 2026

MADA-RL: Multi-Agent Debate-Aware Reinforcement Learning for Parameter-Efficient Reasoning in Compact Models

A controlled study isolates the source of MADA-RL's gains: the counterfactual advantage produces the highest critic improvement rate of any model evaluated, indicating that trained critics learn to correct generator errors rather than to imitate them.

Martino M. L. Pulici, Cuong Xuan Chu, Evgeny Kharlamov et al. · arXiv.org · 0 citations

#machine learning Jul 2026

Off-Context GRPO: Learning to Reason on Hard Problems using Privileged Information

Off-Context GRPO (OC-GRPO), a minimally modified variant of GRPO that uses guided rollouts but applies an importance-corrected objective to steer the update back toward the original unguided objective, avoiding the mismatch that destabilizes uncorrected guided training.

Priyank Agrawal, Ankur Samanta, S. Ghasemlou et al. · arXiv.org · 1 citation

Jul 2026

GFlowRL: Scaling Distribution-Matching RL to Large Language Models

GFlowRL, a streamlined GFlowNet-style RL algorithm that removes the auxiliary partition network entirely while preserving the reward-distribution-matching objective, is proposed, and is the first GFlowNet-style RL algorithm to scale stably across both dense and sparse architectures.

Xiaodong Liu, Michael Xu, Jack W. Stokes et al. · arXiv.org · 0 citations

#machine learning Preprint Sep 2026

Gradients Know What Outcomes Don't: Unlocking Reinforcement Learning for LLM Reasoning with Gradient-Aligned Rewards

Reinforcement learning from verifiable rewards (RLVR) drives chain-of-thought reasoning in large language models, yet its binary outcome reward cannot distinguish among correct trajectories. Existing dense reward alternatives, from surface heuristics to process reward models, either ignore the expert solutions already present in training corpora or require expensive offline annotation. We propose Gradient-Aligned Reward (GAR), which operates in the policy's own gradient space: truncated backpropagation through the output projection layer extracts a compact gradient vector for each rollout, and cosine similarity with an expert-anchor gradient yields a dense, reasoning-aware reward with less than 9% wall-clock overhead. We prove that this cosine admits a multiplicative decomposition into prediction-error and activation-pattern factors, providing a concrete characterization of what the alignment signal measures. On Qwen3-4B and Qwen3-8B, GAR consistently improves over GRPO and other baselines on competition-level math benchmarks and transfers to GPQA Diamond and MMLU-Pro without domain-specific data. Code and data are available at https://github.com/LQgdwind/GAR.

Le-Qi Zheng, Jin-Bo Su, Fang Niu et al. · 2 citations

Preprint Aug 2026

Beyond Solvability: Task Learnability as a Static Prior for LLM RL Post-Training

TrajVal, a lightweight probe-based estimator that approximates per-task learnability from a short probe run and two endpoint evaluations, is proposed and it is found that learnability is reproducible across independently sampled training contexts and predictive of downstream utility.

Ting Zhou, Zhenqing Ling, Daoyuan Chen et al. · 0 citations