A Model Decomposition-in-Time of Recurrent Neural Networks: A Feasibility Analysis

In the context of Recurrent Neural Networks, minimization of the Loss Function (LF) causes the most training overhead. Following the Parallel In-Time approaches, we introduce an ab-initio decomposition across time direction. The key point of our approach lies in the innovative definition of local objective functions which allows us to overcome the sequential nature of the network and the management of dependencies between time steps. In particular, we define local RNNs by adding a suitable overlapping operator to the local objective functions which guarantees their matching between adjacent subsequences. In this way, we get to a fully parallelizable decomposition of the RNN whose implementation avoids global synchronizations or pipelining. Nearest neighbours communications guarantee the algorithm’s convergence. We hope that these findings encourage readers to further extend the framework according to their specific application requirements.

1. Introduction

Recurrent Neural Networks (RNNs) are a class of Deep Learning models that are fundamentally designed to handle sequential data. Unlike feedforward neural networks, RNNs possess the unique feature of maintaining a memory of previous inputs by using their internal state to process sequences of inputs. This makes them ideally suited for applications where the order of data points is crucial [1]. With the increase in the size of the dataset and the improvement of the complexity of the model, the demand for computational strength and storage space in the training process also increases proportionally. The present work is placed in the context of the design of parallel algorithms for RNNs in large-scale applications where parallelizing RNNs can be challenging due to their sequential nature. We focus on RNNs with an additive loss function consisting of a large number of component functions, let us say f_i, such as

$\text{L}\left(\theta \right)=min{\displaystyle \sum _i{f}_i\left(\theta \right)}\text{ (1)}$

where each term corresponds to the error between some data and the output of a parametric model, with being the vector of parameters. An example is linear least squares problems, where f_i has a quadratic structure, except for a regularization function. A more general class of additive cost problems is nonlinear least squares [2,3].

Minimization of the Loss Function (LF) causes the most training overhead. The structure of the additive cost function has facilitated the use of parallel computing approaches that are well-suited for the incremental approach [4]. Very often, besides the exploitation of Graphics Processing Units' acceleration [5], concurrency is introduced inside the operations of each minimization step, at the cost of an all-to-one data synchronization at the end of each step (this is the simplest way to achieve what is defined in the context of deep learning as data parallelism or the mini-batch gradient descent approach). In the alternative, the incremental computing of the values and subgradients of the components f_i could be performed in a distributed manner. This is defined as model parallelism. Naive model parallelism is relatively simple. It’s straightforward to split a large model into chunks of consecutive layers. However, there’s a sequential dependency between inputs and outputs of layers, so a naive implementation can lead to a large amount of idle time while a worker waits for outputs from the previous machine to be used as its inputs. We can reuse the ideas from data parallelism by having each worker only process a subset of data elements at one time, allowing us to cleverly overlap new computations with wait time. Pipeline parallelism allows the execution of a model to be partitioned such that multiple micro-batches can execute different parts of the model code concurrently [6-11].

In the presence of evolutionary problems, in the last decades, Parallel-In-Time methods have been investigated for reducing the temporal dimensionality of such problems. Since Nievergelt, in 1964, proposed for the first time the decomposition algorithm for finding the parallel solutions of evolutionary ordinary differential equations [12], the methods of time-parallel time integration have been extensively expanded and several relevant works can be found in the literature. An extensive and updated literature list can be found on the website [13] collecting information about people, methods, and software in the field of parallel-in-time integration methods. By relying on results we have obtained on the Kalman Filter and on variational methods on the uncertainty quantification models [14-16], in this work we present the main idea underlying an ab-initio model decomposition in time of RNNs. This approach implies both the time series and the objective function decomposition among computing resources. Local functionals are suitably modified, by imposing a regularization constraint in order to enforce the matching of their solutions between adjacent sub-sequences. Finally, according to the Additive Schwarz Method (ASM), synchronization of local solutions is imposed iteratively. Such synchronization guarantees the model convergence [16].

2. Model decomposition-in-time of RNNs

