# Adaptive Inference Batching using Policy Gradients

arXiv: https://arxiv.org/abs/2607.05272 · HTML: https://mlsanity.com/paper/2607.05272/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in RL for energy, networks and operations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 20 of its 48 references that are in the map.
- First posted on arXiv 6 Jul 2026.

## Placed in

- Deep reinforcement learning methods / [RL for energy, networks and operations](https://mlsanity.com/deeprl/rl-for-energy-networks-and-operations/)

## Builds on (most cited first, 20 of 20)

- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Learning both Weights and Connections for Efficient Neural Networks](https://mlsanity.com/paper/1506.02626/) (2015)
- [How transferable are features in deep neural networks?](https://mlsanity.com/paper/1411.1792/) (2014)
- [Characterizing possible failure modes in physics-informed neural networks](https://mlsanity.com/paper/2109.01050/) (2021)
- [When and why PINNs fail to train: A neural tangent kernel perspective](https://mlsanity.com/paper/2007.14527/) (2020)
- [Scientific Machine Learning through Physics-Informed Neural Networks: Where we are and What's next](https://mlsanity.com/paper/2201.05624/) (2022)
- [Importance Estimation for Neural Network Pruning](https://mlsanity.com/paper/1906.10771/) (2019)
- [Movement Pruning: Adaptive Sparsity by Fine-Tuning](https://mlsanity.com/paper/2005.07683/) (2020)
- [Soft Filter Pruning for Accelerating Deep Convolutional Neural Networks](https://mlsanity.com/paper/1808.06866/) (2018)
- [Characterizing and Avoiding Negative Transfer](https://mlsanity.com/paper/1811.09751/) (2018)
- [Explicit Inductive Bias for Transfer Learning with Convolutional Networks](https://mlsanity.com/paper/1802.01483/) (2018)
- [Gradient-enhanced physics-informed neural networks for forward and inverse PDE problems](https://mlsanity.com/paper/2111.02801/) (2021)
- [Soft Threshold Weight Reparameterization for Learnable Sparsity](https://mlsanity.com/paper/2002.03231/) (2020)
- [Finite Basis Physics-Informed Neural Networks (FBPINNs): a scalable domain decomposition approach for solving differential equations](https://mlsanity.com/paper/2107.07871/) (2021)
- [Challenges in Training PINNs: A Loss Landscape Perspective](https://mlsanity.com/paper/2402.01868/) (2024)
- [Physics-informed neural network for ultrasound nondestructive quantification of surface breaking cracks](https://mlsanity.com/paper/2005.03596/) (2020)
- [Can Physics-Informed Neural Networks beat the Finite Element Method?](https://mlsanity.com/paper/2302.04107/) (2023)
- [Enhanced Physics-Informed Neural Networks with Augmented Lagrangian Relaxation Method (AL-PINNs)](https://mlsanity.com/paper/2205.01059/) (2022)
- [Transfer Learning in Physics-Informed Neural Networks: Full Fine-Tuning, Lightweight Fine-Tuning, and Low-Rank Adaptation](https://mlsanity.com/paper/2502.00782/) (2025)
- [Unlearning Noise in PINNs: A Selective Pruning Framework for PDE Inverse Problems](https://mlsanity.com/paper/2602.19967/) (2026)

## Abstract (arXiv)

Inference serving systems must balance throughput and latency under bursty, heterogeneous workloads, yet the industry standard remains static batching policies that require manual tuning and cannot adapt to shifting traffic. We investigate whether reinforcement learning (RL) can learn adaptive batching and routing policies that outperform these heuristics, training REINFORCE and PPO agents on a discrete-event simulator validated against queuing theory and production traces (Azure Functions, BurstGPT). We formulate the problem as an MDP over queue state, request type and GPU availability, evaluating across standard Poisson traffic, extreme bursts, real-world traces and heterogeneous multi-GPU routing. Our central finding is a clear boundary condition for RL's value in systems problems. In single-GPU settings, a well-tuned static batching policy is already near-optimal under Poisson-like arrivals and RL offers only marginal gains (+0.1% to +1.0%). In multi-GPU heterogeneous routing, however, where fast and slow requests compete for shared resources, the agent discovers a workload-segregation policy that eliminates Head-of-Line blocking, yielding a 3.5x (348%) improvement over Round-Robin and a 48% improvement over the strongest heuristic baseline (Shortest-Queue), with 60% higher throughput and 25% lower latency while respecting SLA constraints. The policy generalizes to unseen bursty and real-world traffic despite training only on synthetic Poisson arrivals and an attention-augmented policy network converges roughly 20% faster than an MLP baseline. These results suggest RL's advantage over engineered heuristics concentrates in combinatorial, multi-resource decisions rather than single-resource temporal scheduling, a practical distinction for deciding where learned policies justify their engineering cost in production inference infrastructure.