SortedNet: A Scalable and Generalized Framework for Training Modular Deep Neural Networks

Suppose $\hat{f}$ is a sub-network within a larger neural network architecture, and $f$ represents an identical network architecture trained independently. We aim to understand the relationship between the parameters of these two networks, $\theta$ for $\hat{f}$ and $\phi$ for $f$, as they are trained under identical conditions.

We assume that the gradients of the loss functions for $\hat{f}$ and $f$, denoted as $\mathcal{L}_{\hat{f}}$ and $\mathcal{L}_{f}$ respectively, are $L$-Lipschitz continuous. This implies that:

$$\|\nabla\mathcal{L}_{\hat{f}}(\theta)-\nabla\mathcal{L}_{\hat{f}}(\theta^{\prime})\|\leq L\|\theta-\theta^{\prime}\|$$

$$\|\nabla\mathcal{L}_{f}(\phi)-\nabla\mathcal{L}_{f}(\phi^{\prime})\|\leq L\|\phi-\phi^{\prime}\|$$

for all $\theta,\theta^{\prime}$ and $\phi,\phi^{\prime}$ in the parameter space.

The parameters of the networks are updated via gradient descent as follows:

• •

  For $\hat{f}$:

  $$\theta_{t+1}=\theta_{t}-\eta\nabla\mathcal{L}_{\hat{f}}(\theta_{t})$$

• •

  For $f$:

  $$\phi_{t+1}=\phi_{t}-\eta\nabla\mathcal{L}_{f}(\phi_{t})$$

We derive a bound on the difference in parameters between $\hat{f}$ and $f$ after each training iteration:

$$\|\theta_{t+1}-\phi_{t+1}\|=\|\theta_{t}-\phi_{t}-\eta(\nabla\mathcal{L}_{\hat{f}}(\theta_{t})-\nabla\mathcal{L}_{f}(\phi_{t}))\|$$

Applying the triangle inequality and the Lipschitz continuity of the gradients, we obtain:

$$\|\theta_{t+1}-\phi_{t+1}\|\leq(1+\eta L)\|\theta_{t}-\phi_{t}\|+\eta C$$

where $C$ is a constant that bounds the difference between the gradients of the loss functions of the networks.

This bound quantifies the evolution of the difference in parameters between $\hat{f}$ and $f$ across training iterations, incorporating the impact of the Lipschitz constant $L$, the learning rate $\eta$, and the constant $C$ that bounds the inherent difference in gradients.

Given that $\hat{f}$ and $f$ are trained under perfectly identical conditions (same data, initialization, and hyperparameters), the difference in their gradients can be considered negligible, leading us to conclude that $C$ is practically zero. Under this assumption, the bound simplifies significantly:

$$\|\theta_{t+1}-\phi_{t+1}\|\leq(1+\eta L)\|\theta_{t}-\phi_{t}\|$$

This indicates that the difference in the parameters of $\hat{f}$ and $f$ is governed by the Lipschitz constant $L$ and the learning rate $\eta$, suggesting that the parameters should remain close throughout the training process, especially when $C$ is negligible.

We would like to find a performance bound between a trained sub-model (with optimized parameters $\theta^{*}$) and its corresponding individual model (with optimized parameters $\phi^{*}$). Let’s assume that ${\phi}^{*}={\theta^{*}}+\Delta\theta$. Then, the performance bound can be calculated as $\Delta{f}={f}(x;{\phi^{*}})-\hat{f}(x;{\theta^{*}})$ in the function value from its optimal value due to a parameter perturbation

Step 1: Second-Order Taylor Expansion Applying the second order taylor expansion to the function $f(x;\phi)$ around ($\phi=\phi^{*}$), we get:

	$\displaystyle{f}(x;\phi^{*})={f}(x;\theta^{*}+\Delta\theta)\approx$
	$\displaystyle{f}(x;\theta^{*})+\nabla_{\theta}{f}(x;\theta^{*})^{T}\Delta\theta+\frac{1}{2}\Delta\theta^{T}H(x;\theta^{*})\Delta\theta=$
	$\displaystyle\hat{f}(x;\theta^{*})+\nabla_{\theta}\hat{f}(x;\theta^{*})^{T}\Delta\theta+\frac{1}{2}\Delta\theta^{T}\hat{H}(x;\theta^{*})\Delta\theta.$

