Advancing Foundational AI: Amii Research at ICML 2026

Published

Jul 6, 2026

Categories

Insights

Subject Matter

Research

The forty-third annual International Conference on Machine Learning (ICML) is underway this week in Seoul, South Korea. Amii is proud to showcase the diverse and high-imact research our Fellows, Canada CIFAR AI Chairs, and students are presenting this year.

This year, Amii’s are offering groundbreaking advancements in areas like the robust reinforcement learning better able to handle real-world problems, exposing flaws in how we evaluate LLMs and other models, and establishing new approaches to trust and safety in generative AI and protecting sensitive data in AI.

*- denotes Amii affiliation

Accepted Papers

Intentional Updates for Streaming Reinforcement Learning

Arsalan Sharifnassab ⋅ Mohamed Elsayed ⋅ Kris De Asis ⋅ Rupam Mahmood* ⋅ Rich Sutton*

LINK TO PAPER

In gradient-based learning, a step size chosen in parameter units does not produce a predictable per-step change in the function output. This may lead to instability in the streaming setting (i.e., batch size=1), where stochasticity is not averaged out and update magnitudes can momentarily become arbitrarily big or small. Instead, we propose \emph{intentional updates}: first specify the \emph{intended outcome} of an update and then solve for the step size that approximately achieves it. This strategy has precedent in online supervised linear regression via normalized LMS, which selects a step size to yield a specified change in the function output proportional to the current error. We extend this principle to streaming reinforcement learning by defining appropriate intended outcomes: \emph{Intentional TD} aims for a fixed fractional reduction of the current TD error relative to the momentary bootstrap target, and \emph{Intentional Policy Gradient} aims for a bounded per-step change in the policy, limiting local KL divergence. We develop practical implementations integrating eligibility traces and diagonal scaling; our experiments show that these methods yield state-of-the-art streaming performance often comparable to batch and replay-buffer learning.

A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshots

Vincent Guan ⋅ Lazar Atanackovic* ⋅ Kirill Neklyudov

LINK TO PAPER

The population dynamics of molecules, cells, and organisms are governed by a number of unknown internal and external forces. In the last decade, population dynamics have predominately been modeled with Wasserstein gradient flows. However, since gradient flows minimize free energy, they fail to capture important dynamical properties, such as periodicity. In this work, we propose a change in perspective by considering population dynamics that minimize Wasserstein Lagrangian action, rather than free energy. As our main theoretical contributions, we derive the Hamiltonian equations of motion from the principle of least population-level action and we show that these mechanics encompass classical mechanics, quantum mechanics, and gradient flows. We further leverage the Hamiltonian perspective to propose an algorithm that learns the population mechanics from observed marginals, without specifying the Lagrangian. We demonstrate that by directly learning the population mechanics, our method forecasts and interpolates unseen marginals without a reference process, and outperforms gradient flow and flow matching methods across a wide range of real and simulated experiments.

ReVSI: Rebuilding Visual Spatial Intelligence Evaluation for Accurate Assessment of VLM 3D Reasoning

Yiming Zhang ⋅ Jiacheng Chen ⋅ Jiaqi Tan ⋅ Yongsen Mao ⋅ Wenhu Chen ⋅ Angel X Chang*

LINK TO PAPER

Current evaluations of spatial intelligence can be systematically invalid under modern vision-language model (VLM) settings. First, many benchmarks derive question-answer (QA) pairs from point-cloud-based 3D annotations originally curated for traditional 3D perception. When such annotations are treated as ground truth for video-based evaluation, reconstruction and annotation artifacts can miss objects that are clearly visible in the video, mislabel object identities, or corrupt geometry-dependent answers (e.g., size), yielding incorrect or ambiguous QA pairs. Second, evaluations often assume full-scene access, while many VLMs operate on sparsely sampled frames (e.g., 16-64), making many questions effectively unanswerable under the actual model inputs. We improve evaluation validity by introducing ReVSI, a benchmark and protocol that ensures each QA pair is answerable and correct under the model's actual inputs. To this end, we re-annotate object labels and geometry across 413 scenes from 5 datasets to improve data quality, and regenerate all QA pairs with rigorous bias mitigation and human verification using professional 3D visualization and annotation tools. We further enhance evaluation controllability by providing variants across multiple frame budgets (16/32/64/all) and fine-grained object visibility metadata, enabling controlled diagnostic analyses. Evaluations of general and domain-specific VLMs on ReVSI reveal systematic failure modes that are obscured by prior benchmarks, yielding a more reliable and diagnostic assessment of spatial intelligence.

