Jiawei Zhang 0007

dblp:10/239-7 · DBLP profile ↗
← Back
15ranked-venue papers
3as first author
12since 2021 · last 2025
—ORCID · conflict

Domains — the database's venue-derived domains; a paper can count in several

Artificial intelligence and machine learning · 10 · 2 first-author · 9 since 2021Graphics, computer vision, multimedia, augmented reality and games · 3 · 1 first-author · 1 since 2021Computer networks · 1 · 1 since 2021Applied, interdisciplinary, general and emerging computing · 1 · 1 since 2021
YearPublicationVenuePosition
2025 Contextual Optimization Under Model Misspecification: A Tractable and Generalizable Approach
abstract
Contextual optimization problems are prevalent in decision-making applications where historical data and contextual features are used to learn predictive models that inform optimal actions. However, practical applications often suffer from model misspecification due to incomplete knowledge of the underlying data-generating process, leading to suboptimal decisions. Existing approaches primarily address the well-specified case, leaving a critical gap in handling misspecified models. In this paper, we propose a novel Integrated Learning and Optimization (ILO) framework that explicitly accounts for model misspecification by introducing a tractable surrogate loss function with strong theoretical guarantees on generalizability, tractability, and optimality. Our surrogate loss aligns with the true decision performance objective, ensuring robustness to misspecification without imposing restrictive assumptions. The proposed approach effectively mitigates the challenges of non-convexity and non-smoothness in the target loss function, leading to efficient optimization procedures. We provide rigorous theoretical analysis and experimental validation, demonstrating superior performance compared to state-of-the-art methods. Our work offers a principled solution to the practically relevant challenge of model misspecification in contextual optimization.
Omar Bennouna, Jiawei Zhang 0007, Saurabh Amin, Asuman E. Ozdaglar
ICML2
2025 Beamforming Optimization for Robust Sensing and Communication in Dynamic mmWave MIMO Networks
abstract
Acquiring accurate channel state information (CSI) at low overhead is crucial for millimeter wave MIMO communications but is challenging in dynamic environments. In this work, we exploit the emerging integrated sensing and communication (ISAC) beamforming technique for concurrent CSI sensing and data transmission. Despite its low overhead, the corresponding ISAC transmit beamforming design faces a complex trade-off between CSI sensing accuracy and communication interference management. To address this, we formulate the beamforming design as an optimization problem minimizing the maximum Cramér-Rao bound (CRB) of CSI sensing errors subject to the users’ worst-case communication rates under CSI errors. To efficiently solve the problem, we step-by-step propose three algorithms. The first algorithm is based on the semidefinite relaxation and successive convex optimization techniques, which can serve as a benchmark algorithm but suffers high computational complexity. To efficiently handle the worst-case objective and rate constraints, we propose a complexity-reduced algorithm based on the primal-dual optimization method and first-order min-max algorithm. Furthermore, we dismiss SDR and employ the block coordinate descent method combined with cheap gradient descent steps to achieve a low-complexity algorithm. Extensive simulations show the proposed ISAC beamforming design and low-complexity algorithms can provide robust communication performance and significantly outperform existing schemes.
Lei Li 0030, Jiawei Zhang 0007, Tsung-Hui Chang
IEEE J. Sel. Areas Commun.2
2024 A Smoothed Bregman Proximal Gradient Algorithm for Decentralized Nonconvex Optimization
abstract
Decentralized computation has received considerable research interest lately, due to its wide applications in information processing systems. However, one key requirement to establish convergence for almost all decentralized algorithms, for convex and non-convex problems alike, is that the loss function has Lipschitz-continuous gradient (LipGrad). This is a strong assumption, which does not hold for many practical problems, such as matrix/tensor factorization, neural network training, etc. On the contrary, in the centralized setting, one can utilize techniques such as the Bregman proximal gradient (BPG) method to deal with the lack of LipGrad. This work fills the gap between centralized and decentralized cases by developing a novel smoothed decentralized BPG algorithm to deal with a class of nonconvex decentralized problem, where the local problems do not have LipGrad objective functions. By leveraging the recent notion of relative smoothness and primal-dual error bounds, we show that the proposed algorithm achieves a certain ε-stationary solution by using $\mathcal{O}\left( {{\varepsilon ^{ - 2}}} \right)$ iterations, matching the rate of the centralized Bregman proximal gradient method. To our knowledge, this is the first decentralized algorithm that matches the centralized convergence rate bounds under the class of considered problems. Our numerical results on the decentralized quadratic regression example demonstrate the effectiveness of proposed algorithm.
Wenqiang Pu, Jiawei Zhang 0007, Rui Zhou 0016, Xiao Fu 0001, Mingyi Hong 0001
ICASSP2
2024 A Unified Linear Programming Framework for Offline Reward Learning from Human Demonstrations and Feedback
abstract
Inverse Reinforcement Learning (IRL) and Reinforcement Learning from Human Feedback (RLHF) are pivotal methodologies in reward learning, which involve inferring and shaping the underlying reward function of sequential decision-making problems based on observed human demonstrations and feedback. Most prior work in reward learning has relied on prior knowledge or assumptions about decision or preference models, potentially leading to robustness issues. In response, this paper introduces a novel linear programming (LP) framework tailored for offline reward learning. Utilizing pre-collected trajectories without online exploration, this framework estimates a feasible reward set from the primal-dual optimality conditions of a suitably designed LP, and offers an optimality guarantee with provable sample efficiency. Our LP framework also enables aligning the reward functions with human feedback, such as pairwise trajectory comparison data, while maintaining computational tractability and sample efficiency. We demonstrate that our framework potentially achieves better performance compared to the conventional maximum likelihood estimation (MLE) approach through analytical examples and numerical experiments.
Kihyun Kim 0001, Jiawei Zhang 0007, Asuman E. Ozdaglar, Pablo A. Parrilo
ICML2
2024 Uniformly Stable Algorithms for Adversarial Training and Beyond
abstract
In adversarial machine learning, neural networks suffer from a significant issue known as robust overfitting, where the robust test accuracy decreases over epochs (Rice et al., 2020). Recent research conducted by Xing et al., 2021;Xiao et al., 2022 has focused on studying the uniform stability of adversarial training. Their investigations revealed that SGD-based adversarial training fails to exhibit uniform stability, and the derived stability bounds align with the observed phenomenon of robust overfitting in experiments. This finding motivates us to develop uniformly stable algorithms specifically tailored for adversarial training. To this aim, we introduce Moreau envelope-$\mathcal{A}$ (ME-$\mathcal{A}$), a variant of the Moreau Envelope-type algorithm. We employ a Moreau envelope function to reframe the original problem as a min-min problem, separating the non-strong convexity and non-smoothness of the adversarial loss. Then, this approach alternates between solving the inner and outer minimization problems to achieve uniform stability without incurring additional computational overhead. In practical scenarios, we demonstrate the efficacy of ME-$\mathcal{A}$ in mitigating the issue of robust overfitting. Beyond its application in adversarial training, this represents a fundamental result in uniform stability analysis, as ME-$\mathcal{A}$ is the first algorithm to exhibit uniform stability for weakly-convex, non-smooth problems.
Jiancong Xiao, Jiawei Zhang 0007, Zhi-Quan Luo, Asuman E. Ozdaglar
ICML2
2023 Pruning Deep Neural Networks from a Sparsity Perspective
Enmao Diao, Ganghua Wang, Jiawei Zhang 0007, Yuhong Yang 0002, Jie Ding 0002, Vahid Tarokh
ICLR3
2023 Linearly Constrained Bilevel Optimization: A Smoothed Implicit Gradient Approach
abstract
This work develops analysis and algorithms for solving a class of bilevel optimization problems where the lower-level (LL) problems have linear constraints. Most of the existing approaches for constrained bilevel problems rely on value function-based approximate reformulations, which suffer from issues such as non-convex and non-differentiable constraints. In contrast, in this work, we develop an implicit gradient-based approach, which is easy to implement, and is suitable for machine learning applications. We first provide an in-depth understanding of the problem, by showing that the implicit objective for such problems is in general non-differentiable. However, if we add some small (linear) perturbation to the LL objective, the resulting implicit objective becomes differentiable almost surely. This key observation opens the door for developing (deterministic and stochastic) gradient-based algorithms similar to the state-of-the-art ones for unconstrained bi-level problems. We show that when the implicit function is assumed to be strongly-convex, convex, and weakly-convex, the resulting algorithms converge with guaranteed rate. Finally, we experimentally corroborate the theoretical findings and evaluate the performance of the proposed framework on numerical and adversarial learning problems.
Prashant Khanduri, Ioannis C. Tsaknakis, Jia Liu 0002, Sijia Liu 0001, Jiawei Zhang 0007, Mingyi Hong 0001
ICML6
2023 Revisiting the Linear-Programming Framework for Offline RL with General Function Approximation
abstract
Offline reinforcement learning (RL) aims to find an optimal policy for sequential decision-making using a pre-collected dataset, without further interaction with the environment. Recent theoretical progress has focused on developing sample-efficient offline RL algorithms with various relaxed assumptions on data coverage and function approximators, especially to handle the case with excessively large state-action spaces. Among them, the framework based on the linear-programming (LP) reformulation of Markov decision processes has shown promise: it enables sample-efficient offline RL with function approximation, under only partial data coverage and realizability assumptions on the function classes, with favorable computational tractability. In this work, we revisit the LP framework for offline RL, and provide a new reformulation that advances the existing results in several aspects, relaxing certain assumptions and achieving optimal statistical rates in terms of sample size. Our key enabler is to introduce proper constraints in the reformulation, instead of using any regularization as in the literature, also with careful choices of the function classes and initial state distributions. We hope our insights bring into light the use of LP formulations and the induced primal-dual minimax optimization, in offline RL.
Asuman E. Ozdaglar, Sarath Pattathil, Jiawei Zhang 0007, Kaiqing Zhang
ICML3
2023 Assisted Unsupervised Domain Adaptation
abstract
Unsupervised domain adaptation (UDA) is a popular machine learning technique that allows one to train models over diverse data collected from different domains. However, this technique requires the learner to collect a large number of properly labeled data samples, which can be costly and unrealistic in many applications. In this work, we propose a decentralized assisted learning framework for UDA. In this framework, a learner has only a limited number of labeled data samples collected from a certain source domain and aims to train a classifier for the target domain. To improve domain adaptation performance, it seeks assistance by interacting with an external service provider, who possesses many labeled data samples collected from a related source domain. We develop an assisted UDA algorithm that avoids data sharing and can significantly improve the learner’s domain adaptation performance within a few rounds of interaction. Experiments using deep neural networks on benchmark datasets demonstrate the effectiveness of this algorithm.
Jiawei Zhang 0007, Jie Ding 0002, Yi Zhou 0017
ISIT2
2022 What is a Good Metric to Study Generalization of Minimax Learners?
abstract
Minimax optimization has served as the backbone of many machine learning problems. Although the convergence behavior of optimization algorithms has been extensively studied in minimax settings, their generalization guarantees, i.e., how the model trained on empirical data performs on the unseen testing data, have been relatively under-explored. A fundamental question remains elusive: What is a good metric to study generalization of minimax learners? In this paper, we aim to answer this question by first showing that primal risk, a universal metric to study generalization in minimization problems, fails in simple examples of minimax problems. Furthermore, another popular metric, the primal-dual risk, also fails to characterize the generalization behavior for minimax problems with nonconvexity, due to non-existence of saddle points. We thus propose a new metric to study generalization of minimax learners: the primal gap, to circumvent these issues. Next, we derive generalization bounds for the primal gap in nonconvex-concave settings. As byproducts of our analysis, we also solve two open questions: establishing generalization bounds for primal risk and primal-dual risk in this setting, and in the strong sense, i.e., without assuming that the maximization and expectation can be interchanged. Finally, we leverage this new metric to compare the generalization behavior of two popular algorithms - gradient descent-ascent (GDA) and gradient descent-max (GDMax) in minimax optimization.
Asuman E. Ozdaglar, Sarath Pattathil, Jiawei Zhang 0007, Kaiqing Zhang
NeurIPS3
2021 Communication Efficient Primal-Dual Algorithm for Nonconvex Nonsmooth Distributed Optimization
abstract
Decentralized optimization problems frequently appear in the large scale machine learning problems. However, few works work on the difficult nonconvex nonsmooth case. In this paper, we propose a decentralized primal-dual algorithm to solve this type of problem in a decentralized manner and the proposed algorithm can achieve an $\mathcal{O}(1/\epsilon^2)$ iteration complexity to attain an $\epsilon-$solution, which is the well-known lower iteration complexity bound for nonconvex optimization. To our knowledge, it is the first algorithm achieving this rate under a nonconvex, nonsmooth decentralized setting. Furthermore, to reduce communication overhead, we also modifying our algorithm by compressing the vectors exchanged between agents. The iteration complexity of the algorithm with compression is still $\mathcal{O}(1/\epsilon^2)$. Besides, we apply the proposed algorithm to solve nonconvex linear regression problem and train deep learning model, both of which demonstrate the efficiency and efficacy of the proposed algorithm.
Congliang Chen, Jiawei Zhang 0007, Li Shen 0005, Peilin Zhao, Zhi-Quan Luo
AISTATS2
2021 When Expressivity Meets Trainability: Fewer than $n$ Neurons Can Work
abstract
Modern neural networks are often quite wide, causing large memory and computation costs. It is thus of great interest to train a narrower network. However, training narrow neural nets remains a challenging task. We ask two theoretical questions: Can narrow networks have as strong expressivity as wide ones? If so, does the loss function exhibit a benign optimization landscape? In this work, we provide partially affirmative answers to both questions for 1-hidden-layer networks with fewer than $n$ (sample size) neurons when the activation is smooth. First, we prove that as long as the width $m \geq 2n/d$ (where $d$ is the input dimension), its expressivity is strong, i.e., there exists at least one global minimizer with zero training loss. Second, we identify a nice local region with no local-min or saddle points. Nevertheless, it is not clear whether gradient descent can stay in this nice region. Third, we consider a constrained optimization formulation where the feasible region is the nice local region, and prove that every KKT point is a nearly global minimizer. It is expected that projected gradient methods converge to KKT points under mild technical conditions, but we leave the rigorous convergence analysis to future work. Thorough numerical results show that projected gradient methods on this constrained formulation significantly outperform SGD for training narrow neural nets.
Jiawei Zhang 0007, Yushun Zhang, Mingyi Hong 0001, Ruoyu Sun 0001, Zhi-Quan Luo
NeurIPS1
2020 A Proximal Dual Consensus Method for Linearly Coupled Multi-Agent Non-Convex Optimization
abstract
Motivated by large-scale signal processing and machine learning applications, this paper considers the distributed multi-agent optimization problem for a linearly constrained non-convex problem. Each of the agents owns a local cost function and local variable, but are coupled with each other due to the linear constraint. Most of the existing methods are either applicable for convex problems only or are developed under the non-convex setting subject to a specific type of linear constraint. There still lacks a distributed method for solving the linear constrained problem under the general and non-convex setting. In this paper, we propose such a method, called the proximal dual consensus (PDC) method, that combines a proximal technique and the dual consensus method. Theoretical analysis shows that the proposed PDC method can yield a Karush-Kuhn-Tucker solution of the linearly constrained non-convex problem and it has an O(1/ε) iteration complexity, where ε is a solution accuracy. The practical behavior of the proposed method is examined by numerical results.
Jiawei Zhang 0007, Songyang Ge, Tsung-Hui Chang, Zhi-Quan Luo
ICASSP1
2020 A Single-Loop Smoothed Gradient Descent-Ascent Algorithm for Nonconvex-Concave Min-Max Problems
abstract
Nonconvex-concave min-max problem arises in many machine learning applications including minimizing a pointwise maximum of a set of nonconvex functions and robust adversarial training of neural networks. A popular approach to solve this problem is the gradient descent-ascent (GDA) algorithm which unfortunately can exhibit oscillation in case of nonconvexity. In this paper, we introduce a ``smoothing" scheme which can be combined with GDA to stabilize the oscillation and ensure convergence to a stationary solution. We prove that the stabilized GDA algorithm can achieve an $O(1/\epsilon^2)$ iteration complexity for minimizing the pointwise maximum of a finite collection of nonconvex functions. Moreover, the smoothed GDA algorithm achieves an $O(1/\epsilon^4)$ iteration complexity for general nonconvex-concave problems. Extensions of this stabilized GDA algorithm to multi-block cases are presented. To the best of our knowledge, this is the first algorithm to achieve $O(1/\epsilon^2)$ for a class of nonconvex-concave problem. We illustrate the practical efficiency of the stabilized GDA algorithm on robust training.
Jiawei Zhang 0007, Peijun Xiao, Ruoyu Sun 0001, Zhi-Quan Luo
NeurIPS1
2019 Scalable Gaussian Process Using Inexact Admm for Big Data
abstract
Gaussian process (GP) for machine learning has been well studied over the past two decades and is now widely used in many sectors. However, the design of low-complexity GP models still remains a challenging research problem. In this paper, we propose a novel scalable GP regression model for processing big datasets, using a large number of parallel computation units. In contrast to the existing methods, we solve the classic maximum likelihood based hyper-parameter optimization problem by a carefully designed distributed alternating direction method of multipliers (ADMM). The proposed method is parallelizable over a large number of computation units. Simulation results confirm the benefits of the proposed scalable GP model over the state-of-the-art distributed methods.
Feng Yin 0001, Jiawei Zhang 0007, Wenjun Xu 0001, Shuguang Cui, Zhi-Quan Luo
ICASSP3