ICML 2026 Papers

International Conference on Machine Learning

Decision breakdown

6,341papers

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

  • Spotlight
    5368.5%
  • Poster
    5,80591.5%

Browse all papers

Topic: Social aspects · accountability transparency and interpretability · 541 papers · Reset filters

$\texttt{PRISM}$:A 3D Probabilistic Neural Representation for Interpretable Shape Modeling

Yining Jiao, Sreekalyani Bhamidi, Carlton Jude ZDANSKI, Julia S Kimbell, Andrew Prince, Cameron P Worden, Samuel Kirse, Christopher Rutter, Benjamin H Shields, William Alexander Dunn, Jisan Mahmud, Marc Niethammer

Expand abstract

Understanding how anatomical shapes evolve in response to developmental covariates—and quantifying their spatially varying uncertainties—is critical in healthcare research. Existing approaches typically rely on global time-warping formulations that ignore spatially heterogeneous dynamics. We introduce \texttt{PRISM}, a novel framework that bridges implicit neural representations with uncertainty-aware statistical shape analysis. \texttt{PRISM} models the conditional distribution of shapes given covariates, providing spatially continuous estimates of both the population mean and covariate-dependent uncertainty at arbitrary locations. A key theoretical contribution is a closed-form Fisher Information metric that enables efficient, analytically tractable local temporal uncertainty quantification via automatic differentiation. Experiments on three synthetic datasets and one clinical dataset demonstrate \texttt{PRISM}'s strong performance across diverse tasks—from modeling shape evolution to anomaly detection—within a unified framework, while providing interpretable and clinically meaningful uncertainty estimates.

Poster
Rotation-Invariant Spherical Watermarking via Third-Order SO(3) Representation Coupling

Pengzhen Chen, Yanwei Liu, Xiaoyan Gu, Antonios Argyriou, Wu Liu, Weiping Wang

Expand abstract

Reliable watermarking of panoramic imagery is fundamentally challenged by arbitrary 3D rotations. As panoramas are defined on the sphere, they naturally transform under the action of $SO(3)$, rendering conventional planar representations and augmentation-based robustness strategies inadequate and devoid of theoretical guarantees. To address this, we formulate panoramas as spherical signals and leverage $SO(3)$ representation theory to derive provably rotation-invariant descriptors. While spherical harmonic coefficients transform equivariantly under rotations, the natural invariant constructions are typically limited to zeroth-order statistics which eliminate directional information and severely constrain embedding capacity. In this work, we introduce a principled third-order invariant construction by coupling higher-order $SO(3)$ irreducible representations via tensor products and projecting onto the trivial representation. This yields a spherical invariant bispectrum that preserves phase information while remaining strictly rotation-invariant. Leveraging this property, we embed watermarks into higher-order spherical harmonic coefficients and recover them from invariant bispectral scalars, enabling reliable extraction under arbitrary 3D rotations. We provide a theoretical proof of $SO(3)$ invariance for it and demonstrate experimentally its near-perfect robustness to continuous rotations while maintaining high visual fidelity.

Poster
Rationality Measurement and Theory for Reinforcement Learning Agents

Kejiang Qian, Amos Storkey, Fengxiang He

Expand abstract

