Title: Learning a Single Neuron Robustly to Distributional Shifts and Adversarial Label Noise
Venue: NeurIPS 2024
Accepted: yes
Keywords: Distributionally-Robust Optimization

Nonconvex Optimization

Primal-Dual Algorithm
Abstract: We study the problem of learning a single neuron with respect to the $L_2^2$-loss in the presence of adversarial distribution shifts, where the labels can be arbitrary, and the goal is to find a "best-fit" function.
More precisely, given training samples from a reference distribution $p_0$, 
the goal is to approximate the vector $\mathbf{w}^*$
which minimizes the squared loss with respect to the worst-case distribution 
that is close in $\chi^2$-divergence to $p_{0}$.
We design a computationally efficient algorithm that recovers a vector $ \hat{\mathbf{w}}$
satisfying 
$\mathbb{E}\_{p^*} (\sigma(\hat{\mathbf{w}} \cdot \mathbf{x}) - y)^2 \leq C \hspace{0.2em}  \mathbb{E}\_{p^*} (\sigma(\mathbf{w}^* \cdot \mathbf{x}) - y)^2 + \epsilon$, where $C>1$ is a dimension-independent constant and $(\mathbf{w}^*, p^*)$ is the witness attaining the min-max risk
$\min_{\mathbf{w}:\|\mathbf{w}\| \leq W} \max\_{p} \mathbb{E}\_{(\mathbf{x}, y) \sim p} (\sigma(\mathbf{w} \cdot \mathbf{x}) - y)^2 - \nu \chi^2(p, p_0)$.
Our algorithm follows the primal-dual framework and is 
designed by directly bounding the risk with respect to the original, nonconvex $L_2^2$ loss.
From an optimization standpoint, our work opens new avenues for the design of primal-dual algorithms under structured nonconvexity.

Learning a Single Neuron Robustly to Distributional Shifts and Adversarial Label Noise

Abstract

We study the problem of learning a single neuron with respect to the L 2 2 -loss in the presence of adversarial distribution shifts, where the labels can be arbitrary, and the goal is to find a 'best-fit' function. More precisely, given training samples from a reference distribution p 0 , the goal is to approximate the vector w ∗ which minimizes the squared loss with respect to the worst-case distribution that is close in χ 2 -divergence to p 0 . We design a computationally efficient algorithm that recovers a vector ˆ w satisfying E p ∗ ( σ ( ˆ w · x ) -y ) 2 ≤ C E p ∗ ( σ ( w ∗ · x ) -y ) 2 + ϵ , where C > 1 is a dimension-independent constant and ( w ∗ , p ∗ ) is the witness attaining the min-max risk min w : ∥ w ∥≤ W max p E ( x ,y ) ∼ p ( σ ( w · x ) -y ) 2 -νχ 2 ( p , p 0 ) . Our algorithm follows a primal-dual framework and is designed by directly bounding the risk with respect to the original, nonconvex L 2 2 loss. From an optimization standpoint, our work opens new avenues for the design of primal-dual algorithms under structured nonconvexity.

1 Introduction

The problem of learning a single neuron from randomly drawn labeled examples is a fundamental problem extensively studied in the machine learning literature. Given labeled examples { ( x i , y i ) : ( x i , y i ) ∈ R d × R } N i =1 drawn from a reference distribution p 0 , the goal in this context is to recover a parameter vector w ∗ 0 that minimizes the squared loss Λ σ, p 0 ( w ) over a ball of radius W > 0 :

where σ : R → R is a known (typically non-linear) non-decreasing activation function (e.g., the ReLU activation σ ( t ) = max(0 , t ) ) and we denote by OPT 0 = min w : ∥ w ∥ 2 ≤ W Λ σ, p 0 ( w ) the minimum squared loss. In the realizable setting - where y = σ ( w ∗ 0 · x ) and thus OPT 0 = 0 -this problem is well-understood and by now part of the folklore (see, e.g., [KS09; Kak+11; Sol17; YS20]). The results for the realizable setting also naturally extend to zero-mean bounded-variance label noise.

The more realistic agnostic model [Hau92; KSS92] (a.k.a. adversarial label noise) aims to identify the best-fitting neuron for a reference distribution of the examples, without any assumptions on label structure. However, it is known that in this setting finding a parameter vector with square loss OPT 0 + ϵ requires d poly (1 /ϵ ) time, even if the x -marginal distribution is Gaussian [GKK19; DKZ20; GGK20; Dia+21; DKR23]. Even if we relax our goal to achieve error O (OPT 0 ) + ϵ , efficient

algorithms only exist under strong distributional assumptions. In fact, without such assumptions, this problem is NP-hard [Sim02; MR18]. Recent work has also shown that (under cryptographic assumptions) no polynomial-time constant-factor improper learner exists even for distributions supported on the unit ball [Dia+22b]. Given these intractability results, recent work has focused on developing efficient constant-factor approximate learners under minimal distributional assumptions (see, e.g.,[Dia+20; FCG20; Dia+22a; ATV23; Gol+23; Wan+23a; Zar+24]).

This recent progress notwithstanding, prior work primarily focused on the setting where only the labels might be corrupted, without considering possible distributional shifts or heterogeneity of the data. Such distributional corruptions are frequently observed in practice and have motivated a long line of research in areas such as domain adaptation and (related to it) distributionally robust optimization (DRO); see e.g., [BEN09; ND16; RM22; Bla+24] and references therein. Thus, the main question motivating our work is:

How do adversarial changes in the underlying distribution impact the learnability of a neuron?

We study this question within the DRO framework, where the goal is to minimize the model's loss on a worst-case distribution from a set of distributions close to the reference distribution. 2 This set of distributions, known as the ambiguity set, models possible distributional shifts of the data. In addition to being interesting on its own merits, the DRO framework arises in diverse contexts, including algorithmic fairness [Has+18b] and class imbalance [Xu+20]. Moreover, it has recently found a range of applications in reinforcement learning [Kal+22; Liu+22; Lot+23; Wan+23b; Yan+23; Yu+23], robotics [Sha+20], language modeling [Liu+21], sparse neural network training [Sap+23], and defense against model extraction [Wan+23c].

Despite a range of impressive results in the DRO literature (see, e.g., recent surveys [Kuh+19; CP20; RM22; Bla+24] and references therein), algorithmic results with rigorous approximation guarantees for the loss have almost exclusively been obtained under fairly strong assumptions about the loss function involving both convexity and either smoothness or Lipschitzness, with linear regression being the prototypical example; see, e.g., [CP18; BMN21; DN21]. Unfortunately, this vanilla setting does not capture a range of machine learning applications, where a typical loss function is nonconvex. In particular, even the simplest ReLU learning problem in the realizable setting (with noise-free labels) is nonconvex. Further, existing DRO approaches for nonconvex loss functions such as [SND18; Qi+21] only guarantee convergence to a stationary point, which is insufficient for learning a ReLU neuron even without distributional ambiguity [YS20]. Motivated by this gap in our understanding, in this work we initiate a rigorous algorithmic investigation of learning a neuron (arguably the simplest non-convex problem) in the DRO setting. We hope that this work will stimulate future research in this direction, potentially addressing more complex models in a principled manner.

Due to space constraints, we defer further discussion of related work to Appendix A.

1.1 Problem Setup

To formally define our setting, we recall the definition of χ 2 -divergence between distributions p and p ′ , given by χ 2 ( p , p ′ ) := ∫ ( d p d p ′ -1 ) 2 d p ′ . We focus on the class of monotone unbounded activations introduced in [Dia+22a], for which we additionally assume convexity. Example activations in this class include the ReLU, leaky ReLU, exponential linear unit (ELU), and normalized 3 SoftPlus.

Definition 1.1 (Unbounded [Dia+22a] + Convex Activation) . Let σ : R → R be a non-decreasing convex function, and let α, β > 0 . We say σ is ( α, β ) -unbounded if it satisfies the following: ( i ) σ is β -Lipschitz; ( ii ) σ ( t 1 ) -σ ( t 2 ) ≥ α ( t 1 -t 2 ) for all t 1 ≥ t 2 ≥ 0 , and ( iii ) σ (0) = 0 .

To formally state the problem, we further define the loss, risk, and optimal value (denoted by OPT ).

Definition 1.2 (Loss, Risk, and OPT ) . Given a regularization parameter ν and a reference distribution p 0 , let P = P ( p 0 ) denote the set of all distributions that are absolutely continuous with respect to

p 0 and B ( W ) := { w : ∥ w ∥ 2 ≤ W } . We define the following:

We say that L σ ( w , p ; p 0 ) is the regularized square loss function of a vector w and a distribution p ∈ P ; and R ( w ; p 0 ) is the DRO risk of w with respect to p 0 . We call p ∗ the target distribution.

The minimization of the DRO risk as defined above corresponds to the regularized/penalized DRO formulation studied in prior work; see, e.g., [SND18; Wan+23c; MDH24]. An alternate formulation would have been to instead optimize over a restricted domain. The two are equivalent because of Lagrangian duality. We show in Claim E.1 a concrete relation between our regularization parameter ν and the chi-squared distance between the population distribution p 0 and the target distribution p ∗ . We further require that ν is sufficiently large to ensure that the resulting χ 2 ( p ∗ , p 0 ) is smaller than an absolute constant, which is in line with the DRO being used for not too large ambiguity sets [RM22].

Empirical Version If the reference distribution is the uniform distribution on N labeled examples ( x i , y i ) ∈ R d × R drawn from p 0 , we call it ̂ p 0 = ̂ p 0 ( N ) , and similarly define ̂ p ∈ P ( ̂ p 0 ) . Note that R ( w ∗ ; ̂ p 0 ) = max ̂ p ∈ P ( ̂ p 0 ) E ( x ,y ) ∼ ̂ p ( σ ( w ∗ · x ) -y ) 2 -νχ 2 ( ̂ p , ̂ p 0 ) ; if we let ̂ p ∗ denote the distribution that achieves the maximum, ̂ p ∗ has the same support as ̂ p 0 and can be interpreted as the reweighting of the samples that maximizes the regularized loss.

Formally, our goal is to solve the following learning problem.

Problem 1.3 (Robustly Learning a Single Neuron Under Distributional Shifts) . Given error parameters ϵ, δ ∈ (0 , 1) , regularization parameter ν > 0 , set radius W > 0 , and sample access to labeled examples ( x , y ) drawn i.i.d. from an unknown reference distribution p 0 , output a parameter vector ˆ w ∈ B ( W ) that is competitive with the DRO risk minimizer w ∗ = arg min w ∈ B ( W ) R ( w ; p 0 ) in the sense that with probability at least 1 -δ, ∥ ˆ w -w ∗ ∥ 2 2 ≤ C OPT+ ϵ for an absolute constant C.

While the stated goal is expressed in terms of ∥ ˆ w -w ∗ ∥ 2 , under mild distributional assumptions that we make on the reference and target distributions, this guarantee implies being competitive with the best-fit function on p ∗ in terms of both the square loss and the risk, namely Λ σ, p ∗ ( ˆ w ) = O (OPT)+ ϵ and R ( w , p 0 ) -min w ∈ B ( W ) R ( ˆ w , p 0 ) ≤ O (OPT) + ϵ . Further, our algorithm is primal-dual and it outputs a distribution ̂ p that is close to ̂ p ∗ in the chi-squared divergence.

Since the solution to Problem 1.3 has an error of O (OPT) + ϵ , when we use the term 'convergence' in our paper, we refer to the following weaker notion: the iterates of our algorithm converge to the (set of) solutions such that asymptotically all iterates lie within the set of O (OPT) + ϵ solutions, which are the target solutions, as stated in Problem 1.3.

1.2 Main Result

Our main contribution is the first polynomial sample and time algorithm for learning a neuron in a distributionally robust setting for a broad class of activations (Definition 1.1) and under mild distributional assumption on the target distribution (Assumptions 2.1 and 2.2 in Section 2.1).

Theorem 1.4 (Main Theorem - Informal) . Suppose that the learner has access to N = ˜ Ω( d/ϵ 2 ) samples drawn from the reference distribution p 0 . If all samples are bounded and the distribution p ∗ satisfies the 'margin-like' condition and concentration (Assumptions 2.1 and 2.2 in Section 2.1), then after ˜ O ( d log(1 /ϵ )) iterations, each running in sample near-linear time, with high probability Algorithm 1 recovers ˆ w such that ∥ ˆ w -w ∗ ∥ 2 2 ≤ C OPT+ ϵ , for an absolute constant C .

We emphasize that Theorem 1.4 simultaneously addresses two types of robustness: firstly, robustness concerning labels ( y ); and secondly, robustness due to shifts in the distribution ( p 0 being perturbed). This result is new even when specialized to any nontrivial activation like ReLU, realizable case (where OPT = 0 ), and the simplest Gaussian x -marginal distribution. Without distributional robustness, existing approaches, as previously discussed, yield an error of O (OPT)+ ϵ under certain

x -marginal conditions. We demonstrate that this error rate can be also achieved with respect to p ∗ in a distributionally robust context, as long as p ∗ meets the same conditions specified in [Wan+23a] among the mildest in the literature addressing non-distributionally robust agnostic setting.

1.3 Technical Overview

Our technical approach relies on three main components, described below:

Local Error Bounds Our work is inspired by optimization-theory local error bounds ('sharpness') obtained for learning a single neuron with monotone unbounded activations under structured distributions without considering distributional shift or ambiguity [MBM18; Wan+23a]. These bounds are crucial as they quantify growth of a loss function outside the set of target solutions, essentially acting as a 'signal' to guide algorithms toward target solutions in our learning problems. Concretely, under distributional assumptions on p ∗ from [Wan+23a], the following sharpness property can be established: there is an absolute constant c 1 > 0 such that ∀ w ∈ B (2 ∥ w ∗ ∥ 2 ) ,

