UAI CDM
Causal Bayesian Optimization via Causal Bandits: Open Questions
Arpan Mukherjee, Zirui Yan, Ali Tajer
Proc. Conference on Uncertainty in Artificial Intelligence Workshop on Causality for Decision Making (UAI CDM), 2026.
Abstract
Causal bandits (CBs) provide a principled approach to the sequential design of interventions (experiments) on causal systems to optimize any desired utility of interest. Implicitly, CB algorithms learn the causal model while performing experiments, to the extent that the model's information is needed for optimizing the desired utility. This paper provides a perspective and framework for viewing and analyzing causal Bayesian optimization (CBO) through the lens of CBs. Despite their distinctions, CBO and CB problems share the same principle: causal structure can be used to share information across interventions and avoid treating actions as independent alternatives. Specifically, this paper first compares the two lines of literature on CBs and CBOs with respect to the roles of graph structures, intervention models, and function-class assumptions, distinguishing graph-theoretic sources of complexity from those arising from posterior surrogate-model complexity. Finally, it identifies open problems in CBO that can be addressed using recent advances in CBs.
@inproceedings{mukherjee2026causal,
title={Causal {Bayesian} Optimization via Causal Bandits: Open Questions},
author={Mukherjee, Arpan and Yan, Zirui and Tajer, Ali},
booktitle={Proc. UAI Workshop on Causality for Decision Making},
year={2026},
month={August},
address={Amsterdam, Netherlands}
}
ACL
Multi-component Causal Tracing in Large Language Models
Zirui Yan, Dennis Wei, Dmitriy A. Katz, Prasanna Sattigeri, Ali Tajer
Proc. Annual Meeting of the Association for Computational Linguistics (ACL), 2026.
Oral
Abstract
Causal tracing systematically intervenes on a large language model’s (LLM’s) internal representations to uncover and quantify the causal pathways linking specific inputs or computations to specific metrics of interest, quantifying the LLM’s behavior. Building on previous single-component or single-layer studies, this paper presents a unified framework for causally tracing multiple components simultaneously. This framework systematically identifies the subsets of components (e.g., attention heads and multi-layer perceptron neurons) most critical to a desired target performance metric (e.g., accuracy and fairness). This is achieved by incorporating flexible interventions applied to a wide range of desired metrics. To address the combinatorial complexity of the multi-component problem, an efficient algorithm is designed that leverages soft interventions and a carefully designed metric transformation, converting the combinatorial search problem into a continuous one that can be solved efficiently under proper constraints, thereby generating proper binary decisions for selecting components. Experimental results demonstrate that the proposed method efficiently identifies subsets of the model’s components that have a high impact on the target metric, outperforming existing baseline approaches.
@inproceedings{yan2026multi,
title={Multi-component Causal Tracing in Large Language Models},
author={Yan, Zirui and Wei, Dennis and Katz, Dmitriy A and Sattigeri, Prasanna and Tajer, Ali},
booktitle={Proc. Annual Meeting of the Association for Computational Linguistics},
year={2026},
month={July},
address={San Diego, CA}
}
NeurIPS
Reward-oriented Causal Representation Learning
Zirui Yan*, Emre Acartürk*, Ali Tajer
Proc. Neural Information Processing Systems (NeurIPS), 2025.
Abstract
Causal representation learning (CRL) is the process of disentangling the latent low-dimensional causally-related generating factors underlying high-dimensional observable data. Extensive recent studies have characterized CRL identifiability and perfect recovery of the latent variables and their attendant causal graph. This paper introduces the notion of reward-oriented CRL, the purpose of which is to move away from perfectly learning the latent representation and instead learning it to the extent needed for optimizing a desired downstream task (reward). In reward-oriented CRL, perfectly learning the latent representation can be excessive; instead, it must be learned at the coarsest level sufficient for optimizing the desired task. Reward-oriented CRL is formalized as the optimization of a desired function of the observable data over the space of all possible interventions and focuses on linear causal and transformation models. To sequentially identify the optimal subset of interventions, an adaptive exploration algorithm is designed that learns the latent causal graph and the variables needed to identify the best intervention. It is shown that for an -dimensional latent space and a -dimensional observation space, over a horizon the algorithm's regret scales as , where measures total uncertainty in the graph estimates. Furthermore, an almost-matching lower bound is shown to scale as , in which is replaced by that counts the number of causal paths in the graph.
@inproceedings{yan2025reward,
title={Reward-oriented Causal Representation Learning},
author={Yan, Zirui and Acart{\"u}rk, Emre and Tajer, Ali},
booktitle={Proc. Advances in Neural Information Processing Systems},
year={2025},
month={December},
address={San Diego, CA}
}
NeurIPS
Linear Causal Bandits: Unknown Graph and Soft Interventions
Zirui Yan, Ali Tajer
Proc. Neural Information Processing Systems (NeurIPS), 2024.
Abstract
Designing causal bandit algorithms depends on two central categories of assumptions: (i) the extent of information about the underlying causal graphs and (ii) the extent of information about interventional statistical models. There have been extensive recent advances in dispensing with assumptions on either category. These include assuming known graphs but unknown interventional distributions, and the converse setting of assuming unknown graphs but access to restrictive hard/ interventions, which removes the stochasticity and ancestral dependencies. Nevertheless, the problem in its general form, i.e., unknown graph and unknown stochastic intervention models, remains open. This paper addresses this problem and establishes that in a graph with nodes, maximum in-degree and maximum causal path length , after interaction rounds the regret upper bound scales as where is a constant and is a measure of intervention power. A universal minimax lower bound is also established, which scales as . Importantly, the graph size has a diminishing effect on the regret as grows. These bounds have matching behavior in , exponential dependence on , and polynomial dependence on (with the gap ). On the algorithmic aspect, the paper presents a novel way of designing a computationally efficient CB algorithm, addressing a challenge that the existing CB algorithms using soft interventions face.
@inproceedings{yan2024linear,
title={Linear Causal Bandits: Unknown Graph and Soft Interventions},
author={Yan, Zirui and Tajer, Ali},
booktitle={Proc. Advances in Neural Information Processing Systems},
year={2024},
month={December},
address={Vancouver, Canada}
}
ISIT
Improved Bound for Robust Causal Bandits with Linear Models
Zirui Yan, Arpan Mukherjee, Burak Varıcı, Ali Tajer
Proc. IEEE International Symposium on Information Theory (ISIT), 2024.
Abstract
This paper investigates the robustness of causal bandits (CBs) in the face of temporal model fluctuations. This setting deviates from the existing literature’s widely-adopted assumption of constant causal models. The focus is on causal systems with linear structural equation models (SEMs). The SEMs and the time-varying pre- and post-interventional statistical models are all unknown and subject to variations over time. The goal is to design a sequence of interventions that incur the smallest cumulative regret compared to an oracle aware of the entire causal model and its fluctuations. A robust CB algorithm is proposed, and its cumulative regret is analyzed by establishing both upper and lower bounds on the regret. It is shown that in a graph with maximum in-degree , length of the largest causal path , and an aggregate model deviation , the regret is upper bounded by and lower bounded by . The proposed algorithm achieves nearly optimal regret when is , maintaining sub-linear regret for a broad range of .
@inproceedings{yan2024improved,
title={Improved Bound for Robust Causal Bandits with Linear Models},
author={Yan, Zirui and Mukherjee, Arpan and Varici, Burak and Tajer, Ali},
booktitle={Proc. IEEE International Symposium on Information Theory},
year={2024},
month={July},
address={Athens, Greece}
}
AISTATS
Causal Bandits with General Causal Models and Interventions
Zirui Yan, Dennis Wei, Dmitriy A. Katz, Prasanna Sattigeri, Ali Tajer
Proc. International Conference on Artificial Intelligence and Statistics (AISTATS), 2024.
Abstract
This paper considers causal bandits (CBs) for the sequential design of interventions in a causal system. The objective is to optimize a reward function via minimizing a measure of cumulative regret with respect to the best sequence of interventions in hindsight. The paper advances the results on CBs in three directions. First, the structural causal models (SCMs) are assumed to be unknown and drawn arbitrarily from a general class of Lipschitz-continuous functions. Existing results are often focused on (generalized) linear SCMs. Second, the interventions are assumed to be generalized soft with any desired level of granularity, resulting in an infinite number of possible interventions. The existing literature, in contrast, generally adopts atomic and hard interventions. Third, we provide general upper and lower bounds on regret. The upper bounds subsume (and improve) known bounds for special cases. The lower bounds are generally hitherto unknown. These bounds are characterized as functions of the (i) graph parameters, (ii) eluder dimension of the space of SCMs, denoted by , and (iii) the covering number of the function space, denoted by . Specifically, the cumulative achievable regret over horizon is , where is related to the Lipschitz constants, is the graph’s maximum in-degree, and is the length of the longest causal path. The upper bound is further refined for special classes of SCMs (neural network, polynomial, and linear), and their corresponding lower bounds are provided.
@inproceedings{yan2024causal,
title={Causal Bandits with General Causal Models and Interventions},
author={Yan, Zirui and Wei, Dennis and Katz-Rogozhnikov, Dmitriy and Sattigeri, Prasanna and Tajer, Ali},
booktitle={Proc. International Conference on Artificial Intelligence and Statistics},
year={2024},
month={May},
address={Valencia, Spain}
}
JSAIT
Robust Causal Bandits for Linear Models
Zirui Yan, Arpan Mukherjee, Burak Varıcı, Ali Tajer
IEEE Journal on Selected Areas in Information Theory (JSAIT), 2024.
Abstract
The sequential design of experiments for optimizing a reward function in causal systems can be effectively modeled by the sequential design of interventions in causal bandits (CBs). In the existing literature on CBs, a critical assumption is that the causal models remain constant over time. However, this assumption does not necessarily hold in complex systems, which constantly undergo temporal model fluctuations. This paper addresses the robustness of CBs to such model fluctuations. The focus is on causal systems with linear structural equation models (SEMs). The SEMs and the time-varying pre- and post-interventional statistical models are all unknown. Cumulative regret is adopted as the design criteria, based on which the objective is to design a sequence of interventions that incur the smallest cumulative regret with respect to an oracle aware of the entire causal model and its fluctuations. First, it is established that the existing approaches fail to maintain regret sub-linearity with even a few instances of model deviation. Specifically, when the number of instances with model deviation is as few as , where is the time horizon and is the length of the longest causal path in the graph, the existing algorithms will have linear regret in . For instance, when and , model deviations in out of instances result in a linear regret. Next, a robust CB algorithm is designed, and its regret is analyzed, where upper and information-theoretic lower bounds on the regret are established. Specifically, in a graph with nodes and maximum degree , under a general measure of model deviation , the cumulative regret is upper bounded by and lower bounded by . Comparing these bounds establishes that the proposed algorithm achieves nearly optimal regret when is and maintains sub-linear regret for a broader range of .
@article{yan2024robust,
title={Robust Causal Bandits for Linear Models},
author={Yan, Zirui and Mukherjee, Arpan and Varici, Burak and Tajer, Ali},
journal={IEEE Journal on Selected Areas in Information Theory},
volume={5},
pages={78--93},
year={2024}
}
TON
Optimizing Parameter Mixing under Constrained Communications in Parallel Federated Learning
Xuezheng Liu*, Zirui Yan*, Yipeng Zhou , Di Wu , Xu Chen , Jessie Hui Wang (* equal contributor)
IEEE/ACM Trans. on Networking (TON), 2023.
Abstract
In vanilla Federated Learning (FL) systems, a centralized parameter server (PS) is responsible for collecting, aggregating and distributing model parameters with decentralized clients. However, the communication link of a single PS can be easily overloaded by concurrent communications with a massive number of clients. To overcome this drawback, multiple PSes can be deployed to form a parallel FL (PFL) system, in which each PS only communicates with a subset of clients and its neighbor PSes. On one hand, each PS conducts iterations with clients in its subset. On the other hand, PSes communicate with each other periodically to mix their parameters so that they can finally reach a consensus. In this paper, we propose a novel parallel federated learning algorithm called Fed-PMA, which optimizes such parallel FL under constrained communications by conducting parallel parameter mixing and averaging with theoretic guarantees. We formally analyze the convergence rate of Fed-PMA with convex loss, and further derive the optimal number of times each PS should mix with its neighbor PSes so as to maximize the final model accuracy within a fixed span of training time. Theoretical study manifests that PSes should mix their parameters more frequently if the connection between PSes is sparse or the time cost of mixing is low. Inspired by our analysis, we propose the Fed-APMA algorithm that can adaptively determine the near-optimal number of mixing times with non-convex loss under dynamic communication conditions. Extensive experiments with realistic datasets are carried out to demonstrate that both Fed-PMA and its adaptive version Fed-APMA significantly outperform the state-of-the-art baselines.
@article{liu2023optimizing,
title={Optimizing Parameter Mixing Under Constrained Communications in Parallel Federated Learning},
author={Liu, Xuezheng and Yan, Zirui and Zhou, Yipeng and Wu, Di and Chen, Xu and Wang, Jessie Hui},
journal={IEEE/ACM Transactions on Networking},
volume={31},
number={6},
pages={2640--2652},
year={2023}
}
ICASSP
Federated Multi-armed Bandit via Uncoordinated Exploration
Zirui Yan, Quan Xiao, Tianyi Chen, Ali Tajer
Proceedings of IEEE International Conference on Acoustics, Speech, and Signal Processing (ICASSP), 2022.
Abstract
A wide range of multi-agent decision-making problems can be abstracted as a federated multi-armed bandit (FMAB) problem. A key challenge of the FMAB problem is that the exploration-exploitation dichotomy inherited from the multi-armed bandit aspect is compounded with data heterogeneity in federated learning. This renders the exploration and exploitation of different agents inherently entangled. This paper focuses on overcoming the difficulty of exploration in FMAB problems, and it proposes a novel federated upper confidence bound (UCB) algorithm that requires uncoordinated exploration (UE) decisions by the agents. The major distinction of this algorithm, referred to as FedUCB-UE, with the existing FMAB algorithms is that it allows the agents to explore the non-optimal arms and make personalized arm-selection decisions without coordination. While such uncoordinated exploration makes the regret analysis non-trivial, it comes with both the theoretical and empirical benefit of diversity in explorations. Under certain mild assumptions, this paper establishes that FedUCB-UE has a regret bound. Furthermore, experiments performed on synthetic datasets show that FedUCB-UE outperforms the state-of-the-art algorithms.
@inproceedings{yan2022federated,
title={Federated Multi-Armed Bandit Via Uncoordinated Exploration},
author={Yan, Zirui and Xiao, Quan and Chen, Tianyi and Tajer, Ali},
booktitle={Proc. IEEE International Conference on Acoustics, Speech and Signal Processing},
year={2022},
month={May},
address={Singapore}
}
WACV
Image denoising via K-SVD with primal-dual active set algorithm
Quan Xiao, Canhong Wen, Zirui Yan
Proc. Winter Conference on Applications of Computer Vision (WACV), 2020.
Abstract
K-SVD algorithm has been successfully applied to image denoising tasks dozens of years but the big bottleneck in speed and accuracy still needs attention to break. For the sparse coding stage in K-SVD, which involves l0 constraint, prevailing methods usually seek approximate solutions greedily but are less effective once the noise level is high. The alternative l1 optimization is proved to be powerful than l0, however, the time consumption prevents it from the implementation. In this paper, we propose a new K-SVD framework called K-SVDp by applying the Primal-dual active set (PDAS) algorithm to it. Different from the greedy algorithms based K-SVD, the K-SVDp algorithm develops a selection strategy motivated by KKT (Karush-Kuhn-Tucker) condition and yields to an efficient update in the sparse coding stage. Since the K-SVDp algorithm seeks for an equivalent solution to the dual problem iteratively with simple explicit expression in this denoising problem, speed and quality of denoising can be reached simultaneously. Experiments are carried out and demonstrate the comparable denoising performance of our K-SVDp with state-of-the-art methods.
@inproceedings{xiao2020image,
title={Image denoising via {K-SVD} with primal-dual active set algorithm},
author={Xiao, Quan and Wen, Canhong and Yan, Zirui},
booktitle={Proc. IEEE/CVF Winter Conference on Applications of Computer Vision},
year={2020},
month={March},
address={Snowmass, CO}
}
