SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\nabla_ {x} h (x) - \nabla_ {x} h ^ {\prime} (x) = \left\{ \begin{array}{l l} w u - w ^ {\prime} u ^ {\prime} & \text {w h e n} u ^ {\top} x > 0, (u ^ {\prime}) ^ {\top} x > 0 \\ w u & \text {w h e n} u ^ {\top} x > 0, (u ^ {\prime}) ^ {\top} x <   0 \\ - w ^ {\prime} u ^ {\prime} & \text {w h e n} u ^ {\top} x <   0, (u ^ {\prime}) ^ {\top} x > 0 \\ 0 & \text {w h e n} u ^ {\top} x <   0, (u ^ {\prime}) ^ {\top} x <   0 \end{array} \right.
$$
If we assume that each region has probability mass greater than 0 then our constraint  $\| \nabla_{x}h(x) - \nabla_{x}h^{\prime}(x)\|_{2}\leq \tau$  implies that  $|w| = |w||u|\leq \tau ,|w^{\prime}| = |w^{\prime}||u^{\prime}||\leq \tau ,\| wu - w^{\prime}u^{\prime}\| \leq \tau .$
Case 2:  $\theta (u,u^{\prime}) = 0$
This implies that the boundary of  $\nabla_{x}h(x)$  and  $\nabla_{x}h^{\prime}(x)$  are the same. Then,  $\nabla_{x}h(x) - \nabla_{x}h^{\prime}(x)$  is a piecewise constant over two regions
$$
\nabla_ {x} h (x) - \nabla_ {x} h ^ {\prime} (x) = \left\{ \begin{array}{l l} w u - w ^ {\prime} u ^ {\prime} & \text {w h e n} u ^ {\top} x > 0 \\ 0 & \text {w h e n} u ^ {\top} x <   0 \end{array} \right.
$$
This gives us that  $|w - w'| = \|wu - w'u'\| \leq \tau$ .
Case 3:  $\theta (u,u^{\prime}) = \pi$
Here, we have that the decision boundaries of  $\nabla_{x}h(x)$  and  $\nabla_{x}h^{\prime}(x)$  are the same but the gradients are non-zero on different sides. Then,  $\nabla_{x}h(x) - \nabla_{x}h^{\prime}(x)$  is a piecewise constant on two regions
$$
\nabla_ {x} h (x) - \nabla_ {x} h ^ {\prime} (x) = \left\{ \begin{array}{l l} w u & \text {w h e n} u ^ {\top} x > 0 \\ - w ^ {\prime} u ^ {\prime} & \text {w h e n} u ^ {\top} x <   0 \end{array} \right.
$$
This gives us that  $|w|\leq \tau$  and  $|w^{\prime}|\leq \tau$
These different cases tell us that the constraint  $\| \nabla_{x}h(x) - \nabla_{x}h^{\prime}(x)\| \leq \tau$  reduces  $\mathcal{H}$  into a class of models follows either
1.  $u = u'$  and  $|w - w'| < \tau$ .
2.  $u \neq u'$  and  $|w| < \tau$ . However, this case only possible when  $|w'| < \tau$ .
If  $|w'| > \tau$ , we know that we must only have the first case. Now, we can calculate the Rademacher complexity of our restricted class  $\mathcal{H}_{\phi, \tau}$ . We will again do this in separate cases.
Case 1:  $|w'| > \tau$
For any  $\dot{h} \in \mathcal{H}_{\phi, \tau}$ , we have that  $u = u'$  and  $|w - w'| < \tau$ . For a sample  $S = \{x_1, \dots, x_n\}$ ,
$$
\begin{array}{l} R _ {s} \left(\mathcal {H} _ {\phi , \tau}\right) = \frac {1}{n} \mathbb {E} _ {\sigma} \left[ \sup  _ {h \in \mathcal {H} _ {\phi , \tau}} \sum_ {i = 1} ^ {n} h \left(x _ {i}\right) \sigma_ {i} \right] \\ = \frac {1}{n} \mathbb {E} _ {\sigma} \left[ \sup  _ {w} \sum_ {i = 1} ^ {n} w \sigma \left(\left(u ^ {\prime}\right) ^ {\top} x _ {i}\right) \sigma_ {i} \right] \quad (\text {a s} u = u ^ {\prime}) \\ = \frac {1}{n} \mathbb {E} _ {\sigma} \left[ \sup  _ {w} w \left(\sum_ {i = 1} ^ {n} \sigma \left(\left(u ^ {\prime}\right) ^ {\top} x _ {i}\right) \sigma_ {i}\right) \right]. \\ \end{array}
$$
Since,  $|w - w'| < \tau$ ,
$$
w ^ {\prime} - \tau <   w <   w ^ {\prime} + \tau
$$
Then, we can compute the supremum over  $w$  as
$$
w = \left\{ \begin{array}{l l} w ^ {\prime} - \tau & \text {i f} \left(\sum_ {i = 1} ^ {n} \sigma ((u ^ {\prime}) ^ {\top} x _ {i}) \sigma_ {i}\right) <   0 \\ w ^ {\prime} + \tau & \text {i f} \left(\sum_ {i = 1} ^ {n} \sigma ((u ^ {\prime}) ^ {\top} x _ {i}) \sigma_ {i}\right) \geq 0 \end{array} \right.
$$
$$
\sup  _ {w} w \left(\sum_ {i = 1} ^ {n} \sigma ((u ^ {\prime}) ^ {\top} x _ {i}) \sigma_ {i}\right) = \left(w ^ {\prime} + \tau \operatorname {s i g n} \left(\sum_ {i = 1} ^ {n} \sigma ((u ^ {\prime}) ^ {\top} x _ {i}) \sigma_ {i}\right)\right) \cdot \left(\sum_ {i = 1} ^ {n} \sigma ((u ^ {\prime}) ^ {\top} x _ {i}) \sigma_ {i}\right).
$$
Now, we can calculate the Rademacher complexity as
$$
\begin{array}{l} R _ {S} \left(\mathcal {H} _ {\phi , \tau}\right) = \frac {1}{n} \mathbb {E} _ {\sigma} \left[ \right. w ^ {\prime} \left(\sum_ {i = 1} ^ {n} \sigma \left(\left(u ^ {\prime}\right) ^ {\top} x _ {i}\right) \sigma_ {i}\right) + \tau \mid \sum_ {i = 1} ^ {n} \sigma \left(\left(u ^ {\prime}\right) ^ {\top} x _ {i}\right) \sigma_ {i} \left. \right\rvert\left. \right] \\ = \frac {\tau}{n} \mathbb {E} _ {\sigma} \left[ \left| \sum_ {i = 1} ^ {n} \sigma \left(\left(u ^ {\prime}\right) ^ {\top} x _ {i}\right) \sigma_ {i} \right| \right] \\ \leq \frac {\tau}{n} \sqrt {\mathbb {E} _ {\sigma} \left[ \| \sum_ {i = 1} ^ {n} \sigma ((u ^ {\prime}) ^ {\top} x _ {i}) \sigma_ {i} \| ^ {2} \right]} \quad (\text {J e s s e n ' s i n e q u a l i t y}) \\ = \frac {\tau}{n} \sqrt {\mathbb {E} _ {\sigma} \left[ \sum_ {i = 1} ^ {n} \sigma \left(\left(u ^ {\prime}\right) ^ {\top} x _ {i}\right) ^ {2} \sigma_ {i} ^ {2} \right]} \quad (\text {s i n c e} \sigma_ {i}, \sigma_ {j} \text {a r e i n d e p e n d e n t w i t h m e a n 0}) \\ \leq \frac {\tau}{n} \sqrt {\sum_ {i = 1} ^ {n} ((u ^ {\prime}) ^ {\top} x _ {i}) ^ {2}} \\ \leq \frac {\tau}{n} \sqrt {\sum_ {i = 1} ^ {n} \| x _ {i} \| ^ {2}}. \\ \end{array}
$$
Combining this with the fact that  $\mathbb{E}\left[\| x\| ^2\right]\leq C^2$  , we have
$$
\begin{array}{l} R _ {n} \left(\mathcal {H} _ {\phi , \tau}\right) = \mathbb {E} \left[ R _ {S} \left(\mathcal {H} _ {\phi , \tau}\right) \right] \\ \leq \frac {\tau}{n} \mathbb {E} [ \sqrt {\sum_ {i = 1} ^ {n} \| x _ {i} \| ^ {2}} ] \\ \leq \frac {\tau}{n} \sqrt {\mathbb {E} \left[ \sum_ {i = 1} ^ {n} \| x _ {i} \| ^ {2} \right]} \quad (\text {J e s s e n ' s i n e q u a l i t y}) \\ \leq \frac {\tau C}{\sqrt {n}}. \\ \end{array}
$$
Case 2:  $|w'| \| u' \| < \tau$ .
We know that  $\mathcal{H}_{\phi,\tau} = \mathcal{H}_{\phi,\tau}^{(1)} \cup \mathcal{H}_{\phi,\tau}^{(2)}$  when
$$
\begin{array}{l} \mathcal {H} _ {\phi , \tau} ^ {(1)} = \{h \in \mathcal {H} | h: x \rightarrow w \sigma \left(u ^ {\top} x\right), u = u ^ {\prime}, \left| w - w ^ {\prime} \right| <   \tau \} \\ \mathcal {H} _ {\phi , \tau} ^ {(2)} = \{h \in \mathcal {H} | h: x \rightarrow w \sigma (u ^ {\top} x), \| u \| = 1, u \neq u ^ {\prime}, | w | <   \tau \} \\ \end{array}
$$
$$
\begin{array}{l} R _ {S} \left(\mathcal {H} _ {\phi , \tau}\right) = \frac {1}{n} \mathbb {E} _ {\sigma} \left[ \sup  _ {h \in \mathcal {H} _ {\phi , \tau}} \sum_ {i = 1} ^ {n} h \left(x _ {i}\right) \sigma_ {i} \right] \\ \leq \frac {1}{n} \mathbb {E} _ {\sigma} \left[ \sup  _ {h \in \mathcal {H} _ {\phi , \tau} ^ {(1)}} \sum_ {i = 1} ^ {n} h (x _ {i}) \sigma_ {i} + \sup  _ {h \in \mathcal {H} _ {\phi , \tau} ^ {(2)}} \sum_ {i = 1} ^ {n} h (x _ {i}) \sigma_ {i} \right] \\ = R _ {S} \left(\mathcal {H} _ {\phi , \tau} ^ {(1)}\right) + R _ {S} \left(\mathcal {H} _ {\phi , \tau} ^ {(2)}\right) \\ \end{array}
$$
The second line holds as  $\sup_{x\in A\cup B}f(x)\leq \sup_{x\in A}f(x) + \sup_{x\in B}f(x)$  when  $\sup_{x\in A}f(x)\geq 0$  and  $\sup_{x\in B}f(x)\geq 0$ . We know that both of these supremums be greater than zero, as we can recover the value of 0 with  $w = 0$ . From Case 1, we know that
$$
R _ {n} \left(\mathcal {H} _ {\phi , \tau} ^ {(1)}\right) \leq \frac {\tau C}{\sqrt {n}}.
$$
We also note that  $\mathcal{H}_{\phi, \tau}^{(2)}$  is a class of two layer neural networks with weights with norms bounded by  $\tau$ . From Theorem I.2, we have that
$$
R _ {n} (\mathcal {H} _ {\phi , \tau} ^ {(2)}) \leq \frac {2 \tau C}{\sqrt {n}}.
$$
Therefore, in Case 2,
$$
R _ {n} (\mathcal {H} _ {\phi , \tau}) \leq \frac {3 \tau C}{\sqrt {n}}.
$$
Now, we consider in the general setting (i.e., no restriction on  $m$ ).
Theorem I.4 (Rademacher complexity for two layer neural networks with gradient constraints). Let  $\mathcal{X}$  be an instance space and  $\mathcal{D}_{\mathcal{X}}$  be a distribution over  $\mathcal{X}$  with a large enough support. Let  $\mathcal{H} = \{h : x \mapsto \sum_{j=1}^{m} w_j \sigma(u_j^\top x) | w_j \in \mathbb{R}, u_j \in \mathbb{R}^d, \|u_j\|_2 = 1, \sum_{j=1}^{m} |w_j| \leq B\}$ . Assume that there exists some constant  $C > 0$  such that  $\mathbb{E}_{x \sim \mathcal{D}_{\mathcal{X}}}[\|x\|_2^2] \leq C^2$ . Our explanation constraint is given by a constraint on the gradient of our models, where we want the gradient of our learnt model to be close to a particular target function  $h' \in \mathcal{H}$ . Let this be represented by an explanation loss given by
for some  $\tau > 0$ . Then, we have that
$$
R _ {n} (\mathcal {H} _ {\phi , \tau}) \leq \frac {3 \tau m C}{\sqrt {n}}.
$$
To be precise,
$$
R _ {n} (\mathcal {H} _ {\phi , \tau}) \leq \frac {(2 m + q) \tau C}{\sqrt {n}}.
$$
when  $q$  is the number of node  $j$  of  $h^\prime$  such that  $\left|w_j^{\prime}\right| < \tau$ .
We note that this result indeed depends on the number of hidden dimensions  $m$ ; however, we note that in the general case (Theorem I.2), the value of  $B$  is  $O(m)$  as it is a sum over the values of each hidden node. We now present the proof for the more general version of our theorem.
Proof. For simplicity, we first assume that any  $h \in \mathcal{H}$  has that  $\| u_j\| = 1, \forall j$ . Consider  $h \in \mathcal{H}$ , we write  $h = \sum_{j=1}^{m} w_j' \sigma((u_j')^\top x)$  and let  $h'(x) = \sum_{j=1}^{m} w_j' \sigma((u_j')^\top x)$  be a function for our gradient constraint. The gradient of a hypothesis  $h$  is given by
$$
\nabla_ {x} h (x) = \sum_ {j = 1} ^ {m} w _ {j} u _ {j} \cdot 1 \{u _ {j} ^ {\top} x > 0 \},
$$
which is a piecewise constant function over at most  $2^{m}$  regions. Then, we consider that
$$
\nabla_ {x} h (x) - \nabla_ {x} h ^ {\prime} (x) = \sum_ {j = 1} ^ {m} w _ {j} u _ {j} \cdot 1 \{u _ {j} ^ {\top} x > 0 \} - \sum_ {j = 1} ^ {m} w _ {j} ^ {\prime} u _ {j} ^ {\prime} \cdot 1 \{(u _ {j} ^ {\prime}) ^ {\top} x > 0 \},
$$
which is a piecewise constant function over at most  $2^{2m}$  regions. We again make an assumption that each of these regions has a non-zero probability mass. Our choice of  $\phi(h, x)$  guarantees that the norm of the gradient in each region is less than  $\tau$ . Similar to the case with  $m = 1$ , we will show that the gradient constraint leads to a class of functions with the same decision boundary or neural networks that have weights with a small norm.
Assume that among  $u_{1}, \ldots, u_{m}$  there are  $k$  vectors that have the same direction as  $u_{1}', \ldots, u_{m}'$ . Without loss of generality, let  $u_{j} = u_{j}'$  for  $j = 1, \ldots, k$ . In this case, we have that  $\nabla_{x}h(x) - \nabla_{x}h'(x)$  is a piecewise function over  $2^{2m-k}$  regions. As each region has non-zero probability mass, for each  $j \in \{1, \ldots, k\}$ , we know that  $\exists x$  such that
$$
u _ {j} ^ {\top} x = \left(u _ {j} ^ {\prime}\right) ^ {\top} x > 0, \quad u _ {i} ^ {\top} x <   0 \text {f o r} i \neq j, \quad \left(u _ {i} ^ {\prime}\right) ^ {\top} x <   0 \text {f o r} i \neq j.
$$
In other words, we can observe a data point from each region that uniquely defines the value of a particular  $w_{j}, u_{j}$ . In this case, we have that
$$
\begin{array}{l} \nabla_ {x} h (x) - \nabla_ {x} h ^ {\prime} (x) = w _ {j} u _ {j} - w _ {j} ^ {\prime} u _ {j} ^ {\prime} \\ = \left(w _ {j} - w _ {j} ^ {\prime}\right) u _ {j} ^ {\prime}. \\ \end{array}
$$
From our gradient constraint, we know that  $||\nabla_{x}h(x) - \nabla_{x}h^{\prime}(x)||\leq \tau ,\forall x$  , which implies that  $|w_{j} - w_{j}^{\prime}|\leq \tau$  for  $j = 1,\ldots ,k$
On the other hand, for the remaining  $j = k + 1,\dots,m$ , we know that there exists  $x$  such that
$$
u _ {j} ^ {\top} x > 0, \qquad u _ {i} ^ {\top} x <   0 \text {f o r} i \neq j, \qquad \left(u _ {i} ^ {\prime}\right) ^ {\top} x <   0 \text {f o r} i = 1, \dots , m.
$$
Then, we have that  $\nabla_{x}h(x) = w_{j}u_{j}$ , and our constraint implies that  $|w_{j}||u_{j}|| = |w_{j}| \leq \tau$ . Similarly, we have that  $|w_{j}'||u_{j}'|| = |w_{j}'| < \tau$ , for  $j = k + 1, \ldots, m$ . We can conclude that  $\mathcal{H}_{\phi, \tau}$  is a class of two layer neural networks with  $m$  hidden nodes (assuming  $\| u_{i} \| = 1$ ) that for each node  $w_{j}\sigma(u_{j}^{\top}x)$  satisfies
1. There exists  $l \in [m]$  that  $u_{j} = u_{l}^{\prime}$  and  $|w_{j} - w_{l}^{\prime}| < \tau$ .
2.  $|w_j| < \tau$
We further note that for a node  $w_{l}^{\prime}\sigma((u_{l}^{\prime})^{\top}x)$  in  $h^{\prime}(x)$  that has a high weight  $|w_{l}^{\prime}| > \tau$ , there must be a node  $w_{j}\sigma(u_{j}^{\top}x)$  in  $h$  with the same boundary  $u_{j} = u_{l}$ . Otherwise, there is a contradiction with  $|w_{l}^{\prime}| < \tau$  for all nodes in  $h^{\prime}$  without a node in  $h$  with the same boundary. We can utilize this characterization of the restricted class  $\mathcal{H}_{\phi,\tau}$  to bound the Rademacher complexity of the class. Let
$$
\mathcal {H} ^ {\prime} = \left\{h: x \mapsto \sum_ {j = 1} ^ {m} w _ {j} ^ {\prime} \sigma \left(\left(u _ {j} ^ {\prime}\right) ^ {\top} x\right) a _ {j} \mid a _ {j} \in \{0, 1 \} \text {a n d f o r} j \text {t h a t} \left| w _ {j} ^ {\prime} \right| > \tau , a _ {j} = 1 \right\}.
$$
This is a class of two layer neural networks with at most  $m$  nodes such that each node is from  $h'$ . We also have a condition that if the weight of the  $j$ -th node in  $h'$  is greater than  $\tau$ , the  $j$ -th node must be present in any member of this class. Let
$$
\mathcal {H} ^ {(\tau)} = \left\{h: x \mapsto \sum_ {j = 1} ^ {m} w _ {j} \sigma ((u _ {j}) ^ {\top} x) a _ {j} \mid w _ {j} \in \mathbb {R}, u _ {j} \in \mathbb {R} ^ {d}, | w _ {j} | <   \tau , \| u _ {j} \| = 1 \right\}.
$$
be a class of two layer neural networks with  $m$  nodes such that the weight of each node is at most  $\tau$ . We claim that for any  $h \in \mathcal{H}_{\phi, \tau}$  there exists  $h_1 \in \mathcal{H}', h_2 \in \mathcal{H}^{(\tau)}$  that  $h = h_1 + h_2$ . For any  $h \in \mathcal{H}_{\phi, \tau}$ , let  $p_h : [m] \to [m] \cup \{0\}$  be a function that matches a node in  $h$  with the node with the same boundary in  $h'$ . Formally,
$$
p _ {h} (j) = \left\{ \begin{array}{l l} l & \text {w h e n} u _ {j} = u _ {l} ^ {\prime} \\ 0 & \text {o t h e r w i s e .} \end{array} \right.
$$
The function  $p_h$  maps  $j$  to 0 if there is no node in  $h'$  with the same boundary. Let  $w_0' = 0, u_0' = [0, \dots, 0]$ , we can write
$$
\begin{array}{l} h (x) = \sum_ {j = 1} ^ {m} w _ {j} \sigma \left(u _ {j} ^ {\top} x\right) \\ = \sum_ {j = 1} ^ {m} w _ {j} \sigma (u _ {j} ^ {\top} x) - w _ {p _ {h} (j)} ^ {\prime} \sigma ((u ^ {\prime}) _ {p _ {h} (j)} ^ {\top} x) + w _ {p _ {h} (j)} ^ {\prime} \sigma ((u ^ {\prime}) _ {p _ {h} (j)} ^ {\top} x) \\ = \underbrace{\sum_{p_{h}(j)\neq 0}(w_{j} - w^{\prime}_{p_{h}(j)})\sigma((u^{\prime})_{p_{h}(j)}^{\top}x) + \sum_{p_{h}(j) = 0}w_{j}\sigma(u_{j}^{\top}x)}_{\in \mathcal{H}^{(\tau)}} + \underbrace{\sum_{p(j)\neq 0}w^{\prime}_{p_{h}(j)}\sigma((u^{\prime})_{p_{h}(j)}^{\top}x)}_{\in \mathcal{H}^{\prime}}. \\ \end{array}
$$
The first term is a member of  $\mathcal{H}^{(\tau)}$  because we know that  $|w_{j} - w_{p(j)}^{\prime}| < \tau$  or  $|w_{j}| < \tau$ . The second term is also a member of  $\mathcal{H}'$  since for any  $l$  that  $|w_{l}^{\prime}| > \tau$ , there exists  $j$  that  $p_{h}(j) = l$ . Therefore, we can write  $h$  in terms of a sum between a member of  $\mathcal{H}'$  and  $\mathcal{H}^{(\tau)}$ . This implies that
$$
R _ {n} (\mathcal {H} _ {\phi , \tau}) \leq R _ {n} (\mathcal {H} ^ {\prime}) + R _ {n} (\mathcal {H} ^ {(\tau)}).
$$
From Theorem I.2, we have that
$$
R _ {n} \left(\mathcal {H} _ {\phi , \tau} ^ {(\tau)}\right) \leq \frac {2 \tau m C}{\sqrt {n}}.
$$
Now, we will calculate the Rademacher complexity of  $\mathcal{H}'$ . For  $S = \{x_1, \ldots, x_n\}$ ,
$$
\begin{array}{l} R _ {S} \left(\mathcal {H} ^ {\prime}\right) = \frac {1}{n} \mathbb {E} _ {\sigma} \left[ \sup  _ {h \in \mathcal {H} ^ {\prime}} \sum_ {i = 1} ^ {n} h \left(x _ {i}\right) \sigma_ {i} \right] \\ = \frac {1}{n} \mathbb {E} _ {\sigma} \left[ \sup  _ {h \in \mathcal {H} ^ {\prime}} \sum_ {i = 1} ^ {n} \left(\sum_ {j = 1} ^ {m} w _ {j} ^ {\prime} \sigma \left(\left(u _ {j} ^ {\prime}\right) ^ {\top} x _ {i}\right) a _ {j}\right) \sigma_ {i} \right] \\ = \frac {1}{n} \mathbb {E} _ {\sigma} \left[ \right. \sup  _ {h \in \mathcal {H} ^ {\prime}} \sum_ {i = 1} ^ {n} \left( \right.\sum_ {| w _ {j} ^ {\prime} | <   \tau} w _ {j} ^ {\prime} \sigma \left(\left(u _ {j} ^ {\prime}\right) ^ {\top} x _ {i}\right) a _ {j} + \sum_ {| w _ {j} ^ {\prime} | > \tau} w _ {j} ^ {\prime} \sigma \left(\left(u _ {j} ^ {\prime}\right) ^ {\top} x _ {i})\right) \sigma_ {i} \left. \right] \\ = \frac {1}{n} \mathbb {E} _ {\sigma} \left[ \sup  _ {a _ {j} \in \{0, 1 \}} \sum_ {i = 1} ^ {n} \sum_ {| w _ {j} ^ {\prime} | <   \tau} w _ {j} ^ {\prime} \sigma \left(\left(u _ {j} ^ {\prime}\right) ^ {\top} x _ {i}\right) a _ {j} \sigma_ {i} \right] \\ = \frac {1}{n} \mathbb {E} _ {\sigma} \left[ \sup  _ {a _ {j} \in \{0, 1 \}} \sum_ {| w _ {j} ^ {\prime} | <   \tau} a _ {j} \left(w _ {j} ^ {\prime} \sum_ {i = 1} ^ {n} \sigma \left(\left(u _ {j} ^ {\prime}\right) ^ {\top} x _ {i}\right) \sigma_ {i}\right) \right]. \\ \end{array}
$$
To achieve the supremum, if  $w_{j}^{\prime}\sum_{i = 1}^{n}\sigma ((u_{j}^{\prime})^{\top}x_{i})\sigma_{i} > 0$  we need to set  $a_{j} = 1$ , otherwise, we need to set  $a_{j} = 0$ . Therefore,
$$
\begin{array}{l} R _ {S} (\mathcal {H} ^ {\prime}) = \frac {1}{n} \mathbb {E} _ {\sigma} \left[ \sup _ {a _ {j} \in \{0, 1 \}} \sum_ {| w _ {j} ^ {\prime} | <   \tau} a _ {j} (w _ {j} ^ {\prime} \sum_ {i = 1} ^ {n} \sigma ((u _ {j} ^ {\prime}) ^ {\top} x _ {i}) \sigma_ {i}) \right] \\ = \frac {1}{n} \mathbb {E} _ {\sigma} \left[ \sum_ {| w _ {j} ^ {\prime} | <   \tau} \sigma \left(w _ {j} ^ {\prime} \sum_ {i = 1} ^ {n} \sigma \left(\left(u _ {j} ^ {\prime}\right) ^ {\top} x _ {i}\right) \sigma_ {i}\right) \right] \\ = \frac {1}{2 n} \mathbb {E} _ {\sigma} \left[ \sum_ {| w _ {j} ^ {\prime} | <   \tau} (w _ {j} ^ {\prime} \sum_ {i = 1} ^ {n} \sigma ((u _ {j} ^ {\prime}) ^ {\top} x _ {i}) \sigma_ {i}) + | w _ {j} ^ {\prime} \sum_ {i = 1} ^ {n} \sigma ((u _ {j} ^ {\prime}) ^ {\top} x _ {i}) \sigma_ {i} | \right] \quad (\sigma (x) = \frac {x + | x |}{2}) \\ = \frac {1}{2 n} \mathbb {E} _ {\sigma} \left[ \sum_ {| w _ {j} ^ {\prime} | <   \tau} | w _ {j} ^ {\prime} \sum_ {i = 1} ^ {n} \sigma \left(\left(u _ {j} ^ {\prime}\right) ^ {\top} x _ {i}\right) \sigma_ {i} | \right] \\ \leq \frac {1}{2 n} \left(\sum_ {| w _ {j} ^ {\prime} | <   \tau} | w _ {j} ^ {\prime} |\right) \mathbb {E} _ {\sigma} \left[ \sup  _ {\| u \| = 1} \left| \sum_ {i = 1} ^ {n} \sigma (u ^ {\top} x _ {i}) \sigma_ {i} \right| \right] \\ \leq \frac {1}{n} \left(\sum_ {| w _ {j} ^ {\prime} | <   \tau} | w _ {j} ^ {\prime} |\right) \mathbb {E} _ {\sigma} \left[ \sup  _ {\| u \| = 1} \sum_ {i = 1} ^ {n} \sigma \left(u ^ {\top} x _ {i}\right) \sigma_ {i} \right] \tag {Lemma I.1} \\ \leq \left(\sum_ {| w _ {j} ^ {\prime} | <   \tau} | w _ {j} ^ {\prime} |\right) \underbrace {\mathbb {E} _ {\sigma} \left[ \frac {1}{n} \sup  _ {\| u \| = 1} \sum_ {i = 1} ^ {n} u ^ {\top} x _ {i} \sigma_ {i} \right]} _ {\text {E m p i r i c a l R a d e m a c h e r c o m p l e x i t y o f a l i n e a r m o d e l}} \quad (\text {T a l a g r a n d ' s L e m m a}). \\ \end{array}
$$
From Theorem 4.1, we can conclude that
$$
R _ {n} (\mathcal {H} ^ {\prime}) \leq \sum_ {| w _ {j} ^ {\prime} | <   \tau} | w _ {j} ^ {\prime} | \frac {C}{\sqrt {n}} \leq \frac {q \tau C}{\sqrt {n}} \leq \frac {m \tau C}{\sqrt {n}}
$$
when  $q$  is the number of nodes  $j$  of  $h^\prime$  such that  $|w_j^{\prime}| < \tau$ . Therefore,
$$
R _ {n} (\mathcal {H} ^ {\prime}) \leq \frac {(2 m + q) \tau C}{\sqrt {n}} \leq \frac {3 m \tau C}{\sqrt {n}}.
$$
A tighter bound is given by  $\frac{(2m + q)\tau C}{\sqrt{n}}$  when  $q$  is the number of  $w_{j}^{\prime}$  that  $|w_{j}^{\prime}| < \tau$ . As  $\tau \to 0$ , we also have  $q \to 0$ . This implies that we have an upper bound of  $\frac{2m\tau C}{\sqrt{n}}$  if  $\tau$  is small enough. When comparing this to the original bound  $\frac{2BC}{\sqrt{n}}$ , we can do much better if  $\tau \ll \frac{B}{m}$ . We would like to point out that our bound does not depend on the distribution  $\mathcal{D}$  because we choose a strong explanation loss
which guarantees that  $\| \nabla_{x}h(x) - \nabla_{x}h^{\prime}(x)\|_{2}\leq \tau$  almost everywhere. We also assume that we are in a high-dimensional setting  $d\gg m$  and there exists  $x$  with a positive probability density at any partition created by  $\nabla_x h(x)$ .
J Algorithmic Results for Two Layer Neural Networks with a Gradient Constraint
Now that we have provided generalization bounds for the restricted class of two layer neural networks, we also present an algorithm that can identify the parameters of a two layer neural network (up to a permutation of the weights). In practice, we might solve this via our variational objective or other simpler regularized techniques. However, we also provide a theoretical result for the required amount of data (given some assumptions about the data distribution) and runtime for an algorithm to exactly recover the parameters of these networks under gradient constraints.
We again know that the gradient of two layer neural networks with ReLU activations can be written as
$$
\nabla_ {x} f _ {w, U} (x) = \sum_ {i = 1} ^ {m} w _ {i} u _ {i} \cdot 1 \{u _ {i} ^ {T} x > 0 \},
$$
where we consider  $||u_i|| = 1$ . Therefore, an exact gradient constraint given of the form of pairs  $(x,\nabla_{x}f(x))$  produces a system of equations.
Proposition J.1. If the values of  $u_{i}$ 's are known, we can identify the parameters  $w_{i}$  with exactly  $m$  fixed samples.
Proof. We can select  $m$  datapoints, which each achieve value 1 for the indicator value in the gradient of the two layer neural network. This would give us  $m$  equations, which each are of the form
$$
\nabla_ {x} f _ {w, U} (x _ {i}) = w _ {i} u _ {i}.
$$
Therefore, we can easily solve for the values of  $w_{i}$ , given that  $u_{i}$  is known.
To make this more general, we now consider the case where  $u_{i}$ 's are not known but are at least linearly independent.
Proposition J.2. Let the  $u_{i}$ 's be linearly independent. Assume that each region of the data (when partitioned by the values of  $u_{i}$ ) has non-trivial support  $> p$ . Then, with probability  $1 - \delta$ , we can identify the parameters  $w_{i}, u_{i}$  with  $O\left(2^{m} + \frac{m + \log\left(\frac{1}{\delta}\right)}{\log\left(\frac{1}{1 - p}\right)}\right)$  data points and in  $O(2^{2m})$  time.
Proof. Let us partition  $\mathcal{X}$  into regions satisfying unique values of the binary vector  $(1\{u_1^T x > 0\}, \dots, 1\{u_m^T x > 0\})$ , which by our assumption each have at least some probability mass  $p$ . First, we calculate the probability that we observe one data point with an explanation from each region in this partition. This is equivalent to sampling from a multinomial distribution with probabilities  $(p_1, \dots, p_{2^m})$ , where  $p_i \geq p, \forall i$ . Then,
$\operatorname{Pr}(\text{observe all regions in } n \text{ draws}) = 1 - \operatorname{Pr}(\exists i \text{ s.t. we do not observe region } i)$
$$
= 1 - 2 ^ {m} (1 - p) ^ {n}.
$$
Algorithm 1 Algorithm for identifying parameters of a two layer neural network, given exact gradient constraints
1: Input: We are given  $M = \{\sum_{x\in C}x|C\in \mathcal{P}(\{x_1,\dots,x_m\})\}$ , with  $\{x_{1},\ldots ,x_{m}\}$  linearly independent  
2: Output: The set of basis elements  $\{x_{1},\ldots ,x_{m}\}$   
3: function  
4:  $B = \{\} ,S = \{\}$  {Set for basis vectors and set for a current sum of at least 2 elements}  
5: for  $x\in M$  do  
6: if  $x\in S$  then  
7: pass  
8: else  
9:  $B = B\cup \{x\}$   
10: if  $|B| = 2$  then  
11:  $S = \{y_1 + y_2\}$ , where  $B = \{y_{1},y_{2}\}$   
12: else  
13:  $S = S\cup \{y + x|y\in S\}$  {Updating sums from adding  $x\}$   
14: end if  
15:  $O = B\cap S$  {Computing overlap between current basis and sums}  
16:  $B = B\setminus O$  {Removing elements contained in pairwise span}  
17:  $S = \{y - y_o|y\in S,y_o\in O\}$  {Updating sums  $S$  from removing set  $O\}$   
18: end if  
19: end for  
20: return  $B$   
21: end function