Coded Computing: A Learning-Theoretic Framework

The theory of coded computing has been developed to improve the reliability and security of large-scale machine learning platforms, effectively tackling two major challenges: (1) the detrimental impact of slow workers (stragglers) on overall computation efficiency, and (2) the threat of faulty or malicious workers that can compromise data accuracy and integrity. These challenges have been well-documented in the literature, including the seminal work [1] from Google. Furthermore, [2] reported that in a sample set of $3000$ matrix multiplication jobs on AWS Lambda, while the median job time was $40$ seconds, approximately $5\%$ of worker nodes took $100$ seconds to respond, and two nodes took as long as $375$ seconds. Moreover, coded computing has also been instrumental in addressing privacy concerns, a crucial aspect of distributed computing systems [3, 4, 5, 6, 7, 8, 9, 10, 11, 12].

The concept of coded computing has been motivated by the success of coding in communication over unreliable channels, where instead of transmitting raw data, the transmitter sends a (linear) combination of the data, known as coded data. This redundancy in the coded data enables the receiver to recover the raw data even in the presence of errors or missing values. Similarly, coded computing includes three stages [13, 3, 11] (see Figure 1(a)):

1. (1)

   The Encoding Layer in which a master node sends a (linear) combination of data, as coded data, to each worker node.

2. (2)

   The Computing Layer, in which the worker nodes apply a predefined computation to their assigned coded data and send the results back to the master node.

3. (3)

   The Decoding Layer, in which the master node recovers the final results from the computation results over coded data. In this stage, the decoder leverages the coded redundancy in the computation to recover the missing results of the stragglers and detect and correct the adversarial inputs.

The existing coded computing has largely built upon algebraic coding theory, drawing inspiration from the renowned Reed-Solomon code construction in communication [14], with proven straggler and Byzantine resiliency [15]. However, the coding in communication, designed for the exact recovery of the messages, built on a foundation that is inconsistent with the computational requirements of machine learning. Developing a code that preserves its specific construction while composing with computation is extremely challenging, leading to significant restrictions. Firstly, current methods are mainly restricted to specific computation functions, such as polynomials and matrix multiplication [3, 16, 17, 13, 11]. Secondly, rooted in algebraic error correction codes, existing approaches are tailored for finite field computations, leading to numerical instability when dealing with real-valued data [18, 19]. Furthermore, these methods are unsuitable for approximate, fixed-point, or floating-point computing, where exact computation is neither possible nor necessary, such as in machine learning inference or training tasks. Finally, these schemes typically have a recovery threshold, which is the minimum number of samples required to recover results from coded outputs of worker nodes [3, 13]. If the number of workers falls below this threshold, the recovery process fails entirely.

Several works have attempted to mitigate the aforementioned issues and transform the coded computing scheme into a more robust and adaptable one, applicable to a wide range of computation functions. These efforts include approximating non-polynomial functions with polynomial ones [20, 5], refining the coding mechanism to enhance stability [21, 22, 23, 24, 25], and leveraging approximation computing techniques to reduce the recovery threshold and increase recovery flexibility [26, 27, 28, 29]. However, these attempts fail to bridge the existing gap between coded computing and general distributed computing systems. The root cause of these issues lies in the fact that they are grounded in coding theory, based on a foundation that is not compatible with the requirements of large-scale machine learning. Therefore, this paper aims to address the following objective:

In this paper, we establish a learning-theoretic foundation for coded computing, applicable to general computations. We adopt an end-to-end system perspective, that integrates an end-to-end loss function, to find the optimum encoding and decoding functions, focusing on straggler resiliency. We show that the loss function is upper-bounded by the sum of two terms: one characterizing the generalization error of the decoder function and the other capturing the training error of the encoder function. Regularizing the decoder stage, we derive the optimal decoder in the Reproducing Kernel Hilbert space (RKHS) of second-order Sobolev functions. This provides an explicit solution for the optimum decoder function and allows us to characterize the resulting loss of the decoding stage. The decoder loss appears as a regularizing term in optimizing the encoding function and represents the norm in another RKHS. Thus, the optimum solution for the encoding function can be derived, too. We address two noise-free and noisy computation settings, for which we derive the optimal encoder and decoder and corresponding convergence rate. We prove that the proposed framework exhibits a faster convergence rate compared to the state-of-the-art and the numerical evaluations support the theoretical derivations (see Figure 1(b)).