The local error bounds in [MBM18; Wan+23a] assume identical reference and target distributions. Introducing distributional ambiguity - as in our work - invalidates this assumption, and as a result necessary distributional assumptions for sharpness may not apply to all distributions in the ambiguity set. In this work, distributional assumptions are exclusively applied to the target distribution to exploit the sharpness property proved in [Wan+23a]. We also assume that the sample covariates from the reference distribution are polynomially bounded; this assumption, which is without loss of generality, impacts only the sample and computational complexities and is satisfied by standard distributions.

Primal-Dual Algorithm Our algorithm is a principled, primal-dual algorithm leveraging the sharpness property on the target distribution, the structure of the square loss, and properties of chi-squared divergence. We control a 'gap-like' function of the iterates, Gap( ̂ w , ̂ p ; ̂ p 0 ) := L σ ( ̂ w , ̂ p ∗ ; ̂ p 0 ) -L σ ( w ∗ , ̂ p ; ̂ p 0 ) . The idea of approximating a gap and showing it reduces at a rate 1 /A k , where A k is a monotonically increasing function of k , comes from [DO19] and has been extended to primal-dual methods, including DRO settings, in [SWD21; Dia+22c; Son+22; MDH24].

Unlike past work [SWD21; Dia+22c; Son+22; MDH24], our primal problem is nonconvex, even for ReLU activations without distributional ambiguity. Unfortunately, the previously mentioned results relying on convexity do not apply in our setting. Additionally, sharpness - which appears crucial to approximating the target loss - is a local property, applying only to w such that ∥ w ∥ 2 ≤ 2 ∥ w ∗ ∥ 2 , where ∥ w ∗ ∥ 2 is unknown. This condition is trivially met at initialization, but proving it holds for all iterates requires convergence. We address this issue via an inductive argument, effectively coupling convergence analysis with localization of the iterates.

Additionally, standard primal-dual methods [CP11; Cha+18; SWD21; ACW22; Son+22] rely on bilinear coupling between primal and dual variables in L σ ( w , ̂ p ; ̂ p 0 ) . In our case, L σ ( w , ̂ p ; ̂ p 0 ) is nonlinear and nonconvex in the first argument. Recent work [MDH24] handled nonlinearity by linearizing the function using convexity of the loss, which makes the function bounded below by its linear approximation at any point. However, this approach cannot be applied to our problem as the loss is nonconvex. Instead, we control the chi-squared divergence between the target distribution and the algorithm dual iterates to bound L σ ( w , ̂ p ∗ ; ̂ p 0 ) from below, using a key structural result that we establish in Lemma 3.4. The challenges involved in proving this structural result require us to rely on chi-squared regularization and convex activation σ . Generalizing our result to all monotone unbounded activations and other strongly convex divergences like KL would need a similar structural lemma under these broader assumptions.

An interesting aspect of our analysis is that we do not rely on a convex surrogate for our problem. Instead, we constructively bound a quantity related to the DRO risk of the original square loss, justifying our algorithmic choices directly from the analysis. Although we do not consider convex surrogates, the vector field v ( w ; x , y ) , scaled by 2 β , corresponds to the gradient of the convex surrogate loss ∫ w · x 0 ( σ ( t ) -y ) d t , which has been used in prior literature on learning a single neuron under similar settings without distributional ambiguity [Kak+11; Dia+20; Wan+23a]. In our analysis, the vector field v ( w ; x , y ) is naturally motivated by the argument in the proof of Lemma 3.4.

'Concentration' of the Target Distribution To prove that our primal-dual algorithm converges, we need to prove both an upper bound and a lower bound for Gap( ̂ w , ̂ p ; ̂ p 0 ) . The lower bound relies

on sharpness; however, we need it to hold for the empirical target distribution ( ̂ p ∗ ). This requires us to translate distributional assumptions and/or their implications from p ∗ to ̂ p ∗ . Unfortunately, ̂ p ∗ is not the uniform distribution over samples drawn from p ∗ . Rather, it is the maximizing distribution in the empirical DRO risk, defined w.r.t. ̂ p 0 . This means that prior uniform convergence results do not apply. Additionally, minimax risk rates from prior statistical results, such as those in [DN21], relate R ( w ; ̂ p 0 ) and R ( w ; p 0 ) . However, they do not help in our algorithmic analysis since they do not guarantee that the sharpness holds for ̂ p ∗ .

To address these challenges, we prove (in Corollary C.2) that as long as ν is sufficiently large, there is a simple closed-form expression for ̂ p ∗ as a function of ̂ p 0 and an analogous relationship holds between p ∗ and p 0 . This allows us to leverage the fact that expectations of bounded functions with respect to ̂ p 0 closely approximate those with respect to p 0 to show that expectations with respect to ̂ p ∗ and p ∗ are similarly close. This result then implies that the sharpness also holds for ̂ p ∗ (Lemma C.6). Full details are provided in Appendix C.

2 Preliminaries

In this section, we introduce the necessary notation and state basic facts used in our analysis.

Notation Given a positive integer N , [ N ] denotes the set { 1 , 2 , . . . , N } . Given a set E , E c denotes the complement of E when the universe is clear from the context. We use I E to denote the characteristic function of a set E : I E ( x ) = 1 if x ∈ E and I E ( x ) = 0 otherwise. For vectors x and ̂ x from the d -dimensional Euclidean space R d , we use ⟨ x , ̂ x ⟩ and x · ̂ x to denote the standard inner product, while ∥ · ∥ 2 = √ ⟨· , ·⟩ denotes the ℓ 2 norm. We use ( x (1) , x (2) , . . . , x ( d ) ) to denote the entries of x ∈ R d . We write x ≤ ̂ x to indicate x ( j ) ≤ ̂ x ( j ) for all coordinates j . For r > 0 , B ( r ) := { x : ∥ x ∥ 2 ≤ r } denotes the centered ball of radius r . We use ∆ N to denote the probability simplex: ∆ n := { x ∈ R N : ∑ N j =1 x ( j ) = 1 , ∀ j ∈ [ N ] : x ( j ) ≥ 0 } . We denote by I d the identity matrix of size d × d . We write A ⪰ B to indicate that x ⊤ ( A -B ) x ≥ 0 for all x ∈ R d . For two functions f and g , we say f = ˜ O ( g ) if f = O ( g log k ( g )) for some constant k , and similarly define ˜ Ω . Weuse notation ˜ O c ( · ) and ˜ Ω c ( · ) to hide polynomial factors in (typically absolute constant) parameters c. For two distributions p and p ′ , we use p ≪ p ′ to denote that p is absolutely continuous with respect to p ′ , i.e., for all measurable sets A , p ′ ( A ) = 0 implies p ( A ) = 0 . Typically, ˆ p and ˆ q are empirical distributions, and ˆ p ≪ ˆ q is equivalent to the condition that the support of ˆ p is a subset of the support of ˆ q . For p ≪ p ′ , we use d p d p ′ to denote their Radon-Nikodym derivative, which is the quotient of probability mass functions for discrete distributions. We use χ 2 ( p , p ′ ) to denote the chi-squared divergence of p w.r.t. p ′ , i.e., χ 2 ( p , p ′ ) = ∫ ( d p d p ′ -1) 2 d p ′ .

2.1 Distributional Assumptions

Similar to [Wan+23a], we make two assumptions about the target distribution of the covariates ( p ∗ x ). First, we assume that the optimal solution w ∗ satisfies the following 'margin-like' condition:

Assumption 2.1 (Margin) . There exist absolute constants λ, γ ∈ (0 , 1] such that E x ∼ p ∗ x [ xx T I w ∗ · x ≥ γ ∥ w ∗ ∥ 2 ] ⪰ λ I , where p ∗ x is the x -marginal distribution of p ∗ .

We also assume that p ∗ x is subexponential with parameter B , which is an absolute constant.

Assumption 2.2 (Subexponential Concentration) . There exists a parameter B > 0 such that for any u ∈ B (1) and any r ≥ 1 , it holds that Pr x ∼ p ∗ x [ | u · x | ≥ r ] ≤ exp( -Br ) .

Appendix E of [Wan+23a] shows that Assumptions 2.1 and 2.2 are satisfied by several important families of distributions including Gaussians, discrete Gaussians, all isotropic log-concave distributions, the uniform distribution over {-1 , 0 , 1 } d , etc.

For simplicity, we assume the labeled samples ( x ( i ) , y ( i ) ) drawn from the reference distribution are bounded. This assumption, which does not affect the approximation constant for Problem 1.3, only impacts iteration and sample complexities. We state the bound on the covariates below, while a bound on the labels follows from prior work (Fact 2.6 stated in the next subsection).

Assumption 2.3 (Boundedness) . There exists a parameter S such that for any fixed u ∈ B (1) it holds that u · x ≤ S for all sample covariates x in the support of ̂ p 0 .

We also assume without loss of generality that ∥ w ∗ ∥ 2 2 ≥ C OPT+ ϵ for some absolute constant C , since otherwise 0 would be a valid O (OPT) + ϵ solution. Algorithmically, we can first compute the empirical risk (per Corollary C.3) of the output from our algorithm and of ˆ w = 0 and then output the solution with the lower risk to get an O (OPT) + ϵ solution; see Claim E.2 for a detailed discussion.

2.2 Auxiliary Facts

To achieve the claimed guarantees, we leverage structural properties of the loss function on the target distribution, implied by our distributional assumptions (Assumptions 2.1 and 2.2). Specifically, we make use of Lemma 2.2 and Fact C.4 from [Wan+23a], summarized in the fact below.

Fact 2.4 (Sharpness ([Wan+23a])) . Suppose p ∗ and w ∗ satisfy Assumptions 2.1 and 2.2. Let c 0 = γλα 6 B log(20 B/λ 2 ) . For all w ∈ B (2 ∥ w ∗ ∥ ) and u ∈ B (1) ,

Fact 2.4 applies to the population version of the problem. Such a result also holds for the target distribution of the empirical problem, which we state below. Note that this result cannot be obtained by appealing to uniform convergence results for learning a neuron (without distributional robustness).

Lemma 2.5 (Empirical Sharpness; Informal. See Lemma C.6) . Under Assumptions 2.1 to 2.3, for a sufficiently large sample size N as a function of B,W,S,ν,α,γ,λ,d and with high probability, for all w ∈ B (2 ∥ w ∗ ∥ ) with ∥ w -w ∗ ∥ ≥ √ ϵ and u ∈ B (1) ,

As a consequence, for c 1 = c 2 0 / (24 B ) and any w ∈ B ( W ) (where c 0 is defined in Fact 2.4), we have

where the left inequality uses Cauchy-Schwarz and the right inequality uses β -Lipschitzness of σ ( · ) [Wan+23a] also showed that the labels y can be assumed to be bounded without loss of generality.

.

Fact 2.6. Suppose p ∗ and w ∗ satisfy Assumption 2.1 and Assumption 2.2. Let y ′ = sign( y ) max {| y | , M } where for some sufficiently large absolute constant C M we define

We also make use of the following facts from convex analysis. First, let ϕ : R N → R be a differentiable function and the Bregman divergence of ϕ for any x , y ∈ R N be defined by

Fact 2.7. Let ψ ( x ) = ϕ ( x ) + ⟨ a , x ⟩ + b for some a ∈ R N and b ∈ R . Then D ψ ( y , x ) = D ϕ ( y , x ) for all x , y ∈ R N , i.e., the Bregman divergence is blind to the addition of affine terms to function ϕ .

Second, we state the first-order necessary conditions that a local maximizer must satisfy.

Fact 2.8 (First-Order Optimality Condition) . Let Ω be a closed, convex, and nonempty set and let f : Ω → R be continuously differentiable. If x ∗ is a local maximizer of f on Ω , then it holds that

If f is also concave, then Equation (7) implies that x ∗ is a global maximizer of f .

3 Algorithm and Convergence Analysis

In this section, we introduce our algorithm and state our main results, summarized in Theorem 3.1. We highlight the main components of our technical approach, while most of the technical details are deferred to the appendix, due to space constraints.

To facilitate the presentation of results, we introduce the following auxiliary notation: ℓ ( w ; x , y ) := ( σ ( w · x ) -y ) 2 , v ( w ; x , y ) := 2 β ( σ ( w · x ) -y ) x and ̂ OPT = E ( x ,y ) ∼ ̂ p ∗ ℓ ( w ∗ ; x , y ) . We also note that Assumption 2.3 implies that for all samples { x i , y i } , the function w ↦→ v ( w ; x i , y i ) is bounded above by G and κ -Lipschitz for all i ∈ [ N ] and w ∈ B ( W ) , where G = 2 βS √ d ( √ 2 βWS + M ) and κ = 2 β 2 S 2 d (see Lemma B.4 in Appendix B). Starting from this section, we write L ( w , ̂ p ) to denote L σ ( w , ̂ p ; ̂ p 0 ) , hiding the dependence on ̂ p 0 and σ . We also write Gap( w , ̂ p ) for Gap( w , ̂ p ; ̂ p 0 )

Our main algorithm (Algorithm 1) is an iterative primal-dual method with extrapolation on the primal side via g i . The vector E ̂ p i [ v ( w i ; x , y )] equals the (scaled) gradient of a surrogate loss used in prior works [Kak+11; Dia+20; Wan+23a]. In contrast to prior work, we directly bound the original square loss, with E ̂ p i [ v ( w i ; x , y )] naturally arising from our analysis. Both updates w i and ̂ p i are efficiently computable: w i involves a simple projection onto a Euclidean ball, and ̂ p i involves a projection onto a probability simplex, computable in near-linear time [Duc+08].

Algorithm 1: Main algorithm