Bear in mind that $f(x;\theta)=\hat{f}(x;\theta)$ and $H(x;\theta^{*})$ refers to the Hessian matrix of $f$ at $(\theta=\theta^{*})$.

Step 2: Optimum Condition

$$\nabla_{\theta}\hat{f}(x;\theta^{*})=0$$

$$\Rightarrow\hat{f}(x;\phi^{*})\approx\hat{f}(x;\theta^{*})+\frac{1}{2}\Delta\theta^{T}H(x;\theta^{*})\Delta\theta$$

Step 3: Lipschitz Continuity of Gradient

$$\|\nabla_{\theta}\hat{f}(x;\theta)-\nabla_{\theta}\hat{f}(x;\theta^{\prime})\|\leq L\|\theta-\theta^{\prime}\|$$

Step 4: Bounding the Hessian

$$\|\Delta\theta^{T}\hat{H}(x;\theta)\Delta\theta\|\leq L\|\Delta\theta\|^{2}$$

Step 5: Estimating the Deviation in Function Value

	$\displaystyle\Delta{f}$	$\displaystyle={f}(x;\phi^{*})-\hat{f}(x;\theta^{*})$
	$\displaystyle\Delta{f}$	$\displaystyle\approx\frac{1}{2}\Delta\theta^{T}\hat{H}(x;\theta)\Delta\theta\leq\frac{1}{2}L\|\Delta\theta\|^{2}$

The deviation $\Delta{f}$ in the function value from its optimal value due to a parameter perturbation $\Delta\theta$ is bounded by $\frac{1}{2}L\|\Delta\theta\|^{2}$ under the assumption of L-Lipschitz continuity of the gradient. This result implies that the function value’s deviation grows at most quadratically with the size of the parameter perturbation.

It is of interest to explore whether limiting the number of parameter updates is a suitable approach for investigating the influence of gradient accumulation on SortedNet. One possible way to certify this factor is by running SortedNet with different gradient accumulation values while keeping the number of updates fixed. To that end, we consider the same settings as Table 5 and repeat the experiment while fixing the maximum number of training epochs. By fixing this value and increasing gradient accumulation values, we implicitly decrease the number of parameter updates. Table 6 reports the results. Comparing the results of these two tables, it is obvious that the number of updates plays a significant role in the model’s performance. For instance, when considering $g_{acc}=2$, an average performance drop of approximately 2% is observed across all sub-models. This reduction indicates that the underlying model needs more training time for higher values of $g_{acc}$.

This section provides an overview of the hyperparameters and experimental configurations, detailed in Table 7.

To empirically compare the training time between SortedNet and LCS_p, the elapsed time per epoch for five epochs is recorded independently for each method. We then ignore the first epoch to reduce the impact of first-time loading and initialization. Next, for each method we take the average of the remaining elapsed times. We refer to these averaging times (in seconds) by $\bar{T}_{SortedNet}=49.7\pm 2.06$ and $\bar{T}_{LCS\_p}=292.7\pm 3.17$ for simplicity. As it is mentioned in Subsection C.1, SortedNet with $g_{acc}=4$ can be considered as a fair comparison with LCS_p. As a result, each epoch in LCS_p holds four times the significance of SortedNet in terms of the total number of parameter updates. Therefore, we can simply multiply $\bar{T}_{SortedNet}$ by a factor of four to equalize their impacts in term of total number of parameter updates. By doing that, we have $\bar{T}_{SortedNet}=198.8$, which is almost one-third less than $\bar{T}_{LCS\_p}$.

