SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Po-han  $\mathbf{L}\mathbf{i}^{1*}$  Sravan Kumar Ankireddy $^{1*}$  Ruihan Zhao $^{1}$  Hossein Nourkhiz Mahjoub $^{2}$
Ehsan Moradi-Pari² Ufuk Topcu¹ Sandeep Chinchali¹ Hyeji Kim¹
<sup>1</sup>The University of Texas at Austin <sup>2</sup>Honda Research Institute USA
Efficient compression of correlated data is essential to minimize communication overload in multi-sensor networks. In such networks, each sensor independently compresses the data and transmits them to a central node. A decoder at the central node decompresses and passes the data to a pre-trained machine learning-based task model to generate the final output. Due to limited communication bandwidth, it is important for the compressor to learn only the features that are relevant to the task. Additionally, the final performance depends heavily on the total available bandwidth. In practice, it is common to encounter varying availability in bandwidth. Since higher bandwidth results in better performance, it is essential for the compressor to dynamically take advantage of the maximum available bandwidth at any instant. In this work, we propose a novel distributed compression framework composed of independent encoders and a joint decoder, which we call neural distributed principal component analysis (NDPCA). NDPCA flexibly compresses data from multiple sources to any available bandwidth with a single model, reducing compute and storage overhead. NDPCA achieves this by learning low-rank task representations and efficiently distributing bandwidth among sensors, thus providing a graceful trade-off between performance and bandwidth. Experiments show that NDPCA improves the success rate of multi-view robotic arm manipulation by  $9\%$  and the accuracy of object detection tasks on satellite imagery by  $14\%$  compared to an autoencoder with uniform bandwidth allocation.[2]
Efficient data compression plays a pivotal role in multi-sensor networks to minimize communication overload. Due to the limited communication bandwidth of such networks, it is often impractical to transmit all sensor data to a central server, and compression of the data is necessary. It is important for the sensors to compress the respective data independently, to avoid communication overload in the network. Information theory literature refers to this setting as distributed source coding [1], where the goal is to recover the original data with minimal distortion. In many cases, the data collected by the sensors is only processed by a downstream task model, e.g., an object detection model, but not by humans, and hence the original distributed source coding goal of minimizing reconstruction error is no longer applicable. Instead, the goal should be to maximize the performance of the downstream task model. Additionally, in practice, data collected by multi-sensor networks is often correlated e.g. stereo cameras with overlapping fields of view. To improve communication efficiency, it is
* Equal contribution; order decided randomly. Correspondence to {pohanli, sravan.ankireddy} @utexas.edu.  
2 https://github.com/UTAustin-SwarmLab/Task-aware-Distributed-Source-Coding
Figure 1: Task-aware distributed source coding with NDPCA.  $X_{1}, \ldots, X_{k}$  are correlated data sources. Encoders  $E_{1}, \ldots, E_{k}$  independently compress data to latent representations  $Z_{1}, \ldots, Z_{k}$ . Using linear matrices, the DPCA module projects the representations to any lower dimension at the encoder and projects them back to the original data space at the decoder, which allocates the bandwidth of sources based on the importance of the task  $\Phi$ . The goal is to find the optimal encoders and decoder that minimize the final task loss.
important for the compression framework to take advantage of the correlation and avoid transmission of redundant data. Combining both objectives, the final goal of the distributed compression framework is to learn relevant features that maximize the task performance, while avoiding the transmission of redundant features by exploiting the correlation between sources. Together, we refer to the distributed compression of task-relevant features as task-aware distributed source coding.
However, existing compression methods fail to combine the following three aspects: 1. Existing distributed compression methods perform poorly in the presence of a task model. Although neural networks have been shown to be capable of compressing stereo images [2,3] and correlated images [4], existing methods focus on reconstructing image data, but not for downstream tasks. 2. Existing task-aware compression methods cannot take advantage of the correlation of sources. Previous works only consider compressing task-relevant features of single source [5-9], but not multiple correlated sources. 3. All existing methods for 1 & 2, especially those based on neural networks, only compress data to a fixed level of compression but not to multiple levels. Thus, they cannot operate in environments with different demands of compression levels and require a separate model trained for each compression level. Here, we note that we use the term bandwidth to indicate the information bottleneck in the dimension of transmitted data. Based on the choice of quantization, it is straightforward to convert the latent dimension to other popular metrics such as bits per pixel (bpp) in the case of image sources. Additionally, we consider the scenario of total bandwidth constraint for the uplink, which is typical for wireless networks [10].
We design neural distributed principal component analysis (NDPCA)-a distributed compression framework that can transmit task-relevant features at multiple compression levels. We consider the case where a task model at the central node requires data from all sources and the bandwidth in the network is not consistent over time, as shown in Fig. In NDPCA, neural encoders  $E_{1}, E_{2}, \ldots, E_{K}$  first independently compress correlated data  $X_{1}, X_{2}, \ldots, X_{K}$  to latent representations  $Z_{1}, Z_{2}, \ldots, Z_{K}$ . A proposed module called distributed principal component analysis (DPCA) further compresses these representations to any lower dimension according to the current bandwidth and decompresses the data at the central node. Finally, a neural decoder at the central node decodes the representations  $\hat{Z}_{1}, \hat{Z}_{2}, \ldots, \hat{Z}_{k}$  to  $\hat{X}_{1}, \hat{X}_{2}, \ldots, \hat{X}_{K}$  and feeds them into a task. Task-aware compression aims to minimize task loss, defined as the difference in task outputs with and without compression, such as the difference in object detection results. Due to the significant training cost involved, we avoid training the task model, which is usually a large pre-trained neural network.
To highlight our proposed method, NDPCA learns task-relevant representations with a single model at multiple compression levels. It includes a neural autoencoder to generate uncorrelated task-relevant representations in a fixed dimension. It is desirable to learn uncorrelated representations to prevent the transmission of redundant information. It also includes a module for linear projection, DPCA, to allocate the available bandwidth among sources based on the importance of the task, by observing the respective principal components, and then further compressing the representations to any desired dimension. By harmoniously combining the neural autoencoder and the linear DPCA module, NDPCA generates representations that are more compressible in limited bandwidths, providing a graceful trade-off between performance and bandwidth.
Contributions: Our contributions are three-fold: First, we formulate a task-aware distributed source coding problem that optimizes the compression for a given task instead of reconstructing the sources (Sec. 2). Second, we provide a theoretical justification for the framework by analyzing the case of a linear compressor and a linear task (Sec. 3). Finally, we propose a task-aware distributed source coding framework, NDPCA, that learns a single model for different levels of compression to handle any type of source and task (Sec. 4). We validate NDPCA with tasks of CIFAR-10 image denoising, multi-view robotic arm manipulation, and object detection of satellite imagery (Sec. 5). NDPCA results in a  $1.2\mathrm{dB}$  increase in PSNR, a  $9\%$  increase in success rate, and a  $14\%$  increase in accuracy compared to an autoencoder with uniform bandwidth allocation, for the respective experiments mentioned above.
We now define the problem statement more formally. Consider a set of  $K$  correlated sources. Let  $x_{i} \in \mathbb{R}^{n_{i}}$  denote the sample from source  $i$  where  $i \in \{1,2,\dots,K\}$ . Samples from each source  $i$  are compressed independently by encoder  $E_{i}$  to a latent representation  $z_{i} \in \mathbb{R}^{m_{i}}$  such that  $\sum_{i=1}^{K} m_{i} = m$ , where  $m$  is the total bandwidth available. A joint decoder  $D$  receives the representations  $\{z_{1},z_{2},\ldots,z_{k}\}$  and reconstructs the sources  $\{\hat{x}_{1},\hat{x}_{2},\ldots,\hat{x}_{k}\}$ . In the setting without a task, the goal is to find a set of encoders and a decoder to recover the inputs  $\{x_{1},x_{2},\ldots,x_{k}\}$  with minimal loss:
$$
\underset {E _ {1}, E _ {2}, \dots , E _ {k}, D} {\operatorname {a r g m i n}} \sum_ {i = 1} ^ {K} \mathcal {L} _ {\mathrm {r e c}} \left(x _ {i}, \hat {x} _ {i}\right) \quad (T a k - a g n o s t i c d i s t r i b u t e d s o u r c e c o d i n g), \tag {1}
$$
where  $\mathcal{L}_{\mathrm{rec}}$  is the reconstruction loss, e.g., the mean-squared error loss.
In the presence of a task  $\Phi$ , it takes the reconstructed inputs to compute the final output  $\Phi(\hat{x}_1, \hat{x}_2, \dots, \hat{x}_k)$ . The goal is to find a set of encoders and a decoder such that the task loss  $\mathcal{L}_{\mathrm{task}}$  is minimized, where  $\Phi(x_1, x_2, \dots, x_k)$  is the task output computed without compression. We refer to this problem as task-aware distributed source coding, which is the main focus of this paper:
$$
\underset {E _ {1}, E _ {2}, \dots , E _ {k}, D} {\operatorname {a r g m i n}} \mathcal {L} _ {\text {t a s k}} \left(\Phi \left(x _ {1}, x _ {2}, \dots , x _ {k}\right), \Phi \left(D \left(E _ {1} \left(x _ {1}\right), E _ {2} \left(x _ {2}\right), \dots , E _ {k} \left(x _ {k}\right)\right)\right)\right) \tag {2}
$$
(Task-aware distributed source coding),
where  $\mathcal{L}_{\mathrm{task}}$  is the task loss, e.g., the difference of bounding boxes when the task is object detection.
Bandwidth allocation: In the previous formulations, we assume that the output dimensions of encoders are known a priori. However, the dimensions are related to the compression quality of each encoder, which is also a design factor. That is, given the total available bandwidth  $m$ , we first need to obtain the optimal  $m_{i}$  for each source  $i$ , then, we can design the optimal encoders and decoder accordingly. Finding the optimal set of bandwidths for a given task is a long-standing open problem in information theory [11], even for the simple task of a modulo-two sum of two binary sources [12]. Also, existing works [4, 13, 14] largely assume a fixed latent dimension for sources and train different models for different total available bandwidth  $m$ , which is, of course, suboptimal. In this paper, our framework provides heuristics to the underlying key challenge of optimally allocating available bandwidth, i.e., deciding  $m_{i}$ , while adapting to different total bandwidths  $m$  with a single model.
We start with a motivating example of task-aware distributed source coding under the constraint of linear encoders, a decoder, and a linear task. We first solve the linear setting using our proposed method, distributed principal component analysis (DPCA). We then describe how DPCA compresses data to different bandwidths and analyze the performance of DPCA. In this way, we gain insights into combining DPCA with neural autoencoders in later Sec.4
DPCA Formulation: We consider a linear task for two sources, defined by the task matrix  $\Phi \in \mathbb{R}^{p\times (n_1 + n_2)}$ , where the sources  $x_{1}\in \mathbb{R}^{n_{1}}$  and  $x_{2}\in \mathbb{R}^{n_{2}}$  are of dimensions  $n_1$  and  $n_2$ , respectively, and the task output is given by  $y = \Phi x\in \mathbb{R}^p$ , where  $x = [x_1^\top ,x_2^\top ]^\top$ . Without loss of generality, we assume the sources to be zero-mean. Now, we have  $N$  observations of two sources  $X_{1}\in \mathbb{R}^{n_{1}\times N}$  and  $X_{2}\in \mathbb{R}^{n_{2}\times N}$  and their corresponding task outputs  $Y = \Phi (X)\in \mathbb{R}^{p\times N}$ , where  $X = [X_{1}^{\top}X_{2}^{\top}]^{\top}$ .
We aim to design the optimal linear encoding matrices (encoders)  $E_{1} \in \mathbb{R}^{m_{1} \times n_{1}}$ ,  $E_{2} \in \mathbb{R}^{m_{2} \times n_{2}}$ , and the decoding matrix (decoder)  $D \in \mathbb{R}^{(n_{1} + n_{2}) \times (m_{1} + m_{2})}$  that minimizes the task loss defined as the Frobenius norm of  $\Phi(X) - \Phi(\hat{X})$ , where  $\hat{X}$  is the reconstructed  $X$ . We only consider the non-trivial case where the total bandwidth is less than the task dimension,  $m = m_{1} + m_{2} < p$ , i.e., the encoders cannot directly calculate the task output locally and transmit it to the decoder. For now, we assume that  $m_{1}$  and  $m_{2}$  are given, and we discuss the optimal allocation later in this section.
Letting  $Z_{1} = E_{1}X_{1}\in \mathbb{R}^{m_{1}\times N}$  and  $Z_{2} = E_{2}X_{2}\in \mathbb{R}^{m_{2}\times N}$  denote the encoded representations and  $M = \Phi D$  denote the product of the task and decoder matrices, we solve the optimization problem:
$$
E _ {1} ^ {*}, E _ {2} ^ {*}, M ^ {*} = \underset {E _ {1}, E _ {2}, M} {\operatorname {a r g m i n}} \| Y - M Z \| _ {2} ^ {2} \tag {3a}
$$
$$
\text {s . t .} \quad Z = \left[ \begin{array}{l} Z _ {1} \\ Z _ {2} \end{array} \right] = \left[ \begin{array}{l} E _ {1} X _ {1} \\ E _ {2} \tilde {X} _ {2} \end{array} \right], \tag {3b}
$$
$$
Z Z ^ {\top} = \mathbb {I} _ {m}, \tag {3c}
$$
$$
\hat {Y} = \Phi D Z = M Z, \quad Y = \Phi \left[ \begin{array}{l} X _ {1} \\ X _ {2} \end{array} \right]. \tag {3d}
$$
Note that solving  $M$  is identical to solving the decoder  $D$  since we can always convert  $M$  to  $D$  by the generalized inverse of task  $\Phi$ . The encoders  $E_1$  and  $E_2$  project the data to representations  $Z_1$  and  $Z_2$  in (3b). We constrain the representations to be orthonormal vectors in (3c) as in the normalization in principal component analysis (PCA) for the compression of a single source [15]. This constraint lets us decouple the problem into subproblems later in (5). Finally, in (3d), the decoder  $D$  decodes  $Z_1$  and  $Z_2$  to  $\hat{X}_1$  and  $\hat{X}_2$  and passes the reconstructed data to task  $\Phi$ .
Solution: We now solve the optimization problem in (3). For any given  $E_1, E_2$  (thus, a given  $Z$ ), we can optimally obtain  $M^* = YZ^\top(ZZ^\top)^{-1} = YZ^\top$  by linear regression. Now, we are left to find the optimal encoders  $E_1, E_2$ . First, a preprocessing step removes the correlation part of  $X_1$  from  $X_2$  by subtracting the least-square estimator  $\hat{X}_2(X_1)$ :
$$
\tilde {X} _ {2} = X _ {2} - \hat {X} _ {2} (X _ {1}) = X _ {2} - X _ {2} X _ {1} ^ {\top} (X _ {1} X _ {1} ^ {\top}) ^ {- 1} X _ {1}. \tag {4}
$$
The orthogonality principle of least-square estimators [16 p.386] ensures that  $X_{1}\tilde{X}_{2}^{\top} = \mathbf{0}_{n_{1}\times n_{2}}$ . We decouple the objective in (3a) with respect to  $E_{1},E_{2}$  by the orthogonality principle and (3c):
$$
\min  _ {E 1, E 2} \| Y - M ^ {*} Z \| _ {2} ^ {2} = \| Y \| _ {2} ^ {2} - \max  _ {E _ {1}, E _ {2}} \| M ^ {*} \| _ {2} ^ {2} = \| Y \| _ {2} ^ {2} - \max  _ {E _ {1}} \| Y _ {1} X _ {1} ^ {\top} E _ {1} ^ {\top} \| _ {2} ^ {2} - \max  _ {E _ {2}} \| Y _ {2} \tilde {X} _ {2} ^ {\top} E _ {2} ^ {\top} \| _ {2} ^ {2}, \tag {5}
$$
where  $Y = \Phi X = [\Phi_1\Phi_2][X_1^\top X_2^\top ]^\top = Y_1 + Y_2$ . We then have two subproblems from (3):
$$
E _ {1} ^ {*} = \operatorname * {a r g m a x} _ {E _ {1}} \| \Phi_ {1} X _ {1} X _ {1} ^ {\top} E _ {1} ^ {\top} \| _ {2} ^ {2}
$$
$$
\begin{array}{c c} \text {s . t .} & E _ {1} X _ {1} X _ {1} ^ {\top} E _ {1} ^ {\top} = \mathbb {I} _ {m _ {1}}, \end{array}
$$
$$
E _ {2} ^ {*} = \underset {E _ {2}} {\operatorname {a r g m a x}} \quad \| \Phi_ {2} \tilde {X} _ {2} \tilde {X} _ {2} ^ {\top} E _ {2} ^ {\top} \| _ {2} ^ {2} \tag {7}
$$
$$
\mathrm {s . t .} \quad E _ {2} \tilde {X} _ {2} \tilde {X} _ {2} ^ {\top} E _ {2} ^ {\top} = \mathbb {I} _ {m _ {2}}.
$$
The two subproblems are the canonical correlation analysis [17], which can be solved by whitening  $E_{1}X_{1}, E_{2}\tilde{X}_{2}$  and singular value decomposition (see [17] for details).
Dynamic bandwidth: So far, we solved the case for fixed bandwidths  $m_{1}$  and  $m_{2}$ . We now describe ways to determine the optimal bandwidth allocation given a current total bandwidth  $m$ . To do so, DPCA solves (6) and (7) with  $m_{1} = n_{1}$  and  $m_{2} = n_{2}$  and obtains  $E_{1}^{*}$ ,  $E_{2}^{*}$  and all pairs of canonical directions and correlations. Canonical directions and correlations can be analogized to a more general case of singular vectors and values. Similar to PCA, the sums of squares of canonical correlations are the optimal values of (6) and (7), so DPCA sorts all the canonical correlations in descending order and chooses the first  $m$  pairs of canonical correlations and directions. These canonical correlations determine the optimal encoders  $E_{1}^{*}$ ,  $E_{2}^{*}$  and decoder  $D^{*}$ , which indirectly solves  $m_{1}$  and  $m_{2}$ . Intuitively, the canonical correlations indicate the importance of a direction to the task, and we prioritize the transmission of directions by importance. For simplicity, we only consider the case of 2 sources. DPCA can easily compress more sources simply by constraining all  $Zs$  to be independent and thus decoupling the original problem (3) to more subproblems.
Performance analysis of DPCA: When DPCA compresses new data matrices with encoder  $E_1^*$  and  $E_2^*$ , the preprocessing step is invalid as the encoders cannot communicate with each other. So for
DPCA to perform optimally while skipping the step, the two data matrices need to be uncorrelated, namely,  $\hat{X}_2(X_1) = 0$ , because in such case the preprocessing step removes nothing from the data sources. Given that correlated sources lead to suboptimality of DPCA, we characterize the performance between the joint compression, PCA, and the distributed compression, DPCA, under the same bandwidth in Lemma 3.1 with the simplest case of reconstruction, namely,  $\Phi = \mathbb{I}_p$ . In this setting, the canonical correlation analysis is relaxed to the singular value decomposition, which is later used for NDPCA in Sec.4.
Lemma 3.1 (Bounds of DPCA Reconstruction). Given a zero-mean data matrix and its covariance,
$$
X = \left[ \begin{array}{c} X _ {1} \\ X _ {2} \end{array} \right] \in \mathbb {R} ^ {(n _ {1} + n _ {2}) \times N}, X X ^ {\top} = \underbrace {\left[ \begin{array}{c c} \operatorname {C o v} _ {1 1} & \mathbf {0} \\ \mathbf {0} & \operatorname {C o v} _ {2 2} \end{array} \right]} _ {X _ {\text {d i a g}}} + \underbrace {\left[ \begin{array}{c c} \mathbf {0} & \operatorname {C o v} _ {1 2} \\ \operatorname {C o v} _ {2 1} & \mathbf {0} \end{array} \right]} _ {\Delta X},
$$
assume that  $\Delta X$  is relatively smaller than  $XX^{\top}$ , and  $XX^{\top}$  is positive definite with distinct eigenvalues. For PCA's encoding and decoding matrices  $E_{\mathrm{PCA}}$ ,  $D_{\mathrm{PCA}}$  and DPCA's encoding and decoding matrices  $E_{\mathrm{DPCA}}$ ,  $D_{\mathrm{DPCA}}$ , the difference of the reconstruction losses is bounded by
$$
0 \leq \| X - D _ {\mathrm {D P C A}} E _ {\mathrm {D P C A}} (X) \| _ {2} ^ {2} - \| X - D _ {\mathrm {P C A}} E _ {\mathrm {P C A}} (X) \| _ {2} ^ {2} = - \sum_ {i = m + 1} ^ {n _ {1} + n _ {2}} \lambda_ {i} e _ {i} ^ {\top} \Delta X e _ {i}.
$$
where  $\lambda_{i}$  and  $e_i$  are the  $i$ -th largest eigenvalue and eigenvector of  $XX^{\top}$ ,  $\mathrm{Tr}$  is the trace function, and  $m$  is the dimension of the compression bottleneck.
The proof of Lemma 3.1 is in Appendix A.1. Note that  $\Delta X$  is the correlation of sources, so as  $\| \Delta X \|_F$  gets smaller, the difference of PCA and DPCA is closer to 0. That is, as the covariance decreases, DPCA performs more closely to PCA, which is the optimal joint compression.
To summarize, uncorrelated data matrices  $X_{1},\ldots ,X_{K}$  are desired for DPCA. If so, DPCA optimally decides the bandwidths of all sources based on the canonical correlations, representing their importance for the task. One application of DPCA is that encoders can use the remaining unselected canonical directions to improve compression when the available bandwidth is higher later.
4 Neural Distributed Principal Component Analysis
The theoretical analysis in the previous section indicates that DPCA has two drawbacks: it only compresses data optimally if sources are uncorrelated, and it only works for linear tasks. However, DPCA dynamically allocates bandwidth to sources based on their importance. On the other hand, neural autoencoders are shown to be powerful tools for compressing data to a fixed dimension but cannot dynamically allocate bandwidth. This contrast motivates us to harmoniously combine a neural autoencoder to generate representations and then pass them through DPCA to compress and find the bandwidth allocation. We refer to the combination of a neural autoencoder and DPCA as neural distributed principal component analysis (NDPCA). With a single neural autoencoder and a matrix at each encoder and decoder, NDPCA adapts to any available bandwidth and flexibly allocates bandwidth to sources according to their importance to the task.
Outline: NDPCA has two encoding stages, as shown in Fig.1. First, the neural encoder at each  $k$ -th source encodes data  $X_{k}$  to a fixed-dimensional representation  $Z_{k}$  for  $k \in [K]$ . Then the DPCA linear encoder adapts the dimension of  $Z_{k}$  via linear projection according to the available bandwidth and the correlation among the sources as per 6. Similarly, the decoding of NDPCA is also performed in two stages. First, the DPCA linear decoder reconstructs the  $K$  fixed-dimensional representations  $\hat{Z}_{1}, \hat{Z}_{2}, \ldots, \hat{Z}_{K}$ , based on which the joint neural decoder generates the estimate of data  $\hat{X}_{1}, \hat{X}_{2}, \ldots, \hat{X}_{K}$ . These estimates are then passed to the neural task model  $\Phi$  to obtain the final task output  $\hat{Y}$ . Note that since we have a non-linear task model here, we envision that the neural encoders generate non-linear embedding of the sources, while the DPCA mainly adapts the dimension appropriately as needed; the role of the DPCA here is to reliability reconstruct the embedding  $\hat{Z}$ s, which corresponds to the case described in Lemma 3.1 with the task matrix  $\Phi$  as identity.
Training procedure: During the training of NDPCA, the weights of the task are always frozen because it is usually a large-scale pre-trained model that is expensive to re-train. We aim to learn the
$K$  neural encoders and the joint neural decoder which minimize the loss function:
$$
\mathcal {L} _ {\mathrm {t o t}} = \lambda_ {\text {t a s k}} \underbrace {\left\| \hat {Y} - Y \right\| _ {F} ^ {2}} _ {\text {t a s k l o s s}} + \lambda_ {\text {r e c}} \underbrace {\left(\left\| \hat {X} _ {1} - X _ {1} \right\| _ {F} ^ {2} + \left\| \hat {X} _ {2} - X _ {2} \right\| _ {F} ^ {2} + \dots \left\| \hat {X} _ {K} - X _ {K} \right\| _ {F} ^ {2}\right)} _ {\text {r e c o n s t r u c t i o n l o s s}}. \tag {8}
$$
In the task-aware setting when  $\lambda_{\mathrm{rec}} = 0$ , the neural autoencoder fully restores task-relevant features, which is the main focus of this paper. When  $\lambda_{\mathrm{task}} = 0$ , the neural autoencoder learns to reconstruct the data  $X$ , which is the task-agnostic setting later compared in Sec.5.