Explicitly Modeling Censoring Produces Superior Survival Predictors

Shi-ang Qi ⋅ Yakun Yu ⋅ Russell Greiner*

LINK TO PAPER

Likelihood-based training is the dominant paradigm in survival prediction. Under independent censoring, we can factorize the likelihood and optimize only the terms related to event modeling, effectively treating the censoring mechanism as incidental. This is justified when censoring is non-informative, i.e., when the censoring process shares no parameters with the event-time model. However, this may not hold in practice, and ignoring censoring contributions may discard useful signals for learning representations that can help to effectively estimate event distributions. Motivated by this, we argue that explicitly modeling censoring can improve representation learning and time-to-event estimation, particularly when event and censoring processes are coupled. We introduce a latent decomposition view that partitions covariates into four disjoint factors: those affecting only the event process, only the censoring process, both, or neither. We then learn decomposed representations for the first three categories to guide a better estimation of the event distribution. We instantiate our method on four popular deep-learning survival models and evaluate on 10 datasets (2 semi-synthetic and 8 real-world), showing consistent gains over strong baselines and multiple SOTA methods.

Position: Stop Chasing the C-index when Evaluating Survival Analysis Models

Christian Marius Lillelund ⋅ Shi-ang Qi ⋅ Russell Greiner* ⋅ Christian Fischer Pedersen

LINK TO PAPER

The current state of evaluation in survival analysis is plagued by the persistent use of evaluation metrics in ways that are misaligned with the stated modeling objective. In addition, many such evaluations are based on censoring assumptions that are left implicit or unjustified. This means that the reported performance can be misleading and may fail to answer the scientific or modeling question the evaluation was intended to address. In this position paper, we present a critical analysis of evaluation practices in survival analysis and highlight why evaluation in survival analysis fundamentally differs from standard regression or classification due to censoring. We place particular focus on concordance-based measures, such as the C-index, which our findings indicate are heavily overused in the literature. To help identify appropriate metrics, we propose a set of key desiderata and introduce a double-helix ladder, in which valid evaluation requires alignment between metric and modeling assumptions, and we provide empirical evidence that this is effective. We conclude by providing practical guidance on how to evaluate a survival model.

Learn from A Rationalist: Distilling Intermediate Interpretable Rationales

Jiayi Dai ⋅ Randy Goebel*

LINK TO PAPER

Because of the pervasive use of deep neural networks (DNNs), especially in high-stakes domains, the interpretability of DNNs has received increased attention. The general idea of rationale extraction (RE) is to provide an interpretable-by-design framework for DNNs via a select-predict architecture where two neural networks learn jointly to perform feature selection and prediction, respectively. Given only the remote supervision from the final task prediction, the process of learning to select subsets of features (or rationales) requires searching in the space of all possible feature combinations, which is computationally challenging and even harder when the base neural networks are not sufficiently capable. To improve the predictive performance of RE models that are based on less capable or smaller neural networks (i.e., the students), we propose REKD (Rationale Extraction with Knowledge Distillation) where a student RE model learns from the rationales and predictions of a teacher (i.e., a rationalist) in addition to the student's own RE optimization. This structural adjustment to RE aligns well with how humans could learn effectively from interpretable and verifiable knowledge. Because of the neural-model agnostic nature of the method, any black-box neural network could be integrated as a backbone model. To demonstrate the viability of REKD, we conduct experiments with multiple variants of BERT and vision transformer (ViT) models. Our experiments across language and vision classification datasets (i.e., IMDB movie reviews, CIFAR 10 and CIFAR 100) show that REKD significantly improves the predictive performance of the student RE models.

RSA-CP: Efficient Conformal Prediction in Small-Sample Regimes via Random Score Alignment

Pankaj Bhagwat ⋅ Zhixian Yang ⋅ yihao wang ⋅ Bei Jiang* ⋅ Linglong Kong*

LINK TO PAPER

Conformal Prediction (CP) provides rigorous finite-sample coverage guarantees, yet its statistical efficiency hinges critically on the size of the calibration set. In data-scarce regimes, CP often suffers from volatile quantile estimation, leading to overly conservative and wide prediction intervals. To address this, we propose Random Score Alignment-Conformal Prediction (RSA-CP), a simple framework designed to improve sample efficiency in small-sample CP. Instead of requiring the computationally intensive generation of full synthetic datasets, RSA-CP enhances calibration by directly aligning real scores with a high-resolution reference score distribution. By employing an optimal transport mapping, our framework refines "step-like" quantile increments through a globally optimal use of reference information. We provide theoretical guarantees establishing that RSA-CP maintains robust coverage without any distributional assumptions on the reference scores. Empirical evaluations demonstrate that RSA-CP consistently produces shorter and more precise prediction intervals while maintaining finite-sample coverage guarantees. Overall, RSA-CP offers a computationally efficient and theoretically grounded solution for robust uncertainty quantification under limited data.

