Zhi Wang's Homepage

Selected Publications

Zhi Wang, Chicheng Zhang, and Ramya Korlakai Vinayak. Bridging Lifelong and Multi-Task Representation Learning via Algorithm and Complexity Measure. Proceedings of The 37th International Conference on Algorithmic Learning Theory (ALT-2026).
[arXiv]

In lifelong learning, a learner faces a sequence of tasks with shared structure and aims to identify and leverage it to accelerate learning. We study the setting where such structure is captured by a common representation of data. Unlike multi-task learning or learning-to-learn, where tasks are available upfront to learn the representation, lifelong learning requires the learner to make use of its existing knowledge while continually gathering partial information in an online fashion. In this paper, we consider a generalized framework of lifelong representation learning. We propose a simple algorithm that uses multi-task empirical risk minimization as a subroutine and establish a sample complexity bound based on a new notion we introduce—the task-eluder dimension. Our result applies to a wide range of learning problems involving general function classes. As concrete examples, we instantiate our result on classification and regression tasks under noise.

Daiwei Chen, Yi Chen, Aniket Rege, Zhi Wang, and Ramya Korlakai Vinayak. PAL: Sample-Efficient Personalized Reward Modeling for Pluralistic Alignment. The 13th International Conference on Learning Representations (ICLR-2025).
[OpenReview]

Foundation models trained on internet-scale data benefit from extensive alignment to human preferences before deployment. However, existing methods typically assume a homogeneous preference shared by all individuals, overlooking the diversity inherent in human values. In this work, we propose a general reward modeling framework for pluralistic alignment (PAL), which incorporates diverse preferences from the ground up. PAL has a modular design that leverages commonalities across users while catering to individual personalization, enabling efficient few-shot localization of preferences for new users. Extensive empirical evaluation demonstrates that PAL matches or outperforms state-of-the-art methods on both text-to-text and text-to-image tasks: on Reddit TL;DR Summary, PAL is 1.7% more accurate for seen users and 36% more accurate for unseen users compared to the previous best method, with 100× less parameters. On Pick-a-Pic v2, PAL is 2.5% more accurate than the best method with 156× fewer learned parameters. Finally, we provide theoretical analysis for generalization of rewards learned via PAL framework showcasing the reduction in number of samples needed per user.

Thang Duong, Zhi Wang, and Chicheng Zhang. Beyond Task Diversity: Provable Representation Transfer for Sequential Multitask Linear Bandits. Advances in Neural Information Processing Systems 37 (NeurIPS-2024).
[arXiv]

We study lifelong learning in linear bandits, where a learner interacts with a sequence of linear bandit tasks whose parameters lie in an \( m \)-dimensional subspace of \( \mathbb{R}^d \), thereby sharing a low-rank representation. Current literature typically assumes that the tasks are diverse, i.e., their parameters uniformly span the \( m \)-dimensional subspace. This assumption allows the low-rank representation to be learned before all tasks are revealed, which can be unrealistic in real-world applications. In this work, we present the first nontrivial result for sequential multi-task linear bandits without the task diversity assumption. We develop an algorithm that efficiently learns and transfers low-rank representations. When facing \( N \) tasks, each played over \( \tau \) rounds, our algorithm achieves a regret guarantee of \( \tilde{O} (Nm \sqrt{\tau} + N^{\frac{2}{3}} \tau^{\frac{2}{3}} d m^{\frac13} + Nd^2 + \tau m d ) \) under the ellipsoid action set assumption. This result can significantly improve upon the baseline of \( \tilde{O} (Nd \sqrt{\tau} ) \) that does not leverage the low-rank structure when the number of tasks \( N \) is sufficiently large and \( m \ll d \). We also demonstrate empirically on synthetic data that our algorithm outperforms baseline algorithms, which rely on the task diversity assumption.

Zhi Wang, Geelon So, and Ramya Korlakai Vinayak. Metric Learning from Limited Pairwise Preference Comparisons. Proceedings of the 40th Conference on Uncertainty in Artificial Intelligence (UAI-2024).
[arXiv]

We study metric learning from preference comparisons under the ideal point model, in which a user prefers an item over another if it is closer to their latent ideal item. These items are embedded into \( \mathbb{R}^d \) equipped with an unknown Mahalanobis distance shared across users. While recent work shows that it is possible to simultaneously recover the metric and ideal items given \( O(d) \) pairwise comparisons per user, in practice we often have a limited budget of \( o(d) \) comparisons. We study whether the metric can still be recovered, even though it is known that learning individual ideal items is now no longer possible. We show that in general, \( o(d) \) comparisons reveals no information about the metric, even with infinitely many users. However, when comparisons are made over items that exhibit low-dimensional structure, each user can contribute to learning the metric restricted to a low-dimensional subspace so that the metric can be jointly identified. We present a divide-and-conquer approach that achieves this, and provide theoretical recovery guarantees and empirical validation.

Zhi Wang, Chicheng Zhang, and Kamalika Chaudhuri. Thompson Sampling for Robust Transfer in Multi-Task Bandits. Proceedings of the 39th International Conference on Machine Learning (ICML-2022).
[arXiv]

