Parallelizing Over Artificial Neural Network Training Runs with Multigrid

MGRIT solves an evolution problem over a finite number of evolution steps. For time-dependent PDEs, these steps correspond to time-steps. For a more general evolution problem, these steps need not be PDE-based. Regardless, MGRIT defines a fine level and coarse level of evolution steps, based on the coarsening factor $m$. (Here, we use a level independent coarsening factor.) There are $N^{(0)}$ steps on level 0 and $N^{(1)}=N^{(0)}/m$ for level 1. Two such levels for $m=4$ are depicted in Figure 3.³³3For time-dependent PDEs, these levels represent a fine and coarse time step spacing, respectively, with the fine time-step size four times smaller than the coarse level time-step size. Typical values for $m$ are 2, 4, 8, and 16.

Multigrid reduction is similar to an approximate block cyclic reduction strategy, and similarly eliminates block rows and columns in the system (5). If all the block rows corresponding to the black fine points in Figure 3 are eliminated from the system (5), then using the level notation (level 1 is coarse, and level 0 is fine), the coarse system becomes

where $(\Phi^{(0)})^{m}$ represents the $m$th power, i.e., $m$ consecutive applications of $\Phi^{(0)}$. Here, we assume that $\Phi^{(\ell)}$ is the same function over all training steps for simplicity, although the method extends easily to the case where $\Phi^{(\ell)}$ depends on the training step number. The operator $R^{(0)}$ represents “Ideal” restriction and we also define an interpolation $P^{(0)}$, as in [6],

The operator $P^{(0)}$ injects coarse points from the coarse- to fine-level, and then extends those values to fine-points with a $\Phi^{(0)}$ application. This will be equivalent to injection, followed by F-relaxation (defined below).

The operator $A^{(1)}$ can be formed in a typical multigrid fashion with $A^{(1)}=R^{(0)}A^{(0)}P^{(0)}$, and we refer to this as the ideal coarse-level operator because the solution to (6) yields the exact fine-level solution at the coarse points.⁴⁴4The ideal coarse-level operator can also be obtained with the injection operator $R^{(0)}_{I}$, $R^{(0)}AP^{(0)}=R^{(0)}_{I}AP^{(0)}$. For efficiency, we therefore will use injection later on in the algorithm. The interpolation $P$ is also ideal, and uses $\Phi$ to compute the exact solution at the fine points. However, this exact solution process is in general as expensive as solving the original fine-level problem, because of the $(\Phi^{(0)})^{m}$ applications in $A^{(1)}$. MGRIT “fixes” this issue by approximating $A^{(1)}$ with $B^{(1)}$, i.e.,

where $\Phi^{(1)}$ is an approximation of $(\Phi^{(0)})^{m}$. One obvious choice here is to approximate $(\Phi^{(0)})^{m}$ with $\Phi^{(0)}$, i.e., approximate $m$ training steps with one training step. For time-dependent PDEs, $\Phi^{(1)}$ is usually a single time-step using a larger time-step size that approximates $m$ finer time-steps. How well $\Phi^{(1)}$ approximates $(\Phi^{(0)})^{m}$ will govern how fast MGRIT converges, and choosing a $\Phi^{(1)}$ that allows for fast convergence is typically the main topic of research for a given problem area.

Thus, the two-level algorithm uses the coarse level to compute an error correction based on the fine-level residual, $r^{(0)}=g^{(0)}-A^{(0)}(w^{(0)})$, mapped to the coarse level with $R^{(0)}$ (see Step 2 in Algorithm 1). Complementing the coarse level correction is a local relaxation process on the fine-level which resolves fine-scale behavior. Figure 4 (from [16]) depicts F- and C-relaxation on a level with $m=4$. F-relaxation updates the solution at each F-point by propagating the solution forward from the nearest coarse point (C-point). And, C-relaxation updates every C-point with a solution propagation from the nearest F-point. Overall, relaxation is highly parallel, with each interval of F-points being independently updated during F-relaxation, and likewise for C-relaxation. The nonintrusiveness of MGRIT is also apparent because the internals of $\Phi^{(\ell)}$ are opaque to MGRIT. Thus, this approach can be easily applied by wrapping existing code bases.

The above two-level process can be easily to nonlinear problems using full approximation storage (FAS), a nonlinear multigrid scheme [1, 7].⁵⁵5We mention that the F-relaxation two-level MGRIT algorithm is equivalent to the popular parallel-in-time Parareal algorithm [9]. The key addition of FAS to the multigrid process is how the coarse-level equation is Step 3 is computed.