Contributions: The main contributions of this paper are:

• •

  We develop a new foundation for coded computing, based on the theory of learning, rather than the theory of coding. We define the loss function as the square error of the computation estimation, averaged over all possible sets of at most $S$ stragglers (Section 3.1). To be able to find the best encoding and decoding functions, we bound the loss function with the summation of two terms, one characterizing the generalization error of the decoder function and the other capturing the training error of the encoder function (Section 3).

• •

  Assuming that the encoder and decoder functions reside in the Hilbert space of second-order Sobolev functions, we use the theory of RKHSs to find the optimum encoding and decoding functions and characterize the convergence rate for the expected loss in both noise-free and noisy computations (Section 4).

• •

  We have extensively evaluated the proposed scheme across different data points and computing functions including state-of-the-art deep neural networks and demonstrated that our proposed framework considerably outperforms the state-of-the-art in terms of recovery accuracy (Section 5).

Here, we propose a novel straggler-resistant Learning-Theoretic Coded Computing (LeTCC) framework for general distributed computing. As depicted in Figure 2, our framework comprises two encoding and decoding layers, with a computing layer sandwiched between them. The framework operates according to the following steps:

1. (1)

   Encoding Layer: The master node fits an encoder regression function $\mathbf{u}_{\textrm{enc}}:\mathbb{R}\to\mathbb{R}^{d}$ at points $\{(\alpha_{k},\mathbf{x}_{k})\}^{K}_{k=1}$ for fixed, distinct, and ordered values $\alpha_{1}<\alpha_{2}<\dots<\alpha_{K}\in\mathbb{R}$. Then, it computes the encoder function $\mathbf{u}_{\textrm{enc}}(\cdot)$ on another set of fixed, distinct, and ordered values $\{\beta_{n}\}_{n=1}^{N}$ where $\beta_{1}<\beta_{2}<\dots<\beta_{N}\in\mathbb{R}$, with $k\in[K]$ and $n\in[N]$. Subsequently, the master node sends the coded data points $\tilde{\mathbf{x}}_{n}=\mathbf{u}_{\textrm{enc}}(\beta_{n})\in\mathbb{R}^{d}$ to worker $n$ for $n\in[N]$. Note that each coded data point $\tilde{\mathbf{x}}_{n}$ is a combination of all initial points $\{\mathbf{x}_{k}\}^{K}_{k=1}$.

2. (2)

   Computing Layer: Each worker node $n\in[N]$ computes $\mathbf{f}(\tilde{\mathbf{x}}_{n})=\mathbf{f}(\mathbf{u}_{\textrm{enc}}(\beta_{n}))$ on its assigned input and sends the result back to the master node.

3. (3)

   Decoding Layer: The master node receives the results ${\mathbf{f}(\tilde{\mathbf{x}}_{v})}_{v\in\mathcal{F}}$ from the non-straggler worker nodes in the set $\mathcal{F}$. Next, it fits a decoder regression function $\mathbf{u}_{\textrm{dec}}:\mathbb{R}\to\mathbb{R}^{m}$ at points ${(\beta_{v},\mathbf{f}(\tilde{\mathbf{x}}_{v}))}_{v\in\mathcal{F}}={(\beta_{v},\mathbf{f}(\mathbf{u}_{\textrm{enc}}(\beta_{v})))}_{v\in\mathcal{F}}$. Finally, using the function $\mathbf{u}_{\textrm{dec}}(\cdot)$, the master node can compute $\hat{\mathbf{f}}(\mathbf{x}_{k}):=\mathbf{u}_{\textrm{dec}}(\alpha_{k})$ as an approximation of $\mathbf{f}(\mathbf{x}_{k})$ for $k\in[K]$. Recall that $\mathbf{u}_{\textrm{dec}}(\alpha_{k})\approx\mathbf{f}(\mathbf{u}_{\textrm{enc}}(\alpha_{k}))\approx\mathbf{f}(\mathbf{x}_{k})$.