One-Shot Weighted Ensemble Estimation for Federated Quantile Regression: Optimal Statistical Guarantees under Heterogeneous Structured Data

Guang Yang ⋅ Bo Pan ⋅ Chengdi Lian ⋅ Xingcai Zhou ⋅ Linglong Kong* ⋅ Yafei Wang ⋅ Bei Jiang*


LINK TO PAPER

Federated Quantile Regression (FQR) has emerged as a powerful modelling paradigm for estimating conditional quantiles, offering a more comprehensive understanding of response distributions than standard conditional mean regression. However, achieving communication efficiency and optimal statistical guarantees for FQR remains challenging, particularly due to the nonsmooth nature of quantile loss functions and the presence of heterogeneously structured data, where each local agent trains its conditional quantile models with distinct sets of features. In this paper, we propose a data-driven, one-shot weighted ensemble estimator for FQR that incorporates scalable weighting schemes to effectively leverage the partially observed features at each local agent, thereby enjoying both communication efficiency and estimation optimality. Theoretically, we present a unified analysis of the proposed learning procedure, establishing that the resulting estimator exhibits asymptotic normality and attains uniformly minimum variance. Furthermore, we investigate the estimator's sensitivity to perturbations introduced by local agents and derive conditions under which the estimator achieves stability and enjoys strong out-of-sample generalization. Extensive simulations and real data analysis under various scenarios validate the asymptotic normality of our estimator and demonstrate its superior estimation accuracy and uniform convergence compared to several baseline methods across a range of quantile levels.

Private and Stable Test-time Adaptation with Differential Privacy

Zefeng Li ⋅ Qiaoyue Tang ⋅ Mathias Lécuyer* ⋅ Evan Shelhamer

LINK TO PAPER

Test-time adaptation (TTA) can reduce error on new and different data by updating the model on these inputs during inference. However, these updates raise the issue of privacy w.r.t. the testing data, because the model parameters now depend on all past inputs. To control this privacy risk, we cast multiple popular TTA methods (Tent, EATA, SAR, DeYO, and COME) into differential privacy (DP) forms that apply per-sample gradient clipping and Gaussian noise for all updates. On ImageNet-C, our DP-TTA methods provide adequate privacy at small cost to accuracy, and in the low-privacy regime the clipping mechanism of DP even improves the accuracy and stability of adaptation in the continual setting. These improvements to privacy and accuracy come at only only modest computational overhead. These first results on private TTA raise awareness of the issue, inform the development of more private test-time updates, and identify per-sample clipping as a reliable technique for improving the accuracy and stability of adaptation.

Private and Stable Test-time Adaptation with Differential Privacy

Zefeng Li ⋅ Qiaoyue Tang ⋅ Mathias Lécuyer* ⋅ Evan Shelhamer

LINK TO PAPER

Test-time adaptation (TTA) can reduce error on new and different data by updating the model on these inputs during inference. However, these updates raise the issue of privacy w.r.t. the testing data, because the model parameters now depend on all past inputs. To control this privacy risk, we cast multiple popular TTA methods (Tent, EATA, SAR, DeYO, and COME) into differential privacy (DP) forms that apply per-sample gradient clipping and Gaussian noise for all updates. On ImageNet-C, our DP-TTA methods provide adequate privacy at small cost to accuracy, and in the low-privacy regime the clipping mechanism of DP even improves the accuracy and stability of adaptation in the continual setting. These improvements to privacy and accuracy come at only only modest computational overhead. These first results on private TTA raise awareness of the issue, inform the development of more private test-time updates, and identify per-sample clipping as a reliable technique for improving the accuracy and stability of adaptation.

Jake Tuero ⋅ Michael Buro ⋅ Levi Lelis* ⋅ Laurent Orseau

LINK TO PAPER