The two-level FAS version of MGRIT is described in Algorithm 1 (taken from [7]). A recursive application of the algorithm at Step 4 makes the method multilevel. This can be done recursively until a trivially sized level is reached, which is then solved sequentially. Ideal interpolation (7) is represented as two steps (5 and 6). Initial values for the weights $w^{(0)}_{0}$ can in principle be anything, even random values, which is the choice made here for $w^{(0)}_{0}$ as the initial condition. For all other points, $w^{(0)}_{i}=0$ is the initial value.

Multigrid solves can use a variety of cycling strategies. Here, we consider V- and F-cycles, as depicted in Figure 5, which shows the order in which levels are traversed. The algorithm continues to cycle until the convergence check based on the fine level residual is satisfied in line 6.

In this section, we consider the motivation in Section 2.4, and consider increasing $\alpha^{(\ell)}$ on coarser levels. In particular, we consider doubling $\alpha^{(1)}$ so that $\alpha^{(1)}=2\alpha^{(0)}$. Remarkably, this single change to the Naive Solver yields fast, scalable two-level iteration counts.

More generally, we consider defining $\alpha^{(\ell)}$ as

where $\alpha^{(max)}$ is some experimentally determined maximum stable $\alpha^{(\ell)}$ value.⁶⁶6This connection between the size of $\alpha^{(\ell)}$ and stability (i.e., $w^{(\ell)}$ values that do not diverge over multiple $\Phi^{(\ell)}$ applications) shows another similarity between $\alpha^{(\ell)}$ and time-step size. Here, $\alpha^{(max)}=8.0$, i.e., $\alpha^{(\ell)}$ is allowed to increase up to level $3$. Beyond that point, the training becomes “unstable”, and $\alpha^{(\ell)}$ can no longer increase. While the halt to increasing $\alpha^{(\ell)}$ values may negatively impact the algorithm, this halt occurs only on coarser levels, and so its negative impact on MGRIT convergence will hopefully be somewhat mitigated.

We denote this as Solver 1, which uses equation (14) to define $\alpha^{(\ell)}$, $m=2$ and equation (11) to define $z^{(\ell)}_{i}$ as containing all the training data. Together with $\alpha^{(b)}=1.0$ and $\alpha^{(max)}=8.0$, this fully defines the method explored in this section.

Table 3 depicts the results for Solver 1. Initially the results look scalable (i.e., flat iteration counts and stable convergence rates $\rho$). But, eventually a log-type growth in the iterations becomes apparent. A possible culprit is the fact that $\alpha^{(\ell)}$ is capped and cannot continually increase. V-cycle results are quite poor and hence omitted for brevity.

One important parameter that affects the MGRIT convergence rate is the base learning rate $\alpha^{(b)}$. Decreasing this value improves the MGRIT convergence rate, while also requiring more training steps to adequately train the network. Exploration of this trade-off is future research, but we make note of it here. If we take Solver 1, decrease $\alpha^{(b)}$ from 1.0 to 0.5 and keep $\alpha^{(max)}=8.0$, we get the improved F-cycle iterations in Table 4