This paper proposes a suite of rationality measures and associated theory for reinforcement learning agents, a property increasingly critical yet rarely explored. We define an action in deployment to be perfectly rational if it maximises the hidden true value function in the steepest direction. The expected value discrepancy of a policy's actions against their rational counterparts, culminating over the trajectory in deployment, is defined to be expected rational risk; an empirical average version in training is also defined. Their difference, termed as rational risk gap, is decomposed into (1) an extrinsic component caused by environment shifts between training and deployment, and (2) an intrinsic one due to the algorithm's generalisability in a dynamic environment. They are upper bounded by, respectively, (1) the $1$-Wasserstein distance between transition kernels and initial state distributions in training and deployment, and (2) the empirical Rademacher complexity of the value function class. Our theory suggests hypotheses on the benefits from regularisers (including layer normalisation, $\ell_2$ regularisation, and weight normalisation) and domain randomisation, as well as the harm from environment shifts. Experiments are in full agreement with these hypotheses. The code is available at \href{https://github.com/EVIEHub/Rationality}{https://github.com/EVIEHub/Rationality}.

Poster
SecCodePRM: A Process Reward Model for Code Security

Weichen Yu, Ravi Mangal, Yinyi Luo, Kai Hu, Jingxuan He, Corina S. Pasareanu, Matt Fredrikson

Expand abstract

Large Language Models are rapidly becoming core components of modern software development workflows, yet ensuring code security remains challenging. Existing vulnerability detection pipelines either rely on static analyzers or use LLM/GNN-based detectors trained with coarse program-level supervision. Both families often require complete context, provide sparse end-of-completion feedback, and can degrade as code length grows, making them ill-suited for real-time, prefix-level assessment during interactive coding and streaming generation. We propose \textbf{SecCodePRM}, a security-oriented process reward model that assigns a \textbf{context-aware}, \textbf{step-level} security score along a code trajectory. To train the model, we derive step-level supervision labels from static analyzers and expert annotations, allowing the model to attend more precisely to fine-grained regions associated with inter-procedural vulnerabilities. SecCodePRM has three applications: full-code vulnerability detection (VD), partial-code VD, and secure code generation (CG). For VD, SecCodePRM uses risk-sensitive aggregation that emphasizes high-risk steps; for CG, SecCodePRM supports inference-time scaling by ranking candidate continuations and favoring higher cumulative reward. This design yields dense, real-time feedback that scales to long-horizon generation. Empirically, SecCodePRM outperforms prior approaches in all three settings, while preserving code functional correctness, suggesting improved security without a safety–utility tradeoff.

Poster
Alignment Risks from Capability-Seeking RL Training

Yujun Zhou, Yue Huang, Han Bao, Kehan Guo, Zhenwen Liang, Pin-Yu Chen, Tian Gao, Werner Geyer, Nuno Moniz, Nitesh V Chawla, Xiangliang Zhang

Expand abstract

While most AI alignment research focuses on preventing models from generating explicitly harmful content, a more subtle risk arises from capability-seeking RL training in vulnerable environments. We investigate whether language models, when trained with reinforcement learning (RL) in environments with implicit loopholes, can learn to exploit these flaws to maximize reward, even without being explicitly instructed to do so. To test this, we design a suite of four diverse "vulnerability games'', each presenting a structural vulnerability related to context-conditional compliance, proxy metrics, reward tampering, and self-evaluation. Our experiments show that models often learn to exploit these vulnerabilities, discovering opportunistic strategies that increase reward while sometimes preserving or even improving standard task-performance metrics. More critically, we find that these exploitative strategies are not always narrow "tricks'': they can transfer in structured but limited ways, propagate from a capable teacher model to other student models through SFT, and in several cases remain more persistent when learned through RL than when distilled through SFT. Our findings show that alignment risks from capability-seeking RL training can be difficult to detect with standard performance monitoring, suggesting that future AI safety work should extend beyond content moderation to auditing and securing training environments, reward mechanisms, and evaluation channels. Code is available at https://github.com/YujunZhou/Capability-seeking-RL-risk.

Poster
Robust Linear Dueling Bandits with Post-serving Context under Unknown Delays and Adversarial Corruptions

Youngmin Oh

Expand abstract

We study linear dueling bandits in volatile environments characterized by the simultaneous presence of post-serving contexts, delayed feedback, and adversarial corruption. Feedback is subject to unknown stochastic or adversarial delays and a cumulative corruption budget $\mathcal{C}$. To address these challenges, we propose RCDP-UCB, which integrates a learned approximator that predicts post-serving contexts from pre-serving information. It further employs an adaptive weighting strategy that clips feature vectors to mitigate the impact of corrupted and delayed observations simultaneously. Under standard regularity conditions and a parametric post-serving mapping, we rigorously establish that our algorithm is delay-regime-agnostic, achieving a regret upper bound of $\widetilde{\mathcal{O}}(d(\sqrt{T} + \mathcal{C} + \mathcal{D}))$, where $d$ is the total feature dimension and $\mathcal{D}$ encapsulates the delay complexity, scaling with $\sqrt{\Lambda}$ under adversarial delays or $\mu_{\tau}$ under stochastic delays ($\Lambda$: cumulative delay budget; $\mu_{\tau}$: mean of sub-Gaussian delays). We further establish lower bounds that nearly match our upper bounds up to a $\sqrt{d}$ factor for adversarial delays in the absence of post-serving contexts. Code is available at \url{https://github.com/youngmin0oh/rcdp-public}.

Poster
SOPE: Situation-Aware and Statistically Indistinguishable Privacy Exfiltration for MCP-enabled Agents

Ruixiao Lin, Qingming Li, Jiahao Chen, Chunyi Zhou, Shouling Ji

Expand abstract

The Model Context Protocol (MCP) enables Large Language Model (LLM) agents to interact with external tools, but this extensibility introduces significant supply chain vulnerabilities that enable covert privacy exfiltration. Prior studies have revealed privacy leakage in MCP-enabled agents via indirect prompt injection; however, existing attacks are typically misaligned with the agent's tool-usage context and rely on rigid templates, resulting in recognizable patterns that are readily flagged by existing defenses. In this work, we exploit the observation that privacy exposure is inherently scenario-dependent, to associate certain privacy items with specific tools. We introduce `SOPE`, a Scenario-aware and zerO-click Privacy Exfiltration framework that transforms any benign MCP server into its privacy-exfiltrating variants. `SOPE` (1) identifies privacy items that are appropriate to the tool usage, (2) embeds privacy-probing instructions into tool-invocation prompts, and (3) achieves zero-click data transmission via code-level modifications. We evaluate `SOPE` across 27,216 test cases, where 324 `SOPE`-transformed *real-world* servers attacking four benchmark and three commercial agents with *nine* state-of-the-art defenses. Results demonstrate that `SOPE` remains highly effective and robust, highlighting critical protocol-level safety gaps in the agent ecosystem.

Poster
A Decision-Theoretic View of Test-Time Training: When, How Far, and Which Directions to Adapt

Tomoya Wakayama

Expand abstract

Test-time training (TTT) adapts a pretrained model to each prompt via parameter updates, improving accuracy under pretraining-to-test distribution shifts. Yet, its performance often suffers from instability and sensitivity to hyperparameters such as update steps and subspace. We explain this behavior through a decision-theoretic lens, treating TTT as implicit Bayesian inference in the kernel regime. Under a Gaussian process benchmark, we show that TTT reduces prediction error when updates are spectrally matched to the prompt's signal-to-noise ratio and aligned with query-relevant eigen-directions. This perspective underpins the following results: (1) we show when fixed update steps and subspaces fail under distribution shifts, motivating adaptive strategies; (2) we prove that selecting update steps via prompt evidence admits a PAC-Bayes guarantee against overfitting; and (3) we characterize the Bayes-optimal update subspace under a linear-Gaussian correction model, yielding a scoring rule for selecting Transformer blocks and heads. Our theory helps explain the empirical instability of TTT, taking a step toward principled guidance for when, how far, and which directions to adapt.

Poster
Gradient Inversion Attacks Beyond SGD

Guangnian Wan, Gongfan Fang, Xinyin Ma, Xinchao Wang

Expand abstract

Gradient Inversion Attack (GIA) poses a significant threat to federated learning, enabling adversaries to reconstruct private training data from the information shared during training. Prior research has predominantly focused on the vanilla SGD, where the server or an eavesdropper can directly observe true gradients. In practical deployments, however, models may be trained with adaptive optimizers (e.g., Adam, RMSProp, and AdaGrad), for which the observable signal is not raw gradients but momentum-based parameter updates. This setting remains underexplored and undermines traditional gradient-matching strategies, which struggle to recover labels and images from non-gradient updates. To address this gap, this paper explores attacks tailored to modern adaptive optimizers. We present an analytical rule for recovering labels from optimizer updates and propose an update-matching objective that optimizes dummy inputs to reproduce the observed updates. The proposed approach is general and can be directly applied to various optimizers such as Adam, AdaGrad, and RMSProp. Furthermore, we find that, despite being introduced for adaptive optimizers, the proposed objective function also yields stronger attacks in the standard SGD setting. Experiments on datasets such as ImageNet and PACS highlight the effectiveness of our method over existing gradient matching techniques.

Poster
DB-KSVD: Scalable Alternating Optimization for Disentangling High-Dimensional Embedding Spaces

Romeo Valentin, Sydney Michelle Katz, Vincent Vanhoucke, Mykel Kochenderfer

Expand abstract

Dictionary learning has recently emerged as a promising approach for mechanistic interpretability of large transformer models. Disentangling high-dimensional transformer embeddings requires algorithms that scale to high-dimensional data with large sample sizes. Recent work has explored sparse autoencoders (SAEs) for this problem. However, SAEs use a simple linear encoder to solve the sparse encoding subproblem, which is known to be NP-hard. It is therefore interesting to understand whether this approach is sufficient to find good solutions to the dictionary learning problem or if a more sophisticated algorithm could find better solutions. In this work, we propose Double-Batch KSVD (DB-KSVD), a scalable dictionary learning algorithm that adapts the classic KSVD algorithm. DB-KSVD is informed by the rich theoretical foundations of KSVD but scales to datasets with millions of samples and thousands of dimensions. We demonstrate the efficacy of DB-KSVD by disentangling text embeddings of the Gemma-2-2B and Pythia-160M models and evaluating on six metrics from the SAEBench benchmark, where we achieve competitive results when compared to established approaches based on SAEs. We further show similar results when disentangling image embeddings obtained from the DINOv2-S and DINOv2-B models, solidifying our findings. By matching SAE performance with an entirely different optimization approach, our results suggest that (i) SAEs do find strong solutions to the dictionary learning problem and (ii) traditional optimization approaches can be scaled to the required problem sizes, offering a promising avenue for further research. We make an implementation of DB-KSVD available.

Poster
AdLift: Lifting Adversarial Perturbations to Safeguard 3D Gaussian Splatting Assets Against Instruction-Driven Editing

Ziming Hong, Tianyu Huang, Runnan Chen, Shanshan Ye, Mingming Gong, Bo Han, Tongliang Liu

Expand abstract

Recent studies have extended instruction-driven 2D editing pipelines to 3D Gaussian Splatting (3DGS), enabling faithful 3DGS asset manipulation for advanced content creation. However, it also exposes 3DGS assets to serious risks of unauthorized editing and malicious tampering. Although adversarial perturbations against editing models have proven effective for protecting 2D images, applying them to 3DGS encounters two major challenges: *view-generalizable protection* and *balancing invisibility with protection capability*. In this work, we propose AdLift, a novel editing safeguard for 3DGS that prevents instruction-driven editing across arbitrary views and dimensions by lifting strictly bounded 2D adversarial perturbations into 3D Gaussian-represented safeguard. To ensure both *protective effectiveness* and *invisibility*, these safeguard Gaussians are progressively optimized across training views using a tailored Lifted PGD, which first conducts *gradient truncation* during back-propagation from the editing model to the rendered image and applies projected gradient updates to strictly bound image-level perturbations. Then, the resulting perturbation is backpropagated to the safeguard Gaussian parameters via *image-to-Gaussian fitting*. We alternate these two steps, yielding effective and imperceptible protection that generalizes across both training and novel views. Empirically, qualitative and quantitative results demonstrate that the proposed AdLift effectively protects against state-of-the-art instruction-driven 2D and 3DGS editing.

Spotlight
DPsurv: Dual-Prototype Evidential Fusion for Uncertainty-Aware and Interpretable Whole Slide Image Survival Prediction

Yucheng Xing, Ling Huang, Jingying Ma, Ruping Hong, Jiangdong Qiu, Pei Liu, Kai He, Huazhu Fu, Mengling Feng

Expand abstract

Whole-slide images (WSIs) are widely used for cancer survival analysis because of their comprehensive histopathological information at both cellular and tissue levels, enabling quantitative, large-scale, and prognostically rich tumor feature analysis. However, most existing WSI survival analysis methods struggle with limited interpretability and often overlook predictive uncertainty in heterogeneous slide images. In this paper, we propose DPsurv, a dual-prototype whole-slide image evidential fusion network that outputs uncertainty-aware survival intervals, and enables interpretable survival results through patch prototype distribution assignment, component prototype evidence reasoning, and component-wise relative risk aggregation. Experiments on five publicly available datasets demonstrate strong discriminative performance and well-calibrated predictions, validating its effectiveness and reliability. The interpretation of survival results provides transparency at the feature, reasoning, and decision levels, thereby enhancing the trustworthiness and interpretability of DPsurv.

Poster
Dynamic Linear Attention

Xin Wang, Hui Shen, Boyuan Zheng, Xueshen Liu, Minkyoung Cho, Zhongwei Wan, Zesen Zhao, Zhuoqing Mao, Shen Yan, Mi Zhang

Expand abstract

The scalability of Large Language Models (LLMs) to long contexts is fundamentally constrained by the quadratic complexity of standard attention, motivating the adoption of linear attention mechanisms with sub-quadratic cost. To improve representation capacity under long contexts, recent approaches organize memory in a multi-state manner. However, existing multi-state linear attention methods rely on fixed state merging policies that cannot adapt to dynamically varying token importance, irreversibly obscuring critical tokens and causing severe error accumulation over long sequences. To address this limitation, we propose DLA, a dynamic memory modeling framework for multi-state linear attention. DLA introduces (i) Information-Aware Dynamic State Merging, which adaptively determines state boundaries based on token-level information variation, preserving high-resolution representations around semantic transitions while aggressively summarizing stable regions, and (ii) Capacity-Bounded Memory Modeling, which maintains a fixed-size, chronologically ordered state cache by selectively merging adjacent low-information states to control memory growth with minimal information loss. We pre-train DLA on two different linear attention models and evaluate on 10 datasets from three different aspects. Experimental results demonstrate the superiority of DLA over state-of-the-art.

Poster
MoCo-EA: Exploiting Adversarial Mode Connectivity for Efficient Evolutionary Attacks

Hyo Seo Kim, Gang Luo, Can Chen, Binghui Wang, Yue Duan, Ren Wang

Expand abstract

Evolutionary algorithms for adversarial attacks leverage population-based search to discover perturbations without gradient information, but suffer from inefficient crossover operations that destroy adversarial properties through discrete interpolation. We introduce Mode Connectivity Evolutionary Attack (MoCo-EA), which replaces traditional crossover with a novel Bézier crossover operator that optimizes perturbations along a continuous Bézier curve between parent perturbations. Our key insight is that adversarial examples lie on connected manifolds where intermediate points maintain and often enhance attack effectiveness. We demonstrate three findings: (1) Successful adversarial perturbations exhibit mode connectivity; (2) Intermediate points along optimized paths achieve higher transferability than endpoints; (3) Bézier crossover dramatically outperforms discrete genetic operations while reducing convergence time and query requirements. By exploiting the geometric structure of adversarial space through path optimization, MoCo-EA provides an efficient and reliable method. Our work challenges the traditional view of adversarial examples as isolated points and opens new directions for both attack generation and defense research.

Poster
FedEBA+: Towards Fair and Effective Federated Learning via Entropy-Based Model

Zhichao Wang, Lin Wang, Ye Shi, Sai Praneeth Karimireddy, Xiaoying Tang

Expand abstract

Federated Learning (FL) often suffers from sacrificing global model accuracy when improving client-level fairness due to data heterogeneity, which often leads to inconsistent performance of the globally trained models, resulting in unfair outcomes among users. Existing fair FL algorithms face a bottleneck: they either sacrifice global model accuracy to promote fairness or fall short of achieving optimal fairness. In this paper, we propose a novel framework that effectively improves fairness while preserving global accuracy by integrating information-theoretic principles with model alignment. Specifically, we leverage the Maximum Entropy Principle to derive an analytic, closed-form solution for fair aggregation weights, ensuring significant fairness enhancements. We further employ a step-wise model alignment strategy that synchronizes gradient directions across heterogeneous clients, effectively mitigating the drift induced by local updates. Theoretical analysis proves that our method guarantees convergence even in non-convex settings. Importantly, we push the theoretical frontier of federated fairness by extending performance variance analysis to generalized regression, providing broader guarantees. Extensive experiments on five datasets demonstrate that our approach consistently outperforms state-of-the-art methods, achieving superior fairness without sacrificing global accuracy. Our code is available at \href{https://github.com/T-Lab-CUHKSZ/FedEBA-Plus}{https://github.com/T-Lab-CUHKSZ/FedEBA-Plus}.

Poster
Toward Stable Value Alignment: Introducing Independent Modules for Consistent Value Guidance

Wenhao Chen, Sirui Sun, Shengyuan Bai, Guojie Song

Expand abstract

Aligning large language models (LLMs) with human values typically relies on post-training or inference-time steering that directly manipulates the backbone’s parameters or representation space. However, a critical gap exists: the model’s residual stream is highly dynamic, in which values exist as fragile, low-dimensional properties, inherently incompatible with the stability required for consistent value expression. In this paper, we propose the Stable Value Guidance Transformer (SVGT), which addresses this gap through an independent value module incorporating two key designs: (1) *independent value modeling*, maintaining normative representations in a dedicated value space isolated from the backbone, and (2) *explicit behavioral guidance*, transducing these stable signals into learnable latent Bridge Tokens. These tokens serve as dynamic value anchors to explicitly steer the generative trajectory, ensuring robust adherence across diverse contexts without disrupting the backbone’s internal representations. Experiments across multiple backbones and safety benchmarks show that SVGT generally reduces harmful scores by over 70\% while maintaining generation fluency, demonstrating the efficacy of architecturally grounded value modeling.

Spotlight
PlugGuard: A Streaming Safeguard for Large Models via Latent Dynamics-Guided Risk Detection

Xiaodan Li, Mengjie Wu, Yao Zhu, Yunna Lv, YueFeng Chen, Cen Chen, Jianmei Guo, Hui Xue

Expand abstract

Large models (LMs) are powerful content generators, yet their open‑ended nature can also introduce potential risks, such as generating harmful or biased content. Existing guardrails mostly perform post-hoc detection that may expose unsafe content before it is caught, and the latency constraints further push them toward lightweight models, limiting detection accuracy. In this work, we propose PlugGuard, a novel plug-in framework that enables streaming risk detection within the LM generation pipeline. PlugGuard leverages intermediate LM hidden states through a Streaming Latent Dynamics Head (SLD), which models the temporal evolution of risk across the generated sequence for more accurate real-time risk detection. To achieve reliable streaming moderation in real applications, we introduce an Anchored Temporal Consistency (ATC) loss, ensuring that risk assessments remain consistent with a strict stop-if-harmful policy. Besides, for a rigorous evaluation of streaming guardrails, we also present StreamGuardBench—a model-grounded benchmark featuring on-the-fly responses from each protected model, reflecting real-world streaming scenarios in both text and vision–language tasks. Across diverse models and datasets, PlugGuard consistently outperforms state-of-the-art streaming guardrails (achieving a 22.80% F1 score gain), while using only 20M parameters and adding less than 0.5 ms of per-token latency. The code and StreamGuardBench are released at PlugGuard to facilitate research on streaming guardrails.

Poster
Differentially Private Cross-Silo Recommendation from Implicit Feedback

Xun Ran, Qingqing Ye, Xin Huang, Jianliang Xu, Haibo Hu

Expand abstract

Cross-silo recommendation from implicit feedback is a key task in modern recommender systems, where user-item interaction data are distributed across multiple parties and cannot be centrally collected. Unlike explicit feedback, which provides fully observed real-valued ratings, implicit feedback is one-class and extremely sparse, recording only users' actions or inactions (e.g., clicks, visits, or bookmarks), yet it is far more prevalent in real-world applications. Such behavioral data are often highly sensitive, raising significant privacy concerns when used for collaborative model training. Although differential privacy (DP) has been widely applied to explicit feedback-based models, the resulting utility degradation makes it difficult to apply DP effectively to implicit feedback learning. In this work, we propose DPIMF, a differentially private implicit matrix factorization framework for cross-silo recommendation based on objective perturbation. To improve utility, we redesign the loss function and adopt an importance sampling scheme to reduce the noise scale required for privacy preservation. We further provide formal utility guarantees for the proposed techniques and characterize the conditions under which utility improvements are maximized. Experiments on three benchmark datasets validate our theoretical analysis and demonstrate that DPIMF achieves a better privacy-utility trade-off than state-of-the-art methods.

Poster
Discovering Interpretable Algorithms by Decompiling Transformers to RASP

Xinting Huang, Aleksandra Bakalova, Satwik Bhattamishra, William Merrill, Michael Hahn

Expand abstract

Recent work has shown that the computations of Transformers can be simulated in the RASP family of programming languages. These findings have enabled improved understanding of the expressive capacity and generalization abilities of Transformers. In particular, Transformers have been suggested to length-generalize exactly on problems that have simple RASP programs. However, it remains open whether trained models actually implement simple interpretable programs. In this paper, we present a general method to extract such programs from trained Transformers. The idea is to faithfully re-parameterize a Transformer as a RASP program and then apply causal interventions to discover a small sufficient sub-program. In experiments on small Transformers trained on algorithmic and formal language tasks, we show that our method often recovers simple and interpretable RASP programs from length-generalizing transformers. Our results provide the most direct evidence so far that Transformers internally implement simple RASP programs.

Poster
Multi-agent imitation learning with function approximation: linear Markov games and beyond

Luca Viano, Till Freihaut, Emanuele Nevali, Volkan Cevher, Matthieu Geist, Giorgia Ramponi

Expand abstract

In this work, we present the first theoretical analysis of multi-agent imitation learning (MAIL) in linear Markov games where both the transition dynamics and each agent's reward function are linear in some given features. We demonstrate that by leveraging this structure, it is possible to replace the state-action level \emph{all policy deviation concentrability coefficient} (Freihaut et al., 2025b) with a concentrability coefficient defined at the feature level which can be much smaller than the state-action analog when the features are informative about \emph{states' similarity}. Furthermore, to circumvent the need for any concentrability coefficient, we turn to the interactive setting. We provide the first, computationally efficient, interactive MAIL algorithm for linear Markov games and show that its sample complexity depends only on the dimension of the feature map $d$. Building on these theoretical findings, we propose a deep MAIL interactive algorithm which clearly outperforms BC on games such as Tic-Tac-Toe and Connect4.

Poster