Subgoal-based policy tree search, which uses a policy to guide search, is effective for complex single-agent deterministic problems but often relies on explicit subgoal generation that can incur substantial overhead and hinders scalability. In this paper, we overcome these limitations by using a learned ``rerooter'' through the recently-introduced algorithm. A rerooter implicitly decomposes the problem into soft subtasks. While previous work focused on the formal guarantees for given or handcrafted rerooters, in this work we propose three rerooter designs: (i) a clustering-based rerooter that exploits global state-space structure, (ii) a heuristic-based rerooter that leverages learned cost-to-go estimates, and (iii) a hybrid that combines both signals. Our framework avoids having to explicitly reconstruct and reason over generated subgoals, thereby enabling scalable allocation of search effort with significantly lower computational overhead. Empirically, our rerooting-based methods scale to complex environments where subgoal-based policy tree search fails, and achieve state-of-the-art online training efficiency on the domains tested.

Revisiting OOD Generalization in Programmatic RL

Amirhossein Rajabpour ⋅ Kiarash Aghakasiri ⋅ Sandra Zilles* ⋅ Levi Lelis*

LINK TO PAPER

Programmatic policies are often reported to generalize better than neural policies in reinforcement learning (RL) benchmarks. We revisit some of these claims and show that much of the observed gap arises from uncontrolled experimental factors rather than intrinsic representational reasons. Re-evaluating three core benchmarks used in influential papers---TORCS, Karel, and Parking---we find that neural policies, when trained with a few modifications, such as sparse observations and cautious intrinsic reward functions, can match or exceed the out-of-distribution (OOD) generalization of programmatic policies. We argue that a representation enables OOD generalization if (i) the policy space it induces includes a generalizing policy and (ii) the search algorithm can find it. The neural and programmatic policies in prior work are comparable in OOD generalization because the domain-specific languages used induce policy spaces similar to those of neural networks, and our modifications help the gradient search find generalizing solutions. By disentangling representational factors from experimental confounds, we advance our understanding of what makes a representation succeed or fail at OOD generalization.

Reasoning Models Are Test Exploiters: Rethinking Multiple Choice

Narun Raman ⋅ Taylor Lundy ⋅ Kevin Leyton-Brown*

LINK TO PAPER

When evaluating Large Language Models (LLMs) in question-answering domains, multiple-choice question answering (MCQA) is widely used because it enables automatic grading. However, MCQA also exposes models to answer options that can be exploited in ways that inflate reasoning ability. We study this phenomenon across question-answering benchmarks and 27 LLMs by systematically varying how and when models are exposed to answer options. For non-reasoning LLMs, MCQA can remain a good proxy for free-text performance when any chain-of-thought is produced only before the options are revealed. However, this "decoupled" format is not realizable for most reasoning models: they are designed to emit reasoning tokens whenever they are prompted, so if options are present they inevitably "reason over" the options. In practice, this makes reasoning models particularly effective at extracting signal from options, and can create large, misleading gains over free-text baselines. To characterize how models exploit MCQA, we introduce diagnostic probes that isolate option-only and question-plus-option exploitation pathways, and we quantify how design choices such as distractor strength and "none-of-the-above" answers effect exploitability. Finally, we examined the practice of multiple choice as an error diagnostic: inferring a model's mistake from the wrong option it picks. On benchmarks where reasoning can be expressed as code, we ask models to output code, we then executed it varying the inputs, and compared the resulting input–output behavior, revealing failure modes that MCQA diagnostics obscure. Lastly, we offer practical guidelines when analyzing results from MCQA that better reflect LLMs' genuine reasoning capabilities.

Show less

Position: RL Researchers Need to Distinguish Between Solving Simulators and Using Simulators as a Proxy

Matthew Vandergrift ⋅ Esraa Elelimy ⋅ Martha White*

LINK TO PAPER

One goal in reinforcement learning (RL) research is to understand general purpose sequential decision-making, using benchmark simulators as a proxy for learning in a deployment setting. When running experiments, however, the goal of achieving high performance in the simulator can mutate into focusing exclusively on solving the simulator. To achieve high scores researchers may adopt solutions exclusively meant for solving simulators, rather than learning while the agent is deployed outside of a simulator. Solving simulators is also worthy of investigation, but is a fundamentally different RL research question. In this paper we argue that RL researchers need to distinguish between two uses cases of simulators: solving simulators and using simulators as a proxy for learning in deployment. We first discuss how these two use-cases are importantly different, in terms of constraints on how the agent can use the simulator, which algorithms are appropriate and which evaluation metrics are appropriate. We then highlight several issues and misleading conclusions that can occur by not making the distinction between these two settings clear, supported with examples and simple experiments. This work is a call to the community to begin clearly distinguishing how they are using simulators in their work, hopefully sparking further discussion on which empirical practices work best in each setting.