Theorem 3.1 (Main Theorem) . Under Assumptions 2.1 to 2.3, suppose the sample size is such that N = ˜ Ω B,S,β,α,γ,λ ( W 4 ϵ 2 ( 1 + W 4 ν 2 ) ( d + W 4 log(1 /δ )) ) and ν ≥ 8 β 2 √ 6 B √ OPT (2) + ϵ/c 1 , where OPT (2) = E p ∗ [ ℓ ( w ∗ ; x , y ) 2 ] and c 1 is defined in Lemma 2.5. With probability at least 1 -δ , for all iterates w k , ̂ p k , it holds that

where D 0 = 1 2 ∥ w ∗ -w 0 ∥ 2 2 + ν 0 χ 2 ( ̂ p ∗ , ̂ p 0 ) and χ 2 ( ̂ p ∗ , ̂ p 0 ) ≤ c 1 / (1536 β 4 B ) (and therefore D 0 does not depend on the sample size N).

In particular, after at most k = ˜ O ( max { κ,G } min { ν,c 1 } log( D 0 ϵ )) iterations, it holds that

We focus on the convergence of iterates w i as claimed in Equation (8); the loss bound (Equation (9)) follows directly from the iterate convergence, while the risk bound (Equation (10)) requires a more involved analysis. Complete details for Equations (9) and (10) are provided in Appendix F.

Our strategy for the convergence analysis is as follows. Consider { a i } , a sequence of positive step sizes, and define A i as their cumulative sum ∑ i j =1 a j . Our algorithm produces a sequence of primal-dual pairs w i , ̂ p i , tracking a quantity related to the primal-dual gap, defined by:

We view ( L ( w i , ̂ p ∗ ) -L ( w ∗ , ̂ p ∗ )) as the 'primal gap' and ( L ( w ∗ , ̂ p ∗ ) -L ( w ∗ , ̂ p i )) as the 'dual gap.' Since the squared loss for ReLU and similar activations is nonconvex, L ( w , ̂ p ∗ ) is nonconvex in its first argument. Note that this gap function is not trivially non-negative (see Remark B.7), requiring an explicit lower bound proof.

Our strategy consists of deriving 'sandwiching' inequalities for the (weighted) cumulative gap ∑ k i =1 a i Gap( w i , ̂ p i ) and deducing convergence guarantees for the algorithm iterates from them. A combination of these two inequalities leads to the statement of Theorem 3.1, from which we can deduce that unless we already have an O (OPT) + ϵ solution, the iterates must be converging to the target solutions at rate 1 /A k , which we argue can be made geometrically fast.

Organization The rest of this section is organized as follows - under the standard assumptions we state in this paper, in Lemma 3.2, we prove a lower bound on Gap( w , ˆ p ) for any choice of w and ˆ p . This can be used to get a corresponding lower bound on the weighted sum ∑ k i =1 a i Gap( w i , ˆ p i ) . In Lemma 3.3 we then state an upper bound on ∑ k i =1 a i Gap( w i , ˆ p i ) ; the proof of this technical argument is deferred to Appendix D. These two bounds together give us the first inequality in Theorem 3.1. Claim B.6 then bounds below the convergence rate for our choice of a i in Algorithm 1; and indicates that it is geometric. Finally, we put everything together to prove Theorem 3.1.

To simplify the notation, we use ϕ ( ̂ p ) := χ 2 ( ̂ p , ̂ p 0 ) throughout this section. Note that D ϕ ( ̂ p , ̂ q ) = D ϕ ( ̂ q , ̂ p ) = ∑ N i =1 ( ̂ q ( i ) -̂ p ( i ) ) 2 ( i ) for any ̂ p and ̂ q in the domain.

3.1 Lower Bound on the Gap Function

We begin the convergence analysis by demonstrating a lower bound on Gap( w , p ) .

i ̂ i

Lemma 3.2 (Gap Lower Bound) . Under the setting in which Lemma 2.5 holds, for all w ∈ B (2 ∥ w ∗ ∥ 2 ) , Gap( w , ̂ p ) ≥ -12 β 2 B c 1 ̂ OPT+ c 1 2 ∥ w -w ∗ ∥ 2 2 + νD ϕ ( ̂ p ∗ , ̂ p ) .

Proof. Writing ( σ ( w · x ) -y ) 2 = ( ( σ ( w · x ) -σ ( w ∗ · x )) + ( σ ( w ∗ · x ) -y ) ) 2 and expanding the square, we have

By the Cauchy-Schwarz inequality, we further have that

where in the second inequality we used the definition of ̂ OPT and E x ∼ ̂ p ∗ x [( σ ( w · x ) -σ ( w ∗ · x )) 2 ] ≤ 6 Bβ 2 ∥ w -w ∗ ∥ 2 2 from the right inequality in Equation (5).

On the other hand, by the left inequality in Equation (5), we also have

Thus, combining Equation (11) and Equation (12), we get

where the last inequality is by 2 β √ 6 B ∥ w -w ∗ ∥ 2 √ ̂ OPT ≤ 4 β 2 6 B 2 c 1 ̂ OPT+ c 1 2 ∥ w -w ∗ ∥ 2 2 , which comes from an application of Young's inequality (Fact B.1).

Finally, we use the optimality of ̂ p ∗ , which achieves the maximum over all ̂ p ∈ P for L ( w ∗ , ̂ p ) . By the definition of a Bregman divergence, Fact 2.7, and first-order necessary condition in Fact 2.8:

Summing up Equation (13) and Equation (14) completes the proof.

3.2 Upper Bound on the Gap Function

Having obtained a lower bound on the gap function, we now show an upper bound, leveraging our algorithmic choices. The proof is rather technical and involves individually bounding L ( w i , ̂ p ∗ ) above and bounding L ( w ∗ , ̂ p i ) below to obtain an upper bound on the gap function, which equals L ( w i , ̂ p ∗ ) -L ( w ∗ , ̂ p i ) . We state this result in the next lemma, while the proof is in Appendix D.

Lemma 3.3 (Gap Upper Bound) . Let w i , ̂ p i , a i , A i evolve according to Algorithm 1, where we take, by convention, a -1 = A -1 = a 0 = A 0 = 0 and w -1 = w 0 , ̂ p -1 = ̂ p 0 . Assuming Lemma 2.5 applies, then, for all k ≥ 1 , ∑ k i =1 a i Gap( w i , ̂ p i ) is bounded above by

A critical technical component in the proof of Lemma 3.3 is how we handle issues related to nonconvexity. A key technical result that we prove and use is the following.

Lemma 3.4. Let S i := E ̂ p i [( σ ( w ∗ · x ) -σ ( w i · x )) 2 ]+ E ̂ p i [2( σ ( w i · x ) -y )( σ ( w ∗ · x ) -σ ( w i · x ))] , w i evolve according to Line 6 in Algorithm 1 and suppose we are in the setting where Lemma 2.5 holds. Then, S i ≥ E ̂ p i [ ⟨ v ( w ; x , y ) , w ∗ -w i ⟩ ] -E i where

This bound is precisely what forces us to choose chi-squared as the measure of divergence between distributions and introduce a dependence on ̂ OPT (2) . One pathway to generalize our results to other divergences would be to find a corresponding generalization to Lemma 3.4.

3.3 Proof of Main Theorem

Combining Lemma 3.2 and Lemma 3.3, we are now ready to prove our main result.

Proof of Theorem 3.1. Combining the lower bound on the gap function from Lemma 3.2 with the upper bound from Lemma 3.3 and rearranging, whenever ∥ w i ∥ 2 ≤ 2 ∥ w ∗ ∥ 2 for all i ≤ k so that Lemma 2.5 applies, we get that

To reach the first claim of the theorem, we first argue that ∑ k i =1 a i ( (4 β 2 √ 6 B √ ̂ OPT (2) /c 1 ) χ 2 ( ̂ p i , ̂ p ∗ ) -νD ϕ ( ̂ p ∗ , ̂ p i ) ) ≤ 0 . This follows from (1) Corollary C.2, by which we have ̂ p ∗ ( j ) ≥ ̂ p ( j ) 0 / 2 for all j ∈ [ N ] , hence

and (2) our choice of ν, which ensures, with high probability, that ν ≥ 8 β 2 √ 6 B √ OPT (2) + ϵ/c 1 ≥ 8 β 2 √ 6 B √ ̂ OPT (2) /c 1 , where the last inequality is because for the specified sample size, we have that ̂ OPT (2) + ϵ ≥ OPT (2) by Corollary C.9.

Second, we similarly have that with probability 1 -δ, ̂ OPT ≤ OPT+ ϵ. Hence, since Bregman divergence of a convex function is non-negative, whenever ∥ w i ∥ 2 ≤ 2 ∥ w ∗ ∥ 2 for all i ≤ k , we have

The bound χ 2 ( ̂ p ∗ , ̂ p 0 ) ≤ c 1 / (1536 β 4 B ) is proved in Claim E.1. Finally, in Appendix E, we inductively prove the following claim so that assumptions in Lemma 2.5 are satisfied.

Claim 3.5. For all iterations k ≥ 0 , ∥ w k ∥ 2 ≤ 2 ∥ w ∗ ∥ 2 .

The bound on the growth of A k follows by standard arguments and is provided as Claim B.6. Since A k grows exponentially with (1 + η ) k where η = min { ν,c 1 / 8 } 2 max { κ,G } and since D 0 (1 + η ) -k ≤ ϵ can be enforced by setting k = (1 + 1 /η ) log( D 0 /ϵ ) ≥ log( D 0 /ϵ ) / log(1 + η ) , we have that after ˜ O ( max { κ,G } min { ν,c 1 } log( D 0 /ϵ )) iterations either ∥ w k -w ∗ ∥ 2 ≤ √ ϵ or ∥ w i -w ∗ ∥ 2 ≤ C 3 √ OPT .

4 Conclusion

In this paper, we study the problem of learning a single neuron in the distributionally robust setting, with the square loss regularized by the chi-squared distance between the reference and target distributions. Our results serve as a preliminary exploration in this area, paving the way for several potential extensions. Future work includes generalizing our approach to single index models with unknown activations, expanding to neural networks comprising multiple neurons, and considering alternative ambiguity sets such as those based on the Wasserstein distance or Kullback-Leibler divergence.

References

[ACW22] A. Alacaoglu, V. Cevher, and S. J. Wright. 'On the complexity of a practical primal-dual coordinate method'. In: arXiv preprint arXiv:2201.07684 (2022).

[AHW95] P. Auer, M. Herbster, and M. K. K. Warmuth. 'Exponentially many local minima for single neurons'. In: Advances in Neural Information Processing Systems . 1995.

[ATV23] P. Awasthi, A. Tang, and A. Vijayaraghavan. 'Agnostic Learning of General ReLU Activation Using Gradient Descent'. In: The Eleventh International Conference on Learning Representations, ICLR . 2023.

[Bec17] A. Beck. First-order methods in optimization . SIAM, 2017.

[Ben+10] S. Ben-David, J. Blitzer, K. Crammer, A. Kulesza, F. Pereira, and J. W. Vaughan. 'A theory of learning from different domains'. In: Machine learning 79 (2010).

[BEN09] A. Ben-Tal, L. El Ghaoui, and A. Nemirovski. Robust optimization . Vol. 28. Princeton university press, 2009.

[BBS07] S. Bickel, M. Brückner, and T. Scheffer. 'Discriminative learning for differing training and test distributions'. In: Proceedings of the 24th international conference on Machine learning . 2007.

[BMN21] J. Blanchet, K. Murthy, and V. A. Nguyen. 'Statistical analysis of Wasserstein distributionally robust estimators'. In: Tutorials in Operations Research: Emerging Optimization Methods and Modeling Techniques with Applications . INFORMS, 2021.

[Bla+24] J. Blanchet, J. Li, S. Lin, and X. Zhang. 'Distributionally Robust Optimization and Robust Statistics'. In: 2024.

[Cha+18] A. Chambolle, M. J. Ehrhardt, P. Richtárik, and C.-B. Schonlieb. 'Stochastic primaldual hybrid gradient algorithm with arbitrary sampling and imaging applications'. In: SIAM Journal on Optimization 28.4 (2018).

[CP11] A. Chambolle and T. Pock. 'A first-order primal-dual algorithm for convex problems with applications to imaging'. In: Journal of mathematical imaging and vision 40 (2011).

[CP18] R. Chen and I. C. Paschalidis. 'A robust learning approach for regression models based on distributionally robust optimization'. In: Journal of Machine Learning Research 19.13 (2018).

[CP20] R. Chen and I. C. Paschalidis. 'Distributionally robust learning'. In: Foundations and Trends® in Optimization 4.1-2 (2020), pp. 1-243.

[Che+20] S. Chen, F. Koehler, A. Moitra, and M. Yau. 'Classification under misspecification: Halfspaces, generalized linear models, and evolvability'. In: Advances in Neural Information Processing Systems 33 (2020).

[Dia+20] I. Diakonikolas, S. Goel, S. Karmalkar, A. R. Klivans, and M. Soltanolkotabi. 'Approximation Schemes for ReLU regression'. In: Conference on Learning Theory, COLT 2020 . 2020.

[DKZ20] I. Diakonikolas, D. Kane, and N. Zarifis. 'Near-Optimal SQ Lower Bounds for Agnostically Learning Halfspaces and ReLUs under Gaussian Marginals'. In: Annual Conference on Neural Information Processing Systems . 2020.

[Dia+21] I. Diakonikolas, D. M. Kane, T. Pittas, and N. Zarifis. 'The Optimality of Polynomial Regression for Agnostic Learning under Gaussian Marginals in the SQ Model'. In: Proceedings of Machine Learning Research vol 134 (2021).

[Dia+22a] I. Diakonikolas, V. Kontonis, C. Tzamos, and N. Zarifis. 'Learning a Single Neuron with Adversarial Label Noise via Gradient Descent'. In: Conference on Learning Theory (COLT) . 2022, pp. 4313-4361.