As mentioned above, the master node selects and fixes the regression points, $\{\alpha_{k}\}^{K}_{k=1}$ and $\{\beta_{n}\}^{N}_{n=1}$, which remain constant throughout the entire process. The encoder and decoder functions are the only components subject to optimization.

Note that the computational efficiency of the encoding and decoding layers is crucial. This includes the fitting process of the encoder and decoder regression functions, as well as the computation of these regression functions at points $\{\beta_{v}\}_{v\in\mathcal{F}}$ and $\{\alpha_{k}\}^{K}_{k=1}$. If the master node’s computation time is not substantially decreased compared to computing $\{\mathbf{f}(\mathbf{x}_{k})\}^{K}_{k=1}$ by itself, then adopting this framework would not provide any benefits for the master node.

In this section, we investigate the proposed framework from a theoretical standpoint. We provide a detailed explanation of how to design the decoder and encoder functions and subsequently analyze the convergence rate. For simplicity, we first present the results for a one-dimensional function $f:\mathbb{R}\to\mathbb{R}$. All the results have been generalized to the case where $f:\mathbb{R}^{d}\to\mathbb{R}^{m}$ as stated in Appendix C.

Suppose the regression points, $\{\alpha_{k}\}^{K}_{k=1},\{\beta_{n}\}^{N}_{n=1}$, are confined to the interval $\Omega:=(-1,1)$ and $u_{\textrm{enc}},u_{\textrm{dec}}\in\mathcal{H}^{2}\left(\Omega;\mathbb{R}\right)$ where $\mathcal{H}^{2}\left(\Omega;\mathbb{R}\right)$ is the RKHS of second-order Sobolev functions on interval $\Omega$. We review the definition and properties of Sobolev spaces in Appendix A.1.

Decoder Design: Since $\mathcal{L}_{\textrm{dec}}(\mathbf{\hat{f}})$ in the decomposition (2) characterizes the generalization error of the decoder function, we propose a regularized objective function for the decoder:

The first term in (4) corresponds to the mean squared error, while the second term characterizes the smoothness of the decoder function on the interval $\Omega$. Equation (4) represents a Kernel Ridge Regression problem (KRR). It can be shown that the solution of (4) has the following form [30]:

where $\phi$ is the kernel of $\mathcal{H}^{2}\left(\Omega;\mathbb{R}\right)$, $\mathbf{d}=[d_{0},d_{1}]^{T}\in\mathbb{R}^{2}$, and $\mathbf{c}\in\mathbb{R}^{|\mathcal{F}|}$. Substituting (5) into the main objective (4), the coefficients $\mathbf{c}$ and $\mathbf{d}$ can be efficiently computed by optimizing a quadratic equation [31]. This solution is known as the second-order smoothing spline function. We review the theoretical properties of smoothing splines in Appendix A.2.

Let us define the following variables, which characterize the maximum and minimum distances between consecutive regression points $\mathcal{F}:=\left\{\beta_{i_{1}},\dots,\beta_{i_{|\mathcal{F}|}}\right\}$, and $\beta_{i_{1}}<\dots,<\beta_{i_{|\mathcal{F}|}}$. These are the points available to the decoder after removing the missing ones due to the presence of stragglers.

with $\beta_{0}:=-1$ and $\beta_{N+1}:=1$. The following theorems provide crucial insights for designing the encoder function as well as deriving the convergence rates.

The proof and formula for the constant $C_{1}$ are provided in Appendix B.1.

The proof of Theorem 2 and the detailed expressions for $C_{2}$ and $C_{3}$, can be found in Appendix B.2. Some results in approximation theory suggest utilizing Chebyseb points of the first kind, $\alpha_{k}=\cos(\frac{(2k-1)\pi}{2K})$, as regression points, as it is used in [29]. To cover those cases, we present the following corollary.