In the following, we summarize the key components of the hybrid decomposition of RNNs. In particular, we introduce the operators defining data reduction and localization then, finally, we give the mathematics underlying decomposition in time of RNNs. For simplicity, we consider the simplest RNN, with a self-connected hidden state, through which information cycles across time steps [17]. The decomposition approach could be extended to other types of RNN.

2.1 Recurrent neural networks

RNNs are a class of deep learning models that are fundamentally designed to handle sequential data. At each time step, t = 1,...,q, the RNN takes an input vector, $x_t\in {\Re }^{m_0}$ and returns the output vector $y_t\in {\Re }^{m_1}$ , s.t.

$y_t\mathrm{<mo>=</mo>}{\sigma }_y\mathrm{<mo\; stretchy="false">(</mo>}{W}_{ht}{h}_t+b_t\mathrm{<mo\; stretchy="false">)</mo>}\text{ (2)}$

where

$\begin{array}{l}{h}_t\mathrm{<mo>=</mo>}{\sigma }_h\mathrm{<mo\; stretchy="false">(</mo>}{W}_{xh}{x}_t+W_{hh}{h}_{t-1}+b_h\mathrm{<mo\; stretchy="false">)</mo>}t\mathrm{<mo>></mo>0,}\\ h_0\mathrm{<mo>=</mo>0}t\mathrm{<mo>=</mo>0}\mathrm{,}\text{ (3)}\end{array}$

is the hidden state vector that captures information about previous inputs, W_xh is the weight matrix between the input and hidden layer, W_hh is the weight matrix for the recurrent connection (hidden-to-hidden), b_t and b_h are the bias vectors, and finally σy and σh are the activation functions. The weight matrices and the bias functions characterize the unknown model. These parameters are obtained by minimizing a predefined objective Function.

For convenience, at each time t = 1,..., q we concatenate vectors x_t and h_t−1 to form a m₀ + m₁ dimensional vector a_t. Then, if we pose

$W_I\mathrm{<mo>=</mo><mo\; stretchy="false">(</mo>}{W}_{xh}\mathrm{<mo\; stretchy="false">)</mo>,}{W}_L\mathrm{<mo>=</mo><mo\; stretchy="false">(</mo>}{W}_{hh}\mathrm{<mo\; stretchy="false">)</mo>}$

where $W_I\in {ℝ}^{m_1\times m_0}\mathrm{,}{W}_L\in {ℝ}^{m_1\times m_1}\mathrm{,}{x}_t\in {ℝ}^{m_0}\mathrm{,}{h}_{t-1}\in {ℝ}^{m_1}$ , we write all the weights together in the matrix

$W\mathrm{<mo>=</mo><mo\; stretchy="false">(</mo>}{W}_I\mathrm{,}{W}_L\mathrm{<mo\; stretchy="false">)</mo>}\in {\Re }^{m_1\times \mathrm{<mo\; stretchy="false">(</mo>}{m}_0+m_1\mathrm{<mo\; stretchy="false">)</mo>}}$

It results that equations (2) and (3) can be written as

$y_t\mathrm{<mo>=</mo>}{\sigma }_t\mathrm{<mo\; stretchy="false">(</mo>}W{a}_t\mathrm{<mo\; stretchy="false">)</mo>}t\mathrm{<mo>=</mo>1,}\dots \mathrm{,}q\mathrm{.}\text{ (4)}$

Let $G_t\in {\Re }^{m_1\times \mathrm{<mo\; stretchy="false">(</mo>}{m}_0+m_1\mathrm{<mo\; stretchy="false">)</mo>}}$ be the matrix obtained by each r.h.s. of equation in (4); $g\mathrm{<mo>=</mo>}vec\mathrm{<mo\; stretchy="false">(</mo>}{G}_t\mathrm{<mo\; stretchy="false">)</mo>}\in {ℝ}^{\mathrm{<mo\; stretchy="false">(</mo>}{m}_0+m_1\mathrm{<mo\; stretchy="false">)</mo>}\times \mathrm{<mo\; stretchy="false">(</mo>}{m}_1\times q\mathrm{<mo\; stretchy="false">)</mo>}}$ be the vector obtained by stacking the t r.h.s of the equation in (4) on top of one another; let $z\in {\Re }^{m_1\times q}$ be the column vector obtained collecting the vector y_t (output of the RNN) q times. Finally, let $A\in {\Re }^{q\times \mathrm{<mo\; stretchy="false">(</mo>}{m}_0+m\mathrm{1<mo\; stretchy="false">)</mo>}}$ be the matrix whose t-th row is obtained by the vectors a_t, t = 1...,q. The equation in (4) can be written as the overdetermined linear system [18,19].