[DPT21] I. Diakonikolas, J. H. Park, and C. Tzamos. 'ReLU Regression with Massart Noise'. In: Advances in Neural Information Processing Systems 34 (2021).

[Dia+22b] I. Diakonikolas, D. Kane, P. Manurangsi, and L. Ren. 'Hardness of Learning a Single Neuron with Adversarial Label Noise'. In: Proceedings of The 25th International Conference on Artificial Intelligence and Statistics . 2022.

[DKR23] I. Diakonikolas, D. Kane, and L. Ren. 'Near-optimal cryptographic hardness of agnostically learning halfspaces and relu regression under gaussian marginals'. In: International Conference on Machine Learning . PMLR. 2023.

[Dia+22c] J. Diakonikolas, C. Li, S. Padmanabhan, and C. Song. 'A fast scale-invariant algorithm for non-negative least squares with non-negative data'. In: Advances in Neural Information Processing Systems (2022).

[DO19] J. Diakonikolas and L. Orecchia. 'The Approximate Duality Gap Technique: A Unified Theory of First-Order Methods'. In: SIAM Journal on Optimization 29.1 (2019), pp. 660-689.

[DN19] J. Duchi and H. Namkoong. 'Variance-based regularization with convex objectives'. In: Journal of Machine Learning Research (2019).

[Duc+08] J. Duchi, S. Shalev-Shwartz, Y. Singer, and T. Chandra. 'Efficient projections onto the ℓ 1 -ball for learning in high dimensions'. In: Proceedings of the 25th international conference on Machine learning . 2008.

[DGN21] J. C. Duchi, P. W. Glynn, and H. Namkoong. 'Statistics of robust optimization: A generalized empirical likelihood approach'. In: Mathematics of Operations Research 46.3 (2021).

[DN21] J. C. Duchi and H. Namkoong. 'Learning models with uniform performance via distributionally robust optimization'. In: The Annals of Statistics 49.3 (2021).

[Dwo+12] C. Dwork, M. Hardt, T. Pitassi, O. Reingold, and R. Zemel. 'Fairness through awareness'. In: Proceedings of the 3rd Innovations in Theoretical Computer Science Conference . 2012.

[FCG20] S. Frei, Y . Cao, and Q. Gu. 'Agnostic learning of a single neuron with gradient descent'. In: Advances in Neural Information Processing Systems 33 (2020).

[GGK20] S. Goel, A. Gollakota, and A. R. Klivans. 'Statistical-Query Lower Bounds via Functional Gradients'. In: Annual Conference on Neural Information Processing Systems . 2020.

[GKK19] S. Goel, S. Karmalkar, and A. R. Klivans. 'Time/Accuracy Tradeoffs for Learning a ReLU with respect to Gaussian Marginals'. In: Advances in Neural Information Processing Systems 32 . 2019.

[Gol+23] A. Gollakota, P. Gopalan, A. R. Klivans, and K. Stavropoulos. 'Agnostically Learning Single-Index Models using Omnipredictors'. In: Thirty-seventh Conference on Neural Information Processing Systems . 2023.

[GSS14] I. J. Goodfellow, J. Shlens, and C. Szegedy. 'Explaining and harnessing adversarial examples'. In: arXiv preprint arXiv:1412.6572 (2014).

[HM13] M. Hardt and A. Moitra. 'Understanding Alternating Minimization for Matrix Completion'. In: 2013.

[Has+18a] T. Hashimoto, M. Srivastava, H. Namkoong, and P. Liang. 'Fairness Without Demographics in Repeated Loss Minimization'. In: ICML . 2018.

[Has+18b] T. Hashimoto, M. Srivastava, H. Namkoong, and P. Liang. 'Fairness without demographics in repeated loss minimization'. In: International Conference on Machine Learning . PMLR. 2018.

[Hau92] D. Haussler. 'Decision theoretic generalizations of the PAC model for neural net and other learning applications'. In: Information and Computation 100 (1992), pp. 78-150.

[Hu+18] W. Hu, G. Niu, I. Sato, and M. Sugiyama. 'Does distributionally robust supervised learning give robust classifiers?' In: ICML . 2018.

[Hua+06] J. Huang, A. Gretton, K. Borgwardt, B. Schölkopf, and A. Smola. 'Correcting sample selection bias by unlabeled data'. In: Advances in neural information processing systems 19 (2006).

[Kak+11] S. M. Kakade, V. Kanade, O. Shamir, and A. Kalai. 'Efficient learning of generalized linear and single index models with isotonic regression'. In: Advances in Neural Information Processing Systems 24 (2011).

[KS09] A. T. Kalai and R. Sastry. 'The Isotron Algorithm: High-Dimensional Isotonic Regression.' In: COLT . 2009.

[KSA19] S. M. M. Kalan, M. Soltanolkotabi, and S. Avestimehr. 'Fitting relus via sgd and quantized sgd'. In: 2019 IEEE International Symposium on Information Theory (ISIT) . IEEE. 2019.

[Kal+22] N. Kallus, X. Mao, K. Wang, and Z. Zhou. 'Doubly Robust Distributionally Robust Off-Policy Evaluation and Learning'. In: ICML . 2022.

[KMM20] S. Karmakar, A. Mukherjee, and R. Muthukumar. 'A Study of Neural Training with Iterative Non-Gradient Methods'. In: arXiv e-prints (2020), arXiv-2005.

[KSS92] M. J. Kearns, R. E. Schapire, and L. M. Sellie. 'Toward efficient agnostic learning'. In: Proceedings of the fifth annual workshop on Computational learning theory . 1992.

[KKM17] A. R. Klivans, P. K. Kothari, and R. Meka. 'Learning Halfspaces and Neural Networks with Random Initialization'. In: Proceedings of the 30th Annual Conference on Learning Theory . 2017.

[Kuh+19] D. Kuhn, P. M. Esfahani, V. A. Nguyen, and S. Shafieezadeh-Abadeh. 'Wasserstein distributionally robust optimization: Theory and applications in machine learning'. In: Operations research & management science in the age of analytics . Informs, 2019, pp. 130-166.

[Lam13] H. Lam. 'Robust Sensitivity Analysis for Stochastic Systems'. In: Mathematics of Operations Research (2013).

[Liu+21] E. Z. Liu, B. Haghgoo, A. S. Chen, A. Raghunathan, P. W. Koh, S. Sagawa, P. Liang, and C. Finn. 'Just Train Twice: Improving Group Robustness without Training Group Information'. In: ICML . 2021.

[Liu+22] Z. Liu, Q. Bai, J. Blanchet, P. Dong, W. Xu, Z. Zhou, and Z. Zhou. 'Distributionally Robust Q -Learning'. In: ICML . 2022.

[Lot+23] K. Lotidis, N. Bambos, J. Blanchet, and J. Li. 'Wasserstein Distributionally Robust Linear-Quadratic Estimation under Martingale Constraints'. In: AISTATS . 2023.

[MMR09] Y. Mansour, M. Mohri, and A. Rostamizadeh. 'Domain adaptation: Learning bounds and algorithms'. In: arXiv preprint arXiv:0902.3430 (2009).

[MR18] P. Manurangsi and D. Reichman. 'The computational complexity of training relu (s)'. In: arXiv preprint arXiv:1810.04207 (2018).

[MDH24] R. Mehta, J. Diakonikolas, and Z. Harchaoui. 'A Primal-Dual Algorithm for Faster Distributionally Robust Optimization'. In: arXiv preprint arXiv:2403.10763 (2024).

[MBM18] S. Mei, Y. Bai, and A. Montanari. 'THE LANDSCAPE OF EMPIRICAL RISK FOR NONCONVEX LOSSES'. In: The Annals of Statistics 46.6A (2018), pp. 2747-2774.

[ND16] H. Namkoong and J. C. Duchi. 'Stochastic gradient methods for distributionally robust optimization with f-divergences'. In: Advances in neural information processing systems 29 (2016).

[NW72] J. A. Nelder and R. W. M. Wedderburn. 'Generalized linear models'. In: Journal of the Royal Statistical Society: Series A (General) 135.3 (1972).

[Ore+19] Y. Oren, S. Sagawa, T. B. Hashimoto, and P. Liang. 'Distributionally robust language modeling'. In: arXiv preprint arXiv:1909.02060 (2019).

[PY09] S. J. Pan and Q. Yang. 'A survey on transfer learning'. In: IEEE Transactions on knowledge and data engineering 22.10 (2009).

[Pat+15] V. M. Patel, R. Gopalan, R. Li, and R. Chellappa. 'Visual domain adaptation: A survey of recent advances'. In: IEEE signal processing magazine 32.3 (2015).

[Qi+21] Q. Qi, Z. Guo, Y. Xu, R. Jin, and T. Yang. 'An online method for a class of distributionally robust optimization with non-convex objectives'. In: Advances in Neural Information Processing Systems 34 (2021), pp. 10067-10080.

[RM22] H. Rahimian and S. Mehrotra. 'Frameworks and results in distributionally robust optimization'. In: Open Journal of Mathematical Optimization (2022).

[Sap+23] H. Sapkota, D. Wang, Z. Tao, and Q. Yu. 'Distributionally Robust Ensemble of Lottery Tickets Towards Calibrated Sparse Network Training'. In: NeurIPS . 2023.

[Sha17] A. Shapiro. 'Distributionally robust stochastic programming'. In: SIAM Journal on Optimization 27.4 (2017).

[Sha+20] V. D. Sharma, M. Toubeh, L. Zhou, and P. Tokekar. 'Risk-Aware Planning and Assignment for Ground Vehicles using Uncertain Perception from Aerial Vehicles'. In: 2020 IEEE/RSJ (IROS) . 2020.

[Shi00] H. Shimodaira. 'Improving predictive inference under covariate shift by weighting the log-likelihood function'. In: Journal of statistical planning and inference 90.2 (2000).

[Sim02] J. Sima. 'Training a single sigmoidal neuron is hard'. In: Neural computation 14.11 (2002).

[SND18] A. Sinha, H. Namkoong, and J. Duchi. 'Certifying Some Distributional Robustness with Principled Adversarial Training'. In: International Conference on Learning Representations . 2018.

[Sol17] M. Soltanolkotabi. 'Learning ReLUs via gradient descent'. In: Advances in neural information processing systems 30 (2017).

[Son+22] C. Song, C. Y. Lin, S. Wright, and J. Diakonikolas. 'Coordinate Linear Variance Reduction for Generalized Linear Programming'. In: NeurIPS . 2022.

[SWD21] C. Song, S. J. Wright, and J. Diakonikolas. 'Variance Reduction via Primal-Dual Accelerated Dual Averaging for Nonsmooth Convex Finite-Sums'. In: Proc. ICML'21 . 2021.

[SJ19] M. Staib and S. Jegelka. 'Distributionally robust optimization and generalization in kernel methods'. In: Advances in Neural Information Processing Systems 32 (2019).

[Tan+18] C. Tan, F. Sun, T. Kong, W. Zhang, C. Yang, and C. Liu. 'A survey on deep transfer learning'. In: Artificial Neural Networks and Machine Learning-ICANN 2018 . Springer. 2018.

[Wan+23a] P. Wang, N. Zarifis, I. Diakonikolas, and J. Diakonikolas. 'Robustly Learning a Single Neuron via Sharpness'. In: Proceedings of the 40th International Conference on Machine Learning . Vol. 202. Proceedings of Machine Learning Research. PMLR, July 2023, pp. 36541-36577.

[Wan+23b] S. Wang, N. Si, J. Blanchet, and Z. Zhou. 'A Finite Sample Complexity Bound for Distributionally Robust Q-learning'. In: AISTATS . 2023.

[Wan+23c] Z. Wang, L. Shen, T. Liu, T. Duan, Y. Zhu, D. Zhan, D. S. Doermann, and M. Gao. 'Defending against Data-Free Model Extraction by Distributionally Robust Defensive Training'. In: NIPS . 2023.

[Xu+20] Z. Xu, C. Dan, J. Khim, and P. Ravikumar. 'Class-weighted classification: Trade-offs and robust approaches'. In: International conference on machine learning . PMLR. 2020.

[Yan+23] Z. Yang, Y. Guo, P. Xu, A. Liu, and A. Anandkumar. 'Distributionally Robust Policy Gradient for Offline Contextual Bandits'. In: AISTATS . 2023.

[YS20] G. Yehudai and O. Shamir. 'Learning a single neuron with gradient methods'. In: Conference on Learning Theory . 2020, pp. 3756-3786.

[Yu+23] Z. Yu, L. Dai, S. Xu, S. Gao, and C. P. Ho. 'Fast Bellman Updates for Wasserstein Distributionally Robust MDPs'. In: NeurIPS . Curran Associates, Inc., 2023.

[Zar+24] N. Zarifis, P. Wang, I. Diakonikolas, and J. Diakonikolas. 'Robustly Learning SingleIndex Models via Alignment Sharpness'. In: 41th International Conference on Machine Learning, arXiv preprint: 2402.17756 (2024).

[Zha+21] R. Zhai, C. Dan, Z. Kolter, and P. Ravikumar. 'DORO: Distributional and Outlier Robust Optimization'. In: Proceedings of the 38th International Conference on Machine Learning . Vol. 139. Proceedings of Machine Learning Research. PMLR, July 2021, pp. 12345-12355.

[Zhu+20] J.-J. Zhu, W. Jitkrittum, M. Diehl, and B. Schölkopf. 'Kernel distributionally robust optimization'. In: arXiv preprint arXiv:2006.06981 (2020).

Supplementary Material

Organization In Appendix A we briefly discuss related work. In Appendix B we set up some additional preliminaries for the rest of the appendix. In Appendix C we show that expectations of some important functions with respect to ̂ p ∗ are close to their expectation with respect to p ∗ . In Appendix D we give a detailed proof of an upper bound on the gap of the iterates our algorithm generates (i.e. Lemma 3.3). Finally, in Appendix F we show that the estimate of w ∗ our algorithm returns is a constant factor approximation to the squared loss of w ∗ with respect to the target distribution.