In Figure 2-a, as mentioned before, we adjusted the performance of the classifiers for the baseline in the experiment but not for the SortedNet. Therefore, as an additional experiment, we wanted to analyze the impact of adjusting the classifier over the performance of our proposed method as well. Same as the previous experiment, we trained the classifier layer for 5 epochs for each sub-model and reported the performance. As shown in Figure 6, the gain is much higher for very smaller networks than the large ones. The SortedNet shared classifier already doing a good job without additional computation overhead for all sub-models but further adjustments might be beneficial as shown.

In this section, we are interested to investigate whether SortedNet is applicable to more complex dimensions other than width and depth. For example, can we utilize the SortedNet for sorting the Attention Heads Vaswani et al. [2017]. To achieve this, we conducted an experiment over BERT-large Devlin et al. [2019] which we tried to sort the information across multiple dimensions at once including, number of layers, hidden dimension, and number of attention heads. In other words, we tried to sort information over Bert-large and Bert-base as Bert-base can be seen as a subset of the Bert-large and therefore respect the nested property. As reported in Table 8, in addition to the reported performance of Bert-base and Bert-large according to the original paper Devlin et al. [2019], we reported the performance of these models in the paper experimental setting. The performance of randomly initialized Bert-base has been reported as well. We also extracted a Bert-base from a Bert-large model, and we reported the performance of such model in the same table. Additionally, we highlighted the number of training updates with respect to each objective function in front of each model. For example, in the last row (Sorted $BERT_{LARGE}$), we approximately trained our Sorted model half of the times ($\sim 3Epochs$) over the objective function of Bert-base ($\mathcal{L}_{B}$) and the other half of the times over the objective function of Bert-large ($\mathcal{L}_{L}$) in an iterative random manner as introduced in the section 3. The learned Bert-base performance with these methods is still around 10% behind a pre-trained base but we argue that this is the value of pre-training. To investigate the impact, one should apply the SortedNet during pre-training which we will leave for future research. However, the performance of the learned Bert-large is on-par with an individual Bert-large which suggests sharing the weights does not necessarily have a negative impact over learning. It seems, however, the secret sauce to achieve a similar performance is that we should keep the number of updates for each objective the same as the individual training of Bert-large and Bert-base.

In order to better understand the impact of sorting information, we designed an experiment that compare the dependency order of all the neurons in a sorted network. To keep the experiment simple, we designed a one layer neural network with 10 (hidden dimension) $\times$ 2 (input dimension) neurons as the hidden layer and a classifier layer on top of that which map the hidden dimension to predict the probabilities of a 4 class problem. The task is to predict whether a 2d point belong to a specific class on a synthetic generated dataset. We trained both Sorted Network and the ordinary one for 10 epochs and optimize the networks using Adam optimizer Kingma and Ba [2017] with the learning rate of 0.01 and batch size of 16.

As can be seen in Table 9, the performance of different orders in the original neural network training paradigm can be different and unfortunately there is no specific pattern in it. Therefore, if one search the whole space of different orders (from neuron 1 to neuron n, from neuron n to neuron 1, or even select a subset of neurons by different strategies i.e. for the half size network activate every other neurons like XOXOXOXOXO.) might find better alternatives that work even better than the desirable target order. In this example, the reverse order in average perform better than the target order (86.67% versus 82.22%). However, with the proposed method, we can clearly see that the target order performance consistently is much better than the reverse order (89.25% versus 59.38%). This means, we have been able to enforce the desirable target order as we wanted using our proposed method. For example, neuron 2 is more dependent to neuron 1 in SortedNet in comparison with the ordinary training. In another example, the last 5 neurons are more dependent to the first 5 neurons than other way around. As shown, the performance of the first five neurons is 93.86% while the performance of the last five neurons is only 66.06% in SortedNet. In other words, the gain of adding the last five neurons is quite marginal and most probably prunnable, while the first 5 neurons contains most of the valuable information. It is of interest to further investigate the dependency of neurons to one another and with other metrics which we will leave for future research.