Ag = z

(Least Square Minimization in RNNs) The least-square minimization problem refers to the computation of g such that [20]:

$\widehat{g}\mathrm{<mo>=</mo>}argmi{n}_{g}J\mathrm{<mo\; stretchy="false">(</mo>}g\mathrm{<mo\; stretchy="false">)</mo>}\text{ (5)}$

with objective function

$J\mathrm{<mo\; stretchy="false">(</mo>}g\mathrm{<mo\; stretchy="false">)</mo><mo>=</mo><mo>|</mo><mo>|</mo>}Ag-z{\mathrm{<mo>|</mo><mo>|</mo>}}_2\text{ (6)}$

2.2 Model decomposition in time

The Model Decomposition in Time direction implies the reduction of both the time series and the objective function among computing resources. This is obtained by data reduction operation and functional localization, respectively. Local functions are obtained as a suitable modification of the original objective function. In particular, local functions are composed of two parts: the first part is the restriction of the original objective function to the sub-interval, then a regularization constraint is added in order to enforce the matching of their solutions between adjacent sub-sequences.

(Data Reduction) Let w = [w_t w_t+1 ... w_n]^T ∈ R^s be a vector with t ≥ 1 , n > 0, s = n − t and I_r = {1,...,r}, r > n and n > t. The extension of w to I_r is:

$E{O}_{I_r}\mathrm{<mo>:</mo>}w\in {ℝ}^s\to E{O}_{I_r}\mathrm{<mo\; stretchy="false">(</mo>}w\mathrm{<mo\; stretchy="false">)</mo><mo>=</mo><mo\; stretchy="false">[</mo>}{\overline{w}}_1{\overline{w}}_2\mathrm{...}{\overline{w}}_r{\mathrm{<mo\; stretchy="false">]</mo>}}^T\in {ℝ}^r\mathrm{,}\text{ (7)}$

where for i = 1,...,r