A Related Work

Learning Noisy Neurons The problem of learning noisy neurons has been extensively explored in recent decades; notable works include [Kak+11; KS09; KKM17; NW72]. In the recent past, the focus has shifted towards specific activation functions such as ReLUs, under both easy noise models such as realizable/random additive noise [KSA19; Sol17; YS20] and more difficult ones like adversarial label noise [Dia+20; DKZ20; Dia+21; Dia+22a; Dia+22b; GGK20; GKK19; Wan+23a].

Even with clean labels, this problem has exponentially many local minima when using squared loss [AHW95]. Unfortunately, directly minimizing the squared loss using (S)GD on a bounded distribution does not converge to the global optimum with probability 1 [YS20]. Even so, gradient based methods can achieve suboptimal rates in the agnostic setting ( √ OPT 0 vs OPT 0 ) for distributions with mild distributional assumptions [FCG20]. Making slightly stronger assumptions on the marginal does allow us to get efficient constant factor approximations. [Dia+20] developed an efficient learning method that is able to handle this in the presence of adversarial label noise and for isotropic logconcave distributions of the covariates. This was later extended to broader classes of activation functions and under weaker distributional assumptions by [Dia+22a; Wan+23a]. Without specific distributional assumptions, learning remains computationally difficult [Dia+22b; HM13]. The challenges extend to distribution-free scenarios with semi-random label noise, where methods like those in [DPT21] address bounded noise, and [KMM20] and [Che+20] explore stricter forms of Massart noise in learning a neuron with noise. In this paper, we consider the harder setting of distributionally robust optimization, where an adversary is allowed to impose not only errors in the labels, but also adversarial shifts in the underlying distribution of the covariates.

Distributionally Robust Optimization Distributional mismatches in data have been extensively studied in the context of learning from noisy data. This includes covariate shift, where the marginal distributions might be perturbed, [BBS07; Hua+06; Shi00], and changes in label proportions [Dwo+12; Xu+20]. This research also extends to domain adaptation and transfer learning [Ben+10; MMR09; PY09; Pat+15; Tan+18]. Distributionally robust optimization (DRO) has a rich history in optimization [BEN09; Sha17] and has gained traction in machine learning [DGN21; DN21; Kuh+19; ND16; SJ19; Zhu+20], showing mixed success across applications like language modeling [Ore+19], class imbalance correction [Xu+20], and group fairness [Has+18b].

Specifically this has also been studied in the context of linear regression and other function approximation [BMN21; CP18; DN21]. Typically, DRO is often very sensitive to additional sources of noise, such as outliers ([Has+18a; Hu+18; Zha+21]). However, prior work makes strong assumptions on the label noise as well as requiring convexity of the loss. We study the problem of learning a neuron where the labels have no guaranteed structure, effectively studying the setting for a combination of two notions of robustness - agnostic learning as well as covariate shift.

B Supplementary Preliminaries

B.1 Additional Notation

Given an m × n matrix A , the operator norm of A is defined in the usual way as ∥ A ∥ op = sup {∥ Ax ∥ 2 : x ∈ R n , ∥ x ∥ 2 ≤ 1 } . For problems ( P ) and ( P ′ ) , we use ( P ) ≡ ( P ′ ) to denote the equivalence of ( P ) and ( P ′ ) . For a vector space E , we use E ∗ to denote its dual space.

B.2 Standard Facts and Proofs

