SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathbf {W} _ {l k} ^ {\prime} = \mathbf {W} _ {l k}, \forall l \in [ L - 2 ], \mathbf {w} _ {(L - 1) k j _ {L - 1}} ^ {\prime} = \frac {\mathbf {w} _ {(L - 1) k j _ {L - 1}}}{r _ {k j _ {L - 1}}}, \forall l \in [ L - 1 ], w _ {L k j _ {L - 1}} ^ {\prime} = r _ {k j _ {L - 1}} w _ {L k j _ {L - 1}}.
$$
We now note that
$$
\sum_ {j _ {1}, \dots , j _ {L - 2}} \left(\| \mathbf {w} _ {1 k j _ {1}} ^ {\prime} \| _ {2} ^ {2} \prod_ {l = 2} ^ {L - 1} w _ {l k j _ {l - 1} j _ {l}} ^ {j _ {l} ^ {2}}\right) = \frac {1}{r _ {k j _ {L - 1}} ^ {2}} \sum_ {j _ {1}, \dots , j _ {L - 2}} \left(\| \mathbf {w} _ {1 k j _ {1}} \| _ {2} ^ {2} \prod_ {l = 2} ^ {L - 1} w _ {l k j _ {l - 1} j _ {l}} ^ {2}\right) = 1.
$$
Then, (2) can be restated as follows
$$
\begin{array}{l} p _ {L} ^ {*} = \min  _ {\{\{\mathbf {W} _ {l k} \} _ {l = 1} ^ {L} \} _ {k = 1} ^ {K}} \mathcal {L} \left(\sum_ {k = 1} ^ {K} \left((\mathbf {X} \mathbf {W} _ {1 k}) _ {+} \dots \mathbf {W} _ {(L - 1) k}\right) _ {+} \mathbf {w} _ {L k}, \mathbf {y}\right) + \beta \sum_ {k = 1} ^ {K} \sqrt {\sum_ {j _ {1} , \ldots , j _ {L}} \left(\| \mathbf {w} _ {1 k j _ {1}} \| _ {2} ^ {2} \prod_ {l = 2} ^ {L} w _ {l k j _ {l - 1} j _ {l}} ^ {2}\right)} \\ = \min_{\substack{\{\{\mathbf{W}_{lk}\}_{l = 1}^{L - 1}\}_{k = 1}^{K}\\ \{\mathbf{w}_{Lk}^{\prime}\}_{k = 1}^{K}}}\mathcal{L}\left(\sum_{k = 1}^{K}\sum_{j_{L - 1}}\frac{\big((\mathbf{X}\mathbf{W}_{1k})_{+}\cdots\mathbf{w}_{(L - 1)kj_{L - 1}}\big)_{+}}{r_{kj_{L - 1}}}w_{Lkj_{L - 1}}^{\prime},\mathbf{y}\right) \\ + \beta \sum_ {k = 1} ^ {K} \sqrt {\sum_ {j _ {L - 1}} w _ {L k j _ {L - 1}} ^ {\prime 2} \frac {\sum_ {j _ {1} , \dots , j _ {L - 2}} \left(\| \mathbf {w} _ {1 k j _ {1}} \| _ {2} ^ {2} \prod_ {l = 2} ^ {L - 1} w _ {l k j _ {l - 1} j _ {l}} ^ {2}\right)}{r _ {k j _ {L - 1}} ^ {2}}} \\ = \min_{\substack{\{\{\mathbf{W}_{lk}\}_{l = 1}^{L - 1}\}_{k = 1}^{K}\\ \{\mathbf{w}_{Lk}^{\prime}\}_{k = 1}^{K}}}\mathcal{L}\left(\sum_{k = 1}^{K}\sum_{j_{L - 1}}\left((\mathbf{X}\mathbf{W}_{1k})_{+}\ldots r_{kj_{L - 1}}^{-1}\mathbf{w}_{(L - 1)kj_{L - 1}}\right)_{+}w_{Lkj_{L - 1}}^{\prime},\mathbf{y}\right) + \beta \sum_{k = 1}^{K}\| \mathbf{w}_{Lk}^{\prime}\|_{2} \\ = \min_{\substack{\{\{\mathbf{W}_{lk}^{\prime}\}_{l = 1}^{L}\}_{k = 1}^{K}\\ (\mathbf{W}_{1k}^{\prime},\ldots ,\mathbf{W}_{(L - 1)k}^{\prime})\in \Theta_{s},\forall k}}\mathcal{L}\left(\sum_{k = 1}^{K}\left(\left(\mathbf{X}\mathbf{W}_{1k}^{\prime}\right)_{+}\dots \mathbf{W}_{(L - 1)k}^{\prime}\right)_{+}\mathbf{w}_{Lk}^{\prime},\mathbf{y}\right) + \beta \sum_{k = 1}^{K}\| \mathbf{w}_{Lk}^{\prime}\|_{2}, \\ \end{array}
$$
where  $\Theta_s := \left\{(\mathbf{W}_1, \ldots, \mathbf{W}_{L-1}) : \sum_{j_1, \ldots, j_{L-2}} \left( \| \mathbf{w}_{1j_1} \|_2^2 \prod_{l=2}^{L-1} w_{lj_{l-1}j_l}^2 \right) = 1, \forall j_{L-1} \in [m_{L-1}] \right\}$ .
We also note that one can relax the equality constraint as an inequality constraint without loss of generality. This is basically due to the fact that if a constraint is not tight, i.e., strictly less than one, at the optimum then we can remove that constraint and make the corresponding output layer weight arbitrarily small via a simple scaling to make the objective value smaller. However, this would lead to a contradiction since this scaling further reduces the objective, which means that the initial set of layer weights (that yields a strict inequality in the constraints) are not optimal.
To obtain the bidual problem of (4), we first utilize semi-infinite duality theory as follows. We first compute the dual of (6) with respect to the dual parameter  $\mathbf{v}$  to get
$$
p _ {\infty} ^ {*} := \min  _ {\| \mathbf {w} _ {3} \| _ {2} \leq 1} \min  _ {\boldsymbol {\mu}} \mathcal {L} \left(\int_ {\theta \in \Theta_ {s}} \left((\mathbf {X W} _ {1}) _ {+} \mathbf {W} _ {2}\right) _ {+} \mathbf {w} _ {3} d \mu (\theta), \mathbf {y}\right) + \beta \| \boldsymbol {\mu} \| _ {T V}, \tag {12}
$$
where  $\| \pmb{\mu}\|_{TV}$  denotes the total variation norm of the signed measure  $\pmb{\mu}$ . Notice that (12) is an infinite-dimensional neural network training problem similar to the one studied in [38]. More importantly, this problem is convex since the model is linear with respect to the measure  $\pmb{\mu}$  and the loss and regularization functions are convex [38]. Thus, we have no duality gap, i.e.,  $d_3^* = p_\infty^*$ . In addition to this, even though (12) is an infinite-dimensional convex optimization problem, it reduces to a problem
with at most  $n + 1$  neurons at the optimum due to Caratheodory's theorem [35]. Therefore, (12) can be equivalently stated as the following finite-size convex optimization problem
$$
\begin{array}{l} p _ {\infty} ^ {*} = \min  _ {\theta \in \Theta_ {s}} \mathcal {L} \left(\sum_ {k = 1} ^ {K ^ {*}} \left((\mathbf {X W} _ {1 k}) _ {+} \mathbf {W} _ {2 k}\right) _ {+} \mathbf {w} _ {3 k}, \mathbf {y}\right) + \beta \sum_ {k = 1} ^ {K ^ {*}} \| \mathbf {w} _ {3 k} \| _ {2} \\ = \min  _ {\theta \in \Theta_ {s}} \mathcal {L} \left(f _ {\theta} (\mathbf {X}), \mathbf {y}\right) + \beta \sum_ {k = 1} ^ {K ^ {*}} \| \mathbf {w} _ {3 k} \| _ {2}, \tag {13} \\ \end{array}
$$
where  $K^{*} \leq n + 1$ . We further remark that given  $K \geq K^{*}$ , (13) and (5) are the same problems, which also proves strong duality as  $p_3^* = p_\infty^* = d_3^*$ . In the remainder of the proof, we show that using an alternative representation for the ReLU activation, we can achieve a finite-dimensional convex bidual formulation.
Now we restate the dual problem (6) as
$$
d _ {3} ^ {*} = \max  _ {\mathbf {v}} - \mathcal {L} ^ {*} (\mathbf {v}) \text {s . t .} \max  _ {\theta \in \Theta_ {s}} \left\| \mathbf {v} ^ {T} \left(\left(\mathbf {X} \mathbf {W} _ {1}\right) _ {+} \mathbf {W} _ {2}\right) _ {+} \right\| _ {2} \leq \beta . \tag {14}
$$
We first note that using the representation in (7), the dual constraint in (14) can be written as
$$
\begin{array}{l} \max  _ {\theta \in \Theta_ {s}} \left\| \mathbf {v} ^ {T} \left((\mathbf {X W} _ {1}) _ {+} \mathbf {W} _ {2}\right) _ {+} \right\| _ {2} \leq \beta \iff \max  _ {\theta \in \Theta_ {s}} \sqrt {\sum_ {j _ {2} = 1} ^ {m _ {2}} \left(\mathbf {v} ^ {T} \left((\mathbf {X W} _ {1}) _ {+} \mathbf {w} _ {2 j _ {2}}\right) _ {+}\right) ^ {2}} \leq \beta \\ \Longleftrightarrow \max  _ {\theta \in \Theta_ {s}} \sqrt {m _ {2} \left(\mathbf {v} ^ {T} \left((\mathbf {X} \mathbf {W} _ {1}) _ {+} \mathbf {w} _ {2}\right) _ {+}\right) ^ {2}} \leq \beta \\ \Longleftrightarrow \max  _ {\mathcal {I} _ {j _ {1}} \in \{\pm 1 \}} \max  _ {\theta \in \Theta_ {s}} \sqrt {m _ {2} \left(\mathbf {v} ^ {T} \left(\sum_ {j _ {1} = 1} ^ {m _ {1}} \mathcal {I} _ {j _ {1}} \left(\mathbf {X} \mathbf {w} _ {1 j _ {1}} | w _ {2 j _ {1}} |\right) _ {+}\right) _ {+}\right) ^ {2}} \leq \beta \\ \Longleftrightarrow \max  _ {\substack {i \in [ P _ {1} ] \\ l \in [ P _ {2} ]}} \max  _ {\mathcal {I} _ {j _ {1}} \in \{\pm 1 \}} \max  _ {\left\{\mathbf {w} _ {j _ {1}} \right\} _ {j _ {1}} \in \mathcal {C} _ {i l}} \sqrt {m _ {2}} \left| \mathbf {v} ^ {T} \mathbf {D} _ {2 l} \sum_ {j _ {1} = 1} ^ {m _ {1}} \mathcal {I} _ {j _ {1}} \mathbf {D} _ {1 i j _ {1}} \mathbf {X} \mathbf {w} _ {j _ {1}} \right| \leq \beta \tag{15} \\ \end{array}
$$
where  $\Theta_s = \{(\mathbf{W}_1, \mathbf{W}_2) : \sum_{j_1=1}^{m_1} \| \mathbf{w}_{1j_1} \|^2 2 w_{2j_1j_2}^2 \leq 1, \forall j_2 \in [m_2]\}$ , we apply a variable change as  $\mathbf{w}_{j_1} = |w_{2j_1}| \mathbf{w}_{1j_1}$  and define the set  $\mathcal{C}_{il}$  as
$$
\mathcal {C} _ {i l} := \{\left\{\mathbf {w} _ {j _ {1}} \right\} _ {j _ {1}}: (2 \mathbf {D} _ {2 l} - \mathbf {I} _ {n}) \sum_ {j _ {1} = 1} ^ {m _ {1}} \mathcal {I} _ {j _ {1}} \mathbf {D} _ {1 i j _ {1}} \mathbf {X} \mathbf {w} _ {j _ {1}} \geq 0, (2 \mathbf {D} _ {1 i j _ {1}} - \mathbf {I} _ {n}) \mathbf {X} \mathbf {w} _ {j _ {1}} \geq 0, \forall j _ {1} \in [ m _ {1} ], \sum_ {j _ {1} = 1} ^ {m _ {1}} \| \mathbf {w} _ {j _ {1}} \| _ {2} ^ {2} \leq 1 \}.
$$
We also note that  $P_{1}$  and  $P_{2}$  denote the number of possible hyperplane arrangements for the first and second ReLU layer (see Appendix A.10 for details).
$$
\begin{array}{l} \max_{\theta \in \Theta_{s}}\left\| \mathbf{v}^{T}\left((\mathbf{X}\mathbf{W}_{1})_{+}\mathbf{W}_{2}\right)_{+}\right\|_{2}\leq \beta \iff \max_{\substack{i\in [P_{1}]\\ l\in [P_{2}]\\ \mathcal{T}_{j_{1}}\in \{\pm 1\}}}\max_{\{\mathbf{w}_{j_{1}}\}_{j_{1}}\in \mathcal{C}_{il}}\sqrt{m_{2}}\left|\mathbf{v}^{T}\mathbf{D}_{2l}\sum_{j_{1} = 1}^{m_{1}}\mathcal{I}_{j_{1}}\mathbf{D}_{1ij_{1}}\mathbf{X}\mathbf{w}_{j_{1}}\right|\leq \beta , \\ \max  _ {\left\{\mathbf {w} _ {i j _ {1} l} ^ {s} \right\} _ {j _ {1}} \in \mathcal {C} _ {i l} ^ {s}} \sqrt {m _ {2}} \mathbf {v} ^ {T} \mathbf {D} _ {2 l} \sum_ {j _ {1} = 1} ^ {m _ {1}} \mathcal {I} _ {i j _ {1} l} ^ {s} \mathbf {D} _ {1 i j _ {1}} \mathbf {X} \mathbf {w} _ {i j _ {1} l} ^ {s} \leq \beta , \\ \iff \max  _ {\{\mathbf {w} _ {i j _ {1} l} ^ {s ^ {\prime}} \} _ {j _ {1} \in C _ {i l} ^ {s}}} - \sqrt {m _ {2}} \mathbf {v} ^ {T} \mathbf {D} _ {2 l} \sum_ {j _ {1} = 1} ^ {m _ {1}} \mathcal {I} _ {i j _ {1} l} ^ {s ^ {\prime}} \mathbf {D} _ {1 i j _ {1}} \mathbf {X} \mathbf {w} _ {i j _ {1} l} ^ {s ^ {\prime}} \leq \beta , \\ \end{array}
$$
where we use  $M := |\{\pm 1\}^{m_1}| = 2^{m_1}$  to enumerate all possible sign patterns  $\{\mathcal{I}_{j_1}\}_{j_1=1}^{m_1}$  of size  $m_1$ . Using the equivalent representation above, we rewrite the dual (14) as
$$
\begin{array}{l} \max  _ {\mathbf {v}} - \mathcal {L} ^ {*} (\mathbf {v}) \text {s . t .} \max  _ {\left. \mathbf {w} _ {i j _ {1} l} ^ {s} \right\rbrace_ {j _ {1}} \in \mathcal {C} _ {i l} ^ {s}} \sqrt {m _ {2}} \mathbf {v} ^ {T} \mathbf {D} _ {2 l} \sum_ {j _ {1} = 1} ^ {m _ {1}} \mathcal {I} _ {i j _ {1} l} ^ {s} \mathbf {D} _ {1 i j _ {1}} \mathbf {X} \mathbf {w} _ {i j _ {1} l} ^ {s} \leq \beta , \forall i, l, s (16) \\ \max  _ {\left\{\mathbf {w} _ {i j _ {1} l} ^ {s ^ {\prime}} \right\} _ {j _ {1}} \in \mathcal {C} _ {i l} ^ {s}} - \sqrt {m _ {2}} \mathbf {v} ^ {T} \mathbf {D} _ {2 l} \sum_ {j _ {1} = 1} ^ {m _ {1}} \mathcal {I} _ {i j _ {1} l} ^ {s ^ {\prime}} \mathbf {D} _ {1 i j _ {1}} \mathbf {X} \mathbf {w} _ {i j _ {1} l} ^ {s ^ {\prime}} \leq \beta , \forall i, l, s. (17) \\ \end{array}
$$
Since the problem above is convex and satisfies the Slater's condition when all the parameters are set to zero, we have strong duality [36], and thus we can state (16) as
$$
\begin{array}{l} \min  _ {\gamma_ {i l} ^ {s} \geq 0} \max  _ {\mathbf {v}} \min  _ {\substack {\left. \right.\left. \right.\left. \right.\left. \right.\left. \right.\left. \right.\left. \right.\left. \right.\left. \right.\left. \right.\left. \right.\left.\left.\left.\left.\left.\left.\left.\left.\left.\left.\left.\left.\left.\left.\left.\left.\left.\left.\left.\left.\left.\left.\left. \text {v} ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s} \right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}\right) ^ {s}}} - \mathcal {L} (\mathbf {v}) ^ {*} + \sum_ {s = 1} ^ {M} \sum_ {i = 1} ^ {P _ {1}} \sum_ {l = 1} ^ {P _ {2}} \gamma_ {i l} ^ {s} \left(\beta - \sqrt {m _ {2}} \mathbf {v} ^ {T} \mathbf {D} _ {2 l} \sum_ {j _ {1} = 1} ^ {m _ {1}} \mathcal {I} _ {i j _ {1} l} ^ {s} \mathbf {D} _ {1 i j _ {1}} \mathbf {X} \mathbf {w} _ {i j _ {1} l} ^ {s}\right) (18) \\ + \sum_ {s = 1} ^ {M} \sum_ {i = 1} ^ {P _ {1}} \sum_ {l = 1} ^ {P _ {2}} \gamma_ {i l} ^ {s ^ {\prime}} \left(\beta + \sqrt {m _ {2}} \mathbf {v} ^ {T} \mathbf {D} _ {2 l} \sum_ {j _ {1} = 1} ^ {m _ {1}} \mathcal {I} _ {i j _ {1} l} ^ {s ^ {\prime}} \mathbf {D} _ {1 i j _ {1}} \mathbf {X} \mathbf {w} _ {i j _ {1} l} ^ {s ^ {\prime}}\right). (19) \\ \end{array}
$$
Due to Sion's minimax theorem [57], we can change the order the inner minimization and maximization to obtain closed-form solutions for the maximization over the variable  $\mathbf{v}$ . This yields the following problem
$$
\begin{array}{l} \min_{\substack{\gamma_{il}^{su}\geq 0\\ \{\mathbf{w}_{ij_{1}l}^{s}\}_{j_{1}}\in \mathcal{C}_{il}^{s}\\ \{\mathbf{w}_{ij_{1}l}^{s^{\prime}}\}_{j_{1}}\in \mathcal{C}_{il}^{s}}}\min_{\substack{\mathbf{s}_{1}\in \mathcal{C}_{il}^{s}\\ \mathbf{s}_{1}\in \mathcal{C}_{il}^{s^{\prime}}}}\mathcal{L}\left(\sum_{s = 1}^{M}\sum_{l = 1}^{P_{2}}\sum_{i = 1}^{P_{1}}\sum_{j_{1} = 1}^{m_{1}}\sqrt{m_{2}}\mathbf{D}_{2l}\mathbf{D}_{1ij_{1}}\mathbf{X}(\mathcal{I}_{ij_{1}l}^{s}\gamma_{il}^{s}\mathbf{w}_{ij_{1}l}^{s} - \mathcal{I}_{ij_{1}l}^{s^{\prime}}\gamma_{il}^{s}\mathbf{w}_{ij_{1}l}^{s^{\prime}}),\mathbf{y}\right) \\ + \beta \sum_ {s = 1} ^ {M} \sum_ {i = 1} ^ {P _ {1}} \sum_ {l = 1} ^ {P _ {2}} \left(\gamma_ {i l} ^ {s} + \gamma_ {i l} ^ {s ^ {\prime}}\right). \tag {20} \\ \end{array}
$$
Finally, we introduce a set of variable changes changes as  $\mathbf{z}_{ij_1l}^s = \sqrt{m_2}\gamma_{il}^s\mathcal{I}_{ij_1l}^s\mathbf{w}_{ij_1l}^s$  and  $\mathbf{z}_{ij_1l}^{s'} = \sqrt{m_2}\gamma_{il}^{s'}\mathcal{I}_{ij_1l}^{s'}\mathbf{w}_{ij_1l}^{s'}$  such that (20) can be cast as the following convex problem
$$
\begin{array}{l} \min  _ {\left\{\mathbf {z} _ {i j _ {1} l} ^ {s} \right\} _ {j _ {1}} \in \mathcal {C} _ {i l} ^ {s ^ {\prime}}} \mathcal {L} \left(\sum_ {s = 1} ^ {M} \sum_ {l = 1} ^ {P _ {2}} \sum_ {i = 1} ^ {P _ {1}} \sum_ {j _ {1} = 1} ^ {m _ {1}} \mathbf {D} _ {2 l} \mathbf {D} _ {1 i j _ {1}} \mathbf {X} \left(\mathbf {z} _ {i j _ {1} l} ^ {s} - \mathbf {z} _ {i j _ {1} l} ^ {s ^ {\prime}}\right), \mathbf {y}\right) \\ \{\mathbf {z} _ {i j _ {1} l} ^ {s ^ {\prime}} \} _ {j _ {1}} \in \mathcal {C} _ {i l} ^ {s ^ {\prime}} \\ + \frac {\beta}{\sqrt {m _ {2}}} \sum_ {s = 1} ^ {M} \sum_ {i = 1} ^ {P _ {1}} \sum_ {l = 1} ^ {P _ {2}} \left(\sqrt {\sum_ {j _ {1} = 1} ^ {m _ {1}} \left\| \mathbf {z} _ {i j _ {1} l} ^ {s} \right\| _ {2} ^ {2}} + \sqrt {\sum_ {j _ {1} = 1} ^ {m _ {1}} \left\| \mathbf {z} _ {i j _ {1} l} ^ {s ^ {\prime}} \right\| _ {2} ^ {2}}\right), \tag {21} \\ \end{array}
$$
where the constraint set  $\mathcal{C}_{il}^{s}$  are defined as
$$
\mathcal {C} _ {i l} ^ {s ^ {\prime}} := \left\{\left\{\mathbf {z} _ {j _ {1}} \right\} _ {j _ {1}}: (2 \mathbf {D} _ {2 l} - \mathbf {I} _ {n}) \sum_ {j _ {1} = 1} ^ {m _ {1}} \mathbf {D} _ {1 i j _ {1}} \mathbf {X} \mathbf {z} _ {j _ {1}} \geq 0, (2 \mathbf {D} _ {1 i j _ {1}} - \mathbf {I} _ {n}) \mathbf {X} \mathcal {T} _ {i j _ {1} l} ^ {s} \mathbf {z} _ {j _ {1}} \geq 0, \forall j _ {1} \in [ m _ {1} ] \right\}.
$$
Notice that (21) is a constrained convex optimization problem with  $2dm_{1}MP_{1}P_{2}$  variables and  $2n(m_{1} + 1)MP_{1}P_{2}$  constraints in the set  $C_{il}^{s}$ .
In this section, we prove that once the convex program in (21) is globally optimized to obtain a set of optimal solutions  $\{\mathbf{z}_{ij_1l}^{s^*},\mathbf{z}_{ij_1l}^{s'^*}\}_{i,j_1,l,s}$ , one can recover an optimal solution to the non-convex training
problem (4) via a simple closed-form mapping as detailed below
$$
\mathbf {W} _ {1 k} ^ {*} = \left\{ \begin{array}{l l} \frac {1}{m _ {2}} \left[ \begin{array}{c c c c} \mathcal {I} _ {i 1 l} ^ {s} \mathbf {z} _ {i 1 l} ^ {s ^ {*}} & \mathcal {I} _ {i 2 l} ^ {s} \mathbf {z} _ {i 2 l} ^ {s ^ {*}} & \ldots & \mathcal {I} _ {i m _ {1} l} ^ {s} \mathbf {z} _ {i m _ {1} l} ^ {s ^ {*}} \end{array} \right] & \text {i f} 1 \leq k \leq M P _ {1} P _ {2} \\ \frac {1}{m _ {2}} \left[ \begin{array}{c c c c} \mathcal {I} _ {i 1 l} ^ {s ^ {\prime}} \mathbf {z} _ {i 1 l} ^ {s ^ {\prime \prime}} & \mathcal {I} _ {i 2 l} ^ {s ^ {\prime}} \mathbf {z} _ {i 2 l} ^ {s ^ {\prime \prime}} & \ldots & \mathcal {I} _ {i m _ {1} l} ^ {s ^ {\prime}} \mathbf {z} _ {i m _ {1} l} ^ {s ^ {\prime \prime}} \end{array} \right] & \text {i f} M P _ {1} P _ {2} + 1 \leq k \leq 2 M P _ {1} P _ {2} \end{array} \right.
$$
$$
\mathbf {W} _ {2 k} ^ {*} = \left\{ \begin{array}{l l} \left[ \begin{array}{c c c c} \mathcal {I} _ {i 1 l} ^ {s} & \mathcal {I} _ {i 1 l} ^ {s} & \ldots & \mathcal {I} _ {1 1 l} ^ {s} \\ \vdots & \vdots & \ldots & \vdots \\ \mathcal {I} _ {i m _ {1} l} ^ {s} & \mathcal {I} _ {i m _ {1} l} ^ {s} & \ldots & \mathcal {I} _ {1 m _ {1} l} ^ {s} \end{array} \right] & \text {i f 1 \leq k \leq M P _ {1} P _ {2}} \\ \left[ \begin{array}{c c c c} \mathcal {I} _ {i 1 l} ^ {s ^ {\prime}} & \mathcal {I} _ {i 1 l} ^ {s ^ {\prime}} & \ldots & \mathcal {I} _ {1 1 l} ^ {s ^ {\prime}} \\ \vdots & \vdots & \ldots & \vdots \\ \mathcal {I} _ {i m _ {1} l} ^ {s ^ {\prime}} & \mathcal {I} _ {i m _ {1} l} ^ {s ^ {\prime}} & \ldots & \mathcal {I} _ {1 m _ {1} l} ^ {s ^ {\prime}} \end{array} \right] & \text {i f M P _ {1} P _ {2} + 1 \leq k \leq 2 M P _ {1} P _ {2}} \end{array} \right.
$$
$$
\mathbf {w} _ {3 k} ^ {*} = \left\{ \begin{array}{l l} \left[ \begin{array}{c c c c} 1 & 1 & \ldots & 1 \end{array} \right] ^ {T} & \text {i f} 1 \leq k \leq M P _ {1} P _ {2} \\ \left[ \begin{array}{c c c c} - 1 & - 1 & \ldots & - 1 \end{array} \right] ^ {T} & \text {i f} M P _ {1} P _ {2} + 1 \leq k \leq 2 M P _ {1} P _ {2} \end{array} \right.,
$$
$$
(s, l, i) = \left\{ \begin{array}{l l} \left(\left\lfloor \frac {k - 1}{P _ {1} P _ {2}} \right\rfloor + 1, \left\lfloor \frac {k - 1 - (s - 1) P _ {1} P _ {2}}{P _ {1}} \right\rfloor + 1, k - (s - 1) P _ {1} P _ {2} - (l - 1) P _ {1}\right) & \text {i f} 1 \leq k \leq M P _ {1} P _ {2} \\ \left(\left\lfloor \frac {k ^ {\prime} - 1}{P _ {1} P _ {2}} \right\rfloor + 1, \left\lfloor \frac {k ^ {\prime} - 1 - (s - 1) P _ {1} P _ {2}}{P _ {1}} \right\rfloor + 1, k ^ {\prime} - (s - 1) P _ {1} P _ {2} - (l - 1) P _ {1}\right) & \text {e l s e} \end{array} \right.
$$
with  $k' = k - MP_1P_2$ . Hence, we achieve an optimal solution to the original non-convex training problem (4) as  $\{\mathbf{W}_{1k}^*,\mathbf{W}_{2k}^*,\mathbf{w}_{3k}^*\}_{k = 1}^{2MP_1P_2}$ , where  $\mathbf{W}_{1k}^* \in \mathbb{R}^{d\times m_1}$ ,  $\mathbf{W}_{2k}^* \in \mathbb{R}^{m_1\times m_2}$ , and  $\mathbf{w}_{3k}^* \in \mathbb{R}^{m_2}$  respectively. Next, we confirm that the proposed set of layer weights are indeed optimal by plugging them back to both the convex and non-convex objectives.
We first verify that both the optimal convex and non-convex layer weights give the same network output as follows (8), i.e.,
$$
\sum_ {k = 1} ^ {2 M P _ {1} P _ {2}} \left(\left(\mathbf {X} \mathbf {W} _ {1 k} ^ {*}\right) _ {+} \mathbf {W} _ {2 k} ^ {*}\right) _ {+} \mathbf {w} _ {3 k} ^ {*} = \sum_ {s = 1} ^ {M} \sum_ {l = 1} ^ {P _ {2}} \sum_ {i = 1} ^ {P _ {1}} \sum_ {j = 1} ^ {m _ {1}} \mathbf {D} _ {2 l} \mathbf {D} _ {1 i j _ {1}} \mathbf {X} \left(\mathbf {z} _ {i j _ {1} l} ^ {s ^ {*}} - \mathbf {z} _ {i j _ {1} l} ^ {s ^ {\prime *}}\right).
$$
Now, we show that the proposed set of weight matrices for the non-convex problem achieves the same regularization cost with (21), i.e.,
$$
\sum_ {k = 1} ^ {2 M P _ {1} P _ {2}} \sqrt {\sum_ {j _ {2} = 1} ^ {m _ {2}} \sum_ {j _ {1} = 1} ^ {m _ {1}} \| \mathbf {w} _ {1 k j _ {1}} ^ {*} \| _ {2} ^ {2} w _ {2 k j _ {1} j _ {2}} ^ {* 2} w _ {3 k j _ {2}} ^ {* 2}} = \frac {1}{\sqrt {m _ {2}}} \sum_ {s = 1} ^ {M} \sum_ {i = 1} ^ {P _ {1}} \sum_ {l = 1} ^ {P _ {2}} \left(\sqrt {\sum_ {j _ {1} = 1} ^ {m _ {1}} \| \mathbf {z} _ {i j _ {1} l} ^ {s ^ {*}} \| _ {2}} + \sqrt {\sum_ {j = 1} ^ {m _ {1}} \| \mathbf {z} _ {i j _ {1} l} ^ {s ^ {\prime *} ^ {*}} \| _ {2}}\right).
$$
Since  $\{\mathbf{W}_{1k}^{*},\mathbf{W}_{2k}^{*},\mathbf{w}_{3k}^{*}\}_{k = 1}^{2MP_{1}P_{2}}$  yields the same network output and regularization cost with the optimal parameters of the convex program in (21), we conclude that the proposed set parameters for the non-convex problem also achieves the optimal objective value  $p_3^*$  , i.e.,
$$
p _ {3} ^ {*} = \mathcal {L} \left(\sum_ {k = 1} ^ {2 M P _ {1} P _ {2}} \left(\left(\mathbf {X W} _ {1 k} ^ {*}\right) _ {+} \mathbf {W} _ {2 k} ^ {*}\right) _ {+} \mathbf {w} _ {3 k} ^ {*}, \mathbf {y}\right) + \beta \sum_ {k = 1} ^ {2 M P _ {1} P _ {2}} \sqrt {\sum_ {j _ {2} = 1} ^ {m _ {2}} \sum_ {j _ {1} = 1} ^ {m _ {1}} \left\| \mathbf {w} _ {k j _ {1}} ^ {*} \right\| _ {2} ^ {2} w _ {2 k j _ {1} j _ {2}} ^ {* 2} w _ {3 k j _ {2}} ^ {* 2}}.
$$
We start with defining the optimal parameters for the original and rank- $k$  approximation of the rescaled problem in (5) as
$$
\left\{\left(\mathbf {W} _ {1 k} ^ {*}, \mathbf {W} _ {2 k} ^ {*}, \mathbf {w} _ {3 k} ^ {*}\right) \right\} _ {k = 1} ^ {K} := \operatorname * {a r g m i n} _ {\theta \in \Theta_ {s}} \mathcal {L} \left(\sum_ {k = 1} ^ {K} \left(\left(\mathbf {X} \mathbf {W} _ {1 k}\right) _ {+} \mathbf {W} _ {2 k}\right) _ {+} \mathbf {w} _ {3 k}, \mathbf {y}\right) + \beta \sum_ {k = 1} ^ {K} \| \mathbf {w} _ {3 k} \| _ {2}
$$
$$
\left\{(\hat {\mathbf {W}} _ {1 k}, \hat {\mathbf {W}} _ {2 k}, \hat {\mathbf {w}} _ {3 k}) \right\} _ {k = 1} ^ {K} := \underset {\theta \in \Theta_ {s}} {\operatorname {a r g m i n}} \mathcal {L} \left(\sum_ {k = 1} ^ {K} \left(\left(\hat {\mathbf {X}} _ {r} \mathbf {W} _ {1 k}\right) _ {+} \mathbf {W} _ {2 k}\right) _ {+} \mathbf {w} _ {3 k}, \mathbf {y}\right) + \beta \sum_ {k = 1} ^ {K} \| \mathbf {w} _ {3 k} \| _ {2} \tag {22}
$$
and the objective value achieved by the parameters trained using  $\hat{\mathbf{X}}_r$  as
$$
p _ {r} := \mathcal {L} \left(\sum_ {k = 1} ^ {K} \left(\left(\mathbf {X} \hat {\mathbf {W}} _ {1 k}\right) _ {+} \hat {\mathbf {W}} _ {2 k}\right) _ {+} \hat {\mathbf {w}} _ {3 k}, \mathbf {y}\right) + \beta \sum_ {k = 1} ^ {K} \| \hat {\mathbf {w}} _ {3 k} \| _ {2}.
$$
$$
\begin{array}{l} p _ {3} ^ {*} = \mathcal {L} \left(\sum_ {k = 1} ^ {K} \left(\left(\mathbf {X} \mathbf {W} _ {1 k} ^ {*}\right) _ {+} \left. \mathbf {W} _ {2 k} ^ {*}\right) _ {+} \mathbf {w} _ {3 k} ^ {*}, \mathbf {y}\right) + \beta \sum_ {k = 1} ^ {K} \| \mathbf {w} _ {3 k} ^ {*} \| _ {2} \right. \\ \stackrel {(i)} {\leq} \mathcal {L} \left(\sum_ {k = 1} ^ {K} \left(\left(\mathbf {X} \hat {\mathbf {W}} _ {1 k}\right) _ {+} \hat {\mathbf {W}} _ {2 k}\right) _ {+} \hat {\mathbf {w}} _ {3 k}, \mathbf {y}\right) + \beta \sum_ {k = 1} ^ {K} \| \hat {\mathbf {w}} _ {3 k} \| _ {2} = p _ {r} \\ \stackrel {(i i)} {\leq} \mathcal {L} \left(\sum_ {k = 1} ^ {K} \left(\left(\hat {\mathbf {X}} _ {r} \hat {\mathbf {W}} _ {1 k}\right) _ {+} \hat {\mathbf {W}} _ {2 k}\right) _ {+} \hat {\mathbf {w}} _ {3 k}, \mathbf {y}\right) + (\beta + \sqrt {m _ {1} m _ {2}} R \sigma_ {r + 1}) \sum_ {k = 1} ^ {K} \| \hat {\mathbf {w}} _ {3 k} \| _ {2} \\ \leq \left(\mathcal {L} \left(\sum_ {k = 1} ^ {K} \left(\left(\hat {\mathbf {X}} _ {r} \hat {\mathbf {W}} _ {1 k}\right) _ {+} \hat {\mathbf {W}} _ {2 k}\right) _ {+} \hat {\mathbf {w}} _ {3 k}, \mathbf {y}\right) + \beta \sum_ {k = 1} ^ {K} \| \hat {\mathbf {w}} _ {3 k} \| _ {2}\right) \left(1 + \frac {\sqrt {m _ {1} m _ {2}} R \sigma_ {r + 1}}{\beta}\right) \\ \stackrel {(i i i)} {\leq} \left(\mathcal {L} \left(\sum_ {k = 1} ^ {K} \left(\left(\hat {\mathbf {X}} _ {r} \mathbf {W} _ {1 k} ^ {*}\right) _ {+} \mathbf {W} _ {2 k} ^ {*}\right) _ {+} \mathbf {w} _ {3 k} ^ {*}, \mathbf {y}\right) + \beta \sum_ {k = 1} ^ {K} \| \mathbf {w} _ {3 k} ^ {*} \| _ {2}\right)\left(1 + \frac {\sqrt {m _ {1} m _ {2}} R \sigma_ {r + 1}}{\beta}\right)\left. \right. \\ \stackrel {(i v)} {\leq} \left(\mathcal {L} \left(\sum_ {k = 1} ^ {K} \left((\mathbf {X W} _ {1 k} ^ {*}) _ {+} \mathbf {W} _ {2 k} ^ {*}\right) _ {+} \mathbf {w} _ {3 k} ^ {*}, \mathbf {y}\right) + \beta \sum_ {k = 1} ^ {K} \| \mathbf {w} _ {3 k} ^ {*} \| _ {2}\right) \left(1 + \frac {\sqrt {m _ {1} m _ {2}} R \sigma_ {r + 1}}{\beta}\right) ^ {2} \\ = p _ {3} ^ {*} \left(1 + \frac {\sqrt {m _ {1} m _ {2}} R \sigma_ {r + 1}}{\beta}\right) ^ {2}, \\ \end{array}
$$
where  $(i)$  and  $(iii)$  follow from the optimality definitions of the original and approximated problems in (22). In addition,  $(ii)$  and  $(iv)$  follow from the relations below