${\overline{w}}_i\mathrm{<mo>=</mo>}\{\begin{array}{l}{w}_{i}ift\le i\le n\\ 0ifi\mathrm{<mo>></mo>}nandi\mathrm{<mo><</mo>}t\end{array}\text{ (8)}$

Let B = [B¹ B² ... Bⁿ] ∈ R^m×n be such that B^j is the j −th column of B. For i < j (i = 1,...,n − 1 and j = 2,...,n) we define the set I_i,j = {i,...,j}. The restriction of B to the set I_j is:

${\mathrm{<mo>|</mo>}}_{I_j}\mathrm{<mo>:</mo>}B\in {ℝ}^{m\times n}\to B{\mathrm{<mo>|</mo>}}_{I_j}\mathrm{<mo>=</mo><mo\; stretchy="false">[</mo>}{B}^1{B}^2\mathrm{...}{B}^j\mathrm{<mo\; stretchy="false">]</mo>}\in {ℝ}^{m\times j}\mathrm{,}\text{ (9)}$

and to I_i,j is:

${\mathrm{<mo>|</mo>}}_{I_{i\mathrm{,}j}}\mathrm{<mo>:</mo>}B\in {ℝ}^{m\times n}\to B{\mathrm{<mo>|</mo>}}_{I_{i\mathrm{,}j}}\mathrm{<mo>=</mo><mo\; stretchy="false">[</mo>}{B}^i{B}^{i+1}\mathrm{...}{B}^j\mathrm{<mo\; stretchy="false">]</mo>}\in {ℝ}^{m\times j-i}\mathrm{.}\text{ (10)}$

(Objective Function Reduction) Let n₁ ,n₂ > 0 two natural numbers less than q. Let

$J{\mathrm{<mo>|</mo>}}_{\mathrm{<mo\; stretchy="false">(</mo>}{I}_{n_i}\mathrm{,}{I}_{n_j}\mathrm{<mo\; stretchy="false">)</mo>}}\mathrm{<mo>:</mo><mo\; stretchy="false">(</mo>}u{\mathrm{<mo>|</mo>}}_{I_{n_i}}\mathrm{,}u{\mathrm{<mo>|</mo>}}_{I_{n_j}}\mathrm{<mo\; stretchy="false">)</mo>}\mapsto J{\mathrm{<mo>|</mo>}}_{\mathrm{<mo\; stretchy="false">(</mo>}{I}_{n_i}\mathrm{,}{I}_{n_j}\mathrm{<mo\; stretchy="false">)</mo>}}\mathrm{<mo\; stretchy="false">(</mo>}u{\mathrm{<mo>|</mo>}}_{I_{n_i}}\mathrm{,}u{\mathrm{<mo>|</mo>}}_{I_{n_j}}\mathrm{<mo\; stretchy="false">)</mo>}\forall i\mathrm{,}j\mathrm{<mo>=</mo>1,2}\text{ (11)}$

denote the restriction of J defined in (6).

For simplicity of notations, we let J_i,jm ≡ J|_(Ii,Ij) with i, j = 1,2. We consider the decomposition of RNNs in two subsets only.

(In Time DD-RNN setup) Let I = {1,...,q} be the index set of rows of A. DD-RNN setup consists in decomposing I into 2 sets:

$I_{n_1}\mathrm{<mo>=</mo><mo>\{</mo>1,...,}{n}_1\mathrm{<mo>\}</mo>,}{I}_{n_2}\mathrm{<mo>=</mo><mo>\{</mo>}{n}_1-s+\mathrm{1,...,}q\mathrm{<mo>\}</mo>,}\text{ (12)}$

where $\mathrm{<mo>|</mo>}{I}_1\mathrm{<mo>|</mo><mo>=</mo>}{n}_1\mathrm{<mo>></mo>0}$ , $\mathrm{<mo>|</mo>}{I}_2\mathrm{<mo>|</mo><mo>=</mo>}{n}_2\mathrm{<mo>></mo>0}$ , and

$I_{\mathrm{1,2}}\mathrm{<mo>=</mo><mo>\{</mo>}{n}_1-s+\mathrm{1,...,}{n}_1\mathrm{<mo>\}</mo>,}\text{ (13)}$

is the overlap set. If s = 0, then I₁ ∩ I₂ = ∅ and I_1,2 ≠ ∅. Restrictions of A to I₁ and I₂ are:

$A_1\mathrm{<mo>=</mo>}A{\mathrm{<mo>|</mo>}}_{I_1}\in {ℝ}^{n_1\times \mathrm{<mo\; stretchy="false">(</mo>}{m}_0+m_1\mathrm{<mo\; stretchy="false">)</mo>}}\mathrm{,}{A}_2\mathrm{<mo>=</mo>}A{\mathrm{<mo>|</mo>}}_{I_2}\in {ℝ}^{n_2\times \mathrm{<mo\; stretchy="false">(</mo>}{m}_0+m_1\mathrm{<mo\; stretchy="false">)</mo>}}\mathrm{,}\text{ (14)}$

(In Time DD-RNN Algorithm) According to the ASM in [21], we introduce the loop for k = 0,1,2,..., the model decomposition-in-time leads to two RNNs solving iteratively the following reduced least square minimization problems:

$\begin{array}{ll}{P}_1^{k+1}\mathrm{<mo>:</mo>}{\widehat{g}}_1^{k+1}\hfill & \mathrm{<mo>=</mo>}argmi{n}_{g_1^{k+1}\in {ℝ}^{n_1}}{J}_1\mathrm{<mo\; stretchy="false">(</mo>}{g}_1^{k+1}\mathrm{,}{g}_2^k\mathrm{<mo\; stretchy="false">)</mo>}\hfill \\ \hfill & \mathrm{<mo>=</mo>}argmi{n}_{g_1^{k+1}\in {ℝ}^{k_1}}\left[J{\mathrm{<mo>|</mo>}}_{\mathrm{<mo\; stretchy="false">(</mo>}{I}_1\mathrm{,}{I}_2\mathrm{<mo\; stretchy="false">)</mo>}}\mathrm{<mo\; stretchy="false">(</mo>}{g}_1^{k+1}\mathrm{,}{g}_2^k\mathrm{<mo\; stretchy="false">)</mo>}+\mu \cdot {\mathcal{O}}_{\mathrm{1,2}}\mathrm{<mo\; stretchy="false">(</mo>}{g}_1^{k+1}\mathrm{,}{g}_2^k\mathrm{<mo\; stretchy="false">)</mo>}\right]\hfill \end{array}\text{ (15)}$

$\begin{array}{ll}{P}_2^{n+1}\mathrm{<mo>:</mo>}{\widehat{g}}_2^{k+1}\hfill & \mathrm{<mo>=</mo>}argmi{n}_{g_2^{k+1}\in {ℝ}^{n_2}}{J}_2\mathrm{<mo\; stretchy="false">(</mo>}{g}_2^{k+1}\mathrm{,}{g}_1^{k+1}\mathrm{<mo\; stretchy="false">)</mo>}\hfill \\ \hfill & \mathrm{<mo>=</mo>}argmi{n}_{g_2^{k+1}\in {ℝ}^{k_2}}\left[J{\mathrm{<mo>|</mo>}}_{\mathrm{<mo\; stretchy="false">(</mo>}{I}_2\mathrm{,}{I}_1\mathrm{<mo\; stretchy="false">)</mo>}}\mathrm{<mo\; stretchy="false">(</mo>}{g}_2^{k+1}\mathrm{,}{g}_1^{k+1}\mathrm{<mo\; stretchy="false">)</mo>}+\mu \cdot {\mathcal{O}}_{\mathrm{1,2}}\mathrm{<mo\; stretchy="false">(</mo>}{g}_2^{k+1}\mathrm{,}{g}_1^k\mathrm{<mo\; stretchy="false">)</mo>}\right]\hfill \end{array}\text{ (16)}$

O_1,2 is the overlapping operator representing the data exchange on the overlap set I_1,2, and µ > 0 is the regularization parameter.

In particular, we pose

${\mathcal{O}}_{\mathrm{1,2}}\mathrm{<mo\; stretchy="false">(</mo>}{g}_i\mathrm{,}{g}_j\mathrm{<mo\; stretchy="false">)</mo><mo>=</mo><mo>|</mo><mo>|</mo>}E{O}_{I_i}\mathrm{<mo\; stretchy="false">(</mo>}{g}_i{\mathrm{<mo>|</mo>}}_{I_{\mathrm{1,2}}}\mathrm{<mo\; stretchy="false">)</mo>}-E{O}_{I_i}\mathrm{<mo\; stretchy="false">(</mo>}{g}_j{\mathrm{<mo>|</mo>}}_{I_{\mathrm{1,2}}}\mathrm{<mo\; stretchy="false">)</mo><mo>|</mo><mo>|</mo>,}i\mathrm{,}j\mathrm{<mo>=</mo>1,2}\text{ (17)}$

with $E{O}_{I_i}{g}_1{\mathrm{<mo>|</mo>}}_{I_{\mathrm{1,2}}}$ , $E{O}_{I_i}{g}_2{\mathrm{<mo>|</mo>}}_{I_{\mathrm{1,2}}}$ be the extension to I_i, of restriction to I_1,2 in (13) of $g_1\in {ℝ}^{n_1}$ and $g_2\in {ℝ}^{n_2}$ , respectively.

As a solution, we get the sequence ${\mathrm{<mo>\{</mo>}{g}^{k+1}\mathrm{<mo>\}</mo>}}_{k\in {\mathbb{N}}_0}$ such that:

$g^{k+1}\mathrm{<mo>=</mo>}\{\begin{array}{ll}{\widehat{g}}_1^{k+1}{\mathrm{<mo>|</mo>}}_{I_1\backslash I_{\mathrm{1,2}}}\hfill & on{I}_1\backslash I_{\mathrm{1,2}}\hfill \\ \frac{\mu }{2}\mathrm{<mo\; stretchy="false">(</mo>}{\widehat{g}}_2^{k+1}{\mathrm{<mo>|</mo>}}_{I_{\mathrm{1,2}}}+{\widehat{g}}_1^{k+1}{\mathrm{<mo>|</mo>}}_{I_{\mathrm{1,2}}}\mathrm{<mo\; stretchy="false">)</mo>}\hfill & on{I}_{\mathrm{1,2}}\hfill \\ {\widehat{g}}_2^{k+1}{\mathrm{<mo>|</mo>}}_{I_2\backslash I_{\mathrm{1,2}}}\hfill & on{I}_2\backslash I_{\mathrm{1,2}}\hfill \\ \hfill & \hfill \end{array}\mathrm{,}\text{ (18)}$

with the sets I₁, I₂ defined in (12) and I_1,2 in (13).

Algorithm implementing local RNN on I_n1 and I_n2,

1: repeat

2: k:= k + 1

3: Call Loc-RNN

4: Exchange x^k_i between adjacent subdomains and compute O_1,2

5: until $\parallel g_i^k-g_i^{k-1}\parallel \mathrm{<mo><</mo>}eps$

3. Discussion

The aim of this work is only to illustrate the framework underlying such innovative ideas. Many aspects have been left intentionally unexplained. For example, regarding the size of the overlapping interval, we note that the majority of parallelization techniques first decompose the domain in nonoverlapping subdomains and then extend each subdomain with halo regions to enable stencil operations in a parallel context. Instead, the algorithm allows one to reduce communications among adjacent subdomains to only those nodes lying on the interfaces. In particular, if the widths of the halo regions are equal to one grid node wide, the halo nodes are actually the nodes of the overlapping region. Otherwise, halo regions contain more nodes than overlapping nodes and this means more communications among adjacent subdomains. Hence, the algorithm reduces communication times which depends on the number of inner nodes in overlap regions. In addition, the extra work performed on the overlapped region with an increasing size can be seen as the effect of a preconditioner on the overlapping region which can overestimate the solution [21].

Regarding the regularization parameter, we observe that the regularization parameter plays an important trade-off role. As is evident, when the regularization parameter overfitting occurs larger values are expected to produce reliable solutions even if the DD-step number increases. The choice of a good regularization parameter is one of the most important issues in solving unconstrained minimization problems and there exists a significant amount of research in the literature on the development of appropriate strategies for selecting regularization parameters. Parameter choice methods can be classified according to the input they require. There are two basic types: a-priori methods, requiring information about the noise level on data. The discrepancy principle, developed and analyzed by Morozov is the oldest one of them; data-driven methods, require no extra information. Generalized Cross-Validation (GCV), due to Wahba, is one of the most popular methods [22].

Another issue is the employment of a dynamic load balancing scheme based on adaptive and dynamic redefining of initial decomposition. Specifically, in order to optimally choose the domain decomposition configuration, the partitioning into subdomains must satisfy certain conditions. First, the computational load assigned to subdomains must be equally distributed. Good quality partitioning also requires the volume of communication during calculation to be kept at its minimum. In [15] the authors employed a dynamic load balancing scheme based on adaptive and dynamic redefining of initial decomposition, aimed to balance load between processors according to data location. In particular, the authors focused on the introduction of a dynamic redefining of initial DD in order to deal with problems where the observations are non uniformly distributed and generally sparse.

The most significant finding is that we can solve several smaller problems improving the accuracy-per-parameter metric. Most importantly, subproblems can be solved in parallel, leading to a scalable algorithm where the workers locally exchange parameter updates via a nearest-neighbourhood communication scheme, which does not require a fully connected network. In contrast to other decomposition-in-time approaches, in our approach local solvers run concurrently from the beginning. Overall, the employment of the framework on hybrid high-performance computing systems seems to be a fruitful research area [23].