The proof can be found in Corollary 7 in Appendix B.
Encoder Design: The upper bounds established in Theorems 1, 2 hold for all $u_{\textrm{enc}}\in\mathcal{H}^{2}\left(\Omega;\mathbb{R}\right)$. However, they do not directly lead to a design for $u_{\textrm{enc}}(\cdot)$. To address this, we present the following theorem which bounds the $\left\lVert f\circ u_{\textrm{enc}}\right\rVert^{2}_{{L}^{2}\left(\Omega;\mathbb{R}\right)}$, enabling us to construct $u_{\textrm{enc}}(\cdot)$ without compromising the convergence rate.

The norm $\left\lVert\cdot\right\rVert_{\widetilde{\mathcal{H}}^{2}\left(\Omega;\mathbb{R}\right)}$ is an RKHS and an equivalent norm on Sobolev space introduced by [32] (see (22) in Appendix A.1). The proof can be found in Appendix B.3. By applying the representer theorem [33], we can deduce that the optimal encoder $u_{\textrm{enc}}(.)$, which minimizes the right-hand side of (10) takes the form $u_{\textrm{enc}}(.)=\sum^{K}_{k=1}z_{k}\phi(\alpha_{k},.)$, where $\mathbf{z}\in\mathbb{R}^{K}$, and $\phi$ is the kernel function of $\widetilde{\mathcal{H}}^{2}\left(\Omega;\mathbb{R}\right)$, identical to one in (5). However, due to the non-linearity of $g(\cdot)$, calculating the values of the coefficients $\mathbf{z}$ is challenging. Nevertheless, we demonstrate that the coefficients can be efficiently derived under certain mild assumptions.

See Appendix B.3.1 for the proof. Corollary 2 states that under the assumptions specified in Corollary 2, smoothing spline minimizes the upper bounds (11).

Convergence Rate: Using Theorems 1, 2, and 3, we can derive the convergence rate of the proposed scheme as stated in the following theorem.

Theorem 4 states that for a fixed number of stragglers $S$, the error converges at a rate of $\mathcal{O}(N^{-3})$ in the noiseless setting, which strictly outperforms the state-of-the-art Berrut coded computing approach [29] (see Figure 1). Note that [29] does not analyze the noisy setting.

In this section, we extensively evaluate the proposed scheme across various scenarios. Our assessments involve examining multiple deep neural networks as computing functions and exploring the impact of different numbers of stragglers on the scheme’s efficiency. The experiments are run using PyTorch [34] in a single GPU machine. We evaluate the performance of the LeTCC scheme in three different model architectures:

• •

  Shallow model: We choose LeNet5 [35] architecture as a known shallow network with approximately $6\times 10^{4}$ parameters, trained on the MNIST [36].

• •

  Deep model with low-dimensional output: In this scenario, we evaluate the proposed scheme when the function is a deep neural network trained on color images in CIFAR-10 [37] dataset. We use the recently introduced RepVGG [38] network with around $26$ million parameters which was trained on CIFAR-10¹¹ 1 The pre-trained weights can be found here..

• •

  Deep model with high-dimensional output: Finally, we demonstrate the performance of the LeTCC scheme in a scenario where the input and output of the computing function are high-dimensional, and the function is a relatively large neural network. We consider the Vision Transformer (ViT) [39] as one of the state-of-the-art base neural networks in computer vision for our prediction model, with more than $80$ million parameters (in the base version). The network was trained and fine-tuned on the ImageNet-1K dataset [40]²² 2 We use the official PyTorch pre-trained ViT network from here..

We use the output of the last softmax layer of each model as the output.