Accelerating Q-learning through Efficient Value-sharing across Actions

Prabhat Nagarajan ⋅ Brett Daley ⋅ Martha White*⋅ Marlos C. Machado*

LINK TO PAPER

Learning action-values efficiently is central to reinforcement learning (RL), as they underpin many control algorithms such as Q-learning. However, action-value learning can be slow, requiring many updates to move values from their initialization, typically near zero, to their true values, which may be far from zero. Moreover, action-value learning algorithms typically update each state–action pair independently, without learning shared value structure across actions within a state. In this paper, we address these inefficiencies by introducing the mean-expansion transformation, which accelerates action-value learning by sharing values across actions within a state and by changing the problem from directly learning potentially large action-values to learning a lower-norm representation of them. In deep RL, this transformation can be applied as a parameter-free modification to Q-network architectures without altering the underlying algorithm. Empirically, we show that it improves DQN's performance in aggregate across 57 Atari games while increasing action gaps and dramatically reducing value overestimation.

Laplacian Representations for Decision-Time Planning

Dikshant Shehmar ⋅ Matthew Schlegel ⋅ Matthew Taylor* ⋅ Marlos C. Machado*

LINK TO PAPER

Planning with a learned model remains a key challenge in model-based reinforcement learning~(RL). In decision-time planning, state representations are critical as they must support local cost computation while preserving long-horizon structure. In this paper, we show that the Laplacian representation provides an effective latent space for planning by capturing state-space distances at multiple time scales. This representation preserves meaningful distances and naturally decomposes long-horizon problems into subgoals, also mitigating the compounding errors that arise over long prediction horizons. Building on these properties, we introduce ALPS, a hierarchical planning algorithm, and demonstrate that it outperforms commonly used baselines on a selection of offline goal-conditioned RL tasks from OGBench, a benchmark previously dominated by model-free methods.

The Safety-Aware Denoiser for Text Diffusion Models

Amman Yusuf ⋅ Zhejun Jiang ⋅ Mi Jung Park*

LINK TO PAPER

Recent work on text diffusion models offers a promising alternative to autoregressive generation, but controlling their safety remains underexplored. Existing safety approaches are geared toward autoregressive models and typically rely on post-hoc filtering or inference-time interventions. These are inadequate for effectively addressing safety risks in text diffusion models. We propose the Safety-Aware Denoiser (SAD), a safety-guidance framework in text diffusion models. The SAD modifies the iterative denoising process such that the text sample at the final denoising step is steered toward provably safe regions of the text space. This inference-time method can integrate safety constraints into the denoiser, avoiding computationally expensive retraining of the underlying diffusion model and enabling flexible, lightweight safety guidance. We evaluate the safety of the generated text using the SAD, with respect to hazard taxonomy, memorization, and jailbreak. Experimental results show that SAD substantially reduces unsafe generations while preserving generation quality, diversity, and fluency, outperforming existing methods. These results demonstrate that our safety guidance during denoising provides an effective and scalable mechanism for enforcing safety in text diffusion models.

Towards Parameter-Free Temporal Difference Learning

Yunxiang LI ⋅ Mark Schmidt* ⋅ Reza Babanezhad ⋅ Sharan Vaswani

LINK TO PAPER

Temporal difference (TD) learning is a fundamental algorithm for estimating value functions in reinforcement learning. Recent finite-time analyses of TD with linear function approximation quantify its theoretical convergence rate. However, they often require setting the algorithm parameters using problem-dependent quantities that are difficult to estimate in practice --- such as the minimum eigenvalue of the feature covariance (w) or the mixing time of the underlying Markov chain (Tmix). In addition, some analyses rely on nonstandard and impractical modifications, exacerbating the gap between theory and practice. To address these limitations, we use an exponential step-size schedule with the standard TD(0) algorithm. We analyze the resulting method under two sampling regimes: independent and identically distributed (i.i.d.) sampling from the stationary distribution, and the more practical Markovian sampling along a single trajectory. In the i.i.d. setting, the proposed algorithm does not require the knowledge of problem-dependent quantities such as , and attains the optimal bias-variance trade-off for the last iterate. In the Markovian setting, we propose a regularized TD(0) algorithm with an exponential step-size schedule. The resulting algorithm achieves a comparable convergence rate to prior works, without requiring projections, iterate averaging, or knowledge of or .Temporal difference (TD) learning is a fundamental algorithm for estimating value functions in reinforcement learning. Recent finite-time analyses of TD with linear function approximation quantify its theoretical convergence rate. However, they often require setting the algorithm parameters using problem-dependent quantities that are difficult to estimate in practice --- such as the minimum eigenvalue of the feature covariance () or the mixing time of the underlying Markov chain (). In addition, some analyses rely on nonstandard and impractical modifications, exacerbating the gap between theory and practice. To address these limitations, we use an exponential step-size schedule with the standard TD(0) algorithm. We analyze the resulting method under two sampling regimes: independent and identically distributed (i.i.d.) sampling from the stationary distribution, and the more practical Markovian sampling along a single trajectory. In the i.i.d. setting, the proposed algorithm does not require the knowledge of problem-dependent quantities such as , and attains the optimal bias-variance trade-off for the last iterate. In the Markovian setting, we propose a regularized TD(0) algorithm with an exponential step-size schedule. The resulting algorithm achieves a comparable convergence rate to prior works, without requiring projections, iterate averaging, or knowledge of Tmix or w.