In this section, we note that there is unutilized parallelism in the method. All four training examples are processed simultaneously as a “batch”. That is, $z^{(\ell)}_{i}$ is defined by an $x^{(\ell)}_{i}$ and $y^{(\ell)}_{i}$ with four rows, representing the four training sets (see equation (11). Instead, each of the four training sets can be broken up and processed individually, resulting in four times as many possible steps to parallelize. More precisely, for $K$ training instances, define $x^{(\ell)}_{i}$ and $y^{(\ell)}_{i}$ to correspond to the mod$(i,K)$-th training instance.

Here $K=4$, so $\Phi^{(0)}_{0}$ would train the network with the first training data set with $z^{(0)}_{0}=[x_{(0)},y_{(0)}]$. $\Phi^{(0)}_{1}$ would train the network for the second training data set with $z^{(0)}_{1}=[x_{(1)},y_{(1)}]$, and so on. Once the training step index passes the value $K=4$, the process repeats, so that $\Phi^{(0)}_{4}$ again trains on the first training data set with $z^{(0)}_{4}=[x_{(0)},y_{(0)}]$. This strategy exposes significantly more parallelism than the “unserialized” $\Phi^{(\ell)}_{i}$.

We denote this as Solver 2, which defines

Solver 2 continues using equation (14) to define $\alpha^{(\ell)}$ and $m=2$. Together with $\alpha^{(b)}=1.0$ and $\alpha^{(max)}=8.0$, this fully defines the method explored in this section.

Table 5 depicts two-level results using Solver 2, which show a scalable convergence rate, but the rate is slower than that in Table 2 for Solver 1.

However, Table 6 depicts Solver 2 results for F-cycles, and this shows an unscalable increase in iteration counts.

One possible contributing factor the unscalable iteration growth in Table 6 is the fact that $\alpha^{(\ell)}$ is allowed to only grow on the first four levels. To investigate this possibility, Table 7 runs the same experiment as in Table 6, but with with $\alpha^{(b)}=0.1$ and $\alpha^{(max)}=30.0$. We see that this combination of a smaller base $\alpha^{(b)}$, which allows $\alpha^{(\ell)}$ to grow up to level $\ell=8$ dramatically improves the scalability of the F-cycles. The downside is the smaller learning rate. We leave a further investigation of this for future work, because a deeper theoretical understanding is required to meaningfully guide our tests.

While the proposed method does allow for parallelism over training steps, we need to quantify possible parallel speedups for the proposed approach. To do that, we use a simple computational model that counts the dominant sequential component of the MGRIT process and compares that to the dominant sequential component of the traditional serial training approach. The dominant sequential component in both cases is the required number of consecutive $\Phi^{(\ell)}_{i}$ applications, because this is the main computational kernel.

For sequential time stepping, the number of consecutive $\Phi^{(0)}_{i}$ applications is simply $N^{(0)}$. However for MGRIT, we must count the number of consecutive $\Phi^{(\ell)}_{i}$ applications as the method traverses levels. Consider a single level during a V-cycle with a coarsening factor of 2. FCF-relaxation contains 3 sequential $\Phi^{(\ell)}_{i}$ applications, the sweeps over F- then C- then F-points. Additionally, each level must compute an FAS residual (see Algorithm 1), which requires two matrix-vector products. In parallel at the strong scaling limit, each process in time will own two block rows that will likely be processed sequentially, resulting in 4 total sequential $\Phi^{(\ell)}_{i}$ applications for the computation of the two matrix-vector products. Last, on all levels, but level 0, there is an additional $\Phi^{(\ell)}_{i}$ application to fill in values at F-points for the interpolation. The coarsest level is an exception and only carries out a sequential solve no longer than the maximum allowed coarsest level size (here 10).

Taken together, this results in 7 sequential $\Phi^{(\ell)}_{i}$ applications on level 0, and 8 on all coarser levels, except the coarsest level which has approximately 10. Thus, the dominant sequential component for a V-cycle is modeled as

where niter is the number of MGRIT iterations, or V-cycles, required for convergence.

The dominant sequential component of an F-cycle is essentially the concatenation of successive V-cycles starting on ever finer levels (see Figure 5), i.e.,

The reader will note that there is no final V-cycle done on level 0 in equation (17).

Thus, we define the potential speedup for V-cycles as $N^{(0)}/\sigma_{V}$ and $N^{(0)}/\sigma_{F}$ for F-cycles. This is a useful speedup measure, but it measures speedup relative to an equivalent sequential training run. For example, for the case of the double $\Phi^{(\ell)}$ applications, this measure will compute a speedup against a sequential run that also does a double $\Phi^{(\ell)}$ application.

Figure 6 depicts the potential speedup for Solver 1 from Tables 3, 4 and Solver 3 from Table 7. Although this is a simple model problem, these results indicate that the observed convergence behavior offers promise for a useful speedup on more real-world examples. Also, the trend lines indicate that larger problem sizes would offer improved potential speedups.

In this section, we experiment with Solver 2, which serializes over the training data, and continue to use $\alpha^{(b)}=0.025$ and $\alpha^{(max)}=0.2$. In other words, the first 500 training steps will train over the first 500 training data sets; training step 501 will loop back to the first training data set; and so on. Because of the smaller training data set size per step, we are able to run out to larger problem sizes.

Table 13 depicts two-level results for Solver 2. The iteration counts are good, but not quite scalable, showing some growth.

Table 14 depicts F-cycle results for Solver 2. They are also not scalable, showing a slow growth, but the iteration counts are modest enough to potentially be useful. Table 15 depicts results for Solver 2 for V-cycles. The deterioration for larger problem sizes and F-cycles in Table 14 is visible when compared to the two-level results in Table 13. To investigate possible remedies, we again explore the possibility of allowing $\alpha^{(\ell)}$ to grow on more levels. Table 16 carries out such a experiment, except this time we compute $\alpha^{(\ell)}$ with

as opposed to equation (14), which uses powers of 2. This allows $\alpha^{(\ell)}$ to grow on all levels for all experiments. The F-cycle results now show no deterioration from the two-level results in Table 13 for the large problem sizes. Unfortunately, this choice of $\alpha^{(\ell)}$ did not show a significant improvement for V-cycles. We leave a further investigation of optimal coarse level $\alpha^{(\ell)}$ as future work, because a deeper theoretical understanding is required to meaningfully guide our tests.