NeurIPS 2026 Papers

Conference on Neural Information Processing Systems

Decision breakdown

6,437papers

Counted from the official paper list; may differ slightly from the number the conference announced.

  • Oral
    991.5%
  • Spotlight
    2403.7%
  • Poster
    6,09894.7%

Browse all papers

ENGINE: Endogenous Variational MultiScale Optimization for Zeroth-Order LLM Fine-Tuning

Zhuoli Ouyang, Changxi Chi, Siyuan Li, Tailin Wu

Expand abstract

Zeroth-order (ZO) optimization enables memory-efficient large language model (LLM) fine-tuning, with recent methods employing curvature-aware strategies. However, in the massive LLM parameter space, these methods inherently suffer from exploding gradient variance and prohibitive costs. Conversely, confining updates to a low-dimensional subspace reduces variance but discards critical high-dimensional information, degrading performance. This creates a fundamental dilemma: full-space optimization is informative but high-variance, while subspace optimization is low-variance but blind. To break this trade-off, we propose \textbf{\Engine} (\textbf{En}do\textbf{g}enous Var\textbf{i}atio\textbf{n}al MultiScal\textbf{e}). Inspired by computational mechanics, we adapt static Variational MultiScale theory to dynamic optimization, treating low- and high-dimensional spaces as resolved coarse and unresolved fine scales. Through a multigrid-style V-cycle, \Engine\ alternates between spaces, dynamically probing fine-grained information in the high-dimensional space and injecting it to correct the efficient low-dimensional sampling process. Evaluations on standard LLM benchmarks show that \Engine\ achieves higher accuracy and faster convergence than existing ZO methods, achieving up to a $\mathbf{2.61\times}$ \textbf{speed-up in forward passes} with zero additional peak memory overhead. Code is available at \url{https://anonymous.4open.science/r/ENGINE-D0A4}.

Poster
EMA-Nesterov: Stabilizing Nesterov's Lookahead for Accelerated Deep Learning Optimization

Chung-Yiu Yau, Dawei Li, Athanasios Glentis, Valentyn Boreiko, Hoi To Wai, Mingyi Hong

Expand abstract

Lookahead-based acceleration methods, such as Nesterov’s momentum, are widely used in optimization, but they often become unreliable in deep learning training due to stochastic gradient noise and non-convex loss landscapes. Standard lookahead relies on short-horizon update signals (e.g., differences between consecutive iterates), which are inherently noisy and can lead to unstable extrapolation directions. This work revisits Nesterov's acceleration from a trajectory perspective and argues that effective acceleration in deep learning should follow the low-frequency trend of optimization trajectories rather than extrapolating noisy one-step updates. Leveraging on this insight, we propose EMA-Nesterov, a simple modification that replaces the standard Nesterov's lookahead direction with an exponential moving average of parameter updates. This yields a stabilized lookahead direction that approximates long-horizon trajectory trends while retaining a lightweight single-loop implementation. We show that EMA-Nesterov retains the same theoretical accelerated convergence rate in convex problems. Furthermore, we show that EMA-Nesterov consistently improves optimization performance across a range of optimizers, including Adam, SOAP, and Muon, in deep learning tasks such as language model pre-training. Compared to prior lookahead methods, EMA-Nesterov achieves better performance by avoiding the instability of short-horizon lookahead and inefficiency of multi-step lookahead.

Poster
EDISCO: Equivariant Discrete Diffusion for Euclidean Combinatorial Optimization

Ruogu Chen, Jie Han

Expand abstract

Euclidean combinatorial optimization problems (ECOPs), such as the Traveling Salesman Problem (TSP) and Capacitated Vehicle Routing Problem (CVRP), possess inherent symmetries under the two-dimensional Euclidean group E(2), including rotations, reflections, and translations. Existing learning-based methods, including recent diffusion-based methods, rely on data augmentation or regularization to approximate E(2)-equivariance. This paper presents EDISCO, the first discrete diffusion model for ECOPs with exact E(2)-invariant generative distributions over node-index solutions. EDISCO introduces an E(2)-equivariant edge-score network coupled with a categorical continuous-time Markov chain over discrete edge variables, and exact posterior sampling provides efficient multi-step inference. This design gives EDISCO a local geometric inductive bias: edge neighborhoods with the same relative geometry and combinatorial context are represented consistently regardless of absolute position or orientation, making learning more efficient and inference more robust than non-equivariant methods. EDISCO outperforms previous learning-based state-of-the-art solvers on synthetic TSP from 100 to 10000 nodes and CVRP from 50 to 2000 customers, while using only 33--50\% of the training instances. Trained only on uniform synthetic data, EDISCO also outperforms competing learning-based baselines under spatial distribution shift and CVRP constraint-tightness shift. Code is available at https://anonymous.4open.science/r/EDISCO-7F54.

Poster
E-GEO: A Testbed for Generative Engine Optimization in E-Commerce

Puneet S Bagga, Vivek Farias, Tamar Korkotashvili, Tianyi Peng, Yuhang Wu

Expand abstract

With the rise of large language models (LLMs), generative engines have become powerful alternatives to traditional search, reshaping retrieval tasks. In e-commerce, for instance, conversational shopping agents now guide consumers to relevant products. This shift has created the need for generative engine optimization (GEO)---improving content visibility and relevance for generative engines. Despite its growing importance, current GEO practices are largely ad hoc, and their impacts remain poorly understood, especially in the e-commerce setting. We address this gap by introducing E-GEO, the first dataset built specifically for e-commerce GEO. E-GEO contains 13,747 realistic, multi-sentence consumer product queries, each paired with 10 relevant Amazon listings, capturing rich intent, constraints, preferences, and shopping contexts that existing datasets miss. Using this dataset, we conduct the first large-scale empirical study of e-commerce GEO across five frontier generative engines, seven popular LLM rewriters with a simple prompt, and fifteen hand-crafted rewriting heuristics. We further formulate GEO as a tractable optimization problem and develop a lightweight prompt meta-optimization algorithm that significantly improves over heuristic baselines. Notably, the optimized prompts reveal a stable, domain-agnostic pattern, suggesting the existence of a "universally effective" GEO strategy. Finally, we red-team the GEO system through both heuristic and optimization-based attacks and show that, under a simple in-prompt defense, gains from GEO reflect genuine content improvement rather than manipulation, anchoring GEO as a substantive, well-defined optimization problem.

Poster
DrPO: Drifting Preference Optimization for One-Step Generative Models

Zhou Jiang, Yandong Wen, Zhen Liu

Expand abstract

While one-step text-to-image generators, a fast and deployment-friendly class of visual generative models, can produce samples with a single forward pass, existing preference-finetuning methods fail to simultaneously achieve efficient adaptation, reward-model agnosticism, and diverse generation. In this work, we propose Drifting Preference Optimization (DrPO), an online preference-finetuning method for one-step generators. Inspired by the recent work of Drifting Models, DrPO first uses the target reward model to construct an on-policy dipole reward model from positive and negative samples. This dipole model induces a preference drifting field in latent feature space, whose gradients are then used to optimize the generator without backpropagating through the target reward model. Empirically, we evaluate DrPO across multiple reward models and SD-Turbo and SDXL-Turbo one-step generators, including HPSv3 and GenEval as representative alignment benchmarks, and validate that DrPO can efficiently and effectively finetune one-step generative models. Preliminary experiments further suggest that this sample-based gradient synthesis can be extended to offline preference finetuning.

Poster
Don't Waste Population: Post-Anneal Refinement for Combinatorial Optimization

Yuma Ichikawa

Expand abstract

Modern GPU-native solvers for combinatorial optimization exploit parallelism by evolving large populations of relaxed candidates; yet, their post-processing collapses this computational effort into a single, rounded solution. We show that the discarded replicas of Parallel Quasi-Quantum Annealing (PQQA) form an implicit elite archive that contains basin-level information recoverable after annealing. We introduce Iterative Path-Relinking Polish (I-PRP), a deterministic, training-free post-anneal operator that leaves the PQQA inner loop unchanged. I-PRP polishes the top rounded elites, traverses cost-guided relinking paths from the current best elite toward the others, re-polishes intermediate states to cross basin boundaries, and iterates only under strict improvement. The operator never degrades the paired upstream polish, terminates after a finite number of updates, and reduces to the original polish when the population collapses to a single basin. Across QUBO and graph-coloring benchmarks, I-PRP preserves saturated cases while improving PQQA on rugged multi-basin instances, revealing that parallel annealing populations contain reusable search structures beyond their best replicas.

Poster
Do We Need Asynchronous SGD? On the Near-Optimality of Synchronous Methods

Grigory Begunov, Alexander Tyurin

Expand abstract

Modern distributed optimization methods mostly rely on traditional synchronous approaches, despite substantial recent progress in asynchronous optimization. We revisit Synchronous SGD and its robust variant, called $m$-Synchronous SGD, which can be interpreted as a method with a backup-worker strategy, and theoretically show that they are nearly optimal in many heterogeneous computation and communication scenarios, which is somewhat unexpected. We analyze the synchronous methods under random times and adversarial partial participation of workers, and prove that their time complexities are optimal in many practical regimes, up to logarithmic factors. While synchronous methods are not universal solutions and there exist tasks where asynchronous methods may be necessary, we show that they are sufficient for many modern heterogeneous computation and communication scenarios.

Oral
Distributionally Robust Token Optimization in RLHF

Yeping Jin, Jiaming Hu, Yannis Paschalidis

Expand abstract

Large Language Models (LLMs) tend to respond correctly to prompts that align well with the data they were trained and fine-tuned on. Yet, small shifts in wording, format, or language can trigger surprisingly large failures, especially on multi-step reasoning problems. To address this problem, we propose a $\textbf{Distributionally Robust Token Optimization (DRTO)}$ approach, which combines token-level Reinforcement Learning from Human Feedback (RLHF) with $\textit{Distributionally Robust Optimization (DRO)}$. DRTO constructs f-divergence ambiguity sets over span-level actor losses, providing a principled way to emphasize difficult response segments during policy optimization. Empirically, DRTO enhances consistency under distribution shifts in multiple reasoning benchmarks among different tasks.

Poster
Distributionally Robust Listwise Preference Optimization

Xudong Wu, Jian Qian, Pangpang Liu, Vaneet Aggarwal, Jiayu Chen

Expand abstract

Existing robust preference optimization for language-model alignment mainly studies pairwise supervision and places robustness at the dataset, prompt, or preference-pair level. We instead study listwise preference optimization under ranking-label uncertainty: given a prompt and a candidate list, the observed ranking over that list may be ambiguous due to annotator inconsistency, near-ties, lossy rankwise feedback, or reward-model noise. We propose a pointwise total-variation robust Plackett--Luce objective that directly robustifies the ranking label conditional on the candidate list. The robust loss admits an exact decomposition into the nominal PL loss plus a worst-case PL correction, and the worst-case ranking is obtained by sorting current implicit scores in ascending order, reducing the inner maximization from $K!$ enumeration to $O(K\log K)$. This tractable structure yields strong offline and online optimization guarantees. In the offline fixed-list setting, the robust objective is convex and projected stochastic subgradient reaches global $\epsilon$-suboptimality with $O(\epsilon^{-2})$ sample complexity. In the online policy-induced setting, where candidate lists are generated by the current policy, we establish weak convexity and $\widetilde O(\epsilon^{-2})$ Moreau-envelope stationarity. Experiments in offline LLM alignment show that the proposed robust correction largely preserves performance under clean labels and improves robustness under noise. In online alignment, it makes reward-model-ranked candidate expansion more reliable and improves both reward-model andexternal GPT-4 judge metrics.

Poster
Distributionally Robust Black Box Optimization-based Bidding Strategy in Auction-based Federated Learning

Xiaoli Tang, Zhuang Qi, Ying-Peng Tang, Xianjie Guo, Han Yu, Jie Zhang, Hengjie Song

Expand abstract

Auction-based Federated Learning (AFL) provides a market-based mechanism to incentivize decentralized data owners (DOs) to join FL model training initiated by data consumers (DCs). However, optimizing the DC's bidding strategy remains fundamentally challenging. Existing methods predominantly rely on Reinforcement Learning (RL), which suffers from severe instability caused by the mismatch between stepwise Markovian rewards and the trajectory-level, non-decomposable feedback inherent in AFL. To address this mismatch, we transition from the stepwise RL formulation to a distributionally robust black-box optimization framework. By treating the entire recruitment-to-training process as a single high-fidelity function evaluation, this transition eliminates the need for explicit reward decomposition. We propose DR-AFL , which accounts for market non-stationarity and data distribution shifts by optimizing policies against worst-case environments within a Wasserstein ambiguity set. Our approach further incorporates an adversarial reweighting scheme to enhance robustness, alongside a dynamic safety projection that leverages real-time cost feedback to enforce budget-aware exploration. We provide theoretical analysis showing that DR-AFL induces a smooth surrogate over the discontinuous auction landscape and guarantees a robust performance lower bound under distributional shifts. Extensive experiments on 6 widely adopted datasets show that DR-AFL consistently outperforms state-of-the-art RL-based baselines, achieving a 2.43% improvement in model accuracy while reducing the number of training cycles required to reach target accuracy by 62.3%.

Poster
Distribution-Adaptive Policy Optimization

Yuxiao He, Ziqi Wang, Xingzhou Lou, Xiaoqian Liu, Junge Zhang

Expand abstract

Reinforcement Learning with Verifiable Rewards empowers Large Language Models to enhance their reasoning capabilities. Most algorithms rely on on-policy training, which is stable but inefficient due to its strict dependency on fresh samples. To overcome this inefficiency, off-policy RL decouples experience generation from policy optimization via Importance Sampling (IS). However, this introduces a policy discrepancy between the behavior policy and the target policy. We observe that mainstream RLVR algorithms struggle to maintain a balance between performance and stability in off-policy scenarios, suffering from either severe performance degradation or training collapse. We identify the root cause as the failure of static trust regions to maintain the bias--variance trade-off associated with IS ratios in off-policy settings. To address this, we propose **DI**stribution-adaptive **P**olicy **O**ptimization (**DIPO**), a novel hyperparameter-free RLVR algorithm employing an adaptive trust region. DIPO leverages the intrinsic statistics of the IS ratio distribution to achieve a robust bias--variance trade-off. Experimental results demonstrate that DIPO enables stable training in extreme off-policy scenarios, achieving performance comparable to or even surpassing the on-policy baseline, and outperforming other off-policy baselines by 13.9%.

Poster
Disentangling Optimization Geometry via Hierarchical Polar Adapters for Class-Incremental Learning

Dat Q Mac, Thanh Hai Dang, Duc-Trong Le, Quynh-Trang Pham Thi

Expand abstract

Deep neural networks suffer from catastrophic forgetting when learning on sequence tasks because standard adaptation mechanisms update a monolithic set of Cartesian weights, inherently entangling feature intensity with spatial displacement and destructively overwriting ancestral representations. Inspired by the brain’s multi-timescale consolidation and frequency-aware memory dynamics, we propose HiPo(Hierarchical Polar Adapters), a novel continual learning framework that fundamentally redefines how neural knowledge is parameterized, stored, and consolidated. First, HiPo projects token representations into a complex orthogonal Fourier basis and explicitly parametrizes the adapter weights in polar coordinates, decoupling feature intensity from geometric correlation to probabilistically isolate task updates. Second, we embed these polar weights within a multi-tiered hierarchical memory stack, functionally separating a highly plastic working memory from a deep, stabilizing long-term memory. Finally, an adaptively thresholded snapshot Fisher mechanism evaluates parameter criticality in the decoupled polar space. By executing mathematically safe, Cartesian-invariant knowledge transfers into the deep memory tiers, HiPo safely hollows out the working memory, paralyzing obsolete gradients and preserving plasticity without inducing structural distortion. Extensive experiments across standard CIL benchmarks demonstrate that HiPo establishes a highly resilient optimization manifold, achieving state-of-the-art performance and exceptional geometric stability, particularly under severe out-of-distribution shifts.

Poster
Discrete Flow Matching: Convergence Guarantees Under Minimal Assumptions

Le-Tuyet-Nhi PHAM, Giovanni Conforti, Zhenjie Ren, Alain Durmus

Expand abstract

Flow Matching has recently emerged as a popular class of generative models for simulating a target distribution $\mu_1$ from samples drawn from a source distribution $\mu_0$. This framework relies on a fixed coupling between $\mu_0$ and $\mu_1$, and on a deterministic or stochastic bridge to define an interpolating process between the two distributions. The time marginals of this process can then be approximately sampled by estimating the transition rates, or more generally the generator, of its Markovian projection. This framework has recently been extended to the case of discrete source and target distributions, under the name *Discrete Flow Matching* (DFM). However, theoretical guarantees for such models remain scarce. In this paper, we study two DFM models on $\mathbb{Z}_m^d = \\{0,\ldots,m-1\\}^d$, sampled through time discretization, and derive non-asymptotic associated bounds for both of them. In contrast to previous work, we establish non-asymptotic bounds in Kullback--Leibler divergence for the early-stopped version of the target distribution. We also derive explicit convergence guarantees in total variation distance with respect to the true target distribution. Importantly, these bounds rely only on an approximation error assumption, relaxing standard score assumptions used in earlier works, while also yielding improved dependence on the vocabulary size $m$ and the dimension $d$.

Poster
DisCoMBO: Steering Expert-in-the-Loop Black Box Optimization via Distributional Conformance

Jonas Seng, Bennet Wittelsbach, Kristian Kersting

Expand abstract

Sequential Model-Based Optimization (SMBO) traditionally relies on Bayesian or ensembling surrogates for uncertainty quantification. While historically treated as fully data-driven, SMBO increasingly integrates external domain expertise to accelerate discovery. To overcome the opaque guidance and diminished integration fidelity of standard acquisition re-weighting, Probabilistic Circuits (PCs) have emerged as a generative surrogate alternative, enabling direct knowledge injection via conditional sampling. However, these generative routines lack the formal exploration-exploitation semantics required for rigorous optimization. We introduce the Distributional Conformance Score (DisCo), a principled metric that unifies the flexibility of PCs with a rigorous uncertainty framework. DisCo provides a bounded, $[0, 1]$-normalized measure of model "surprise" that (1) recovers properties comparable to kernel-based uncertainty known from, e.g., Gaussian Processes while maintaining linear-time inference, and (2) enables principled and accurate assessment of conformance of external knowledge w.r.t. model evidence. We then present DisCoMBO, a framework leveraging these properties for robust, knowledge-aware optimization. We prove that DisCoMBO is a zero-regret algorithm and demonstrate its effectiveness across diverse benchmarks from AutoML, material optimization, and wind park optimization.

Poster
Diffusion-State Policy Optimization for Masked Diffusion Language Models

Daisuke Oba, Hiroki Furuta, Naoaki Okazaki

Expand abstract

Masked diffusion language models generate text through iterative masked-token filling, but terminal-only rewards on final completions provide coarse credit assignment for the intermediate filling decisions that shape the generation process. We propose Diffusion-State Policy Optimization (DiSPO), a plug-in credit-assignment layer that directly optimizes intermediate filling decisions. At selected intermediate masked states, DiSPO branches by resampling the currently masked positions from rollout-cached logits, scores the resulting completions, and updates only the newly filled tokens, requiring no additional multi-step diffusion rollouts or optimizer steps. We formalize a fixed-state objective for branched completions and derive a policy-gradient estimator that reuses the same rollouts as terminal-feedback policy optimization. Experiments on LLaDA-8B-Instruct show that DiSPO consistently improves terminal-feedback baselines, including diffu-GRPO and SPG, on math and planning benchmarks under matched rollout compute and optimizer steps, supporting its use as a general plug-in for masked diffusion policy optimization.

Poster
Diffusion Transformers with Residual Adaptive Layer Normalization

Ge Wu, Minxing Luo, Yikai Ge, Lei Wang, Dandan Zheng, Rui Liu, Libin Wang, Yujie Bao, taihang Hu, JingDong Chen, Xiang Li

Expand abstract

Residual connections with PreNorm are the default design in modern transformers due to their strong optimization stability, but recent research has shown that they weaken the functional differentiation of deeper blocks and increase representational redundancy in later layers. This issue also appears in diffusion models using similar architectures, where the deep block is expected to support progressive denoising refinement rather than repeated computation over similar states. Motivated by this, we revisit cross-layer shortcuts in diffusion transformers and argue that their role should go beyond merely transferring shallow information, instead guiding deeper blocks to improve depth utilization. To this end, we propose \textit{\textbf{Res}idual Adaptive Layer \textbf{Norm}alization} (\textbf{ResNorm}), which routes residual cross-layer shortcuts into the existing adaLN modulation. ResNorm preserves the optimization advantages of PreNorm residual connections, while reusing the native adaLN interface to transmit early-layer representations as modulation signals. This design naturally combines cross-layer history with the model's native condition, such as timestep and class embeddings. ResNorm allows the earlier structural representation to adaptively scale, shift, and gate deeper blocks. This not only produces more differentiated deep representations and reduces cross-layer similarity in later blocks, but also accelerates training convergence and improves performance. On ImageNet 256$\times$256, SiT-XL/2 + ResNorm achieves up to $\textbf{2}\times$ and $\textbf{36}\times$ faster training than SiT-XL/2 + REG and SiT-XL/2 + REPA. Moreover, SiT-L/2 + ResNorm trained for only 400K iterations already outperforms SiT-XL/2 + REG.

Poster
Diff3R: Feed-forward 3D Gaussian Splatting with Uncertainty Aware Differentiable Optimization

Yueh-Cheng Liu, Jozef Hladký, Matthias Niessner, Angela Dai

Expand abstract

Recent advances in 3D Gaussian Splatting (3DGS) present two main directions: feed-forward models offer fast inference in sparse-view settings, while per-scene optimization yields high-quality renderings but is computationally expensive. To combine the benefits of both, we introduce Diff3R, a novel framework that explicitly bridges feed-forward prediction and test-time optimization. By incorporating a differentiable 3DGS optimization layer directly into the training loop, our network learns to predict an optimal initialization for test-time optimization rather than a conventional zero-shot result. To overcome the computational cost of backpropagating through the optimization steps, we propose computing gradients via the Implicit Function Theorem and a scalable, matrix-free PCG solver tailored for 3DGS optimization. Additionally, we incorporate a data-driven uncertainty model into the optimization process by adaptively controlling how much the parameters are allowed to change during optimization. This approach effectively mitigates overfitting in under-constrained regions and increases robustness against input outliers. Since our proposed optimization layer is model-agnostic, we show that it can be seamlessly integrated into existing feed-forward 3DGS architectures for both pose-given and pose-free methods, providing improvements for test-time optimization.

Poster
Deriving Hyperparameter Scaling Laws via Modern Optimization Theory

Egor Shulgin, Jorg Franke, Dimitri von Rütte, Tianyue Zhang, Niccolò Ajroldi, Korbinian Pöppel, Bernhard Scholkopf, Aaron Klein, Peter Richtarik, Antonio Orvieto

Expand abstract

Hyperparameter transfer has become an important component of modern large-scale training recipes. Existing methods, such as muP, primarily focus on transfer between model sizes, with transfer across batch sizes and training horizons often relying on empirical scaling rules informed by insights from timescale preservation, quadratic proxies, and continuous-time approximations. We study hyperparameter scaling laws for modern first-order optimizers through the lens of recent convergence bounds for methods based on the Linear Minimization Oracle (LMO), a framework that includes normalized SGD, signSGD (approximating Adam), and Muon. Treating bounds in recent literature as a proxy and minimizing them across different tuning regimes yields closed-form power-law schedules for learning rate, momentum, and batch size as functions of the iteration or token budget. Our analysis, holding model size fixed, recovers most insights and observations from the literature under a unified and principled perspective, with clear directions open for future research. Our results draw particular attention to the interaction between momentum and batch-size scaling, suggesting that optimal performance may be achieved with several scaling strategies.

Poster
Decoupled Optimization for Teacher-Student Semi-Supervised Learning via a Pioneer Student

Haorong Han, Jidong Yuan, Chixuan Wei, Yongqi Sun

Expand abstract

Semi-supervised learning (SSL) relies on two core mechanisms: self-training under the Teacher-Student (T-S) framework and joint optimization of labeled and unlabeled losses. Despite their effectiveness, we find both mechanisms introduce distinct optimization pathologies. First, parameter coupling enforces strict synchronization between teacher and student, where strong regularization on the student degrades the teacher's fitting ability, thereby limiting the permissible generalization intensity. Second, the imbalance in gradient update consistency between labeled and unlabeled losses drives the shared parameters to prematurely converge to labeled-dominated local minima, creating a bottleneck for global optimization. To address both issues, we propose the Pioneer Student (PiS), an auxiliary branch that operates in an independent parameter space and periodically transfers accumulated knowledge back to the T-S model. Extensive experiments show that PiS is a universal plug-and-play module that consistently improves mainstream SSL methods.

Poster
Decision-Focused Learning in MDPs: An Occupancy Measure Approach

Zihao Zhao, Ashwath K Karunakaram, Ali Eshragh, Yuexing Li, Kai Wang

Expand abstract

In this work, we consider decision-focused learning (DFL) for a Markov decision process (MDP), where existing methods differentiate through the KKT conditions of the Bellman equation and require solving a linear system over all state-action pairs, limiting its scalability. We address this by reformulating the MDP as an occupancy measure-based linear program (LP), whose feasible region is induced by predicted dynamics, and we derive a closed-form gradient by identifying the active constraints in the feasible polyhedron via the pivoting algorithm. This occupancy measure-based LP layer raises two challenges: (1) LP's solution gradient is discontinuous when active constraints change, and (2) the LP backward cost still scales with the state size, which is costly for large or continuous state spaces. We address the challenges with an augmented Lagrangian surrogate and smooth the boundary jumps by random row sketching of the constraints, and a learnable soft state-aggregation layer and its function-approximation generalization that scales the LP to large finite and continuous-state MDPs. Across multiple tasks, our methods reach lower regret than KKT-based DFL and two-stage baselines with significantly lower computation cost. The source code for all experiments is available [here](https://anonymous.4open.science/r/State_Aggregation_Project-3E64/).

Poster