Flatland: The Adventures of Gradient Descent with Large Step Sizes

Leonardo Galli ⋅ Curtis Fox ⋅ Wiebke Bartolomaeus ⋅ Mark Schmidt* ⋅ Holger Rauhut

LINK TO PAPER

The training of neural networks often entails objective functions that are not globally -smooth. For these functions, it is both theoretically and practically difficult to reply to the question: what is the largest possible step size that ensures the convergence of gradient descent (GD)? We address this longstanding open question in deep learning by providing a unifying definition of "large'' step sizes that requires only local Lipschitz (or even Hölder) continuity of the gradient. We design first-order adaptive methods that provably yield large step sizes and show that they operate at the edge of stability (EoS) right from the start of the training. In particular, the loss decreases nonmonotonically and the product between the step size and sharpness, i.e., the largest eigenvalue of the hessian, stays above the EoS threshold of 2 throughout training. Using our method, we are also able to minimize the sharpness all the way down to its global minimum. Contrary to expectation, we find that encountering globally-flat regions too early in the training may both slow down convergence and jeopardize the generalization ability of the network. Exploiting a self-stabilization argument, we allow GD to enter slightly sharper valleys and turn unsuccessful training runs into very successful ones.

Sequential Kernel-based Conditional Independence Testing via Adaptive Betting

Zheng He ⋅ Danica J Sutherland*

LINK TO PAPER

Testing conditional independence is a fundamental yet inherently difficult challenge, as controlling Type I error is impossible in general. The recently popular "Model-X" paradigm offers a solution by relying on a perfectly known conditional distribution. In traditional "one-shot" testing regimes, slight deviations from perfect knowledge are sometimes allowable, but existing work in more realistic online settings has required exact adherence to Model-X. We propose a new approach for sequential testing of conditional independence that is far more robust to estimation errors in the conditional distribution. Our method, based on online optimization of the Kernel Conditional Independence statistic, introduces a novel normalization and "truncate-and-shift" calibration strategy to the testing-by-betting paradigm. This framework greatly improves validity with estimated conditionals while still providing high power across high-dimensional synthetic benchmarks and real-world fairness tasks.

Robust AI Evaluation through Maximal Lotteries

Hadi Khalaf ⋅ Flavio Calmon ⋅ Daniel Halpern ⋅ Ariel Procaccia ⋅ Itai Shapira ⋅ Serena Wang*

LINK TO PAPER

The standard way to evaluate language models on subjective tasks is through pairwise comparisons: an annotator chooses the "better" of two model responses for a given prompt. These comparisons are then aggregated into a single ranking via the Bradley–Terry (BT) framework, forcing heterogeneous preferences into a total order and violating basic social-choice desiderata. In contrast, social choice theory provides an alternative approach called maximal lotteries, which aggregates pairwise preferences without imposing any assumptions on their structure. However, we show that maximal lotteries can be highly sensitive to heterogeneity among annotators and across prompts. We introduce robust lotteries, which optimize worst-case performance under plausible shifts in the preference data. On large-scale preference datasets, robust lotteries achieve more reliable win rate guarantees across the annotator distribution and recover a stable set of top performing models.

Pluralistic Leaderboards

Nika Haghtalab ⋅ Ariel Procaccia ⋅ Han Shao ⋅ Serena Wang* ⋅ Kunhe Yang

LINK TO PAPER