Fact B.1 (Young's inequality) . If a ≥ 0 and b ≥ 0 are nonnegative real numbers and if p > 1 and q > 1 are real numbers such that

then

Equality holds if and only if a p = b q .

Fact B.2 (Hoeffding's Inequality) . Let X 1 , X 2 , . . . , X n be independent random variables such that a i ≤ X i ≤ b i almost surely for all i . Let X = 1 n ∑ n i =1 X i . Then, for any t > 0 ,

Fact 2.7. Let ψ ( x ) = ϕ ( x ) + ⟨ a , x ⟩ + b for some a ∈ R N and b ∈ R . Then D ψ ( y , x ) = D ϕ ( y , x ) for all x , y ∈ R N , i.e., the Bregman divergence is blind to the addition of affine terms to function ϕ .

Proof of Fact 2.7. The Bregman divergence D ϕ and D ψ are defined by:

Since ∇ ψ ( x ) = ∇ ϕ ( x ) + a , substituting in the definition gives:

Thus, the Bregman divergence is blind to the addition of linear terms to the function ϕ .

B.3 Auxiliary Facts

We first state and prove Lemma B.3 to obtain upper bounds on the norm of each point, projections onto vectors of norm at most W , and the loss value at each point.

Lemma B.3 (Boundedness) . Fix w ∈ B ( W ) . For all samples ( x i , y i ) with truncated labels | y i | < M as per Fact 2.6 and bounded covariates as per Assumption 2.3, it holds that

Proof. Equation (18) follows from Assumption 2.3, as

To prove Equation (19), for each coordinate j ∈ [ d ] , we have | x ( j ) | = sign( x ( j ) ) e j · x ≤ S , by again using Assumption 2.3. Therefore, ∥ x ∥ 2 ≤ S √ d .

For Equation (20), we recall Fact 2.6 that for some sufficiently large absolute constant C M , it holds that | y | ≤ M := C M WBβ log( βBW/ϵ ) . Thus, Equation (20) follows from Young's inequality (Fact B.1) and Equation (19), since | σ ( t ) | ≤ β | t | , which follows from β -Lipschitzness of σ and σ (0) = 0 .

Lemma B.4 (Bounds on v ) . Let

Then v is uniformly bounded by G in ℓ 2 -norm and κ -Lipschitz for all samples x , y with truncated labels | y | < M and bounded covariates as per Assumption 2.3, where G = 2 βS √ d ( √ 2 βWS + M ) and κ = 2 β 2 S 2 d .

Proof. We first uniformly upper bound v . An application of Lemma B.3 gives us,

Taking square roots, we get ∥ v ( w ; x , y ) ∥ 2 ≤ 2 βS √ d ( √ 2 βWS + M ) =: G .

We now upper bound the Lipschitz constant κ . We will use the fact that σ is β -Lipschitz.

Corollary B.5. Fix a reference distribution p 0 . Suppose ∥ v ( w ; x , y ) ∥ 2 ≤ G for all w almost surely. Then for all distributions p , q ∈ P ( p 0 ) it holds that

Proof.

where ( i ) is an application of Jensen's inquality and ( ii ) follows from Lemma B.4.

Claim B.6 (Convergence Rate) . For all i ≥ 0 , let a i be defined as in Line 3. Then it holds that 2 G 2 a i 2 1+0 . 5 c 1 A i ≤ ν 0 + νA i -1 and 2 κ 2 a i 2 1+0 . 5 c 1 A i ≤ 1+0 . 5 c 1 A i -1 4 for all i . Moreover, A k = ∑ k i =0 a i = ((1 + min { ν,c 1 / 8 } 2 max { κ,G } ) k -1) min { ν 0 , 1 / 4 } / min { ν, c 1 / 8 } .

Proof. In order for both 2 G 2 a i 2 1+0 . 5 c 1 A i ≤ ν 0 + νA i -1 and 2 κ 2 a i 2 1+0 . 5 c 1 A i ≤ 1+0 . 5 c 1 A i -1 4 to hold for all iterations i , it suffices that

for which it suffices to enforce

where we used A i -1 ≤ A i .

Taking a square root on both sides using a i > 0 , we obtain

Solving this recurrence relation using Mathematica, we compute that for all iterations i and k ,

Remark B.7 . Note that in our case, the gap is not guarenteed to be non-negative as is usually the case for convex-concave min-max problems. Recall that Gap( w , ̂ p ) = ( L ( w , ̂ p ∗ ) -L ( w ∗ , ̂ p ∗ )) + ( L ( w ∗ , ̂ p ∗ ) -L ( w ∗ , ̂ p ))

Consider the following example:

Let p 0 be the uniform distribution over { ( -2 , 2) , (2 , 1 . 5) } , σ ≡ ReLU and ν = 0 . Then, w ∗ = -1 and p ∗ is the distribution which places all its mass on (2 , 1 . 5) . Then, Gap(1 , p ∗ ) = L (1 , p ∗ ) -L ( -1 , p ∗ ) = 0 . 25 -2 . 25 < 0 .

This is why we also need an explicit lower bound on the Gap that we have shown in Lemma 3.2.

C Concentration

Recall that that ̂ q w is not guaranteed to act as an empirical estimate of q w , because we cannot draw samples from the (unknown) distribution q w but only from p 0 . In this section, we show that for certain important functions f , it holds that E ̂ q w [ f ] ≈ E q w [ f ] . We will abuse terminology and say that f 'concentrates' with respect to q w .

Organization: In Appendix C.1 we derive closed-form expressions for q w and ̂ q w in terms of p 0 and ̂ p 0 respectively. Note that bounded functions concentrate with respect to p 0 . In Appendix C.2 we use the closed-form expressions found in Appendix C.1 to translate these concentration properties to ̂ q w . Finally, in Appendix C.3 we show that ̂ p ∗ satisfies sharpness, ̂ OPT ≈ OPT and ̂ OPT (2) ≈ OPT (2) .

C.1 Closed-form expression

The following lemma gives us a closed-form expression for q w and ̂ q w in terms of p 0 and ̂ p 0 , respectively.

We start with an additional definition to Definition 1.2:

Lemma C.1 (Closed-form q w ) . Let p 0 be a fixed distribution. Then, there exists ξ ∈ R such that,

When p 0 is the empirical distribution ̂ p 0 ( N ) , this result implies that there exists ˆ ξ ∈ R such that

The constants ξ and ˆ ξ can be interpreted as normalization that ensures ∫ d q w = ∫ d ̂ q w = 1 .

Proof. Recall that the dual feasible set is given by P = P ( p 0 ) = { p ≪ p 0 : ∫ d p = 1 , p ≥ 0 } = { p ≪ p 0 : ∫ d p d p 0 d p 0 = 1 , d p d p 0 ≥ 0 } and the function p ↦→ L ( w , p ) is strongly concave.

Consider the following optimization problem

By Fact 2.8, the first-order necessary and sufficient condition that corresponds to q w := arg max p ∈ P ( p 0 ) L ( w , p ) is the following: for any p ∈ P ( p 0 ) ,

where we recall that both ∇ p L ( w , q w ) and Radon-Nikodym derivatives d p d p 0 , d q w d p 0 are real-valued measurable functions on R d × R . We will also write ℓ = ℓ ( w ∗ , · , · ) for short.

We claim Equation (21) is satisfied if there exists ξ ∈ R and a bounded measurable function ψ ≥ 0 such that

Indeed, for any p ∈ P ( p 0 ) ,

where (i) is because ∫ d q w d p 0 > 0 ( d q w d p 0 ) d p 0 = ∫ ( d q w d p 0 ) d p 0 .

Observe from the definition of L ( w , q w ) that ∇ p L ( w , q w )( x , y ) = ℓ ( w ; x , y ) -2 ν ( d q w d p 0 ( x , y ) -1) . Plugging this into Equation (22) and rearranging, we have,

For the case where d p ∗ d p 0 > 0 , d p ∗ d p 0 = 2 ν + ℓ -ξ 2 ν , so the condition d q w d p 0 > 0 becomes 2 ν + ℓ -ξ > 0 . On the other hand, if the above condition fails, it has to be the case that d q w d p 0 = 0 . Combining, we have

Instead of using the expression in Lemma C.1, we will set ν to be big enough to ensure that there is no maximum in the expression for q w . This is captured in Corollary C.2.

Corollary C.2 (Simpler Closed-form q w ) . Fix w ∈ R d . If ν ≥ 1 2 E p 0 ℓ ( w ) , then

Similarly, if ν ≥ 1 2 E ̂ p 0 ℓ ( w ) , then q w ( i ) > 0 for all i ∈ [ N ] , and

Furthermore, if ν ≥ E ̂ p 0 ℓ ( w ) , then, in particular, for each coordinate j ∈ [ N ] , we have

Similarly, if ν ≥ E p 0 ℓ ( w ) , then for any non-negative function g , we have

Recall from Definition 1.2 that when w = w ∗ , we define p ∗ = q w ∗ and ̂ p ∗ = ̂ q w ∗ . If ν ≥ 8 β 2 √ 6 B √ OPT (2) + ϵ/c 1 as assumed in Theorem 3.1, then both conditions ν ≥ E ̂ p 0 ℓ ( w ) and ν ≥ E p 0 ℓ ( w ) hold.

Proof. Setting ν ≥ 1 2 E p 0 ℓ ( w ) and ξ = E p 0 ℓ ( w ) in Lemma C.1 implies ℓ ( w ; x , y ) -E p 0 ℓ ( w ; x , y )+ 2 ν > 0 , which, in turn, means

The empirical version follows analogously.

To establish the last claim, we show that 8 β 2 √ 6 B √ ̂ OPT (2) /c 1 ≥ E ̂ p 0 ℓ ( w ) . By Corollary C.9, it holds that

By definition of ̂ p ∗ , we have L ( w ∗ , ̂ p ∗ ) ≥ L ( w ∗ , ̂ p 0 ) = E ̂ p 0 ( σ ( w ∗ · x ) -y ) 2 -νχ 2 ( ̂ p 0 , ̂ p 0 ) = E ̂ p 0 ( σ ( w ∗ · x ) -y ) 2 . Combining, we obtain 8 β 2 √ 6 B √ ̂ OPT (2) /c 1 ≥ E ̂ p 0 ℓ ( w )8 β 2 √ 6 B/c 1 . We conclude by observing 8 β 2 √ 6 B/c 1 ≥ 1 .

Another consequence of Corollary C.2 is a closed form expression for the risk, as a varianceregularized loss, similar to [DN19; Lam13].

Corollary C.3. Fix an arbitrary distribution p 0 . Recall the risk defined in Definition 1.2,

If ν ≥ 1 2 E ̂ p 0 ℓ ( w ) , it holds that

Proof. Both of these follow from Corollary C.2. To see the first equality holds, observe that,

The second equality follows by a similar substitution.

Setting d q w as per Corollary C.2, we get

Finally, an important consequence of Lemma C.1 is that it is possible to efficiently compute the risk of a given vector w with respect to ̂ p 0 . We use this to compare the risk of our final output with the risk that is achieved by the zero vector.

C.2 Concentration

The expression we get in Corollary C.2 for p ∗ in terms of p 0 allows us to translate concentration properties of p 0 to p ∗ . We first state and prove a helper lemma, Lemma C.4, that shows E p 0 ℓ ( w ∗ ) ≈ E ̂ p 0 ℓ ( w ∗ ) . Note that this is for the reference distribution p 0 , and not the target distribution p ∗ .

For ease of notation, we define U := 2 β 2 W 2 ( S 2 + C 2 M B 2 log 2 ( WBβ/ϵ )) , which is the upper bound for the loss value in Equation (20) throughout this section.

Lemma C.4. Suppose p 0 satisfies Assumption 2.3. Then for any fixed w ∈ B ( W ) and all t > 0 , it holds that

with probability at least 1 -2 exp ( -t 2 N 8( β 2 W 2 ( S 2 + C 2 M B 2 log 2 ( WBβ/ϵ ))) 2 ) . In particular, the above inequality holds for w ∗ .

Proof. By Equation (20), ∀ i ∈ [ N ] , 0 ≤ ( σ ( w · x i ) -y i ) 2 ≤ U . Hoeffding's inequality (Fact B.2) now implies that, for all t > 0 ,

Rearranging and plugging the definition of U , we get the lemma.

We now use Lemma C.4 show that bounded Lipschitz functions concentrate with respect to p ∗ .

Lemma C.5. Fix ζ > 0 . Let h = h ( z ; x , y ) : B ( ζ ) × R d × R → R be a measurable function with respect to x , y that satisfies the condition that | h ( z ; · ; · ) | ≤ b almost surely. Then, for N = O B,S,β ( b 2 t 2 ( 1 + W 4 log 4 ( W/ϵ ) ν 2 ) log(1 /δ ) ) samples drawn from the reference distribution p 0 to construct ̂ p 0 ( N ) , for any fixed z ∈ B ( ζ ) , with probability at least 1 -4 δ , it holds that

Moreover, suppose z ↦→ h ( z ; x , y ) is a -Lipschitz. Then, for

with probability at least 1 -4 δ , it holds that for all z ∈ B ( ζ ) ,

Proof. We use Lemma C.1 to change the distribution with respect to which we are taking the expectation,

Lemma C.4 now implies that with probability 1 -2 exp( -2 Nt 2 / ( bU/ 2 ν ) 2 ) ,

We now show that the expectation on the right hand side concentrates. To this end, we will use Hoeffding's inequality (Fact B.2). To apply this, we will need a bound on the quantity in the expectation. We bound this via an application of Equation (20) and the fact that | h | ≤ b to get,

This means, with probability at least 1 -2 exp( -2 t 2 N/ ( b 2 (1 + U/ 2 ν ) 2 )) ,

Since w ↦→ h ( w ) is a -Lipschitz, a standard net argument over exp( O ( d log( ζa/t ))) vectors yields: with probability at least 1 -2 exp( O ( d log( ζa/t ) -t 2 N/ ( b 2 (1 + U/ 2 ν ) 2 ))) , it holds that for all z ∈ B ( ζ ) ,

Putting things together, we see that if we choose

with probability at least 1 -4 δ , for all z ∈ B ( ζ ) ,

C.3 Sharpness and Optimal Loss Value

Finally, as a consequence of Lemma C.5, we can derive that ̂ p ∗ satisfies sharpness, ̂ OPT ≈ OPT and ̂ OPT (2) ≈ OPT (2) .

Lemma C.6 (Shaprness for ̂ p ∗ ) . Suppose Assumptions 2.1 to 2.3 are satisfied, then for large enough N :

with probability at least 1 -4 δ , for all w ∈ B (2 ∥ w ∗ ∥ ) with ∥ w -w ∗ ∥ ≥ √ ϵ and u ∈ B (1) ,

Proof. Fact 2.4 shows that p ∗ the conditions above (with different constants). We need to translate these to ̂ p ∗ . For each of the inequalities above, we will do this via an application of Lemma C.5.

Proof of Equation (25) : Set h in Lemma C.5 to be h ( w ; x , y ) := ( σ ( w · x ) -σ ( w ∗ · x ))( w · x -w ∗ · x ) . We proceed to set the constants a and b that used in Lemma C.5. Equation (18) implies that for all ( x , y ) in the support of ̂ p 0 , | h ( w ; x , y ) | ≤ 4 βW 2 S 2 =: b . Also, w ↦→ h ( w ) is a := 2 WS 2 ( β +1) √ d -Lipschitz as a consequence of Equations (18) and (19).

Lemma C.5 now gives us that for N = ˜ O B,S,β ( W 4 t 2 ( 1 + W 4 log 4 (1 /ϵ ) ν 2 ) ( d +log(1 /δ )) ) with probability at least 1 -4 δ , for all w ∈ B (2 ∥ w ∗ ∥ ) ,

Using the fact that ∥ w -w ∗ ∥ 2 ≥ √ ϵ , we set t = c 0 ϵ/ 2 , giving us the sample complexity

Proof of Equation (26) : This follows analogously to the proof above. Set h ( u ; x , y ) = ( x · u ) τ in Lemma C.5 for τ = 2 , 4 and we proceed to calculate constants a, b . By Equations (18) and (19), it holds that h ( u ) ≤ S 4 =: b and u ↦→ h ( u ) is a := 4 S 4 √ d -Lipschitz. Setting t = B , by Lemma C.5, for N = ˜ O B,S,β (( 1 + W 4 log 4 ( W/ϵ ) ν 2 ) ( d + log(1 /δ )) ) the conclusion follows. Note that this is dominated by Equation (27).

We now show that OPT ≈ ̂ OPT .

Lemma C.7. Suppose Assumptions 2.1 to 2.3 are satisfied and the sample size N is large enough and N = ˜ O B,S,β ( W 4 log 4 (1 /ϵ ) t 2 ( 1 + W 4 log 4 (1 /ϵ ) ν 2 ) log(1 /δ ) ) . Then for any fixed w ∈ B ( W ) and all t > 0 , it holds that

with probability at least 1 -4 δ. In particular, the above inequality holds for w ∗ , i.e., | OPT -̂ OPT | ≤ t .

Proof. Set b := β 2 W 2 (2 S 2 + 2 C 2 M B 2 log 2 ( WBβ/ϵ ) ≥ ∥ ℓ ( w ) ∥ 2 in Lemma C.5, where the inequality is a consequence of Lemma B.3. Then, with N = ˜ O B,S,β ( W 4 log 4 (1 /ϵ ) t 2 ( 1 + W 4 log 4 (1 /ϵ ) ν 2 ) log(1 /δ ) ) samples, with probability 1 -4 δ , | E p ∗ ℓ ( w ; x , y ) -E ̂ p ∗ ℓ ( w ; x , y ) | ≤ t .

Finally, we show OPT (2) ≈ ̂ OPT (2) .

Lemma C.8. Suppose Assumptions 2.1 to 2.3 are satisfied and the sample size N is large enough and N = ˜ O B,S,β ( W 8 log 8 (1 /ϵ ) t 2 ( 1 + W 4 log 4 (1 /ϵ ) ν 2 ) log(1 /δ ) ) . Then for any fixed w ∈ B ( W ) and all t > 0 , it holds that

with probability at least 1 -4 δ. In particular, the above inequality holds for w ∗ , i.e., | OPT (2) -̂ OPT (2) | ≤ t .

Proof. Analogous to the previous proof, observe that ∥ ℓ 2 ( w ) ∥ 2 ≤ 8 β 4 W 4 ( S 4 + C 4 M B 4 log 4 ( WBβ/ϵ ) =: b . By Lemma C.5, with N = ˜ O B,S,β ( W 8 log 8 (1 /ϵ ) t 2 ( 1 + W 4 log 4 (1 /ϵ ) ν 2 ) log(1 /δ ) ) , with probability 1 -4 δ , it holds that | E p ∗ ℓ 2 ( w ; x , y ) -E ̂ p ∗ ℓ 2 ( w ; x , y ) | ≤ t .

We capture properties of OPT and OPT (2) in the following corollary.

Corollary C.9 (Properties of OPT , OPT (2) ) . Suppose Assumptions 2.1 to 2.3 are satisfied and the sample size N is large enough and N = ˜ O B,S,β ( W 8 log 8 (1 /ϵ ) t 2 ( 1 + W 4 log 4 (1 /ϵ ) ν 2 ) log(1 /δ ) ) . Then for all t > 0 , the following hold:

With probability 1 -4 δ , | OPT (2) -̂ OPT (2) | ≤ t .

With probability 1 -4 δ , | OPT -̂ OPT | ≤ t .

Proof. The first two items follow immediately from Lemma C.8 and Lemma C.7. The third item is a consequence of Cauchy-Schwarz.

D Gap Upper Bound

To prove Lemma 3.3, we need to construct an upper bound on the gap Gap( w i , ̂ p i ) = L ( w i , ̂ p ∗ ) -L ( w ∗ , ̂ p i ) . To achieve this, we establish an upper bound on L ( w i , ̂ p ∗ ) , which motivates the update rule for ̂ p i . We also establish a lower bound on L ( w ∗ , ̂ p i ) , which guides the update rule for w i and the construction of g i . Note that the construction of the lower bound is more challenging here, due to the nonconvexity of the square loss. This is where most of the (non-standard) technical work happens. To simplify the notation, we use ϕ ( ̂ p ) := χ 2 ( ̂ p , ̂ p 0 ) throughout this section.

Upper bound on L ( w i , ̂ p ∗ ) . We begin the analysis with the construction of the upper bound, which is used for defining the dual updates. Most of this construction follows a similar argument as used in other primal-dual methods such as [Dia+22c; SWD21].

Lemma D.1 (Upper Bound on a i L ( w i , ̂ p ∗ ) ) . Let ̂ p i evolve as outlined in Line 7. Then, for all i ≥ 1 ,

Proof. Recall that ϕ ( ̂ p ) := χ 2 ( ̂ p , ̂ p 0 ) . Observe that L ( w i , ̂ p ∗ ) as a function of ̂ p ∗ is linear minus the nonlinearity νϕ . We could directly maximize this function and define ̂ p i correspondingly (which would lead to a valid upper bound); however, such an approach appears insufficient for obtaining our results. Instead, adding and subtracting ( ν 0 + νA i -1 ) D ϕ ( ̂ p ∗ , ̂ p i -1 ) , we have

where we define, for notational convenience:

Observe that by the definition of ̂ p i , h ( ̂ p ) is maximized by ̂ p i . Hence, using the definition of a Bregman divergence, we have that

where in the inequality we used that ⟨∇ h ( ̂ p i ) , ̂ p ∗ -̂ p i ⟩ ≤ 0 (as ̂ p i maximizes h ) and D h ( ̂ p ∗ , ̂ p i ) = -( ν 0 + νA i ) D ϕ ( ̂ p ∗ , ̂ p i ) (as h ( ̂ p ) can be expressed as -( ν 0 + νA i ) ϕ ( ̂ p ) plus terms that are either linear in ̂ p or independent of it. See Fact 2.7). Combining with Equation (28) and the definition of h and simplifying, the claimed bound follows.

An important feature of Lemma D.1 is that the first two Bregman divergence terms usefully telescope, while the last one is negative and can be used in controlling the error terms arising from the algorithmic choices.

Lower bound on L ( w ∗ , ̂ p i ) . The more technical part of our analysis concerns the construction of a lower bound on L ( w ∗ , ̂ p i ) , which leads to update rule for w ∗ . In standard, Chambolle-Pock-style primal-dual algorithms [ACW22; CP11; SWD21], where the coupling L ( w , ̂ p ) between the primal and the dual is bilinear , the lower bound would be constructed using an analogue of the upper bound, with a small difference to correct for the fact that w i is updated before ̂ p i and so one cannot use information about ̂ p i in the w i update. This is done using an extrapolation idea, which replaces ̂ p i with an extrapolated value from prior two iterations and controls for the introduced error.

In our case, however, the coupling is not only nonlinear, but also nonconvex because ℓ ( w ; x , y ) = ( σ ( w · x ) -y ) 2 is nonconvex. Nonlinearity is an issue because if we were to follow an analogue of

the construction from Lemma D.1, we would need to assume that we can efficiently minimize over w the sum of L ( w , ̂ p ) and a convex function (e.g., a quadratic), which translates into proximal point updates for the L 2 2 loss for which efficient computation is generally unclear. Nonlinearity alone (but assuming convexity) has been handled in the very recent prior work [MDH24], where this issue is addressed using convexity of the nonlinear function to bound it below by its linear approximation around w i . Unfortunately, as mentioned before, this approach cannot apply here as we do not have convexity. Instead, we use a rather intricate argument that relies on monotonicity and Lipschitzness properties of the activation σ and structural properties of the problem which only hold with respect to the target distribution ̂ p ∗ (and the empirical target distribution p ∗ , due to our results from Lemma 2.5. Handling these issues related to nonconvexity of the loss in the construction of the upper bound is precisely what forces us to choose chi-square as the measure of divergence between distributions; see Lemma D.3 and the discussion therein.

Proposition D.2. Consider the sequence { w i } i evolving as per Line 6. Under the setting in which Lemma 2.5 holds, we have for all i ≥ 1 ,

where is E i is defined by Equation (37) .

Proof. From the definition of L ( w ∗ , ̂ p i ) , we have:

Writing ( σ ( w ∗ · x ) -y ) 2 = (( σ ( w ∗ · x ) -σ ( w i · x )) + ( σ ( w i · x ) -y )) 2 and expanding the square, we have

where for notational convenience we define

Observe that L ( w i , ̂ p i ) on the right-hand side also appears in the upper bound on L ( w i , ̂ p ∗ ) in Lemma D.1 and so it will get cancelled out when L ( w ∗ , ̂ p i ) is subtracted from L ( w i , ̂ p ∗ ) in the gap computation. Thus, we only need to focus on bounding S i . This requires a rather technical argument, which we defer to Lemma D.3 below. Instead, we call on Lemma D.3 to state that

and carry out the rest of the proof under this assumption (which is proved in Lemma D.3).

At this point, we have obtained a 'linearization' that was needed to continue by mimicking the construction of the upper bound. However, v ( w i ; x , y ) depends on w i , and so trying to define w i based on this quantity would lead to an implicitly defined update, which is generally not efficiently computable. Instead, here we use the idea of extrapolation: instead of defining a step w.r.t. w i , we replace E ̂ p i [ v ( w i ; x , y )] by an 'extrapolated gradient' defined by (cf. Line 5 in Algorithm 1):

Combining with the bound on S i from Equation (31) and simplifying, we now have

Let ψ ( w ) = a i ⟨ g i -1 , w ⟩ + 1+0 . 5 c 1 A i -1 2 ∥ w -w i -1 ∥ 2 2 and observe that (by Line 6 in Algorithm 1) w i = arg min w ∈ B ( W ) ψ ( w ) . Then, by a similar argument as in the proof of Lemma D.1, since ψ is minimized by w i and is a quadratic function in w i , we have

On the other hand, by the definition of g i , we have

The first two terms on the right-hand side of Equation (34) telescope, so we focus on bounding the last term. We do so using Young's inequality (Fact B.1) followed by κ -Lipschitzness of v , which leads to

where in ( ii ) we used 2 a i -1 2 κ 2 1+0 . 5 c 1 A i -1 ≤ 1+0 . 5 c 1 A i -2 4 and 2 a i -1 2 G 2 1+0 . 5 c 1 A i -1 ≤ ν 0 + νA i -2 , which both hold by the choice of the step size, while ( i ) follows by boundedness and κ -Lipschitzness of v and Corollary B.5, using

Combining Equations (32) to (35), we now have

To complete the proof, it remains to combine the last inequality with Equation (29).

Lemma D.3. Let S i be defined by Equation (30) . Then, under the setting of Proposition D.2, we have

where

Proof. Define the event G = { ( x , y ) : σ ( w i · x ) -y ≥ 0 } . Then,

where the last inequality uses convexity of σ ( · ) to bound the term that involves I G and concavity of -σ ( · ) to bound the term that involves I G c and where σ ′ denotes any subderivative of σ (guaranteed to exist, due to convexity and Lipschitzness).

Recall that v ( w i ; x , y ) = 2 β ( σ ( w i · x ) -y ) x . Using that σ ′ ( t ) = β +( σ ′ ( t ) -β ) for all t and combining with the inequality above, we see

and so to prove the lemma we only need to focus on bounding the terms in the last two lines.

Recall that σ is assumed to be monotonically increasing and β -Lipschitz, and so 0 ≤ σ ′ ( w ∗ · x ) ≤ β. Thus, we have

where we have used σ ′ ( w ∗ · x ) -β ≤ 0 (by Lipschitzness) and ( σ ( w i · x ) -σ ( w ∗ · x ))( w ∗ · x -w i · x ) ≤ 0 (by monotonicity of σ ). By the same argument,

To complete the proof of the lemma, it remains to bound the expectation of the term in Equation (40). We proceed using that | σ ′ ( w · x ) -β | ≤ β , ∀ w , and thus for w ∈ { w ∗ , w i } :

Taking the expectation on both sides, and combining with Equation (40), we further have

In the last equality, the first integral is just the expectation with respect to ̂ p ∗ , and thus using CauchySchwarz inequality, the definition of OPT , and Lemma C.6, the first term in Equation (41) can be

bounded by

where the last line is by Young's inequality and the second last line uses Lemma 2.5.

For the remaining integral in Equation (41), using the definition of chi-square divergence and CauchySchwarz inequality, we have

where ( i ) is by Cauchy-Schwarz, ( ii ) is by Cauchy-Schwarz again, ( iii ) is by the definition of ̂ OPT (2) and Lemma C.6, and (iv) is by Young's inequality.

To complete the proof, it remains to plug Equations (41) to (43) back into Equation (39), and simplify.

Gap upper bound proof of Lemma 3.3. Combining the upper and lower bounds from Lemma D.1 and Proposition D.2, we are now ready to prove Lemma 3.3, which we restate below.

Lemma 3.3 (Gap Upper Bound) . Let w i , ̂ p i , a i , A i evolve according to Algorithm 1, where we take, by convention, a -1 = A -1 = a 0 = A 0 = 0 and w -1 = w 0 , ̂ p -1 = ̂ p 0 . Assuming Lemma 2.5 applies, then, for all k ≥ 1 , ∑ k i =1 a i Gap( w i , ̂ p i ) is bounded above by

Proof. Combining the upper bound on a i L ( w i , ̂ p ∗ ) from Lemma D.1 with the lower bound on a i L ( w ∗ , ̂ p i ) from Proposition D.2 and recalling that Gap( w i , ̂ p i ) = L ( w i , ̂ p ∗ ) -L ( w ∗ , ̂ p i ) and

A i = A i -1 + a i , we have

Observe that except for the first line on the right-hand side of the above inequality, all remaining terms telescope. Summing over i = 1 , 2 , . . . , k and recalling that, by convention, a 0 = A 0 = a -1 = A -1 = 0 , w -1 = w 0 , and ̂ p -1 = ̂ p 0 , we have

To complete the proof, it remains to bound a k |⟨ E ̂ p k [ v ( w k ; x , y )] -E ̂ p k -1 [ v ( w k -1 ; x , y )] , w ∗ -w k ⟩| , which is done similarly as in the proof of Proposition D.2. In particular,

where ( i ) is by Young's inequality and ( ii ) is by Equation (36), and ( iii ) is by 2 G 2 a k 2 1+0 . 5 c 1 A k ≤ ν 0 + νA k -1 and 2 κ 2 a k 2 1+0 . 5 c 1 A k ≤ 1+0 . 5 c 1 A k -1 4 , which both hold by the choice of the step sizes in Algorithm 1.

To complete the proof, it remains to plug Equation (45) back into Equation (44), use the definition of E i from Equation (37), and simplify.

E Omitted Proofs in Main Theorem

Claim 3.5. For all iterations k ≥ 0 , ∥ w k ∥ 2 ≤ 2 ∥ w ∗ ∥ 2 .

Proof of Claim 3.5. It trivially holds that 0 = w 0 ∈ B (2 ∥ w ∗ ∥ 2 ) . Suppose ∥ w i ∥ 2 ≤ 2 ∥ w ∗ ∥ 2 for all iterations i ≤ t where t ≥ 0 . Then

where we used the gap upper bound Lemma 3.3 again as it does not require w ∈ B ( ∥ w ∗ ∥ 2 ) . We proceed to deduce a different lower bound for Gap( w k +1 , ̂ p k +1 ) . Similar to Lemma 3.2, we break into two terms L ( w ∗ , ̂ p k +1 ) -( -L ( w ∗ , ̂ p ∗ )) ≥ νD ϕ ( ̂ p ∗ , ̂ p k +1 ) and L ( w k +1 , ̂ p ∗ ) -L ( w ∗ , ̂ p ∗ ) = E ̂ p ∗ [( σ ( w k +1 · x ) -y ) 2 -( σ ( w ∗ · x ) -y ) 2 ] ≥ -̂ OPT , where the first term is bounded the same way as in Lemma 3.2. Hence, Gap( w k +1 , ̂ p k +1 ) ≥ -̂ OPT . Therefore, we simplify as before

which implies by nonnegativity of Bregman divergence that:

Rearranging and using 2 + c 1 A k ≥ 2 ,

Our choice of stepsizes a i implies a k +1 2+ c 1 A k ≤ 1 / max { κ, G } ≤ 1 , hence

Claim E.1. For ν ≥ 8 β 2 √ 6 B √ ̂ OPT (2) /c 1 , it holds that

Similarly, for ν ≥ E p 0 ℓ ( w ∗ ) , 8 β 2 √ 6 B √ OPT (2) /c 1 , it holds that

Proof of Claim E.1. By Corollary C.3, χ 2 ( ̂ p ∗ , ̂ p 0 ) = E ̂ p 0 [ ℓ 2 ( w ∗ )] -( E ̂ p 0 [ ℓ ( w ∗ )]) 2 4 ν 2 ≤ E ̂ p 0 [ ℓ 2 ( w ∗ )] 4 ν 2 ≤ ̂ OPT (2) 2 ν 2 ≤ c 1 / (1536 β 4 B ) , where the second last inequality uses ̂ p ∗ ≥ ̂ p 0 / 2 from Corollary C.2 and the last inequality comes from lower bound on ν in the assumption.

The population version follows analogously.

Since D ϕ ( ̂ p ∗ , ̂ p 0 ) = χ 2 ( ̂ p ∗ , ̂ p 0 ) , by choosing ν 0 = 768 β 4 Bϵ/c 1 , we ensure 2 ν 0 D ϕ ( ̂ p ∗ , ̂ p 0 ) ≤ ϵ .

By choosing ν 0 small enough and initialization w 0 = 0 , it holds that

We may assume without loss of generality that 1 max { κ,G } ≪ 288 β 2 B c 2 1 because both κ and G is O ( d ) but the right hand side is an absolute constant. We may also assume without loss of generality that 300 β 2 B c 2 1 ̂ OPT + ϵ ≤ 2 ∥ w ∗ ∥ 2 2 , thus completing the induction step ∥ w ∗ -w k +1 ∥ 2 ≤ 2 ∥ w ∗ ∥ 2 . The reason for the last no loss of generality is the following: otherwise, we can compare, per Claim E.2, the empirical risk of the output from our algorithm and of ˆ w = 0 and output the solution with the lower risk to get an O (OPT) + ϵ solution.

Claim E.2 (Zero-Tester) . In the setting of Theorem 3.1, it is possible to efficiently check if R ( 0 ; ̂ p 0 ) > R ( ̂ w ; ̂ p 0 ) or not; where ̂ w is the output of Algorithm 1.

Proof. Observe that L ( w , ̂ p ) = ∑ N i =1 ̂ p i ( σ ( w · x ) -y ) 2 -νχ 2 ( ̂ p , ̂ p 0 ) is 1 /ν -strongly concave in ( ̂ p (1) , . . . , ̂ p ( N ) ) . Now, since R ( w ; ̂ p 0 ) = max ̂ p L ( w , ̂ p ) , we can estimate the risk at any given w using standard maximization techniques (such as gradient descent).

To test which risk is larger, we estimate R ( 0 ; ̂ p 0 ) and R ( ̂ w ; ̂ p 0 ) to a necessary accuracy and then compare.

F Parameter Estimation to Loss and Risk Approximation

Theorem 3.1 shows that Algorithm 1 recovers a vector ̂ w such that ∥ ̂ w -w ∗ ∥ 2 ≤ √ OPT + √ ϵ , where OPT := E p ∗ ( σ ( w ∗ · x ) -y ) 2 .

F.1 Loss Approximation

In this section we show that this implies that the neuron we recover achieves a constant factor approximation to the optimal squared loss.

Lemma F.1. Let p ∗ satisfy Assumption 2.1 and Assumption 2.2. Suppose ( ̂ w , ̂ p ) is the solution returned by Algorithm 1 when given N = samples drawn from p 0 . Then, E p ∗ ( σ ( ̂ w · x ) -y ) 2 ≤ O β,B (OPT) + ϵ .

Proof. Recall that σ is β -Lipschitz, and Fact 2.4 gives us E p ∗ ( u · x ) 2 ≤ 5 B for all unit vectors u . These together imply,

F.2 Risk Approximation

Fix ˆ w as output by Algorithm 1 and w ∗ as defined in Definition 1.2. Since we are bounding the population risk throughout this subsection, we write R ( w ) = R ( w ; p 0 ) in short. The goal of this subsection is to show

We first introduce some convex analysis results that we rely on in this subsection:

Fact F.2 (Strong convexity of chi-square divergence) . Consider the space P ( p 0 ) = { p : p ≪ p 0 } . For p ∈ P ( p 0 ) , we denote by d p d p 0 the Radon-Nikodym derivative of p with respect to p 0 , and we define ∥ p ∥ ∗ p 0 = √∫ ( d p d p 0 ) 2 d p 0 . Then p ↦→ χ 2 ( p , p 0 ) is 2-strongly convex with respect to ∥·∥ ∗ p 0 .

Fact F.3. Consider the space P ( p 0 ) = { p : p ≪ p 0 } . Denote by ⟨· , ·⟩ p 0 the inner product ⟨ ℓ 1 , ℓ 2 ⟩ p 0 = ∫ ℓ 1 ℓ 2 d p 0 and denote by ∥·∥ p 0 the correponding norm. Then ∥·∥ p 0 is the dual norm of ∥·∥ ∗ p 0 defined in Fact F .2.

Definition F.4 (Convex conjugate) . Given a convex function defined on a vector space E denoted by f : E → R , its convex conjugate is defined as:

for all y ∈ E ∗ where E ∗ is the dual space of E and ⟨ y, x ⟩ denotes the inner product.

Fact F.5 (Conjugate Correspondence Theorem, [Bec17, Theorem 5.26]) . Let ν > 0 . If f : E → R is a ν -strongly convex continuous function, then its convex conjugate f ∗ : E ∗ → R is 1 ν -smooth.

We are then able to state and prove the key technical corollary in this subsection:

Corollary F.6. For any p 0 -measurable function ℓ : R d × R → R , let R ( ℓ ) = max p ≪ p 0 E p 0 ℓ -νχ 2 ( p , p 0 ) . The function R ( · ) is 1 / (2 ν ) -smooth with respect to the norm ∥·∥ p 0 defined in Fact F .3.

Proof. Observe by definition of the convex conjugate that R ( · ) is the convex conjugate of the function νχ 2 ( · , p 0 ) . Since the function νχ 2 ( · , p 0 ) is 2 ν -strongly convex with respect to the norm ∥·∥ ∗ p 0 by Fact F.2, it follows from Fact F.5 that R ( · ) is 1 / (2 ν ) -smooth with respect to the norm ∥·∥ p 0 .

For ease of presentation, we define the following quantities: let ℓ ∗ ( x , y ) = ( σ ( w ∗ · x ) -y ) 2 and ˆ ℓ ( x , y ) = ( σ ( ˆ w · x ) -y ) 2 . We first compute ∇ ℓ R ( ℓ ∗ ) by conjugate subgradient theorem.

Fact F.7 (Conjugate Subgradient Theorem [Bec17, Theorem 4.20]) . Let f : E → R be convex and continuous. The following claims are equivalent for any x ∈ E and y ∈ E ∗ :

Corollary F.8. Let p ∗ be as defined in Definition 1.2. Then p ∗ ∈ ∂ ℓ R ( ℓ ∗ ) .

Proof. We verify that R ( ℓ ∗ ) = max p ≪ p 0 E p 0 [ σ ( w ∗ · x -y ) 2 ] -νχ 2 ( p , p 0 ) = E p ∗ [ σ ( w ∗ · x -y ) 2 ] -νχ 2 ( p ∗ , p 0 ) = E p ∗ ℓ ∗ -νχ 2 ( p ∗ , p 0 ) = ⟨ p ∗ , ℓ ∗ ⟩ -νχ 2 ( p ∗ , p 0 ) , where the second equality is the definition of p ∗ and the third equality is the definition of ℓ ∗ . By Fact F.7 and observing that R ( · ) is the convex conjugate of the function νχ 2 ( · , p 0 ) , we have p ∗ ∈ ∂ ℓ R ( ℓ ∗ ) .

Theorem F.9. Suppose Corollary C.2 holds for both w ∗ and ˆ w with respect to the population distribution. Then

where C 4 = 1 + 2(10 Bβ 2 + c 1 ) C 3 + c 1 5 Bβ 2 C 2 3 . In particular, Corollary C.2 holds for both w ∗ and ˆ w is satisfied under the assumptions in Theorem 3.1.

Proof. By the definition of smoothness, it holds that for any p ∈ ∂ ℓ R ( ℓ ∗ ) , we have

Hence it follows from Corollary F.8 that

We use the following shorthand for ease of presentation.

Then

where the last inequality is the standard inequality ( a + b ) 2 ≤ 2 a 2 +2 b 2 .

From the second- and fourth-moment bounds, we have E p ∗ [∆ τ ] = E p ∗ [( σ ( x · ˆ w ) -σ ( x · w ∗ )) τ ] ≤ β τ E p ∗ [( x · ( ˆ w -w ∗ )) τ ] ≤ 5 Bβ τ ∥ ˆ w -w ∗ ∥ τ 2 for τ = 2 , 4 , where the second last inequality follows from β -Lipschitzness of σ ( · ) . Taking τ = 2 gives us a bound for E p ∗ [∆ 2 ] .

For E p ∗ [∆Ψ] , it follows from Cauchy-Schwarz that E p ∗ [∆Ψ] ≤ √ E p ∗ [∆ 2 ] E p ∗ [Ψ 2 ] ≤ √ 5 Bβ 2 OPT ∥ ˆ w -w ∗ ∥ 2 .

Finally, similarly by Corollary C.2, it follows additionally from Cauchy-Schwarz that E p 0 [∆ 2 Ψ 2 ] ≤ 2 E p ∗ [∆ 2 Ψ 2 ] ≤ 2 √ E p 0 [∆ 4 ] E p ∗ [Ψ 4 ] ≤ 2 ∥ ˆ w -w ∗ ∥ 2 2 √ 5 Bβ 4 OPT (2) . By Theorem 3.1, we have ν ≥ 8 β 2 √ 6 B √ OPT (2) + ϵ/c 1 by assumption, hence 4 E p 0 [∆ 2 Ψ 2 ] /ν ≤ c 1 ∥ ˆ w -w ∗ ∥ 2 2 .

Combining the above four bounds and the guarantee of Theorem 3.1 that ∥ ˆ w -w ∗ ∥ 2 2 ≤ 2 C 3 OPT+2 ϵ , we have

By Corollary C.9, OPT ≤ √ OPT (2) , hence ν ≥ 8 β 2 √ 6 B √ OPT (2) + ϵ/c 1 ≥ 8 β 2 √ 6 B max { OPT , √ ϵ } /c 1 , hence

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract and introduction clearly outline the open problem we provide the first results for, clearly state the assumptions involved (cf. Section 1.1 - Problem Setup and Section 1.2 - Main Result), and give a detailed technical overview of our algorithm (cf. Section 1.3 - Technical Overview).

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]

Justification: The limitations are clearly stated in the statements of each theorem and are discussed in the introduction of the paper and description of the key lemma (Lemma 3.4).

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Each theorem statement provides all the assumptions and we provide complete proofs for all statements that are either in the main body of the paper or in the appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [NA]

Justification: The paper is theoretical in nature and does not include experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [NA]

Justification: The paper is theoretical in nature and does not include experiments.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [NA]

Justification: The paper is theoretical in nature and does not include experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [NA]

Justification: The paper is theoretical in nature and does not include experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [NA]

Justification: The paper is theoretical in nature and does not include experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Our research conforms in every respect with the NeurIPS Code of Ethics. Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: The work is theoretical and we do not see any major or immediate implications

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The work is theoretical.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification: This paper does not use existing assets.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: This paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors consider the problem of fitting a single neuron with respect to distributional uncertainty modeled by a $\\chi^2$ divergence. The authors prove convergence to within a constant factor of the optimal value. The proof of this results proceeds by calculating lower and upper bounds on a duality gap. Unfortunately, prior work shows that finding an algorithm that converges to optimality is in general computationally infeasible, so this is the best result one could hope for.\n\nSTRENGTHS:\nLearning with a single neuron (aka the perceptron algorithm) is a well-studied method in machine learning that sheds light on the behavior of more complicated methods. The authors seek to understand this method in the context of distributional uncertainty modeled by a chi-squared divergence.\n\nWEAKNESSES:\n1. **Main concern 1**.  The steps and intermediate quantities of Algorithm 1 were not sufficiently explained.  For instance: the step size  $a_i$ is chosen to always be larger than 1. Why does the algorithm not diverge with such a choice? What is the \"interpolated quantity\" $g_i$? Can you explain the choices for $v$ and $w$ in lines 4 and 6?\n\n2. **Main concern 2**. The set $\\mathcal P (p_0)$ is not sequentially compact. Thus one cannot conclude the existence of the maximizer $q_w$ defined in the second equation below line 74 just from compactness. The existence of this maximizer seems to be related to properties of the $\\chi^2$ divergence. Could you elaborate?\n\nQUESTIONS:\nDoes your algorithm actually converge? The classical perceptron algorithm is not guaranteed to converge when OPT$\\neq 0$. \n\n\nDo you have a stopping criterion for your algorithm?\n\n\nOther comments:\n\n\nIn the introduction/title/abstract: when you discuss \"robustness\" and \"adversarial\" you should make it clear that you don't mean adversarial robustness\n\n\n1st equation below line 74: The $(=\\Lambda_{\\sigma,\\rho}\\ldots)$ portion is very confusion\n\n2nd equation below line 74: why is there another expectation over $(\\mathbf x ,y)\\sim \\rho)$? This expectation was already computed in the definition of $L_\\sigma$ in the line above\n\nlines 84-88: it seems your algorithm only considers \"reweightings\" of the sampled points for $\\hat \\rho$. Does this suffice for modelling the entire $\\chi^2$ ball around $\\hat p_0$?\n\n\nline 144-145: \"This condition is ..requires convergence\". It's really not obvious that this condition holds at initialization, or that convergence would imply that it holds for all iterates\n\nline 204: What is $h(r)$?\n\nline 212-213: Shouldn't Assumption 2.3 always be true because the sample data set is finite? \n\nline 250: $M$ was not defined\n\nline 273: Why would one consider looking at $\\sum_{i=1}^k a_iGap(w_i,\\hat p_i)$? Can you better motivate this quantity?\n\nlines 274-276: this was hard to understand\n\nline 214-216: this comment was also non obvious\n\nLIMITATIONS:\nYes. The authors acknowledge that their theorems don't prove convergence.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

comments: SUMMARY:
This paper addresses the problem of learning a single neuron in a distributionally robust setting, where both the input distribution and labels can be adversarially perturbed. The authors propose an efficient primal-dual algorithm that achieves a constant-factor approximation to the optimal loss, with respect to the worst-case distribution within a chi-squared divergence neighborhood of the reference distribution.

The proposed algorithm iteratively updates both the weight vector and a reweighting of the empirical distribution. The analysis relies on carefully bounding a gap function related to the primal-dual objective. A key technical lemma bounds the nonconvex parts of the objective, allowing the analysis to go through despite the nonconvexity.

STRENGTHS:
Learning single neuron is an important problem towards understanding deep learning. This paper extends the prior works to the setting of distributionally robust loss. They give concrete primal-dual algorithm to optimize, and provide rigorous theoretical generalization guarantees of the algorithm. The algorithm is computationally efficient, with polynomial ($O(d/\epsilon)$) sample complexity.
The analysis introduces novel techniques for handling nonconvexity in primal-dual methods, which to me, is interesting. Their theory also applies to various activation function, not restricted to ReLU.

WEAKNESSES:
It seems that [Dia+22a] and [Dia+22b] also studied robustly learning single neuron. I am not familiar with those literatures, so I cannot give a very accurate judgement on the degree of technical novelty of this paper.

QUESTIONS:
Please see weanknesses part.

LIMITATIONS:
They discussed the limitations. This is a pure theory paper so I do not see potential negative societal impact.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

comments: SUMMARY:
The authors study the problem of computing a distributionally robust optimization with \ell_2^2 loss functions with sample distribution p, and a chi-square regularizer that ensures that p is close to a given p0 for a single d-dimensional neuron with monotone lipschitz activation function. Their main contribution is an efficient primal dual algorithm which recovers a weight w close to the optimal weight (in \ell_2^2 norm) in O(d) iterations, as long as the optimal distribution p^* satisfies certain sub-exponential concentration properties.

STRENGTHS:
The problem is natural and a practically important one since distribution shifts issues occur in practice and before we study large neural nets, it's important to be able to answer the question for a single neuron. The proofs are fairly detailed, which is often missing in many conference papers.

WEAKNESSES:
The crux of the paper is Theorem 3.1 which bounds the error in each iterate of their primal dual algorithm. The proof of the theorem lies in bounding the primal dual gap and while the argument is interesting, the ideas are elementary and do not seem to be new to the area. I am happy to be corrected if I have missed any significantly new intuition or novelty in their arguments in Lemma 3.2 or 3.3 (which is the crux of their proof).

QUESTIONS:
Can this be extended to two single neuron GANs? The min-max structure might be a simple yet tractable extension to your single neuron setting.

LIMITATIONS:
No societal impact, theory result.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work considers the labels with noise and possible distribution shifts of the data. The authors aim to minimize the model's loss on a worst-case distribution from a set of distributions close to the reference distribution, which is defined as ambiguity set. The activation functions used to produce labels are nonconvex, including ReLU, leaky ReLU, exponential linear unit, and normalized SoftPlus. In such situation, the distributionally robust optimization can achieve a stationary point, which is insufficient for learning a ReLU neuron. To close the gap, this work proposed the first polynomial sample and time algorithm for learning a neuron in a distributionally robust setting for a broad class of activations. \n\nTo this end, the authors make two assumptions about the target distribution covariates. Then, the loss function on the target distribution is claimed to be sharp. The algorithm is proposed to produce a sequence of primal-dual pairs w, p given a sequence of positive step sizes. The difference between optimal w and empirical w is called gap. The gap function is bounded, which further leads to the convergence of the algorithm as shown in Theorem 3.1.\n\nSTRENGTHS:\nThis work considers the labels with noise and possible distribution shifts of the data. Existing works do not consider both of the concerns. \n\nThe problem is with practical interests.\n\nThe authors propose the novel algorithm to solve the problem.\n\nWEAKNESSES:\n1. The variable \\hat{p}_i in Algorithm 1 (line 261 on page 7 ), is not updated in other iterations. \n\n2. The analysis about time complexity and memory cost is missing. \n\n3. It is not clear where the noise of labels is discussed in Theorem 3.1 on page 6. \n\n4. How to control distribution shift is not discussed in Theorem 3.1.\n\nQUESTIONS:\n1. What is the time complexity and memory cost of the algorithm?\n\n3. Where do authors consider label noise in Theorem 3.1?\n\n4. How to control distribution shift is not discussed in Theorem 3.1.\n\nLIMITATIONS:\nThe time and memory cost is not provided. It may limit the applications in real-world applications.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-16"}