We study the problem of online multi-task learning where the tasks are performed within similar but not necessarily identical multi-armed bandit environments. In particular, we study how a learner can improve its overall performance across multiple related tasks through robust transfer of knowledge. While an upper confidence bound (UCB)-based algorithm has recently been shown to achieve nearly-optimal performance guarantees in a setting where all tasks are solved concurrently, it remains unclear whether Thompson sampling (TS) algorithms, which have superior empirical performance in general, share similar theoretical properties. In this work, we present a TS-type algorithm for a more general online multi-task learning protocol, which extends the concurrent setting. We provide its frequentist analysis and prove that it is also nearly-optimal using a novel concentration inequality for multi-task data aggregation at random stopping times. Finally, we evaluate the algorithm on synthetic data and show that the TS-type algorithm enjoys superior empirical performance in comparison with the UCB-based algorithm and a baseline algorithm that performs TS for each individual task without transfer.

Chicheng Zhang and Zhi Wang. Provably Efficient Multi-Task Reinforcement Learning with Model Transfer. Advances in Neural Information Processing Systems 34 (NeurIPS-2021).
[arXiv]

We study multi-task reinforcement learning (RL) in tabular episodic Markov decision processes (MDPs). We formulate a heterogeneous multi-player RL problem, in which a group of players concurrently face similar but not necessarily identical MDPs, with a goal of improving their collective performance through inter-player information sharing. We design and analyze an algorithm based on the idea of model transfer, and provide gap-dependent and gap-independent upper and lower bounds that characterize the intrinsic complexity of the problem.

Zhi Wang*, Chicheng Zhang*, Manish Kumar Singh, Laurel D. Riek, and Kamalika Chaudhuri. Multitask Bandit Learning through Heterogeneous Feedback Aggregation. Proceedings of the 24th International Conference on Artificial Intelligence and Statistics (AISTATS-2021).
[arXiv]

In many real-world applications, multiple agents seek to learn how to perform highly related yet slightly different tasks in an online bandit learning protocol. We formulate this problem as the ε-multi-player multi-armed bandit problem, in which a set of players concurrently interact with a set of arms, and for each arm, the reward distributions for all players are similar but not necessarily identical. We develop an upper confidence bound-based algorithm,
RobustAgg
(ε), that adaptively aggregates rewards collected by different players. In the setting where an upper bound on the pairwise similarities of reward distributions between players is known, we achieve instance-dependent regret guarantees that depend on the amenability of information sharing across players. We complement these upper bounds with nearly matching lower bounds. In the setting where pairwise similarities are unknown, we provide a lower bound, as well as an algorithm that trades off minimax regret guarantees for adaptivity to unknown similarity structure.

Zhi Wang, Liron Cohen, Sven Koenig, and T. K. Satish Kumar. The Factored Shortest Path Problem and Its Applications in Robotics. Proceedings of the 28th International Conference on Automated Planning and Scheduling (ICAPS-2018).
[pdf]
Many real-world combinatorial problems exhibit structure in the way in which their variables interact. Such structure can be exploited in the form of "factors" for representational as well as computational benefits. Factored representations are extensively used in probabilistic reasoning, constraint satisfaction, planning, and decision theory. In this paper, we formulate the factored shortest path problem (FSPP) on a collection of constraints interpreted as factors of a highdimensional map.We show that the FSPP is not only a generalization of the regular shortest path problem but also particularly relevant to robotics.We develop factored-space heuristics for A* and prove that they are admissible and consistent. We provide experimental results on both random and handcrafted instances as well as on an example robotics domain to show that A* with factored-space heuristics outperforms A* with the Manhattan Distance heuristic in many cases.

T. K. Satish Kumar, Zhi Wang, Anoop Kumar, Craig Milo Rogers, and Craig A. Knoblock. Load Scheduling of Simple Temporal Networks Under Dynamic Resource Pricing. Proceedings of the 32nd AAAI Conference on Artificial Intelligence (AAAI-2018).
[pdf]
We study load scheduling of simple temporal networks (STNs) under dynamic pricing of resources. We are given a set of processes and a set of simple temporal constraints between their execution times, i.e., an STN. Each process uses a certain amount of resource for execution. The unit price of the resource is a function of time, \( f(t) \). The goal is to find a schedule of a given STN that trades off makespan minimization against cost minimization within a user-specified suboptimality bound. We provide a polynomial-time algorithm for solving the load scheduling problem when \( f(t) \) is piecewise constant. This has important applications in many real-world domains including the smart home and smart grid domains. We then study the dependency of the unit price of the resource on time as well as the total demand at that time. This leads to a further characterization of tractable, NP-hard, and conjectured tractable cases.

T. K. Satish Kumar, Zhi Wang, Anoop Kumar, Craig Milo Rogers, and Craig A. Knoblock. On the Linear Programming Duals of Temporal Reasoning Problems. The 15th International Symposium on Artificial Intelligence and Mathematics (ISAIM-2018).
[pdf]
Temporal reasoning problems occur in many application domains of Artificial Intelligence; therefore, it is important for us to develop algorithms for solving them efficiently. While some problems like Simple Temporal Problems are known to be tractable, some other problems like Disjunctive Temporal Problems are known to be NP-hard. In this paper, we provide a Linear Programming (LP) duality perspective on temporal reasoning problems. In many cases, we show that their LP duals are the commonly-studied flow problems in Graph Theory. Using the general theory of LP duality, we develop novel algorithms for efficiently solving several temporal reasoning problems. We also show that other previously-known efficient methods in temporal reasoning also fit into this perspective of LP duality.

Zhi Wang

Selected Publications

Teaching

Professional Service

Contact