Recent leaderboard-based evaluations of large language models aggregate user feedback by fitting a Bradley--Terry model to pairwise comparisons, producing a single global ranking based on a latent quality score. While appealing for its simplicity, this approach is incompatible with heterogeneous preferences: when LLMs are used across diverse tasks and use cases, users who favor fundamentally different model behaviors can be systematically misrepresented when collapsed into a single quality score. To address this issue, we study \emph{pluralistic leaderboards} that aim to remain \emph{stable} with respect to heterogeneous user populations. Drawing on ideas from social choice theory, we adapt the notion of \emph{local stability}, which requires that no model outside the top-k positions is collectively preferred to the top-k set by more than O(1/k) fraction of users. Building on techniques from the social choice literature, we design an alternative leaderboard mechanism that satisfies local stability while eliciting only O˜(k) pairwise comparisons per user, where k is the size of the prefix for which stability is guaranteed. Using data from LMArena, we show that standard Bradley--Terry aggregation can violate local stability in practice, whereas our method provides substantially stronger stability guarantees.

REVIS: Sparse Latent Steering to Mitigate Object Hallucination in Large Vision-Language Models

Jialin Wu ⋅ Wei Shi ⋅ Han Shen ⋅ Peigui Qi ⋅ Kunsheng Tang ⋅ Zhicong Huang ⋅ Binghao Wang ⋅ Zhou Yang*

LINK TO PAPER

Despite the advanced capabilities of Large Vision-Language Models (LVLMs), they frequently suffer from object hallucination. One reason is that visual features and pretrained textual representations often become intertwined in the deeper network layers. To address this, we propose REVIS, a training-free framework designed to explicitly re-activate this suppressed visual information. Rooted in latent space geometry, REVIS extracts the pure visual information vector via orthogonal projection and employs a calibrated strategy to perform sparse intervention only at the precise depth where suppression occurs. This surgical approach effectively restores visual information with minimal computational cost. Empirical evaluations on standard benchmarks demonstrate that REVIS reduces object hallucination rates by approximately 19% compared to state-of-the-art baselines, while preserving general reasoning capabilities.

PrivCode++ : Latent-Conditioned Differentially Private Code Generation for Comprehensive Guarantees

Zheng Liu ⋅ Chen GONG ⋅ Terry Yue Zhuo ⋅ Zhou Yang* ⋅ Kecen Li ⋅ Wenlong Meng ⋅ Xinwen Hou ⋅ Yu Liu ⋅ Xiaochen Li

LINK TO PAPER

Large language models fine-tuned on instruction–code pairs may memorize and subsequently leak sensitive training data. Existing differentially private (DP) code generation methods primarily protect code snippets while assuming prompts are public, which fails in realistic scenarios where prompts may also contain sensitive information. When prompts cannot be explicitly learned or used during generation, code synthesis suffers from severe utility degradation and reduced diversity. To address these challenges, we propose PrivCode++, the first work to explore DP code generation under where both prompts and code snippets are considered sensitive in LLM fine-tuning. PrivCode++ introduces a two-stage DP framework with a Privacy-Free Latent Conditioning module, enabling effective DP fine-tuning and data synthesis without direct access to sensitive prompts or code. Extensive experiments show that PrivCode++ achieves substantially higher utility than baselines, remains competitive with the method with relaxing privacy assumptions, and provides stronger privacy guarantees.

Controlled Collaboration Geometry for Personalized Federated Learning

Hongbo Yin ⋅ Wu Jichun ⋅ Zhou Yang* ⋅ Chi Jiang ⋅ Yin Zhang ⋅ Yan Zhang

LINK TO PAPER

In personalized federated learning (PFL), collaboration graphs specify model aggregation among clients. However, without constraints on the collaboration geometry, training can drift into two degenerate regimes: global consensus or spontaneous clustering. This paper provides a unified dynamical analysis: under the same budget of representative models, collaborative PFL is more expressive and achieves higher-order approximation accuracy than clustered PFL. An upper bound on disagreement further reveals two degeneration mechanisms—overly strong collaboration drives consensus (reducing to standard federated learning), while similarity-driven weight updates make the graph nearly reducible and induce self-clustering (collapsing to clustered PFL). Motivated by these findings, we propose pFedCCG. pFedCCG preserves the expressivity advantage via controlled collaboration geometry (CCG): it builds a static similarity-based collaboration template decoupled from training, optimizes a Markovian collaboration matrix with a prescribed stationary distribution via reversible parameterization and Euclidean projection, and schedules collaboration strength to avoid self-clustering. Experiments across diverse heterogeneity settings show consistent personalization gains and markedly reduced collapse and self-clustering. Code will be available at https://anonymous.4open.science/r/pFedCCG-CB88.

