We demonstrate that the sample complexity of (Vanilla) Gradient Descent (GD), with common choices of hyperparameters, has sample complexity that is linear in the dimension. In particular, we show that when the dimension is strictly larger than number of samples then $T=\Omega(1/\epsilon^4)$ iterations are necessary to avoid overfitting. This resolves an open problem by <cit.>,  and improves over previous lower bounds that demonstrated that the sample size must be at least square root of the dimension.

We further show that, when hyper-parameters are chosen to achieve minimal empirical risk at optimal rate, GD has the same worst-case sample complexity of any empirical risk minimizer. That means GD does not demonstrate any algorithmic advantage in terms of generalization. This is in contrast with other algorithms (e.g. SGD) that can outperform the sample complexity of general ERMs.

§ INTRODUCTION

Stochastic Convex Optimization (SCO) is a theoretical model that depicts a learner that minimizes a convex function when given finite noisy observations of the objective <cit.>. While often considered simplistic, in recent years SCO has become a focus of much theoretical research, partly, because of its importance to the study of first-order optimization methods.
But, also, it has become focus of study because it is one of few theoretical settings that exhibit overparameterized learning .

In more detail, classical learning theory often focuses on the tension between number of samples, or training data, and the complexity of the model to be learnt. A common wisdom of classical theories <cit.> is that, to avoid overfitting, the complexity of a model should be adjusted in proportion to the amount of training data. However, recent advances in the Machine Learning have challenged this viewpoint. Evidently <cit.>, novel state of the art algorithms generalize well but without, explicitly, controlling the capacity of the model to be learnt. In turn, today, it is one of the most emerging challenges, for learning theory, to understand learnability when the number of parameters in a learnt model exceeds the number of examples, and when, seemingly, nothing withholds the algorithm from overfitting.

Towards this, we look at SCO. In SCO, <cit.> showed how algorithms can overfit with dimension dependent sample size. But, at the same time, it was known <cit.> that there are algorithms that provably avoid overfitting with far fewer examples than dimensions. As such, SCO became a canonical model to study how a well-designed algorithm can avoid overfitting, and what are the algorithmic trait that enable this phenomena of generalization in an overpaametrized task <cit.>.

Surprisingly, though, sometimes even the sample complexity of the simplest algorithms is not yet fully known. Yet, a prerequisite to answering how algorithms avoid overfitting, is to ask which algorithms avoid overfitting.

Perhaps the simplest algorithm, whose sample complexity is not yet understood, is Vanilla Gradient Descent (GD). And we turn to the following basic question:

How many samples are required when applying Gradient Descent on the empirical risk?

While the question above remained open, there have been several advancements and intermediate answers: The first, dimension independent, generalization bound was given by <cit.> that provided stability bounds <cit.>. The result of  demonstrated that, for some hyperparameter choice, GD has dimension-independent sample complexity rate. However, to achieve that, the number of rounds becomes quadratic in the size of the sample. On the other hand, a standard covering argument shows that linear dependence in the dimension is the worst possible, for any empirical risk minimizer, irrespective of properties such as stability.

In terms of lower bounds, <cit.> were the first to show that GD may have a dimension dependence in the sample complexity. They showed that, with natural hyperparameters, the algorithm must observe number of samples that is at least logarithmic in the dimension. This result was improved by <cit.> that showed that at least square root of the dimension is required. Taken together, so far it was shown that either the algorithm's hyperparameters are tuned to achieve stability, at a cost in running time, or the algorithm must suffer some dimension dependence, linear at worst square root at best.

Here, we close the gap and show that linear dependence is necessary, unless one take care of the stability term, and increases round complexity. Informally, we provide the following generalization error bound, in terms of dimension $d$, sample size, $m$, and hyperparameters of the algorithm, $\eta$ and $T$. We show that when $T$ is at most cubic in the dimension (see <ref> for a formal statement):

\[ \textrm{Generalization gap of GD} = \Omega\left ( \min\left\{ \frac{d}{m},1\right\}\cdot \min \left\{\eta \sqrt{T},1\right\}\right).\]

The first factor in the RHS describes the linear dependence of the generalization error in the dimension, and corresponds to the optimal sample complexity of empirical risk minimizers, as demonstrated by <cit.>.

The second term lower bounds the stability of the algorithm <cit.>, and played a similar role in previous bounds <cit.>. Our improvement is that it lower bounds generalization already when the dimension is linear in sample size. Each factor is optimal for a certain choice of hyper-parameters, and cannot be improved. Specifically, for a standard choice of $\eta=O(1/\sqrt{T})$, the first term is dominant, and the aformentioned lower bound is complemented with the upper bound of <cit.>. Our result implies, then, a sample complexity of $\tilde \Theta(d/m+1/\sqrt{m})$. When $d\ge m$, the second factor is dominant, unless running time is quadratic. This term also governs the stability of the algorithm, hence the result of <cit.> provides a complementary upper bound. for suitable choice of hyperparameters (see further discussion in <ref>).

Arguably, one of the most studied topics in learning, supervised learning, can be roughly described as the the task of fitting a model (or hypothesis) when given a finite amount of past observations, or data, coming out of some stochastic distribution to be learnt. Contemporary theories tend to stress the tension between the complexity of the model to be learnt and the scarcity of the data to be trained upon. Common wisdom suggests that the complexity of a model should be adjusted in proportion to the amount of training data, in order to prevent overfitting. This principle has been formalized and refined over recent years <cit.>, leading to a mature understanding of what can be termed underparameterized learning. On the other hand, recent advances in the development of deep learning algorithms, has led to novel learning algorithms, that evidently <cit.> show capabilities of generalizing even without explicitly trying to avoid overfitting. In turn, it is, today, one of the most emerging challenges in learning theory to provide a theoretical account for learning overparametrized setups.

Towards this goal, a very popular, and useful, theoretical model for learning has turned out to be Stochastic Convex Optimization (SCO) <cit.>. In SCO, the examples can be thought of as stochastic loss function that are convex and Lipschitz when considered as mapping over the parameter space of the model to be learnt. Though this setup may seem simplistic (and in practice, only rarely a learning algorithm may be considered as solving a convex problem), it turns out to be reach enough, and even this simplistic setup challenges the common view that learning should be constrained to an underparametrized setup. More concretely, <cit.> demonstrated how learning can occur in SCO even in the overparameterized regime. In more detail, in SCO, one consider a loss function $f(w,z)$ over a parameter space $\W\subseteq \mathbb{R}^d$, in dimension $d$. Given a distribution $D$ over the space $Z$, a learner that gets to observe finite set of samples $\{z_1,\ldots, z_m\}$ needs to minimize the population loss:

\[ \textrm{minimize}~F(w) = \min_{z\sim D}[ f(w,z)].\]

<cit.> showed that the capacity of the model to be learned indeed scales with the dimension, and they showed that one can overfit the data, unless $\Omega (\log d)$ examples are observed. Later, <cit.> improved the result and showed that in fact $\Omega(d)$ examples are necessary to avoid overfitting. On the other hand, it was known <cit.> that there are algorithms that can find an approximate minimizer to the population loss, even without observing a dimension dependent amount of data. In other words, it seems that overparameterized learning is possible, already in a simple task such as stochastic convex optimization.

§ BACKGROUND

We consider the standard setup of stochastic convex optimization (SCO) as in <cit.>. Set $\W= \{w: \|w\|\le 1\}$, and let $\Z$ be an arbitrary, finite[Our main result is a lower bound on the generalization error, hence there is no loss of generality by assuming $\Z$ is finite], domain. We assume that there exists a function $f(w,z)$ that is convex and $L$-Lipschitz in $w\in \W$ for every choice of $z\in \Z$. Recall that a function $f$ is convex and $L$-Lipschitz if for any $w_1,w_2 \in \W$ and $0\le \lambda \le 1$:

\begin{equation} f(\lambda w_1 + (1-\lambda)w_2)\le \lambda f(w_1) + (1-\lambda) f(w_2),~ \textrm{and},~ |f(w_1)-f(w_2)|\le L\|w_1-w_2\|.\end{equation}

First order optimization Algorithmically we require further assumptions concerning any interaction with the function to be optimized.
Recall that the sub-gradient set of $f(w,z)$ at point $w, z$ is the set:

\[ \partial f(w,z)= \left\{g: f(w',z) \ge f(w,z) +g^\top (w'-w), \forall w'\in \W\right\}.\]

A first order oracle for $f$ is a mapping $\O_z(w)$ such that
$ \O_z(w)\in \partial f(w,z).$
Our underlying assumption is that a learner has a first order oracle access. In other words, given a function $f(w,z)$, we will assume that there is a procedure $\O_z$ that calculates and returns a subgradient at every $w$ for every $z$. Recall <cit.> that When $|\partial f(w,z)|=1$, the function is differentiable, at $w$, and in that case, the unique subdifferential is the gradient $\nabla f(w,z)$.

§.§ Learning

A learning algorithm $A$, in SCO, is any algorithm that receives as input a sample $S=\{z_1,\ldots, z_m\}\in \Z^m$ of $m$ examples, and outputs a parameter $w_S$. An underlying assumption in learning is that there exists a distribution $D$, unknown to the learner $A$, and that the sample $S$ is drawn i.i.d from $D$. The goal of the learner is to minimize the population loss:

\[ F(w) = \E_{z\sim D}[f(w,z)],\]

More concretely, We will say that the learner has sample complexity $m(\epsilon)$ if, assuming $|S|\ge m(\epsilon)$, the w.p.[As we will mostly care about lower bounds, fixing the confidence to $1/2$ will have no loss of generality.] $1/2$:

\begin{equation}\label{eq:learn} F(w_S) - \min_{w\in \W} F(w) \le \epsilon.\end{equation}

Empirical Risk Minimization A natural approach to perform learning is by Empirical Risk Minimization (ERM). Given a sample $S$, the empirical risk is defined to be:
\[F_S(w) = \frac{1}{|S|} \sum_{z\in S} f(w,z).\]
An $\epsilon$-ERM is any algorithm that, given sample $S$, returns a solution $w_S\in W$ that minimizes the empirical risk up to additive error $\epsilon>0$.

Recently, <cit.> showed that any $\epsilon$-ERM algorithm has a sample complexity bound of
\begin{equation}\label{eq:erm} m(\epsilon)= \tilde O\left(\frac{d}{\epsilon}+\frac{1}{\epsilon^2}\right),\end{equation}

The above rate is optimal up to logarithmic factor <cit.>. Namely, there exists a construction and an ERM that will fail, w.p. $1/2$, unless $m=\Omega(d/\epsilon)$ examples are provided[the $\Omega(1/\epsilon^2)$ sample complexity bound is more straightforward and follows from standard information-theoretic arguments]. Importantly, though, there are algorithms that can learn with much smaller sample complexity. In particular SGD <cit.>, stable-GD <cit.> and regularized ERMs <cit.>.

§.§ Gradient Descent

We next depict Gradient Descent whose sample complexity is the focus of this work. GD depends on hyperparameters $T\in \mathbb{N}$ and $\eta\ge 0$ and operates as follows on the empirical risk. The algorithm receives as input a sample $S=\{z_1,\ldots, z_m\}$, defines $w_0=0$, operates for $T$ iterations according to the following recursion:

\begin{equation}\label{eq:gd} w_{t} = \Pi\left[w_{t-1} - \frac{\eta}{|S|}\sum_{z\in S} \O_z(w_t)\right] \Rightarrow~  w_S^{GD} : = \frac{1}{T} \sum_{t=1}^T w_t,
\end{equation}

where $\Pi$ is the projection onto the unit ball, and $\O_z(w_t)$ is a subgradient of the loss function $f(w,z)$. The final output, $w_S^{GD}$, of the algorithm is the averaged iterate.

GD constitutes an $\epsilon$-ERM. Concretely, it is known <cit.> that GD minimizes the empirical risk and its optimization error is given by:

\begin{equation}\label{eq:gd_empirical} F_S(w_S^{GD}) - \min_{w\in W}F_S(w) = \Theta\left(\min\left\{\eta + \frac{1}{\eta T},1\right\}\right).\end{equation}

The above bound is tight irrespective of the dimension[For completeness, we demonstrate the lower bound for $d=1$ at <ref>].
The population loss have also been studied, and <cit.> demonstrated the following learning guarantee:

\begin{equation}\label{eq:gd_generalization} \E_{S\sim D^m}\left[ F(w_S^{GD}) - \min_{w\in W}F(w)\right] = O\left(\frac{1}{\eta T} + \frac{\eta T}{m} + \eta\sqrt{T}\right).\end{equation}

The last two terms in the RHS follow from a stability argument, provided in <cit.>, and the first term follows from the optimization error of GD as depicted in <ref>.

Notice that there is a $O(\eta \sqrt{T})$ gap between the generalization error and empirical error of gradient descent.  <cit.>, later improved by <cit.>, showed that, in high dimension, this gap is unavoidable. In fact this is true to a much larger family of algorithms <cit.>.

[<cit.>] For sufficiently large $T$ and $\eta =O(1/\sqrt{T})$. For $d=\Omega(mT+m^2)$, there exists a distribution $D$ over instances $\Z$ and a convex, differentiable, and $1$-Lipschitz function $f$ such that for GD, w.p. $1/6$:

\begin{equation}\label{gd:previous} F(w_S^{GD})- \min_{w\in \W} F(w) = \Omega\left(\min\left\{\eta \sqrt{T}+\frac{1}{\eta T},1\right\}\right).\end{equation}

One consequence of this gap is that if we run $GD$ with step size $\eta= O(1/\sqrt{T})$ and $T= O(1/\epsilon^2)$, the algorithm will achieve empirical error of $O(\epsilon)$ but generalization error of $\Omega(1)$, meaning it will overfit unless we observe at least $\Omega(\sqrt{d})$ examples. Our main result improves over this and shows that GD may overfit unless it osbserves $\Omega(d)$ examples.

Stochastic Gradient Descent
Our next set of result deals with yet another algorithm known as Stochastic Gradient Descent (SGD):

SGD is yet another first order method that depends on a hyper parameter $\eta\ge 0$ that may depend on the sample size. The algorithm, receives as input a sample $S=\{z_1,\ldots, z_m\}$ and access to a first order oracle $\O_z$ and operates for $T=m$ iterations as follows: we define $w_0 =0 $. Then at iteration $t$ we define:

\[ w_{t} = \Pi\left[w_{t-1} - \O_{z_t}(w_t)\right].\]

The algorithm then outputs

\begin{equation}\label{eq:gd} w_S^{SGD} = \frac{1}{T} \sum_{t=1}^T w_t,\end{equation}

Again, the averaged iterate.
SGD has the following learning guarantee, which improves over <ref>:

\begin{equation}\label{eq:sgd_generalization} \E_{S\sim D^m}[ F(w_S^{SGD}) - \min_{w\in W}F(w)] = O\left(\frac{1}{\eta T} + \eta\right),\end{equation}

which by choice of $\eta = O(1/\sqrt{m})$ yields the optimal sample complexity rate for a learning algorithm of $m(\epsilon)= O(1/\epsilon^2)$. Surprisingly, though, SGD does not fall into the category of empirical risk minimizers. <cit.> showed that, in high dimension, with the above choice of parameters, SGD may have generalization gap and

\[ \E_{S\sim D}[F_S(w_S^{SGD})- \min_{w\in \W} F_S(w_S^{SGD})] = \Omega(1).\]

Again, this result was improved by <cit.> that provided improved dependence on the dimension:

[<cit.>] For sufficiently large $T$ and $\eta =O(1/\sqrt{T})$. For $d=\Omega( m^2)$, there exists a distribution $D$ over instances $\Z$ and a convex, differentiable, and $1$-Lipschitz function $f$ such that for SGD, w.p. $1/6$:

\begin{equation}\label{sgd:previous} F_S(w_S^{SGD})- \min_{w\in \W} F_S(w) = \Omega\left(\min\left\{\eta \sqrt{T}+\frac{1}{\eta T},1\right\}\right).\end{equation}

§ MAIN RESULT

Let $\eta>0$, $T\ge 10$. For every $m,d, T$, there exists a distribution $D$, and a $5$-Lipschitz convex function $f(w,z)$ in $\mathbb{R}^{d+1}$, such that for any first order oracle of $f(w,z)$, with probability $1/2$, if we run GD with $\eta$ as a learning rate then:

\[ F(w_\SS^{GD}) - F(0) \ge \thebound.\]

We now specialize our bound for two interesting regimes. First, we improve previous dependence in the dimension in <cit.> and obtain a generalization error bound for $d=\Omega(m+T^{1/3})$:

Fix $\eta$, and suppose $d =\Omega\left( \max\{m,\frac{1}{4} T^{1/3}\}\right)$. There exists a distribution $D$, and an $O(1)$-Lipschitz convex function $f(w,z)$ in $\mathbb{R}^{d+1}$, such that for any first order oracle of $f(w,z)$, with probability $1/2$, if we run GD for $T$ iterations, then:

\begin{equation}\label{eq:main1} F(\wgd) - F(0) \ge \Omega\left( \min\left\{\eta \sqrt{T}+ \frac{1}{\eta T},1\right\}\right).\end{equation}

The first term follows from <ref>, the second term follows from the optimization error in <ref>. <Ref> does not hold for $d< m$, and the linear improvement over <cit.> is tight. This can be seen from <ref> that shows that GD achieves $\epsilon$ empirical excess error when $\eta= O(1/\sqrt{T})$ and $T=O(1/\epsilon^2)$. <Ref> becomes vacuous for such choice of parameters, but <cit.> showed that the sample complexity of any ERM is bounded by $\tilde O((d+\sqrt{m})/m)$. However, as depicted next, this upper bound becomes tight and GD does not improve over a worst-case ERM:

Suppose $T\le m^{1.5}$ (for example, $T=O(m)$, and $\eta = \Theta(1/\sqrt{T})$. There exists a distribution $D$, and an $O(1)$-Lipschitz convex function $f(w,z)$ in $\mathbb{R}^{d+1}$, such that for any first order oracle of $f$, with probability $1/2$, if we run GD with $\eta$, for $T$ iterations:

\[ F(\wgd) - F(0) \ge  \lbound{\frac{d}{m}+\frac{1}{\sqrt{m}}}.\]

<Ref> complements the upper bound provided by <cit.> to general ERMs, and improves over <cit.> that showed the lower bound for general ERMs. To see that <ref> follows from <ref>, notice that when $d\le \sqrt{m}$, then $d/m < 1/\sqrt{m}$ and the bound is dominated by the second term, which is a well known-information theoretic lower bound for learning. When $d> \sqrt{m}$, and $T<m^{1.5}$ we have that $T\le d^3$, plugging $\eta=O(1/\sqrt{T})$ yields the bound.

§.§ Discussion

<Ref> provides a new generalization error bound, and shows that already in linear dimension a (natural) first order algorithm such as GD can already overfit unless dimension dependent sample is presented, or the algorithm's hyperparameter are tuned correctly, which leads to inefficiency in terms of round complexity. As discussed, the bound is tight in several regimes, nevertheless still there are unresolved open problems.

Stability in low dimension
As discussed, when GD is treated as a naive empirical risk minimizer, and $\eta = O(1/\sqrt{T})$, $T=O(m)$, there is no improvement, using GD, over a worst-case ERM. In the other direction, for dimension that is linear in $m$, one cannot improve over the $\Omega (\eta \sqrt{T})$ term that governs stability. Our bound, though, provide a hope that stability in low dimension can yield an improved bound. In particular, consider the case where $\eta =1/T^{1/4}$ and $d< m$. This is a case where we apply a stable algorithm in small dimensions. Our bound does not negate the possibility of an improved generalization bound. That would mean that, at least at some regime, GD can improve over the worst-case ERM behaviour. We leave it as an open problem for future study

Is there a generalization bound for GD such that

\[ \E_{S\sim D^m} \left[F(w_S^{GD}) \min F(w)\right] = O\left(\frac{d\eta \sqrt{T}}{m}+\frac{1}{\sqrt{m}} \right).\]

Alternatively, can we prove an improved generalization error bound such that:

\[ \E_{S\sim D^m} \left[F(w_S^{GD}) \min F(w)\right] = \Omega\left(\min\left\{\frac{d}{m}, \eta\sqrt{T},1\right\}\right).\]

Late stopping
Another regime where there is a gap between known upper bound and lower bound appears when $T=\Omega(m^2)$. Specifically, the stability upper bound for GD by <cit.> gives

\[ \E_{S\sim D^m} \left[F(w_S^{GD}) \min F(w)\right] =\ubound{\eta\sqrt{T}+ \frac{\eta T}{m} + \frac{1}{\eta T}}.\]

By <ref>, for large enough dimension:

\[ \E_{S\sim D^m} \left[F(w_S^{GD}) \min F(w)\right] = \lbound{\eta\sqrt{T}+  \frac{1}{\eta T}}.\]

When $T=O(m^2)$, the two bounds coincide. Indeed, for the $\eta T/m$ term to dominate the $\eta\sqrt{T}$ term, we must have $T=\Omega(m^2)$. One has to take at least $T=O(m^2)$ iterations in order to generalize with GD (in fact, any full batch method <cit.>), however $T=O(m^2)$ iterations are sufficient. Nevertheless, the above gap does yield the possibility of an unstable GD method that does generalize. Particularly, if we just regulate the term $\eta \sqrt{T}$, but allow $\eta T/m = \Omega(1)$, then this may yield a regime where GD is unstable (and ERM bounds do not apply) and yet generalize.

Are there choices of $\eta$ and $T$ (that depend on $m$) such that $\eta T/m \in \Omega(1)$, but GD has dimension indpendent sample complexity?

Notice that the $\eta T/m$ term also governs stability in the smooth convex optimization setup <cit.>. Recall that a function $f(w,z)$ is said to be $\beta$-smooth if for all $z$, $f(w,z)$ is differentiable, and the gradient is an $\beta$-Lipschitz mapping <cit.>. For smooth optimization, even if $\eta \sqrt{T}=\Omega(1)$, GD is still stable. <cit.> showed that the stability of GD in the smooth case is governed by $O\left(\frac{\eta T}{m}\right)$ for $\eta < 1/\beta$. Therefore, the question of generalization when $\eta T/m \in \Omega(1)$ remains open, even under smoothness assumptions:

Assume that $f(w,z)$ is $\Theta(1)$-smooth. What is the sample complexity of GD, when $\eta$ and $T$ are chosen so that
$\eta +\frac{1}{\eta T} = o(1),$
but $\frac{\eta T}{m}=\Omega(1)$.

§ TECHNICAL OVERVIEW

We next provide a high level overview of our proof technique. For simplicity of exposition we begin with the case $T=m=d$.
We begin by a brief overview of previous construction by <cit.> that demonstrated <ref> when $m=\Omega(\log d)$. The construction in <cit.> can be decomposed into three terms:

\[ f(w,z) = g(w,z)+ N_0(w) + h(w,z).\]

The function $g$ has the property that an ERM may fail to learn, unless dimension dependent sample size is considered. <cit.> incorporated <cit.> construction. Later, <cit.> used Feldman's function <cit.> to construct $g$. The shift from the construction of <cit.> to 's function is the first step that allows to move from logarithmic to polynomial dependence in the dimension.

In both constructions and underlying property of of $g$ is that there exists a distribution $D$ such that, when the sample is small, with significant probability, there exists a $w_S \in \{0,1/\sqrt{d}\}^d$  such that

\begin{equation}\label{eq:berm} \frac{1}{|S|}\sum_{z\in S} g(w_S,z) - \E_{z\sim D} [g(w_S,z)] = \Omega(1).\end{equation}

The challenge is then, to make gradient descent's trajectory move towards the point $w_S$. The idea can be decomposed into two parts:

§.§.§ Simplifying with an adversarial subgradient:

To simplify the problem, let us first ease the challenge and suppose we can choose our subgradient oracle in a way that depends on the observed sample.

Let $N_0$ be a variant of the Nemirovski function <cit.>:

\[ N_0(w) = \max\{- w(i),0)\}.\]

Notice that $N_0$ is not differentiable and the choice of subgradient at certain points is not apriori determined. For example, notice that every standard basis vector $-e_i \in \partial N(0)$. More generally, given a sample $S$, let $I=i_1,\ldots, i_{d'}$ be exactly the set of indices such that $w_S$, from <ref>, $w_{S}(i)\ne 0$. Now assume by induction that $w_t(i)> 0$ exactly for $i=i_1,\ldots i_t$, then one can show that we can define the subgradient oracle of $N$:

\[ \O(w_t) = -e_{i_{t+1}} \in \partial N_0(w_t).\]

In that case $w_{t+1}$ will satisfy our assumption for $i_{t+1}$ and we can continue to follow this dynamic for $T$ steps.

Notice that, in this case, GD will converge to $w_S$ (if $\eta=1/\sqrt{d}$ which we assume now for concreteness). One can also show that the output of GD (the averaged iterate) will overfit.

The caveat is that our subgradient oracle depends on the sample $S$. In reality, the sample is drawn independent of the subgradient oracle. and all previous constructions, as well as ours need to handle this. This is discussed in the next section. But before that, let us review another challenge which is the case of small dimension.

Improving when $d\ll T$
Another challenge we face with the construction above is that it works when we assume that $T\approx d$. That is because, in Nemirovski's function, the number of iterates we can perform is bounded by the dimension. After $d$ iterations we will end up with the vector $v= \sum_{t=1}^d \eta e_{i_t}$. If $T= \omega(d)$ then $\eta = o(1/\sqrt{d})$,
and the dynamic will end up with a too small norm vector to induce a sizeable population loss. This strategy will provide, at best, with a factor of the form $\Omega\left(\eta\sqrt{\min\{d,T\}}\right)$. Such a factor may be unsatisfactory in a very natural setting where, say, $T=O(m)$, $\eta = O(1/\sqrt{m})$, and $d= \Omega(\sqrt{m})$.

To obtain the $d^3$ dependence, we perform the following alternation over the Nemirovski function. Consider the function:

\begin{equation}\label{eq:nemonster} N(w) = \max\{0,\max_{i\le d}\{-w(i)\},\max_{i\le j\le d}\{w(j)-w(i)\}\}.\end{equation}

And suppose that at each iteration we return a subgradient as follows:

* If there is $i\le d$, such that $w(i)=w(i+1)> \eta$, return subgradient $e_{i+1} -e_{i}$ and $w$ is updated by $w_{t+1}= w_t-\eta e_{i+1} + \eta e_{i}$.

* If there is no such $i$, then take the minimal $i$ (if exists) such that $w(i)=0$, and return subgradient $-e_i$ and update $w_{t+1}=w_t + \eta e_i$.

* When non of the above is met, return subgradient $0$.

The dynamic of the above scheme is depicted in <ref>. One can show that GD will run for at least $d^3$ iterations, and will increase each coordinate, on average, by an order of $O(\eta d)$. This is better than the increase of $\eta$, in each coordinate, that we obtain by Nemirovski's function, and that's what help us obtain the improved result of $\eta\sqrt{\min \{d^3,T\}}$.

Depiction of the dynamics induced by <ref> and our choice of sub-differentials

§.§.§ Reduction to advarsarial subgradients:

As discussed, the construction above does not yield a legitimate lower bound as it relies on a subgradient oracle that is dependent on the whole sample. To simulate such and adversarial subgradient oracle, we incorporate the function $h$. The idea is that if we can infer the sample $S$ from the trajectory, i.e. if the state $w_t$ “encodes" the sample, then formally the subgradient is allowed to “decipher" the sample from the point $w_t$, and then act upon it. In that way we achieve this behaviour of sample dependent subgradient oracle. This part becomes challenging and may depend on the way we choose $g$, and $N$. The simplest case, studied by <cit.>, $h$ was a small perturbation function that elevates coordinates in $I$ and inhibits coordinates not in $I$. The function $h$ depends on $z$ and not on $S$, hence it cannot know apriori $I$. But, an important observation is that, in <cit.> construction, if $i\notin I$, there exists $z\in S$ that “certifies" that. In fact, each $z$ can be thought of as a subset of indices, and if an index appears in $z$, then it cannot be in $I$. So we can build the perturbation in a way that every coordinate is elevated, unless $z$ certifies $i\not\in I$: In that case we define $h(w,z)$ so that its gradient will radically inhibit $i$.

The last observation is what becomes challenging in our case. As discussed, to achieve improved rate, we need to use Feldman's function. When using Feldman's function the coordinates cannot be ruled out, or identified, by a single $z$ but one has to look at the whole sample to identify $I$. <cit.> tackles this in one way, and takes a different trajectory.

We take the following approach: For each $z$ assign a random, positive, number $\alpha(z)$. We can think of this $\alpha$ as a hash function. Let us add another coordinate to the vector $w$, $w(d+1)$. Consider the function

\[ h(w,z)= \gamma \max\{w(d+1)+\alpha(z),0\}.\]

Then $\partial h(0,z) = \alpha(z)e_{d+1}$. Write $\alpha(S) = \frac{1}{|S|}\sum_{z\in S} \alpha(z)$ then in turn:
\[
w_1(d+1)= -\partial \frac{1}{|S|}\sum_{z\in S} h(w,z) = -\gamma\alpha(S)e_{d+1}.\]

And since $\partial h(w_1,z)=0$ we have that $w_t(d+1)=w_1(d+1)$.

If $\alpha(z)$ are chosen correctly, $\alpha(S)$ is a one to one mapping from samples to real numbers. Then, we may define the subgradient oracle to be dependent on coordinate $d+1$ which encodes the whole sample. Our final construction will take a different $h$, for reasons next explained:

(Almost) Smoothing the construction Notice that our statement is slightly stronger than what we so far illustrated. <ref> states that, for any subgradient oracle, GD will fail. For that, we need to be a little bit more careful, and we want to replace our function with a function that leads to the same guaranteed trajectory, but at the same time it should be differentiable at visited points. On the one hand, this will ensure a unique derivative, making the construction independent of the choice of (sub)gradient oracle, on the other hand this makes it challenging to depend on a single coordinate in the manner we care.

Towards this goal, we start with the construction depicted so far and consider the set of all values, gradients, and points $\{f_j, g_j, w_j\}_{j\in J}$ that our algorithm may visit (for any possible time step and any possible samples) in the function above. What we want is to interpolate a new function through these triplets. In contrast with our original construction, we require a differentiable function at the designated points. Notice, that such an interpolation will have the exact same behaviour when implementing GD on it (with the added feature that the oracle is well defined and unique).

The problem of convex interpolation is well studied, and <cit.> show sufficient and necessary conditions for interpolation of a smooth function. Our case is slightly easier as we do not care about the smoothness parameter. On the other hand we do require Lipschitzness of the interpolation. We therefore provide an elementary, self-contained, proof to the following easy to prove Lemma, (proof is provided in <ref>)

Let $G=\{f_j,g_j,w_j\}_{j\in J} \subseteq \mathbb{R}\times\mathbb{R}^d\times\mathbb{R}^d$ be a triplet of values, gradients, and points in $\mathbb{R}^d$, such that $\|g_j\|\le L$. Suppose that for every $i,j\in J$:

\begin{equation}\label{eq:ineq} f_i \ge f_j + g_j^\top (w_i-w_j),\end{equation}

and let
\[I_{\textrm{diff}}= \{i: f_i=f_j+ g_j^\top(w_i-w_j) \Rightarrow g_i=g_j\}.\]

Then there exists a convex $L$-Lipschitz function $\hat {f}$ such that for all $j\in J$: $\hat f(w_j)=f_j$, and for all $i\in I_{\textrm{diff}}$, $\hat f$ is differentiable at $w_i$ and:

\[ \nabla f(w_i)= g_i.\]

To achieve the above condition, we choose a slightly different $h$ that not only encodes the sample in the trajectory, but also  encodes the time-step in the coordinate and guarantees a small amount of strong convexity in that coordinate. This will give rise to the strict inequality
in the convexity condition, which we can use by <ref>.

In the next two sections we provide a more formal statements of the two main ingredients: First, we define a setup of optimization with a sample-dependent first order Oracle and state a lower bound for the generalization error in this setup. The second ingredient is a reduction from the standard setup of first order optimization.

§.§ Sample-dependent Oracle

As discussed, the first step in our proof is to consider a slightly weaker setup where the first-order oracle may depend on the whole sample. Let us formally define what we mean by that. Define
\[\S^T_m= \{\SS=(S_1,\ldots, S_t), S_i\in \cup_{i=1}^m \Z^{m}, t\le T\},\]
the set of all subseqences of samples of size at most $m$.

Given a function $f(w,z)$, a sample dependent oracle, $\OSS$, is a finite sequence of first order oracles
\[\OSS= \{\OS{t}(S; w,z)\}_{t=1}^T,\]
that each receive as input a finite sample $S$, as well as $w$ and returns a subgradient:

\[\OS{t}(S, w,z)\in \partial f(w,z).\]

The sequence of samples can be thought of as the past samples that were observed by the algorithm. In the case of full-batch GD these will be the whole sample, and for SGD, each $S$ provided to $\OS{t}$ will be all previously observed samples.

Given $\SS\in \S_m^{T}$ Let us also denote

\begin{equation}\label{eq:ost} \OS{t}(\SS,w) = \frac{1}{|S_{t}|} \sum_{z\in S_t} \OS{t}(S_{1:t-1},w,z) \in \partial \left(\frac{1}{|S_t|} \sum_{z\in S_t} f(w,z)\right),\end{equation}

where we let $\SS_{1:0}=\emptyset$, and $S_{1:t-1}= (S_1,\ldots, S_{t-1})$ is the concatenated subsample of all previously observed samples in the sequence. As discussed, working with a sample-dependent oracle is easier (for lower bounds). And indeed, our first result shows that, if the subgradient can be chosen in a way that depends on the sample, we can provide the desired lower bound.

For fixed and known $\eta>0 $, $T$, a sample dependent first order oracle $\OSS$, and a sequence of samples $\SS=(S_1,S_2,\ldots, S_T)$, define $w_0=0$ and inductively:

\begin{equation}\label{eq:s_dependent_gd} w^\SS_{t} = \Pi\left[w^\SS_{t-1} - \eta\OS{t}(\SS,w^\SS_{t-1})\right] \Rightarrow~  w_\SS^{GD} : = \frac{1}{T} \sum_{t=1}^T w^\SS_t,
\end{equation}

Our first Lemma is then as follows:

For every $m,d,T\ge 18$ and $\eta>0$ there are a distribution $D$, a $3$-Lipschitz convex function $f(w,z)$ in $\mathbb{R}^d$, as well as a sample dependent first order oracle $\OSS$ such that: if $\SS= (S,S,\ldots S)\in \S_m^T$ for $S\sim D^{m}$ i.i.d, then w.p. $1/2$:

\[ F(w_\SS^{GD}) - F(0) \ge \thebound.\]

The proof of <ref> is provided in <ref>. We next move to describe the second ingredient of our proof.

§.§ Reduction to sample-dependent oracles

As discussed, the second ingredient of our proof is a reduction to the sample-dependent setup. Instead of using a perturbation function as in <cit.>, we take a more black box approach and show that, given a sample dependent first order oracle, there exists a function that basically induces the same trajectory. Proof is provided in <ref>.
Suppose $q \in \mathbb{R}^T$, $\|q\|_\infty \le 1$. And suppose that $f(w,z)$ is a convex, $L$-Lipschitz, function over $w\in \mathbb{R}^d$, let $\eta>0$, let $\OSS$ be a sample dependent first order oracle, and for every sequence of samples $\SS=(S_1,S_2,\ldots, S_T)$ define the sequence $\{w_{t}^\SS\}_{t=1}^T$ as in <ref>.

Then, for every $\epsilon>0$ there exists an
$L+1$ Lipschitz convex function[i.e. $w\in \mathbb{R}^d$ and $x\in \mathbb{R}$] $\bar f((w,x),z)$ over $\mathbb{R}^{d+1}$ (that depends on $q,f,T,\eta,m,\OSS,\epsilon$).

such that
for any first order oracle $\O_z$ for $\bar f$, define $u_0=0 \in \mathbb{R}^d$ and $x_0=0 \in \mathbb{R}$, and:

\[ (u_t, x_t)  = (u_{t-1},x_{t-1}) - \frac{\eta}{|S_t|} \sum_{z\in S_t} \O_z((u_t,x_t))\]
then if we define:

\[ u_{\q}= \sum_{t=1}^T \q(t)u_t, ~\textrm{and},\quad x_{\q}= \sum_{t=1}^T \q(t)x_t,~\textrm{and} ~w_{\q}^\SS= \sum_{t=1}^T \q(t) w_t^\SS.\]

then, we have that $u_{\q}=w_{\q}^\SS$ and:

\[ |\bar f((w^\SS_{\q},x_{\q}),z)- f(w_{\q}^\SS,z)| \le \epsilon.\]

\[ |\bar f((0,0),z)- f(0,z)| \le \epsilon.\]

§ PROOF OF <REF>

The proof is an immediate corollary of <ref>, which we prove in <ref> respectively. To see that <ref> indeed follows from these Lemmas, start with $\eta>0$ and $T\ge 4$ and $m$ and $d$. Let $f(w,z)$ be the function and $\OSS$ the sample dependent first order oracle, whose existence follows from <ref>. And let $\bar f$ be the function whose existence follows from <ref> to some arbitrarily small $\epsilon_0$, with $\q(t)=\frac{1}{T}$ for all $t$. It is easy to see that if we apply GD on $\bar f$ and define its output $(u^{GD},x^{GD})$ then: $(u^{GD},x^{GD}) = (u_{\q}, x_{\q})$, and $w^\SS_{\q} = w_\SS^{GD}$.

Then, we have that w.p. $1/2$:

\begin{align*} \bar F(u^{GD}) - \bar F(0) &\ge  F(w_\SS^{GD}) -\epsilon_0 - F(0)
\\
\thebound - \epsilon_0
\\
\thebound

Where in the last equation, we assume $\epsilon_0$ to be sufficiently small.

§.§ Proof of <ref>

We start by recalling Feldman's construction <cit.>: There exists a set $\V \subseteq \{0,1\}^d$, such that:

\begin{equation}\label{eq:V} \left(\forall v_1\ne v_2\in \V,~ v_1\cdot v_2 \le \frac{5d}{16}\right) \textrm{,~and}~\left( \|v\|\ge \frac{7d}{16}\right)\textrm{,~and~}\left( |\V|> e^{d/260}\right).\end{equation}

Indeed, suppose we pick randomly $w\in \{0,1\}^d$ according to probability $P$ where each coordinate $P(w(i))=1$ with probability $1/2$ (independently). Then, by Hoeffding's inequality for two $w_1,w_2\sim P$ independently:

\[ P\left (w_1\cdot w_2 > \frac{d}{4} + \frac{d}{16} \right) \le e^{-\frac{d}{128}},\]

\[ P\left(w_1 \cdot w_1 < \frac{d}{2} - \frac{d}{16}\right) \le e^{-\frac{d}{128}}.\]

Thus, picking $\V$ elements i.i.d according to $P$, randomly, of size $|\V|= e^{d/260}$ we can show by union bound that with positive probability , all $|\V|^2$ pairs in $\V$ will satisfy the requirement. Next,  we define a distribution $D_\epsilon$ supported on subsets of $\V$ such that for a random variable $V\subseteq \V$ we take $v\in V$ w.p. $\epsilon$ independently.

We start by assuming that $T\le d^3$, and let $k$ be such that

\[ d^3 \le \frac{k}{17} T < 2d^3.\]

The case $T> d^3$ is handled at the end of the section.

For every

will assume that:

\[ d^3 <\frac{1}{17} T \le 2d^3.\]

The general case is handled at the end of the section, and follows immediately from this special case (by either, stopping at an early iteration, or reducing the dimension).

Define $\alpha=\min\{\frac{1}{\eta \sqrt{2T}},1\}$, and let:

\[f(w,V)= g(w,V) + \alpha N(w).\]

where $N$ is defined in <ref>, and $g$ is defined to be Feldman's function with a suitable choice of threshold:

\begin{equation}\label{eq:gs} g(w,V) = \frac{1}{\sqrt{d}} \max_{v\in V}\left\{ \frac{45\cdot \eta\alpha  d^2}{2\cdot 16^2}  ,  w\cdot v\right\}.\end{equation}

Next we define a sample dependent oracle. We only define it for samples $\SS$ such that there exists $v^\star \notin V_i$ for all $V_i\in S$ (define it arbitrarily to any other type of sample). Let $I=\{i_1,\ldots, i_{d'}\}$ be a set of $\frac{7d}{16}$
indices such that $v^\star(i_j)\ne 0$.

Let us first define only an oracle for the function $\alpha N$. In other words, let $\OS{t}_{\alpha N}$ be a sample dependent oracle for $\alpha N$.

\[ \OS{1}_{\alpha N}(S,0)=0,\]

Next, we define for $t>1$ as follows: For any $w$ such that $0\notin \partial N(w)$, define it arbitrarily. If $0\in \partial N(w)$ we define it as follows:

* If there exists a coordinate $i_j$ such that $w(i_j)=w(i_{j+1})> \eta\alpha$: Then:

\[ \OS{t}_{\alpha N}(S,w)= \alpha(e_{j+1} - e_{j}).\]

* If there is no such coordinate, and if $j$ is the minimal coordinate such that $w(i_j)=0$, set:

\[ \OS{t}_{\alpha N}(S,w) = - \alpha(e_{i_j}).\]

* If both conditions cannot be met, then:

\[ \OS{t}_{\alpha N}(S,w) = 0.\]

Now we assume that $\{w_t^\SS\}_{t=1}^T$ follows the trajectory depicted in <ref> with that choice of Oracle. It can be seen that the update step is such that $w_t$ is non zero only in coordinates $i\in I$, and that for any subset $I_B\subseteq I$ such that $|I|=B$:

\begin{equation}\label{eq:decreasing} \sum_{i\in I_B} w_{t}^\SS(i) \le \max\left\{ \sum_{i\in I_B} w_T^\SS(i): I_B\subseteq I, |I|=B\right\}\end{equation}

One can also show that after $T'=1+ \sum_{t'=1}^{d'} \sum_{k=1}^{t'} k$ rounds, we will have that:

\begin{equation}\label{eq:nemonster} w_{{T'}}^\SS(i_t) = \alpha\eta(d'+1-t),\end{equation}

and that for $t>T'$: $\OS{t}_{\alpha N}(S,w_{{T'}}^\SS)=0$.

A proof by picture of <ref> is depicted in <ref>.
A formal proof is provided in <ref>.
Notice that:

\[ 1+ \sum_{t'=1}^{d'} \sum_{k=1}^{t'} k < d'^3 \le \frac{1}{17}T.\]

Next, for any $v\ne v^\star$, we have that:

\begin{align*} w_t^\SS \cdot v &\le \eta \sum_{i=1}^{d'}  w_t^\SS(i_t) \mathbf{1}\{v(i_t)=1\}\\
&\le \eta \max \left\{\sum_{i\in I_B} w_{T}^\SS(i): I_B\subseteq I, |I|=\frac{5d}{16}\right\} &\cref{eq:decreasing,eq:V}\\
\sum_{t=1}^{\frac{5d}{16}} \eta\alpha(d'+1-t) &\cref{eq:nemonster}\\
&= \eta\alpha \left( \frac{5d (d'+1)}{16}- \frac{5d(5d+16)}{2\cdot 16^2}\right)
\\
& = \eta\alpha \left( \frac{7d (d'+1)}{2\cdot 16}+\frac{3d (d'+1)}{2\cdot 16}- \frac{5d(5d+16)}{2\cdot 16^2}\right)
\\
&\le \eta\alpha \left( \frac{d'(d'+1)}{2} + \frac{3d(7d+16)}{2\cdot 16^2}- \frac{5d(5d+16)}{2\cdot 16^2}\right)\\
& = \eta\alpha \left( \frac{49d^2+7d}{2\cdot 16^2} - \frac{4d^2 +32d}{2\cdot 16^2}\right)\\
& \le \eta\alpha \left(\frac{45 d^2 -25d}{2\cdot 16^2}\right)\\
&\le\eta\alpha \left( \frac{45 d^2}{2\cdot 16^2}\right)

As such $0 \in \partial g(w_t^\SS,V_i)$ for all $V_i$, and we can define $\OS{t}$ so that

\[\OS{t}(S, w_t^\SS,V_i) = \OS{t}_{\alpha N}(S,w_t^\SS).\]

And we have that, $\{w_t^\SS\}$ is the trajectory obtained from <ref> with respect to this oracle of $g$ also. We also have:

\[ w_{T'}^\SS\cdot v^\star =  \eta\alpha  \sum_{i=1}^{d'}  (d'+1-t)  = \eta\alpha\left( \frac{d'\cdot (d'+1)}{2}\right) \ge \eta\alpha\left(\frac{49d^2}{2\cdot16^2}\right) .\]

And because $T'\le \frac{1}{17}T$:

\[ v^\star \cdot \bar w_T \ge \frac{16}{17} w_{T'}^\SS\cdot v^\star \ge \eta\alpha \left( \frac{46 d^2}{2\cdot 16^2}\right).\]

Recall that we assume $34 d^3 \ge T$:

\begin{equation}\label{eq:genbound} F(\wgd) - F_S(0) \ge
\epsilon \left(\frac{46\alpha\eta d^{1.5}}{2\cdot 16^2} -\frac{45\alpha\eta d^{1.5}}{2\cdot 16^2}\right)\ge \epsilon \frac{\alpha\eta d^{1.5}}{2\cdot 16^2}
\ge \frac{\epsilon}{2\cdot6\cdot 16^2}\min\left\{\eta\sqrt{T}, \right\}\end{equation}

<ref> lower bounds the generalization under the event that there exists $v\in \V$ such that $v\notin V_i$ for every $i=1,\ldots, m$. Now assume $\epsilon = \min\{\frac{d}{520m},\frac{1}{4}\}$, then for every $v$, using the inequality $(1-\epsilon)\le e^{-2\epsilon}$ for $\epsilon<1/2$:

\[ P(v\notin \cap_{i=1}^m  V_i)= (1-\epsilon)^m\ge e^{-2\epsilon\cdot m} \ge e^{-2d/520},\]

\[ P(\exists v, \notin \cap_{i=1}^m  V_i) \ge 1 - (1 - e^{-2d/520})^{|V|} \ge 1 - (1 - e^{-d/260})^{e^{d/260}}\ge 1- e^{-1} \ge 1/2.\]

And we obtain the desired result.

When $T\le d^3$
Then, note that we can allways embed our construction for $d^3= \frac{1}{17}T$ in a larger dimension, hence there is no loss in generality in assuming that $d^3=\floor{\frac{1}{17} T}$.

When $T\ge 34\cdot d^3$
Notice that if we take the construction with $T= 34\cdot d^3$ then:

\[\OS{T}(w^\SS_T)=0,\]

hence we can use the above construction for any $T'>T$, and the subdifferential oracle is defined for every iteration above $T$ as returning $0$, and we obtain a similar analysis to the case that $T=34 d^3$, and our bounds yield as in <ref>:

\begin{equation*} F(\wgd) - F_S(0) \ge \frac{\epsilon}{16^3} \min\left\{\eta \sqrt{34d^3},1\right\}\end{equation*}

§.§ Proof of <ref>

The main ingredient in our proof is the following claim whose proof is provided in <ref>. To state the claim we will need a few notations. First, for any two tuples of sample-sequences in $\S^T_m$, and $t>0$ $(\SS,t)$ and $(\SS',t)$, let us denote $(\SS,t)\equiv (\SS',t')$ if $t=t'$ and for every $S_{1:t} = S'_{1:t}$, namley the prefixes of the sample sequences agree up to point $t'=t$. Second, for a mapping $\alpha: \Z\to [0,1]$, from sample points to real numbers, and a sample $S$, let us denote

\[ \alpha(S)= \sum_{z\in S} \alpha(z).\]

For every $\eta,\gamma$ and  $\alpha$ define for every sequence of samples $\SS= (S_1,\ldots, S_T)$, and $t$ inductively: $x^\SS_0=0$ and

\begin{equation}\label{eq:xts} x^\SS_t = (1-\gamma \eta) x^\SS_{t-1} - \gamma \eta \alpha(S_t).\end{equation}

and define:

\begin{equation}\label{xss} x_{\q}^{\SS}= \sum_{t=1} \\q(t) x^{\SS}_{t} \end{equation}

for any $\epsilon>0$, there is a choice of $\gamma<\epsilon $ and $\alpha$ such that:

*  For $(\SS,t)\not\equiv (\SS',t')$: $ x^\SS_t \ne x^{\SS'}_{t'}.$

*  For every $\SS,t$, such that $t<\max\{t': q(t')\ne 0\}$ and $\SS'$: $x_{t}^\SS \ne \sum_{t=1}^T \q(t) x_t^{\SS'}$.

The $x_t^\SS$ represents the different states the trajectory can be in. $x_q^\SS$ represent the output of the trajectory which can be an aggreagated sum. We require that each provide a signature for the state, and this will allow us to “code" the state of the trajectory along GD.

We now continue with the proof of <ref>. For every sample dependent first order oracle $\OS{t}$, we have that:

\begin{equation}\label{eq:subgradsum} \OS{t}(\SS;w) = \frac{1}{|S_t|}\sum_{z\in S_t} \OS{t}(S_{1:t-1};w,z).\end{equation}

To simplify notations, let us denote:

\[\OS{t}(S_{1:t-1},w_t^\SS,z) = \O_{\SS,t,z},\]

as $\SS,t,z$ completely determine the output. We next define
\[h_z(x) := \frac{1}{2}\left(\gamma (x^2- 2 \alpha(z)x)\right) ,\]

where $\gamma>0$ is arbitrarily small[$\gamma \le \epsilon/(\eta T)$, will suffice].

Observe that

\begin{equation}\label{eq:hdynamic} x_{t}^\SS = x_{t-1}^\SS - \frac{\eta }{|S_{t}|}\sum_{z\in S_{t}}\nabla h_z(x_{t-1}^\SS) = (1-\gamma \eta)x_{t-1}^\SS -\gamma\eta \alpha(S_t).\end{equation}

In particular, by <ref>, if $x_0^\SS=0$:

\begin{equation}\label{eq:clmreiterated} (\SS,t)\not\equiv (\SS',t') \Leftrightarrow ~ x_{t}^\SS \ne x_{t}^{\SS'} ,\end{equation}

and observe that, if $x_t^\SS$ is defined as in <ref>:

\begin{equation}\label{eq:hdynamic} x_{t}^\SS = (1-\gamma \eta)x_{t-1}^\SS-\gamma\eta \alpha(S_t) = x_{t-1}^\SS- \frac{\eta}{|S_t|} \sum_{z\in S_t}\nabla h_z(x_{t-1}^\SS).\end{equation}

To simplify notation, let us denote:

\[ w_{T+1}^\SS = w_q^\SS, ~\textrm{and}~ x_{T+1}^\SS= x_q^\SS,\]

and assume without loss of generality that $\max\{q(i) \ne 0\}= T$ (Otherwise, we look only at the sequence up to point $\max\{q(i) \ne 0\}$). Consider now the sets of triplets:

\[ G(z) = \left\{ (v,g,u)= \left( f(w_t^\SS,z)+h_z(x_t^\SS), (\O_{\SS,t,z},\nabla h_z(x_t^\SS)), (w_t^\SS,x_t^\SS)\right) :  \SS_T^m\in \S, t\le T+1\right\},\]

where  $\O_{\SS,T+1,z} \in \partial (f(w_{\q}^\SS)+h_z(x_{\q}^\SS))$ is chosen arbitrarily.

Convexity of $f+h_z$ ensure that the triplets in $G(z)$ satisfy <ref> for all $t\le T+1$, as in <ref>. To apply the Lemma, we also want to achieve differentiability at points such that $t\le T$. Therefore, take any two triplets
\[(v_i,g_i,u_i) =\left( f(w_{t_i}^\SS,z)+h_z(x_{t_i}^\SS), \left(\O_{\SS,t,z},\nabla h_z(x_{t_i}^\SS)\right),\left(w_{t_i}^\SS,x_{t_i}^\SS\right)\right), i={1,2},\]

where $t_1\le T$ and $t_2\le T+1$, and suppose $g_1\ne g_2$. To simplify notations, let us write $w_{t_i}^\SS = \wz{i}$ and $x_{t_i}^\SS=\xz{i}$.

First, by convexity of $f$ we have that:

\begin{align*}
&\f{1}+\h{1} - \f{2}-\h{2} - \g{2}^\top\left(\u{1}-\u{2}\right)\\
= &\f{1}-\f{2}-\oz{2}^\top\left(\wz{1}-\wz{2}\right)+\h{1}-\h{2} - \hz{2}^\top\left(\xz{1}-\xz{2}\right)\\
\ge&
\h{1}-\h{2} -\hz{2}^\top\left(\xz{1}-\xz{2}\right).
\end{align*}

Next, because $g_1\ne g_2$, either
$\hz{1}\ne \hz{2}$, which implies
\ne x_2$, or $\O_{\SS_{1},t_1,z}\ne \O_{\SS_{2},t_2,z}$ which implies $(\SS_1,t_1)\not\equiv (\SS_2,t_2)$ which again implies $x_1 \ne x_2$ by <ref>. In other words, if $g_1\ne g_2$ then $\xz{1}\ne \xz{2}$:

\begin{align*}
\h{1}-\h{2} -\hz{2}^\top\left(\xz{1}-\xz{2}\right)&= \gamma\left(\xz{1}^2-2\alpha(z)\xz{1}-\xz{2}^2-2\alpha(z)\xz{2}-(2\xz{1}-2\alpha(z))\cdot(\xz{1}-\xz{2})\right)  \\
& = \gamma\left(\xz{1}^2 + \xz{2}^2 - 2\xz{1}\cdot \xz{2} \right)\\
& = \gamma\left(\xz{1}-\xz{2}\right)^2\\
& > 0 &\xz{1}\ne \xz{2}
\end{align*}

We showed then, that $v_1-v_2+g_2^\top (u_2-u_1)=0$, implies $g_1=g_2$.

We obtain, by <ref>, that there exists a function $\bar f((w,x),z)$ such that for all $t$ and $\SS$:

\begin{equation}\label{eq:prv1} \bar f(w_t^\SS,x_t^\SS,z)= f(w_t^\SS,z)+h_z(x_t^\SS)
\end{equation}

and for all $t\le T$,

\begin{equation}\label{eq:prv2}
\nabla  \bar f((w_t^\SS,x_t^\SS)) = (\O_{\SS,t,z},\nabla h_z(x_{t}^\SS))

This proves <ref>.
Indeed. By the Lipschitzness of $h$ in the unit ball, we have that $|x_t|\le \gamma \eta T$. For sufficiently small $\gamma$, from <ref>, since $\{(f(0,z)+h_z(0),(0,\nabla h_z(0)),(0,0))\}\in G(z)$, we obtain that

\[|\bar f((0,0),z)- f(0,z)|= \gamma |h_z(0)|\le \gamma^2 \eta T\le \epsilon.\]

\[|\bar f((w_{\q}^\SS,x_{\q}^\SS),z)- f((w_{\q}^\SS,z))|= \gamma |h_z(x_{\q}^\SS)|\le \epsilon.\]

Further, if we assume by induction that for every $t'\le t-1$, $u_{t'}=w^{\SS}_{t'}$, and $x_{t'}=x_{t'}^\SS$, then from <ref> we have that for any first order oracle:

\begin{align*} (u_t,x_t) &= (w_{t-1}^\SS,x_{t-1}^\SS)-\frac{\eta}{|S_t|}\sum_{z\in S_t} \O_z(w^\SS_{t-1},x_{t-1}^\SS)
\\
&= (w_{t-1}^\SS,x_{t-1}^\SS)-\frac{\eta}{|S_t|}\sum_{z\in S_t}\left(\O_{\SS,t,z},\nabla h_z(x_{t-1}^{\SS},z)\right)
\\
&= (w_{t-1}^\SS,x_{t-1}^\SS)-\frac{\eta}{|S_t|}\sum_{z\in S_t}\left(\OS{t}(S_{t-1},w^\SS_{t-1},z),\nabla h_z(x_{t-1}^{\SS},z)\right)
\\
& = (w_{t-1}^\SS,x_{t-1}^\SS)-\eta \left(\O(\SS,w^\SS_{t-1}),\frac{1}{|S_t|} \sum_{t=1}^T \nabla h_z(x_{t-1}^{\SS},z)\right)
\\
\\
&= \left(w_{t-1}^\SS - \eta O(\SS,w_{t-1}^\SS), x_{t-1}^\SS - \frac{\eta}{|S_t|}\sum \nabla h_z(x_{t-1}^\SS)\right) \\
& = (w_t^{\SS},x_{t}^\SS), &\cref{eq:hdynamic}
\end{align*}

which proves, by linearity, that $u_{\q} = w_{\q}^\SS$.

§ PROOF OF <REF>

We choose

\[ \hat f(w) = \max_{j\in J}\{f_j+ g_j^\top(w-w_j\}.\]

$\hat f$ is indeed convex as it is the maximum of linear functions. Further, it is known <cit.> that at each point $w$:

\[ \partial \hat f(w) = \{g_j: \hat f(w)= f_j+ g_j^\top(w-w_j)\}.\]
It follows that, $\hat f$ is $L$-Lipschitz. Next, for any $w_i$, notice that our assumption implies:

\[ f_i \ge \hat f(w_i).\]
On the other hand,

\[f(w_i) = f_i + g_i^\top(w_i-w_i)\le \hat f(w_i).\]

Hence $\hat f(w_i)= f_i$. Finally, to see the function is differentiable at designated points, take any $i\in I_{diff}$ and consider $w_i$. It is known <cit.> that

\[ \partial \hat f(w) = \mathrm{conv}\{g_j: \hat f(w) = f_j+g_j^\top (w-w_j)\}.\]

So it is enough to show that if $g_j \in \partial f(w_i)$ then $g_i = g_j$, but this clearly follows from our assumption, and the fact that $\hat f(w_i)=f_i$.

§ PROOF OF <REF>

We first prove by induction that $x_0^\SS=0$ and, for $t\ge 1$:

\begin{equation}\label{eq:xst}
x_t^{\SS} =  \gamma \eta\sum_{z\in Z}\alpha(z)\left(\sum_{\{t'\le t,z\in S_{t'}\}} (1-\gamma\eta)^{t-t'}\right) .
\end{equation}

\begin{align*}
x_t^{\SS}&= x_{t-1}^{\SS} -\gamma\eta \left(\eta x_{t-1}^{\SS}- \alpha(\SS_t)\right) \\
&= (1-\gamma \eta) x_{t-1}^\SS + \gamma \eta \alpha(\SS_t)
\\
&= (1-\gamma \eta)\left(\sum_{z \in Z} \alpha(z) \sum_{\{t'<t : z\in S_{t'}\}}(1-\gamma\eta)^{t-1-t'} \eta \gamma \right)+ \sum_{z\in S_t}\gamma \eta \alpha(z)\\
&= \sum_{z \in Z} \alpha(z) \sum_{\{t'\le t-1 : z\in S_{t'}\}}(1-\gamma\eta)^{t-t'} \eta \gamma + \gamma \eta \mathbf{1}[ z\in \SS_t]\\
&=  \eta \gamma  \sum_{z \in Z}\alpha(z)  \sum_{\{t'\le t : z\in S_{t'}\}}(1-\gamma\eta)^{t-t'}
\end{align*}

Now, for every $\SS, t, z$, define a polynomial:

\[ P_{\SS,t,z}(X) = \sum_{n=0}^{t-1} \mathbf{1}[z\in S_{t-n}]X^{n},\]

and let $q$ be a rational point, sufficinetly small, so that $1-q$ is not the root of any polynomial of the form $P_{\SS,t,z}-P_{\SS',t',z'}$ that is distinct from $0$. In other words, we choose $q$ so that

\[ P_{S,t,z}(1-q) = P_{S',t',z'}(1-q) \Leftrightarrow  P_{S,t,z}(X) = P_{S',t',z'}(X).\]

as well as

\[ \frac{1}{T} \sum_{t=1}^T P_{S,t,z}(1-q) = P_{S',t',z}(1-q) \Leftrightarrow  \frac{1}{T} \sum_{t=1}^T P_{S,t,z}(X) = P_{S',t',z'}(X).\]

Notice that there are only finitely many polynomials of the above form, hence we can choose such $q$ in any interval $(0,\epsilon)$ for any $\epsilon>0$.

Rewriting <ref> we have:

\[ x_t^\SS = \gamma\eta \sum_{z\in Z} \alpha(z)P_{\SS,t,z}\left(1-\gamma\eta\right)\]

Now, suppose we choose $\{\alpha(z)\}$ to be reals, independent over the rationals, and suppose we choose $\gamma = q/\eta$.

Proof of <ref>
Assume that $x_t^\SS = x_{t'}^{\SS'}$. Because $\alpha(z)$ are independent over the rationals, and because $P(1-\gamma\eta)$ are always rationals, we have that $P_{\SS',t',z}(1-q) = P_{\SS',t',z}(1-q)$ for every $z$. But then,
\[ \forall z\in Z: P_{\SS,t,z}(X)= P_{\SS',t',z}(X),\]
by choice of $q$.

But then, $t=t'$. Indeed, assume by contradiction and w.l.og assume $t< t'$: then for any $z\in S'_1$ $P_{\SS,t,z}$ is a $t'-1$-degree polynomial, on the other hand $P_{\SS,t,z}$ is at most of degree $t-1< t'-1$.
Next, if $S_i\ne S'_i$, for $i\le t$, then we can assume (w.l.o.g) that there is $z\in S_i$ such that $z\notin S'_i$ it follows that, by looking at the coefficient of the two polymials of the monomial $X^{t-i}$ we have that:

\[ P_{\SS',t',z}(X) \ne  P_{\SS,t,z}(X).\]

Overall then, we obtain that if $x_{t}^\SS = x_{t'}^{\SS'}$ then $(\SS,t)\equiv (\SS',t').$.

Proof of <ref>: The proof of <ref> is similar to how we proved $t=t'$. Notice that:

\[ \sum_{t=1} \q(t) x_t^{\SS} = \gamma\eta \sum_{z\in Z} \alpha(z) \sum_{t=1}^T \q(t) P_{\SS,t,z}(1-\eta\gamma).\]

Hence $x_t^{\SS}= x^{\SS'}$ implies for all $z\in Z$:

\[ \sum_{t=1}^T \q(t) P_{\SS',t,z} = P_{\SS,t,z}.\]

For $z\in S'_1$ we have that $\sum_{t=1}^T \q(t) P_{\SS,t,z}$ is a $\max\{t: \q(t)\ne 0\}$-degree polynomial, but on the other hand we assume that $t<\max\{t: \q(t)\ne 0\}$, hence $P_{\SS,t,z}$ is a lower degree polynomial.

§ DIMENSION INDEPENDENT LOWER BOUND FOR GD

In this section we prove that the optimization error of GD in <ref> is optimal. The lower bound is an optimization error for first order methods and is well established (see <cit.>). The point here is to show that the bound is valid in any dimension, for GD.

For every choice of $\eta, T$, there exists a convex and $1$ Lipschitz function $f(x):\mathbb{R}\to \mathbb{R}$ such that, if we run GD on $f$, then

\[ f(w^{GD}) - f(0) \ge \frac{\eta}{2} + \frac{1}{6\eta T}.\]

We divide the proof into two cases:

case1: $\eta \ge 1/\eta T$:

In this case we choose $f(x)= |x-\gamma|$, where $\gamma>0$ is a arbitrarily small (may depend on $\eta$). It can be seen that for every even iteration, we have that:

\[\nabla f(x_{2t}) = -1,\]

at at every odd iteration

\[ \nabla f(x_{2t+1}) =+1.\]

As such, for every even iteration we have that $x_{2t} = 0$, and at every odd iteration we have that $x_{2t+1}= \eta$. We thus have that $x^{GD}=\frac{\eta}{2}$ and

\[ f(x^{GD})- 0 = \frac{\eta}{2} - 0 \ge \frac{\eta}{4} + \frac{1}{2\eta T}.\]

case1: $\eta \le 1/\eta T\le 1$:
In this case choose $f(x) = \alpha x$, where $\alpha = \frac{1}{(T+1)\eta}\le 1.$ Then, one can show that GD outputs

\[ x^{GD} = -\frac{\eta}{T}\sum t\cdot\alpha = -\frac{(T+1)\eta}{2} \alpha,\]

\[f(x^{GD}) -f(-1) = \alpha - \frac{(T+1)\eta}{2} \alpha^2 \ge  \frac{1}{2(T+1)\eta} \ge \frac{1}{3T\eta} \ge \frac{1}{6T\eta}+\frac{\eta}{2}.\]

§ PROOF OF <REF>

§.§ Proof of <ref>

To avoid cumbersome notations, we will supress dependence on $\SS$, and assume w.l.o.g that $I=\{1,\ldots, d'\}$. Namely, we will write $w_t$ instead of $w_t^\SS$, and $w_t^\SS(j)$ instead of $w_t^\SS(i_j)$, where $i_j\in I$.

As a first step, observe that at each iteration no projection is performed. Indeed, let us show by induction that:

\[ \|w_{t+1}\|^2=\|w_t-\eta \OS{t}(S,w,V)\|^2 \le 2\eta^2\alpha^2 (t+1).\]

The definition of $\alpha$ then implies that $w_t$ are restricted to the unit ball without projections.

To see the above is true, let us consider the case where the first type of update is performed:

\begin{align*}\|w_{t+1}\|^2&= \|w_{t}^\SS-\eta \OS{t}(S,w,V)\|^2\\
&=\|w_{t}^\SS-\eta\alpha e_{i_{j+1}} +\eta\alpha e_{i_j}\|^2\\
&= \sum_{s=1}^{j-1} (w_t(i_s))^2 +
(w_t(i_{j})+\eta\alpha)^2 +(w_t(i_{j+1})-\eta\alpha)^2
+\sum_{s=j+2}^{d'} (w_t(i_s))^2\\
&=\sum_{s=1}^{d'} (w_t(i_s))^2+2\eta\alpha(w_t(i_j)-w_t(i_{j+1}))+2\eta^2\alpha^2\\
&= \sum_{s=1}^{d'} (w_t(i_s))^2+2\eta^2\alpha^2 & w_t(i_{j+1})=w_t(i_{j})\\
&\le 2\eta^2\alpha^2 \cdot t+ \eta^2\alpha^2\\
&=2\eta^2\alpha^2 \cdot (t+1).
\end{align*}

And if the second type of update is performed:

\begin{align*}\|w_{t+1}\|^2&= \|w_{t}^\SS-\eta \OS{t}(S,w,V)\|^2\\
&=\|w_{t}^\SS-\eta\alpha e_j\|^2\\
&= \sum_{s\ne j} w_t(s) + \eta^2\alpha^2 & w_t(s)=0\\
& \le 2\eta^2\alpha^2\cdot t +\eta^2\alpha^2\\
& \le 2\eta^2\alpha^2\cdot (t+1)
\end{align*}

We now move on to prove that <ref> holds by induction. Specifically, we will show that for $d_0 \le d'$ that at time $T_{d_0} = 1+\sum_{d=0}^{d_0} \sum_{k=0}^{d} k$, for any $t\le d_0$:

\[ w_{T_{d_0}}(t) =\alpha\eta \begin{cases}
(d_0 + 1 -t) & t\le  d_0\\
0 & \textrm{o.w.}
\end{cases} .\]

For $d_0=0$, $T_0=1$ and we have, indeed, $w_1=0$. Next assume we proved the statement for $d_0$, and we will prove it for $d_0+1$. Here too, we will use induction, and we prove that, for $d_1\le d_0+1$, at time
\[T_{d_0,d_1}= T_{d_0} + \sum_{k=0}^{d_1-1} (d_0+1-k),\]
we have that:

\[ w_{T_{d_0,d_1}}(t) = \alpha\eta \begin{cases}
(d_0+2-t) & t\le d_1\\
(d_0+1-t) & d_1< t\le  d_0\\
0& \textrm{o.w.}
\end{cases}\]

For the case $d_1$=0, $T_{d_0,d_1}=d_0$, and it follows from our outer-induction step. Assume the statement is true for $d_1$ and we will prove it for $d_1+1$, here, yet again, we use induction. And we will show that for $1 \le d_2< d_0+2-d_1$ we have at time
\[T_{d_0,d_1,d_2}=T_{d_0} + T_{d_1} + d_2,\]
we have that:

\begin{equation}\label{eq:sksd} w_{T_{d_0,d_1,d_2}}(t) = \alpha\eta \begin{cases}
(d_0+2-t) & t \le  d_1\\
(d_0 +1 -t) & d_1 < t < d_0 + 2 -  d_2\\
(d_0+2 -t) & t= d_0 + 2 - d_2 \\
(d_0+1-t) & d_0+2-d_2 < t\le  d_0\\
0 & \textrm{o.w.}
\end{cases}\end{equation}

We start the induction with the case $d_2=1$, in that case notice that $T_{d_0,d_1,d_2}= T_{d_0,d_1}+1$, and by induction hypothesis:

\[ w_{T_{d_0,d_1}}(t) = \alpha\eta \begin{cases}
(d_0+2-t) & t\le d_1\\
(d_0+1-t) & d_1< t\le  d_0\\
0 & \textrm{o.w.}
\end{cases}\]

In this case, note that there are no two consecutive coordinates that are equal, hence our choice of oracle is defined so that $\OS{t} = - e_{d_0+1}$. Hence, by our update rule (and the lack of projections which we proved at the beginning):

\[ w_{T_{d_0,d_1,1}}(t) = \alpha\eta \begin{cases}
(d_0+2-t) & t\le d_1\\
(d_0+1-t) & d_0< t\le  d_0\\
1 & t=d_0+1\\
0 & \textrm{o.w.}
\end{cases}\]

Which satisfies <ref>. Now assume that <ref> holds for $d_2$, and take $d_2+1 < d_0+2-d_1$ (otherwise, we are done). Notice that $T_{d_0,d_1,d_2+1} = T_{d_0,d_1,d_2}+1$. Observe that $w_{T_0,d_1,d_2}(d_0+2-d_2) =w_{T_0,d_1,d_2}(d_0+1-d_2)$ (notice that $d_0+1-d_2>d_1$), and our update rule is such that $\OS{t}= e_{d_0+2-d_2}- e_{d_0+1-d_2}$ and we obtain then:

\[ w_{T_{d_0,d_1,d_2+1}} =w_{T_{d_0,d_1,d_2}} - \eta \alpha e_{d_0+2-d_2} + \eta \alpha e_{d_0+1-d_2}
= \alpha\eta \begin{cases}
(d_0+2-t) & t \le  d_1\\
(d_0 +1 -t) & d_1 < t < d_0 + 2 - ( d_2+1)\\
(d_0+2 -t) & t= d_0 + 2 - (d_2+1) \\
(d_0+1-t) & d_0+2-(d_2+1) < t\le  d_0\\
0 & \textrm{o.w.}
\end{cases}\]

The most inner induction step is now complete. We now notice that $T_{d_0,d_1,d_0+1-d_1}= T_{d_0,d_1+1}$, which proves the middle-induction step. And we notice that $T_{d_0,d_0+1}= T_{d_0+1}$, which proves the whole induction argument.

§.§ Proof of <ref>

We only need to show that the following quantity is increasing in $t$

\[X_t= \max\left\{\sum_{i\in I_B} w_t^\SS(i): I_B\subseteq I, |I| = B\right\}.\]

But, as shown in <ref>, the update rule is such that we don't perform projections. It then follows easily from our update step. Indeed if we increase a coordinate by $\alpha \eta$, then clearly the $X_t$ only increases. Also, if we perform update of the form:

\[ w_t^\SS(i_1)=w_{t-1}^\SS(i_{1})+\alpha \eta,\quad  w_t^\SS(i_2)=w_{2}^\SS(i_{2})-\alpha \eta,\]

for two consecutive and equal coordinates, then for any $I_B$: If $I_B$ includes both $i_1$ and $i_2$ then the magnitude doesn't change. If $I_B$ contains only $i_{1}$ then it increases, and if $I_B$ contains only $i_2$ then consider $I_{B'}$ that swaps $i_{1}$ with $i_{2}$ then we clearly have:

\[ \sum_{i\in I_{B'}} w_{t}^\SS(i)>\sum_{i\in I_{B'}} w_{t-1}^\SS(i) > \sum_{i\in I_{B}} w_t^\SS(i)
\]


# ProDyn0: Inferring calponin homology domain stretching behavior using graph
neural networks

Ali Madani, Cyna Shirazinejad, Jia Rui Ong, Hengameh Shams, Mohammad Mofrad
Molecular Cell Biomechanics Laboratory
University of California, Berkeley
Berkeley, CA 94720, USA
<EMAIL_ADDRESS>
###### Abstract

Graph neural networks are a quickly emerging field for non-Euclidean data that
leverage the inherent graphical structure to predict node, edge, and global-
level properties of a system. Protein properties can not easily be understood
as a simple sum of their parts (i.e. amino acids), therefore, understanding
their dynamical properties in the context of graphs is attractive for
revealing how perturbations to their structure can affect their global
function. To tackle this problem, we generate a database of 2020 mutated
calponin homology (CH) domains undergoing large-scale separation in molecular
dynamics. To predict the mechanosensitive force response, we develop neural
message passing networks and residual gated graph convnets which predict the
protein dependent force separation at 86.63 percent, 81.59 $kJmol^{-1}nm^{-1}$
MAE, 76.99 $psec$ MAE for force mode classification, max force magnitude, max
force time respectively– significantly better than non-graph-based deep
learning techniques. Towards uniting geometric learning techniques and
biophysical observables, we premiere our simulation database as a benchmark
dataset for further development/evaluation of graph neural network
architectures.111https://github.com/a-mad/prodyn

##  1 Introduction

The primary sequence of a well-ordered protein encodes its structural
information as well as its functional properties. Proteins undergo dynamic
shifts in conformation to perform their designated tasks. The growth of
computational resources and simulation algorithms allow for studying molecular
trajectories at atomic resolution commonplace. However, because of the
heterogeneity of bio-molecules and their chemical environments, computational
models can suffer in performance when tasked with predicting molecular
conformations, physical and chemical properties, and protein-protein
interaction networks. Along with the growing excitement in machine learning in
tangent fields, there is emerging interest in data-driven approaches to
studying biological function at scales ranging from biochemical properties to
tissues mechanics (Glazer et al. (2008); Botu & Ramprasad (2015); Schütt et
al. (2017); Gastegger et al. (2017); Madani et al. (2018); Wu et al. (2018);
Ramakrishnan et al. (2015)).

Mechanosensitive proteins are central to a plethora of biological phenomena
(Galkin et al. (2010)). Our goal in this paper is to bridge the primary
sequence of proteins to their intrinsic force observables while undergoing
large-scale conformational changes. To accomplish this, we generated a
molecular dynamics database of trajectories of two calponin homology domain
mutants being pulling apart [Figure 1] by Steered Molecular Dynamics (SMD)
(Isralewitz et al. (2001)). SMD is a technique that allows for biasing the
interaction potential between explicit regions in a simulation; in our case,
two globular domains, CH1 and CH2, are gradually pulled apart during the
simulation. While these simulations are computationally expensive and require
careful consideration for proper physical dynamics, they provide information
at the atomic resolutions about the sensitivity of the roles mutations play in
response to mechanical loads in proteins. Our aim to cast protein structures
as graphical models allows us to take the first imperative steps towards
building relationships between sequences and physical observables that can be
computed with more expensive techniques such as molecular dynamics. We show
that the force between the two CH domains, a global property over a long
simulated trajectory, can accurately be predicted using graph neural networks.

Figure 1: Workflow Schematic: From Biophysics to Graph Neural Networks

##  2 Methods

###  2.1 Generating a database of point-mutated protein dynamics

To generate models to predict protein dynamics, we established the first
ProDyn dataset with high-throughput in silico mutagenesis experiments of
proteins undergoing large-scale conformational changes. The generated
trajectories span 2020 unique point-mutations along two calponin homology
domains, CH1 and CH2, connected by a tether that comprise a total of 224
residues. A full description of the molecular dynamics methods can be found in
Supplementary Materials. The data per simulation was processed to yield a
graphical input representation of residue interactions in terms of node and
edge features, in addition to an output representation for force profile
classification and regression. We then construct and compare conventional and
graph neural network architectures to predict the force characteristics of our
biophysical system.

###  2.2 Data Pre-processing

A molecular dynamics (MD) simulation under a given force field outputs a
trajectory file that includes type, position, and velocity information for
every atom in the system for each time step. There are several prediction
problems of interest that can be encapsulated by our data. In this study, our
specific objective is to predict the force characteristics for a given protein
initial state (i.e. residue information and interactions pre-steering).

The input representation is a directed graph, $G$, with node features, $x_{v}$
and edge features, $e_{vw}$. As shown in Table 1, node features comprise of
the residue type, residue properties, and initial secondary structure type as
calculated by mdtraj (McGibbon et al. (2015)). Edge features currently consist
of pair-wise dihedral angles, Kabsch-Sander (K-S) hydrogen bonds (Kabsch &
Sander (1983)), and the initial distance between steered and fixed residues.

The output representation, a graph-level target of protein mechanosensitivity,
is of two types: regression and classification. Each SMD simulation exhibits a
characteristic force response over time of the simulation. The regression
target is a 2 dimensional vector encoding maximum pulling force magnitude and
maximum pulling force time-point respectively. The classification target is a
one-hot encoded 2-dimensional vector that encodes a force mode category. The
force mode categories were ascertained via spectral k-means (k=2) clustering
on the force-response graphs in a lower dimensional t-SNE (Maaten & Hinton
(2008)) derived space as shown in Figure 2. There emerged two clearly well-
separated clusters which we utilized as categories for graph-level
classification. Lastly, all continous data was scaled to [0,1], and the entire
dataset was split randomly into a 404 sample held-out test set and 1616 sample
80-20 training/validation split. For the classification task, we compute the
accuracy and macro-averaged F1-score as performance metrics. For the
regression task, we calculate the mean absolute error (MAE) for maximum force
magnitude and maximum force time-point.

Table 1 - Input Representation Features
---
Feature | Node/Edge | Type | Index | Examples
Residue Type | Node | One-hot | $[0:20]$ | ALA, GSP, ALY, …
Residue Secondary Structure | Node | One-hot | $[21:26]$ | alpha-helix, turn, …
Residue Properties | Node | Multi-label | $[27:43]$ | acidic, polar, …
Dihedral-$\phi$ | Edge | Continuous | $[0]$ | $0-2\pi$
Dihedral-$\psi$ | Edge | Continuous | $[1]$ | $0-2\pi$
Dihedral-$\omega$ | Edge | Continuous | $[2]$ | $0-2\pi$
Hydrogen Bond | Edge | Continuous | $[3]$ | K-S energy
Steer-to-Fixed Residue Dist | Edge | Continuous | $[4]$ | average distance

###  2.3 Baseline Conventional Deep Learning Models

As a comparison to graph-based neural networks described below, conventional
multi-layer perceptron (MLP) and gated recurrent unit (GRU) were trained (Cho
et al. (2014)). The graph node features $x_{v}$ were provided as input either
by concatenation or sequentially. For the MLP architectures with one vector
input of concatenated features, the number of layers [1-5], activation
functions [ReLU, eLU], nodes per layer [128-512], dropout regularization
[0.0-0.5], and loss functions [L1, L2, logcosh, cross-entropy] were
experimented with. For the GRU+MLP architecture with inputs features passed
sequentially through the GRU, we experimented with different GRU
implementations- varying hidden dimensions [16-64], stacking [1-2], bi-
directionality, regularization [via dropout], and outputs [last layer depth-
wise vs time-wise].

###  2.4 Neural Message Passing Network Model

In line with Gilmer et al. (2017), we developed various neural message passing
network (MPNet) architectures. Our models learn to compute a function of the
entire graph by learning features, invariant to graph isomorphisms, through a
message passing algorithm. We define the hidden state of each node in the
graph by $h^{t}_{v}$ with $h^{0}_{v}=x_{v}$ where v denotes residue/node
number and t denotes the cycle number, and the directed edge from node w to v
as $e_{v\leftarrow w}$. The message passing scheme runs for a total of $T$
cycles which can be interpreted as the message-passing neighborhood size.
During each cycle, messages are passed between a neighboring node and the
current node based on the edge direction. There is a unique message function,
$M$, that computes the message given the current node features, neighboring
node features, and edge features which is then aggregated into one message
update per node, $m_{v}^{t}$. This message is then used to update the hidden
node state by the function, $U$. The final hidden states of all nodes are then
passed through a readout function, $R$, which yields an output layer of the
dimension of the sample target. To formalize, for each message passing cycle
the node states are updated in the following fashion:

$m_{v}^{t}=\frac{1}{|N(v)|}\sum_{w\in
N(v)}M(h_{v}^{t},h_{w}^{t},e_{v\leftarrow w})$ (1)
$h_{v}^{t+1}=U(h_{v}^{t},m_{v}^{t})=h_{v}^{t}+m_{v}^{t}$ (2)

where $N(v)$ denotes the neighbors of v in graph G. The message function, M,
is a 3-layer MLP with 0.4 dropout, 256 hidden layer dimensions, and output
layer of the same dimension as node features. After T cycles, the readout
function R is applied as follows:

$\hat{y}=R(\\{h_{v}^{T}|v\in G\\})$ (3)

We did not see a significant increase in performance for increasing T cycles
and present our MP models for T=1 in Results. For the choice of R, we
experimented with a max-pooling operation over the feature dimension and also
feeding the nodes in primary sequence order to a GRU. The final output time-
wise of the GRU is fed to a 2-layer MLP with 0.5 dropout. The final loss
function, $L(\hat{y},y)$ was either a L1-norm for regression or softmax+cross-
entropy for classification. Models were trained using stochastic gradient
descent with ADAM (Kingma & Ba (2014)) and batch size of 25.

###  2.5 Residual Gated Graph ConvNet Model

Residual Gated Graph ConvNets (RGGCN) is a method introduced by Bresson &
Laurent (2017) that combines the vanilla graph ConvNet (Sukhbaatar et al.
(2016)) and the edge gating mechanism (Marcheggiani & Titov (2017)) with
residual connections for problems involving variable graphs. We extend this
method by incorporating edge features, such that each graph convolution layer
follows the layer-wise propagation rule:

$h_{i}^{\ell+1}=\operatorname{ReLU}\left(U^{\ell}h_{i}^{\ell}+\sum_{j\rightarrow
i}\eta_{j\rightarrow i}\odot(V^{\ell}h_{j}^{\ell}+D^{\ell}e_{j\rightarrow
i})\right)$

where $h_{i}^{\ell}$ denote features of node $i$ at layer $l$,
$e_{j\rightarrow i}$ denote features from edge $j\rightarrow i$, and edge
gates $\eta_{j\rightarrow
i}=\sigma\left(A^{\ell}h_{i}^{\ell}+B^{\ell}h_{j}^{\ell}+C^{\ell}e_{j\rightarrow
i}\right)$. $U$, $V$, $A$, $B$, $C$ and $D$ are learnable parameters.

In our experiments, we used a 4-layer RGGCN and applied max pooling as our
readout function. For classification, our model had 50 hidden dimensions and a
single linear layer. For regression, our model had 256 hidden dimensions with
a 2-layer MLP. Batch normalization was employed, as with residual connections
between successive convolution layers.

Our models were initialized using Glorot initialization (Glorot & Bengio
(2010)). We used the Adam SGD optimizer with an initial learning rate of
0.0005, and computed loss using binary cross-entropy for classification, and
L1 for regression. Dropout rate was set to 0.4 for both node features and edge
gates in classification, and only for the 2-layer MLP in regression.

##  3 Results and Discussion

One of our main achievements is creating a database that is uniquely well-
positioned to benchmark the rapid development and evaluation of novel graph
neural network architectures on both a graph-level classification and
regression task. As shown in Table 2, our baseline non-graph-based deep
learning models– both (1) variations of concatenating features and processing
through an MLP and (2) sequentially feeding primary sequence-ordered nodes
through a GRU – were unable to learn effectively and predicted tightly near
the majority/mean target. Within our data size scale, the complexity of many
biophysical phenomena, through initial states and short/long-range bonded and
non-bonded interactions partially captured in this study, is too difficult to
train effectively unconstrained via an MLP or recurrently via a GRU.

Table 2 - Prediction Performance of Various Architectures
---
Model | Accuracy | F1-Score |  $\text{Force}_{mag}$ MAE |  $\text{Force}_{time}$ MAE
Conventional | 70.54 | 0.414 | 89.68 | 111.39
MPNet + Maxpool | 77.09 | 0.673 | 85.39 | 92.44
MPNet + GRU | 86.63 | 0.839 | 81.59 | 76.99
RGGCN + Maxpool | 86.39 | 0.904 | 89.00 | 100.52

Our goal was to experiment with 2 types of graphical neural networks: our own
designed MPNet and a state-of-the-art graph technique such as the RGGCN. Both
graph neural network architectures, MPNet and RGGCN, capture the graphical
structure, information communication, and structural invariances to learn an
effective node embedding for classification/regression tasks. In end-to-end
training after node embedding, a readout function is applied that is agnostic
to the number of nodes/residues. Among functions invariant to graph
isomorphism, we observed that maxpool outperformed average pool. As proteins
have multiple hierarchies of structure, we can view the MPNet with GRU readout
as encoding the protein structure hierarchy with the GRU applied after message
passing cycle, as to allow message passing interactions to be communicated
between nodes before the recurrent primary structure network. The choice of
GRU readout is likely more effective for protein-like graph systems; the
RGGCN+Maxpool would be more generalizable otherwise.

Figure 2: MPNet + GRU Performance. a) Classification training statistics b)
Regression histograms of predicted vs actual values

As shown in Table 2, all graph neural networks significantly outperform
conventional deep learning techniques. The RGGCN+Maxpool and MPNet+GRU perform
well on the classification task at around 86%. For the regression task, the
MPNet+GRU is able to learn with a MAE of 81.59 $kJmol^{-1}nm^{-1}$ MAE on max
force magnitude prediction and 76.99 $psec$ MAE on max force time prediction.
To put into context the performance to the natural stochasticity of the MD
system, we ran duplicate SMD simulations and observed the mean standard
deviation of the max force mag was 57.94. Also, in Figure 2, we show the well-
behaved convergence of the classification model training in addition to the
resemblant output statistics of the true and predicted targets for regression.

Future work can include learning edge states in addition to node states,
set2vec as readout function (Vinyals et al. (2015)) or GRU as an update
function, and experimenting with unsupervised graph embedding techniques such
as Veličković et al. (2018). From the biophysical perspective, we can expand
the choice of node/edge features, try different protein systems, or formulate
a plethora of pertinent biophysical prediction tasks.

To conclude, the progress of geometric deep learning techniques is
inextricably linked to the quality of benchmark datasets– analagous to the
impact of benchmark datasets for vision and language. We view our work as
laying the groundwork for the intersection of graph neural network and
biophysics researchers for the advancement of both fields.

## References

  * Berendsen et al. (1995) Herman JC Berendsen, David van der Spoel, and Rudi van Drunen.  Gromacs: a message-passing parallel molecular dynamics implementation.  _Computer physics communications_ , 91(1-3):43–56, 1995.
  * Botu & Ramprasad (2015) V. Botu and R. Ramprasad.  Learning scheme to predict atomic forces and accelerate materials simulations.  _Phys. Rev. B_ , 92:094306, Sep 2015.  doi: 10.1103/PhysRevB.92.094306.  URL https://link.aps.org/doi/10.1103/PhysRevB.92.094306.
  * Bresson & Laurent (2017) Xavier Bresson and Thomas Laurent.  Residual gated graph convnets.  _arXiv preprint arXiv:1711.07553_ , 2017.
  * Cho et al. (2014) Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio.  Learning phrase representations using rnn encoder-decoder for statistical machine translation.  _arXiv preprint arXiv:1406.1078_ , 2014.
  * Galkin et al. (2010) Vitold E Galkin, Albina Orlova, Anita Salmazo, Kristina Djinovic-Carugo, and Edward H Egelman.  Opening of tandem calponin homology domains regulates their affinity for f-actin.  _Nature structural & molecular biology_, 17(5):614, 2010.
  * Gastegger et al. (2017) Michael Gastegger, Jörg Behler, and Philipp Marquetand.  Machine learning molecular dynamics for the simulation of infrared spectra.  _Chemical science_ , 8(10):6924–6935, 2017.
  * Gilmer et al. (2017) Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl.  Neural message passing for quantum chemistry.  In _Proceedings of the 34th International Conference on Machine Learning-Volume 70_ , pp. 1263–1272. JMLR. org, 2017.
  * Glazer et al. (2008) Dariya S Glazer, Randall J Radmer, and Russ B Altman.  Combining molecular dynamics and machine learning to improve protein function recognition.  In _Biocomputing 2008_ , pp. 332–343. World Scientific, 2008.
  * Glorot & Bengio (2010) Xavier Glorot and Yoshua Bengio.  Understanding the difficulty of training deep feedforward neural networks.  In _Proceedings of the thirteenth international conference on artificial intelligence and statistics_ , pp. 249–256, 2010.
  * Henikoff & Henikoff (1992) Steven Henikoff and Jorja G Henikoff.  Amino acid substitution matrices from protein blocks.  _Proceedings of the National Academy of Sciences_ , 89(22):10915–10919, 1992.
  * Isralewitz et al. (2001) Barry Isralewitz, Mu Gao, and Klaus Schulten.  Steered molecular dynamics and mechanical functions of proteins.  _Current opinion in structural biology_ , 11(2):224–230, 2001.
  * Kabsch & Sander (1983) Wolfgang Kabsch and Christian Sander.  Dictionary of protein secondary structure: pattern recognition of hydrogen-bonded and geometrical features.  _Biopolymers: Original Research on Biomolecules_ , 22(12):2577–2637, 1983.
  * Kingma & Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Lindahl et al. (2001) Erik Lindahl, Berk Hess, and David Van Der Spoel.  Gromacs 3.0: a package for molecular simulation and trajectory analysis.  _Molecular modeling annual_ , 7(8):306–317, 2001\.
  * Maaten & Hinton (2008) Laurens van der Maaten and Geoffrey Hinton.  Visualizing data using t-sne.  _Journal of machine learning research_ , 9(Nov):2579–2605, 2008.
  * Madani et al. (2018) Ali Madani, Jia Rui Ong, Anshul Tibrewal, and Mohammad RK Mofrad.  Deep echocardiography: data-efficient supervised and semi-supervised deep learning towards automated diagnosis of cardiac disease.  _npj Digital Medicine_ , 1(1):59, 2018.
  * Marcheggiani & Titov (2017) Diego Marcheggiani and Ivan Titov.  Encoding sentences with graph convolutional networks for semantic role labeling.  _arXiv preprint arXiv:1703.04826_ , 2017.
  * McGibbon et al. (2015) Robert T. McGibbon, Kyle A. Beauchamp, Matthew P. Harrigan, Christoph Klein, Jason M. Swails, Carlos X. Hernández, Christian R. Schwantes, Lee-Ping Wang, Thomas J. Lane, and Vijay S. Pande.  Mdtraj: A modern open library for the analysis of molecular dynamics trajectories.  _Biophysical Journal_ , 109(8):1528 – 1532, 2015\.  doi: 10.1016/j.bpj.2015.08.015.
  * Ramakrishnan et al. (2015) Raghunathan Ramakrishnan, Mia Hartmann, Enrico Tapavicza, and O Anatole Von Lilienfeld.  Electronic spectra from tddft and machine learning in chemical space.  _The Journal of chemical physics_ , 143(8):084111, 2015.
  * Šali & Blundell (1993) Andrej Šali and Tom L Blundell.  Comparative protein modelling by satisfaction of spatial restraints.  _Journal of molecular biology_ , 234(3):779–815, 1993.
  * Schütt et al. (2017) Kristof T Schütt, Farhad Arbabzadah, Stefan Chmiela, Klaus R Müller, and Alexandre Tkatchenko.  Quantum-chemical insights from deep tensor neural networks.  _Nature communications_ , 8:13890, 2017.
  * Shams et al. (2016) Hengameh Shams, Javad Golji, Kiavash Garakani, and Mohammad RK Mofrad.  Dynamic regulation of $\alpha$-actinin’s calponin homology domains on f-actin.  _Biophysical journal_ , 110(6):1444–1455, 2016\.
  * Sukhbaatar et al. (2016) Sainbayar Sukhbaatar, Rob Fergus, et al.  Learning multiagent communication with backpropagation.  In _Advances in Neural Information Processing Systems_ , pp. 2244–2252, 2016.
  * Van Der Spoel et al. (2005) David Van Der Spoel, Erik Lindahl, Berk Hess, Gerrit Groenhof, Alan E Mark, and Herman JC Berendsen.  Gromacs: fast, flexible, and free.  _Journal of computational chemistry_ , 26(16):1701–1718, 2005.
  * Veličković et al. (2018) Petar Veličković, William Fedus, William L Hamilton, Pietro Liò, Yoshua Bengio, and R Devon Hjelm.  Deep graph infomax.  _arXiv preprint arXiv:1809.10341_ , 2018.
  * Vinyals et al. (2015) Oriol Vinyals, Samy Bengio, and Manjunath Kudlur.  Order matters: Sequence to sequence for sets.  _arXiv preprint arXiv:1511.06391_ , 2015.
  * Weiss et al. (2000) Gregory A Weiss, Colin K Watanabe, Alan Zhong, Audrey Goddard, and Sachdev S Sidhu.  Rapid mapping of protein functional epitopes by combinatorial alanine scanning.  _Proceedings of the National Academy of Sciences_ , 97(16):8950–8954, 2000.
  * Wu et al. (2018) Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh S Pappu, Karl Leswing, and Vijay Pande.  Moleculenet: a benchmark for molecular machine learning.  _Chemical science_ , 9(2):513–530, 2018.

##  4 Supplementary Materials

The creation of the mutation library was primarily informed by a BLOSUM95
matrix displaying frequencies of swaps between pairs of residues in closely-
related protein homologues (Henikoff & Henikoff (1992)). In addition, an
alanine scan over most of the sequence was made to complement work done by
saturation mutagenesis experiments (Weiss et al. (2000)). To augment our data
set beyond this, a random sample of positions were mutated to random residues
with no further bias.

To generate the database of single point-mutated calponin homology domain
trajectories, all-atom molecular dynamics was performed. Experimental
procedures were largely adapted from previous work documenting the SMD
experimental protocol for the wild-type CH1-CH2 domain (Shams et al. (2016)).
The wild type structure from Shams et al. was modified at single residue
locations using a MODELLER (https://salilab.org/modeller/) script that swaps
residues, optimizes dihedral angles, reduces bad contacts with conjugate
gradient steps, and performs a few steps of molecular dynamics to allow the
swapped residue to relax (Šali & Blundell (1993)).

All subsequent molecular dynamics was performed with GROMACS (Van Der Spoel et
al. (2005); Lindahl et al. (2001); Berendsen et al. (1995)) using the CHARMM27
all-atom force-field. Pulling simulations were performed in a 11.5 x 11.5 x 13
nm3 box solvated with SPC216 water molecules and 100 mM KCL. The system was
minimized with steepest-descent to remove bad contacts until the total energy
was less than 1000.0 kJ/mol/nm. Periodic boundary conditions were set in all 3
directions, a Verlet cutoff scheme truncated non-bonded long-range
interactions at 1.2 nm, and PME electrostatics with Fourier spacing of 0.12 nm
was used. Following minimization, each system was equilibrated for 100 ps in
NVT followed by 100 ps of NPT, both using 2 fs time steps and a target
temperature of 300K and 1 bar, respectively. The equilibration was closely
monitored for each mutation to ensure that the system was properly sampling
the correct equilibrium distribution before steering experiments were
performed.

SMD experiments were performed by anchoring down and pulling on a selected
group of heavy carbon atoms on the interface between the CH1 and CH2. To
maintain this experimental control, residues that include these atoms were not
mutated in our experiments. All other parameters of our simulations is held
constant except for a single point mutation that is unique to each simulation.
The reaction coordinate distance is defined here as the distance between the
center of masses of the fixed and steered residues. Steering was performed by
moving an umbrella potential with stiffness 1255 kJ mol-1 nm-2 at a rate of 2
nm/ps along a reaction coordinate defined by the axis connecting the center of
masses of the fixed and steered atoms. Temperature coupling was performed by a
Nose-Hoover thermostat targeting 310K and coupling was maintained at 1 bar
with an isotropic Parrinello-Rahman piston for the duration of the pulling
experiment. Throughout the 600 ps pulling experiment, readouts of the reaction
coordinate distance and the force between the fixed and steered atoms were
measured. The maximum force applied between the domains as well as the
entirety of the force profile are the observables of interest predicted with
our neural networks. Figure 1 shows the force profiles can be vary
dramatically in shape for different mutations, while remaining relatively
stationary for different initial conditions per mutation.