8 papers
Multi-Reward Hacking: From Aggregation-Induced Drift to Trust-Anchored Defense
Yushu Li, Jiin Woo, Xiaoxiao Li +1
To improve task performance while encouraging desirable response properties, LLM post-training often combines a primary reward (e.g., correctness) with auxiliary rewards (e.g., bre…
Compatibility Matters: A Framework for Token-Level Teacher Selection in On-Policy Distillation
Stanley Wei, Jiin Woo, Kaan Ozkara +1
On-policy distillation (OPD) has recently been introduced as a promising paradigm for continual learning, where a student policy can be supervised at the token-level. Typically, OP…
How to Train Your Little Dragon: Teaching a Small Language Model to Construct and Refine GPU Optimization Scripts
Annabelle Sujun Tang, Wei Tang, Songyao Jiang +7
Compiler transformations interact, making GPU kernel optimization a search for combinations that preserve the computation while improving performance. We introduce HOW TO TRAIN YOU…
Federate the Router: Learning Language Model Routers with Sparse and Decentralized Evaluations
Baris Askin, Shivam Patel, Anupam Nayak +4
Large language models (LLMs) are increasingly accessed as remotely hosted services by edge and enterprise clients that cannot run frontier models locally. Since models vary widely…
Sample Complexity of Average-Reward Q-Learning: From Single-agent to Federated Reinforcement Learning
Yuchen Jiao, Jiin Woo, Gen Li +2
Average-reward reinforcement learning offers a principled framework for long-term decision-making by maximizing the mean reward per time step. Although Q-learning is a widely used…
TritonRL: Training LLMs to Think and Code Triton Without Cheating
Jiin Woo, Shaowei Zhu, Allen Nie +3
The rapid evolution of Large Language Models (LLMs) has driven a growing demand for automated, high-performance system kernels to accelerate machine learning workloads. We introduc…
Large Language Model-Enhanced Reinforcement Learning for Diverse and Novel Recommendations
Jiin Woo, Alireza Bagheri Garakani, Tianchen Zhou +2
In recommendation systems, diversity and novelty are essential for capturing varied user preferences and encouraging exploration, yet many systems prioritize click relevance. While…
Federated Offline Reinforcement Learning: Collaborative Single-Policy Coverage Suffices
Jiin Woo, Laixi Shi, Gauri Joshi +1
Offline reinforcement learning (RL), which seeks to learn an optimal policy using offline data, has garnered significant interest due to its potential in critical applications wher…