Hyper-parameters: The entire encoding and decoding process is the same for different functions, as we adhere to a non-parametric approach. The sole hyper-parameters involved are the two smoothing parameters ($\lambda_{\textrm{enc}},\lambda_{\textrm{dec}}$) which are determined using cross-validation and greed search over different values of the smoothing parameters.
Baseline: We compare LeTCC with the Berrut approximate coded computing (BACC) introduced by [29] as the state-of-the-art coded computing scheme for general computing. The BACC framework is used in [29] for training neural networks and in [41] for inference.
Interpolation Points: We choose Chebyshev points of the first and second kind, $\{\alpha_{i}\}^{K}_{k=1}=\cos\left(\frac{(2k-1)\pi}{2K}\right)$ and $\{\beta_{n}\}^{N}_{n=1}=\cos\left(\frac{(n-1)\pi}{N-1}\right)$, for fair comparison with [29].
Evaluation Metrics: We employ two evaluation metrics to assess the performance of the proposed framework: Relative Accuracy (RelAcc) and Mean Squared Error (MSE). RelAcc is defined as the ratio of the base model’s prediction accuracy to the accuracy of the estimated model on the initial data points. MSE, on the other hand, is our main loss defined in (1) which measures the empirical average of the mean squared difference over multiple batches of and non-straggler set $\mathcal{F}$, $\mathbb{E}_{\mathbf{x}\sim\mathcal{X},\mathcal{F}}\left[\frac{1}{K}\sum^{K}_{k=1}\|\mathbf{f}(\mathbf{x}_{k})-\mathbf{\hat{f}}(\mathbf{x}_{k})\|^{2}_{2}\right]$, where $\mathcal{X}$ is the distribution of input data.
Performance Evaluation: Table 1 presents both MSE and RelAcc metrics side by side. The results demonstrate that LeTCC outperforms BACC across various architectures and configurations, with an average improvement of 15%, 17%, and 9% in MSE for LeNet, RepVGG, and ViT architectures, respectively, and a 2%, 5%, and 4% enhancement in RelAcc.

In the subsequent analysis, we assess the performance of LeTCC in comparison to BACC for a diverse range of stragglers. For a constant number of stragglers, we execute both schemes with the same set of input data points and the same set of stragglers multiple times and record the average RelAcc and MSE metrics. Figure 3 presents the performance of both schemes on three model architectures. As depicted in Figure 3, the proposed scheme consistently outperforms BACC across nearly all values for the number of stragglers.

Coded Points: We also compare the coded points $\{\mathbf{\tilde{x}_{n}}\}_{n=1}^{N}$ sent to the workers in LeTCC and BACC schemes. The results, shown in Figure 4, demonstrate that BACC coded samples exhibit high-frequency noise which causes the scheme to approximate the original prediction worse than LeTCC.

Coded computing was initially introduced to tackle the challenges of distributed computation, particularly the existence of stragglers or slow workers, and also faulty or adversarial nodes. Traditional approaches to deal with stragglers primarily rely on repetition [42, 1, 43, 44, 45, 46], where each task is assigned to multiple workers either proactively or reactively. Recently, coded computing approaches have reduced the overhead of repetition by leveraging coding theory and embedding redundancy in the worker’s input data [29, 3, 13, 26, 47, 48, 49, 50, 51]. This technique, which mainly relies on theory of coding, has been developed for specific types of structured computations, such as polynomial computation and matrix multiplication [13, 3, 52, 17, 53, 48, 12, 7, 54]. Recently, there have been attempts to generalize coded computing for general computations [29, 55, 4, 5]. Towards extending the application of coding computing to machine learning computation, Kosaian et al. [55] suggest training a neural network to predict coded outputs from coded data points. However, the scheme of Kosaian et al. [55] requires a complex training process and tolerates only one straggler. In another work, Jahani-Nezhad and Maddah-Ali [29] proposes BACC, a model-agnostic and numerically stable framework for general computations. They successfully employed BACC to train neural networks on a cluster of workers, while tolerating a larger number of stragglers. Building on the BACC framework, Soleymani et al. [41] introduced ApproxIFER scheme, as a straggler resistance and Byzantine-robust prediction serving system. However, the scheme of BACC uses a reasonable rational interpolation (Berrut interpolation [56]), off the shelf, for encoding and decoding, without considering any end-to-end cost function to optimize. In contrast, we theoretically formalize a new foundation of coded computing grounded in learning theory, which can be naturally used for machine learning applications.

In this paper, we developed a new foundation for coded computing based on learning theory, contrasting with existing works that rely on coding theory and use metrics like minimum distance and recovery threshold for design. This shift in foundations removes barriers to using coded computing for machine learning applications, allows us to design optimal encoding and decoding functions, and achieves convergence rates that outperform the state of the art. Moreover, the experimental evaluations validate the theoretical guarantees. While this work focuses on straggler mitigation, future work will extend our proposed scheme to achieve Byzantine robustness and privacy, offering promising avenues for further research.