Optimization
From Data to Rewards: a Bilevel Optimization Perspective on Maximum Likelihood Estimation
Benechehab, Abdelhakim, Singer, Gabriel, Léger, Corentin, Hili, Youssef Attia El, Paolo, Giuseppe, Thomas, Albert, Filippone, Maurizio, Kégl, Balázs
Generative models form the backbone of modern machine learning, underpinning state-of-the-art systems in text, vision, and multimodal applications. While Maximum Likelihood Estimation has traditionally served as the dominant training paradigm, recent work have highlighted its limitations, particularly in generalization and susceptibility to catastrophic forgetting compared to Reinforcement Learning techniques, such as Policy Gradient methods. However, these approaches depend on explicit reward signals, which are often unavailable in practice, leaving open the fundamental problem of how to align generative models when only high-quality datasets are accessible. In this work, we address this challenge via a Bilevel Optimization framework, where the reward function is treated as the optimization variable of an outer-level problem, while a policy gradient objective defines the inner-level. We then conduct a theoretical analysis of this optimization problem in a tractable setting and extract insights that, as we demonstrate, generalize to applications such as tabular classification and model-based reinforcement learning. We release the code at https://github.com/abenechehab/nll_to_po .
Understanding In-context Learning of Addition via Activation Subspaces
Hu, Xinyan, Yin, Kayo, Jordan, Michael I., Steinhardt, Jacob, Chen, Lijie
To perform few-shot learning, language models extract signals from a few input-label pairs, aggregate these into a learned prediction rule, and apply this rule to new inputs. How is this implemented in the forward pass of modern transformer models? To explore this question, we study a structured family of few-shot learning tasks for which the true prediction rule is to add an integer $k$ to the input. We introduce a novel optimization method that localizes the model's few-shot ability to only a few attention heads. We then perform an in-depth analysis of individual heads, via dimensionality reduction and decomposition. As an example, on Llama-3-8B-instruct, we reduce its mechanism on our tasks to just three attention heads with six-dimensional subspaces, where four dimensions track the unit digit with trigonometric functions at periods $2$, $5$, and $10$, and two dimensions track magnitude with low-frequency components. To deepen our understanding of the mechanism, we also derive a mathematical identity relating ``aggregation'' and ``extraction'' subspaces for attention heads, allowing us to track the flow of information from individual examples to a final aggregated concept. Using this, we identify a self-correction mechanism where mistakes learned from earlier demonstrations are suppressed by later demonstrations. Our results demonstrate how tracking low-dimensional subspaces of localized heads across a forward pass can provide insight into fine-grained computational structures in language models.
Efficient Prompt Optimisation for Legal Text Classification with Proxy Prompt Evaluator
Lee, Hyunji, Li, Kevin Chenhao, Grabmair, Matthias, Xu, Shanshan
Prompt optimization aims to systematically refine prompts to enhance a language model's performance on specific tasks. Fairness detection in Terms of Service (ToS) clauses is a challenging legal NLP task that demands carefully crafted prompts to ensure reliable results. However, existing prompt optimization methods are often computationally expensive due to inefficient search strategies and costly prompt candidate scoring. In this paper, we propose a framework that combines Monte Carlo Tree Search (MCTS) with a proxy prompt evaluator to more effectively explore the prompt space while reducing evaluation costs. Experiments demonstrate that our approach achieves higher classification accuracy and efficiency than baseline methods under a constrained computation budget.
DYNAMIX: RL-based Adaptive Batch Size Optimization in Distributed Machine Learning Systems
Dai, Yuanjun, He, Keqiang, Wang, An
Abstract--Existing batch size selection approaches in distributed machine learning rely on static allocation or simplistic heuristics that fail to adapt to heterogeneous, dynamic computing environments. We present DYNAMIX, a reinforcement learning framework that formulates batch size optimization as a sequential decision-making problem using Proximal Policy Optimization (PPO). Our approach employs a multi-dimensional state representation encompassing network-level metrics, system-level resource utilization, and training statistical efficiency indicators to enable informed decision-making across diverse computational resources. Our approach eliminates the need for explicit system modeling while integrating seamlessly with existing distributed training frameworks. Through evaluations across diverse workloads, hardware configurations, and network conditions, DY - NAMIX achieves up to 6.3% improvement in the final model accuracy and 46% reduction in the total training time. Our scalability experiments demonstrate that DYNAMIX maintains the best performance as cluster size increases to 32 nodes, while policy transfer experiments show that learned policies generalize effectively across related model architectures. Distributed machine learning (DML) has emerged as the predominant paradigm for training increasingly complex models on expansive datasets. As model architectures grow in parameter count and computational demands, practitioners increasingly rely on distributed training across multiple computational nodes to maintain feasible training timelines. Within this paradigm, batch size selection represents a critical hy-perparameter that significantly influences both training efficiency and model convergence properties. While larger batch sizes generally improve hardware utilization through increased parallelism, they may adversely affect statistical efficiency, potentially degrading convergence rates and generalization performance [19], [32]. The optimization complexity intensifies substantially in heterogeneous distributed environments, characterized by variance in computational capabilities, network characteristics, and hardware specifications across training nodes. These heterogeneous configurations arise from several practical considerations: cost optimization through spot instance utilization [12], consolidation of diverse hardware generations within organizational clusters [13], and workload deployment in multi-tenant infrastructure [15]. Under such conditions, the conventional approach of uniform batch size allocation frequently leads to suboptimal resource utilization, as demonstrated by Jia et al. [16], who observed significant throughput degradation due to synchronization barriers in heterogeneous clusters. Existing approaches to batch size optimization in distributed environments fall into several distinct categories, each exhibiting particular limitations.
Reliability of Single-Level Equality-Constrained Inverse Optimal Control
Bečanović, Filip, Jovanović, Kosta, Bonnet, Vincent
Abstract-- Inverse optimal control (IOC) allows the retrieval of optimal cost function weights, or behavioral parameters, from human motion. The literature on IOC uses methods that are either based on a slow bilevel process or a fast but noise-sensitive minimization of optimality condition violation. Assuming equality-constrained optimal control models of human motion, this article presents a faster but robust approach to solving IOC using a single-level reformulation of the bilevel method and yields equivalent results. Through numerical experiments in simulation, we analyze the robustness to noise of the proposed single-level reformulation to the bilevel IOC formulation with a human-like planar reaching task that is used across recent studies. The approach shows resilience to very large levels of noise and reduces the computation time of the IOC on this task by a factor of 15 when compared to a classical bilevel implementation.
Mix- and MoE-DPO: A Variational Inference Approach to Direct Preference Optimization
Bohne, Jason, Polak, Pawel, Rosenberg, David, Bloniarz, Brian, Kazantsev, Gary
Direct Preference Optimization (DPO) has recently emerged as a simple and effective alternative to reinforcement learning from human feedback (RLHF) for aligning large language models (LLMs) with user preferences. However, existing DPO formulations rely on a single monolithic model, which limits their expressivity in multi-task settings and their adaptability to heterogeneous or diverse preference distributions. In this work, we propose Mix- and MoE-DPO, a framework that extends DPO with both soft mixture models and mixture-of-experts (MoE) architectures, using a stochastic variational inference approach. Our method introduces a latent-variable model over expert assignments and optimizes a variational evidence lower bound (ELBO), enabling stable and efficient learning of specialized expert policies from preference data. Mix- and MoE-DPO provides three key advantages over standard DPO: (i) generalization via universal function approximation through mixtures; (ii) reward and policy specialization through expert components tailored to distinct preference modes; and (iii) contextual alignment through input-dependent soft gating that enables user-specific mixture policies. Our framework supports both shared base architectures with expert-specific policy heads and fully independent expert models, allowing flexible trade-offs between parameter efficiency and specialization. We validate our approach on a variety of model sizes and multi-preference datasets, demonstrating that Mix- and MoE-DPO offers a powerful and scalable method for preference-based LLM alignment.
EB-MBD: Emerging-Barrier Model-Based Diffusion for Safe Trajectory Optimization in Highly Constrained Environments
Mishra, Raghav, Manchester, Ian R.
We propose enforcing constraints on Model-Based Diffusion by introducing emerging barrier functions inspired by interior point methods. We show that constraints on Model-Based Diffusion can lead to catastrophic performance degradation, even on simple 2D systems due to sample inefficiency in the Monte Carlo approximation of the score function. We introduce Emerging-Barrier Model-Based Diffusion (EB-MBD) which uses progressively introduced barrier constraints to avoid these problems, significantly improving solution quality, without the need for computationally expensive operations such as projections. We analyze the sampling liveliness of samples each iteration to inform barrier parameter scheduling choice. We demonstrate results for 2D collision avoidance and a 3D underwater manipulator system and show that our method achieves lower cost solutions than Model-Based Diffusion, and requires orders of magnitude less computation time than projection based methods.
LiveThinking: Enabling Real-Time Efficient Reasoning for AI-Powered Livestreaming via Reinforcement Learning
Sun, Yuhan, Huang, Zhiwei, Cui, Wanqing, Xiong, Shaopan, Guo, Yazhi, Jin, Meiguang, Ma, Junfeng
In AI-powered e-commerce livestreaming, digital avatars require real-time responses to drive engagement, a task for which high-latency Large Reasoning Models (LRMs) are ill-suited. We introduce LiveThinking, a practical two-stage optimization framework to bridge this gap. First, we address computational cost by distilling a 670B teacher LRM into a lightweight 30B Mixture-of-Experts (MoE) model (3B active) using Rejection Sampling Fine-Tuning (RFT). This reduces deployment overhead but preserves the teacher's verbose reasoning, causing latency. To solve this, our second stage employs reinforcement learning with Group Relative Policy Optimization (GRPO) to compress the model's reasoning path, guided by a multi-objective reward function balancing correctness, helpfulness, and brevity. LiveThinking achieves a 30-fold reduction in computational cost, enabling sub-second latency. In real-world application on Taobao Live, it improved response correctness by 3.3% and helpfulness by 21.8%. Tested by hundreds of thousands of viewers, our system led to a statistically significant increase in Gross Merchandise Volume (GMV), demonstrating its effectiveness in enhancing user experience and commercial performance in live, interactive settings.
Estimating Fair Graphs from Graph-Stationary Data
Navarro, Madeline, Buciulea, Andrei, Rey, Samuel, Marques, Antonio G., Segarra, Santiago
We estimate fair graphs from graph-stationary nodal observations such that connections are not biased with respect to sensitive attributes. Edges in real-world graphs often exhibit preferences for connecting certain pairs of groups. Biased connections can not only exacerbate but even induce unfair treatment for downstream graph-based tasks. We therefore consider group and individual fairness for graphs corresponding to group- and node-level definitions, respectively. To evaluate the fairness of a given graph, we provide multiple bias metrics, including novel measurements in the spectral domain. Furthermore, we propose Fair Spectral Templates (FairSpecTemp), an optimization-based method with two variants for estimating fair graphs from stationary graph signals, a general model for graph data subsuming many existing ones. One variant of FairSpecTemp exploits commutativity properties of graph stationarity while directly constraining bias, while the other implicitly encourages fair estimates by restricting bias in the graph spectrum and is thus more flexible. Our methods enjoy high probability performance bounds, yielding a conditional tradeoff between fairness and accuracy. In particular, our analysis reveals that accuracy need not be sacrificed to recover fair graphs. We evaluate FairSpecTemp on synthetic and real-world data sets to illustrate its effectiveness and highlight the advantages of both variants of FairSpecTemp.
HJCD-IK: GPU-Accelerated Inverse Kinematics through Batched Hybrid Jacobian Coordinate Descent
Yasutake, Cael, Kingston, Zachary, Plancher, Brian
Inverse Kinematics (IK) is a core problem in robotics, in which joint configurations are found to achieve a desired end-effector pose. Although analytical solvers are fast and efficient, they are limited to systems with low degrees-of-freedom and specific topological structures. Numerical optimization-based approaches are more general, but suffer from high computational costs and frequent convergence to spurious local minima. Recent efforts have explored the use of GPUs to combine sampling and optimization to enhance both the accuracy and speed of IK solvers. We build on this recent literature and introduce HJCD-IK, a GPU-accelerated, sampling-based hybrid solver that combines an orientation-aware greedy coordinate descent initialization scheme with a Jacobian-based polishing routine. This design enables our solver to improve both convergence speed and overall accuracy as compared to the state-of-the-art, consistently finding solutions along the accuracy-latency Pareto frontier and often achieving order-of-magnitude gains. In addition, our method produces a broad distribution of high-quality samples, yielding the lowest maximum mean discrepancy. We release our code open-source for the benefit of the community.