Watermarking LLM Agent Trajectories

Wenlong Meng ⋅ Chen GONG ⋅ Terry Yue Zhuo ⋅ Fan Zhang ⋅ Kecen Li ⋅ Zheng Liu ⋅ Zhou Yang* ⋅ Chengkun Wei ⋅ Wenzhi CHEN

LINK TO PAPER

LLM agents rely heavily on high-quality trajectory data to guide their problem-solving behaviors, yet producing such data requires substantial task design, high-capacity model generation, and manual filtering. Despite the high cost of creating these datasets, existing literature has overlooked copyright protection for LLM agent trajectories. This gap leaves creators vulnerable to data theft and makes it difficult to trace misuse or enforce ownership rights. This paper introduces ActHook, the first watermarking method tailored for agent trajectory datasets. Inspired by hook mechanisms in software engineering, ActHook embeds hook actions that are activated by a secret input key and do not alter the original task outcome. Like software execution, LLM agents operate sequentially, allowing hook actions to be inserted at decision points without disrupting task flow. When the activation key is present, an LLM agent trained on watermarked trajectories can produce these hook actions at a significantly higher rate, enabling reliable black-box detection. Experiments on mathematical reasoning, web searching, and software engineering agents show that ActHook achieves an average detection AUC of 94.3 on Qwen-2.5-Coder-7B while incurring negligible performance degradation.

Workshop Papers

A Systematic Investigation of RL-Jailbreaking in LLMs

Montaser Mohammedalamen* ⋅ Kevin Roice* ⋅ Reginald McLean* ⋅ Alyssa Lefaivre Škopac*

LINK TO PAPER

LLM agents rely heavily on high-quality trajectory data to guide their problem-solving behaviors, yet producing such data requires substantial task design, high-capacity model generation, and manual filtering. Despite the high cost of creating these datasets, existing literature has overlooked copyright protection for LLM agent trajectories. This gap leaves creators vulnerable to data theft and makes it difficult to trace misuse or enforce ownership rights. This paper introduces ActHook, the first watermarking method tailored for agent trajectory datasets. Inspired by hook mechanisms in software engineering, ActHook embeds hook actions that are activated by a secret input key and do not alter the original task outcome. Like software execution, LLM agents operate sequentially, allowing hook actions to be inserted at decision points without disrupting task flow. When the activation key is present, an LLM agent trained on watermarked trajectories can produce these hook actions at a significantly higher rate, enabling reliable black-box detection. Experiments on mathematical reasoning, web searching, and software engineering agents show that ActHook achieves an average detection AUC of 94.3 on Qwen-2.5-Coder-7B while incurring negligible performance degradation.

A Systematic Investigation of RL-Jailbreaking in LLMs

Reginald McLean* ⋅ Tabitha Edith Lee Montaser Mohammedalamen* ⋅ Kevin Roice* ⋅ Glen Berseth Patrick Pilarski*Marlos C. Machado* Alyssa Lefaivre Škopac* ⋅ Benjamin Rosman

LINK TO PAPER

Abstract:

With the rapid advancement and deployment of Agentic AI, our scientific understanding of capabilities and limitations has not kept pace, leading to cases where AI agents cause harm. We argue that many of these safety limitations are not novel problems. Instead, the safety challenges currently facing AI agents can be seen as instances of problems the reinforcement learning (RL) community has studied rigorously for decades. The core of this argument concerns the problem formulation of AI agents. AI agents are designed to solve sequential decision-making problems: problems with long-term objectives in which actions have delayed consequences. To model these types of problem, the problem is set up the problem such that the agent receives observations, feedback on its progress, and then takes actions. This is precisely the formulation of the RL problem. In this paper, we formalize the problem equivalence, which we then leverage to argue that \textbf{AI Agent safety is a reinforcement learning problem: the failure modes currently observed in deployed AI agents are structural instances of problems RL has formalized for decades, and the RL safety literature provides principled tools to diagnose and address them.}. We conclude with a call for deliberate collaboration between the RL and AI agent research communities: AI agent researchers gain access to principled frameworks, while RL researchers gain a class of real-world problems that could expose fundamental gaps in current RL benchmarks and theory.


Want to stay up-to-date on the latest research from the Amii community? Sign up for our monthly newsletter!

Share