# Pipelined Stochastic Gradient Descent with Taylor Expansion
## 3 Proposed Methods
### Pipelined SGD with Taylor Expansion (TaylorPipe)  
To alleviate the delayed gradient problem occurring in ASGD, the gradient at a future point in time may be approximated by applying the Taylor expansion [15]. As gradients are the first-order derivatives of the loss function, the second-order derivatives should be used in applying the Taylor expansion. To predict the gradient of the master model \(\bar{W}_{n}\) that is from the future in time using the current time local model \(W_{n}^{t}\), the following Taylor expansion can be used:  
\[\
abla\mathcal{L}(\bar{W}_{n};x^{t},y^{t})=\
abla\mathcal{L}(W_{n}^{t};x^{t}, y^{t})+\
abla^{2}\mathcal{L}(W_{n}^{t};x^{t},y^{t})(\bar{W}_{n}-W_{n}^{t})+\epsilon, \tag{2}\]  
where \(\epsilon\) is the error caused by the first-order Taylor expansion. The Hessian matrix, which is the second-order derivative of the loss function, can be approximated using an appropriate value of \(\lambda\) as follows [15]:  
\[\
abla^{2}\mathcal{L}(W_{n}^{t};x^{t},y^{t})\approx\lambda\
abla\mathcal{L}(W_ {n}^{t};x^{t},y^{t})\odot\mathcal{L}(W_{n}^{t};x^{t},y^{t}), \tag{3}\]  
where \(\odot\) represents the pair-wise dot product. The proposed TaylorPipe alternately performs the forward pass and backward pass according to the work scheduling in Figure 2. Each computing node calculates the Taylor expansion only for the stage of the model assigned to the computing node, instead of for the entire model, because the pipelined SGD work scheduling and the model replicas described in Section 3.1 allow model parallelism as well as data parallelism. A pseudo-code of TaylorPipe is presented in Algorithm 1, which is executed by all computing nodes simultaneously. In the algorithm, \(t_{\text{f}}\) and \(t_{\text{b}}\) represent the mini-batch indices for the forward and backward passes, respectively, and \(T\) denotes the total number of mini-batches. During the initial and final phases, only the forward and backward passesare executed, respectively. After the \((n+1)\)-th computing node finishes executing the backward pass \(B_{n+1}^{\text{f}_{\text{b}}}\), the \(n\)-th computing node calculates the gradient in the backward pass \(B_{n}^{\text{f}_{\text{b}}}\) using the local model \(W_{n}^{\text{f}_{\text{b}}}\) for the mini-batch of index \(t_{\text{b}}\). Subsequently, it performs the Taylor expansion of the gradient for the master model \(\tilde{W}_{n}\), and the updated master model is copied back to the local model. At a given time point, each stage of the master model in the computing nodes contains the updated model parameters trained on different mini-batches. For example, when \(B_{1}^{1}\) and \(B_{2}^{2}\) are finished, \(\tilde{W}_{1}\) and \(\tilde{W}_{2}\) contain the model parameters trained using up to the 1 st and 2 nd mini-batches, respectively. However, when the pipeline is flushed completely, all stages of the master model in the computing nodes finally contain the model parameters of the same time version. ```
Initialize \(\tilde{W}_{n}\) and copy it to \(W_{n}^{1}\), \(W_{n}^{2},\cdots,W_{n}^{2N-1}\). Set \(t_{\text{f}}\) and \(t_{\text{b}}\) to \(1\). repeat if not final phase then /* forward pass */ if\((1<n)\)then  Wait for \(F_{n-1}^{t_{\text{f}}}\) to be finished endif  Execute \(F_{n}^{t_{\text{f}}}\) if\((n<N)\)then  Send the result of \(F_{n}^{t_{\text{f}}}\) to node \(n+1\) endif \(t_{\text{f}}\gets t_{\text{f}}+1\) endif if not initial phase then /* backward pass */ if\((n<N)\)then  Wait for \(B_{n+1}^{\text{f}_{\text{b}}}\) to be finished endif  Execute \(B_{n}^{\text{f}_{\text{b}}}\) as follows: \(\left| \begin{array}{l}g\leftarrow\
abla\mathcal{L}(W_{n}^{\text{f}_{\text{b}}} ;x^{\text{f}_{\text{b}}},y^{\text{f}_{\text{b}}})\\ \tilde{W}_{n}\leftarrow\tilde{W}_{n}-\eta(g+\lambda g\odot g\odot(\tilde{W}_{n} -W_{n}^{\text{f}_{\text{b}}}))\\ w_{n}^{\text{f}_{\text{b}}}\leftarrow\tilde{W}_{n}\end{array}