Transfer Learning with Reconstruction Loss

In this paper, we propose a novel transfer learning approach to explicitly enforce the learning of transferable features within a specific location in the neural network computation flow¹¹1The code for this paper is available at:
https://github.com/willtop/Transfer_Learning_with_Reconstruction_Loss. Firstly, we establish the following concept for transfer learning:

Common Information²²2The term “common information” has also been used in information theory as a similarity measurement between two correlated random variables [20], which is not to be confused with the definition in this paper.: The information required for specifying and for solving both the source task and the target task.

Although common information may be difficult to identify in general, in this paper, we make a key observation that when the source task and the target task share the same input, the problem input itself always forms a (possibly non-strict) super-set of the common information. Therefore, in this paper, we consider transfer learning for problems that share the same input distribution and propose to use the problem input as a general choice for common information. Moreover, this paper also goes beyond using the input as common information and further considers applications where we can extract specific common information with lower dimensional representations. In this case, we can further adopt the proposed approach for these specific representations.

With the concept of common information established, the proposed transfer learning approach can be described as follows. When training the neural network model on the source task, besides the task-specific loss, we introduce an additional reconstruction loss to be minimized jointly: i.e., we let the neural network reconstruct the common information using features from a specific hidden layer (referred to as the feature layer) and compute the common information reconstruction loss. Through minimizing this reconstruction loss, we encourage the features learned at the feature layer to be informative about all the correlated tasks that take the same input distribution. To perform transfer learning on a target task, we fix the trained model parameters up to this feature layer as the already-optimized feature learning stage for the model and further train the remaining model parameters on the target task.

When the proposed approach utilizes a choice of common information that is generic (e.g. the problem inputs), the features learned in the model can be used for all target tasks that have the same input. Essentially, the proposed transfer learning approach is target-task agnostic. This is in contrast to some of the prior transfer learning works [21, 22, 17, 23] where the training approach is dedicated to a given source-task and target-task pair. We note that several works also explore the similar idea of encouraging input reconstruction from the model’s internal features, e.g., in the field of semi-supervised learning [24, 25], multi-task learning [26], or domain adaptation (a sub-category of transfer learning where the input distribution changes from the source to the target domain) [22, 27, 28, 29, 30, 31]. Nonetheless, these works deal with different problem setups: in semi-supervised learning, the focus is on obtaining quality features and latent representations of the inputs with only limited ground-truth labels; in multi-task learning, the model is trained on multiple task simultaneously using ample training data; in domain adaptation, the task stays the same while the input distribution changes. While the input reconstruction technique in these works all aims to encourage the neural network models to extract salient features, the present paper differs in that we focus on transfer learning and specifically on quick adaptation to new unseen tasks, after the initial training stage, with only limited further adjustments of model parameters. Moreover, this paper goes one step further in expanding the possibilities for the reconstruction target beyond the problem inputs, through introducing the concept of common information.

For numerical simulations, we first demonstrate the proposed approach through a classical machine learning application: the MNIST handwritten digits classification [32]. We formulate the source task and the target task as correlated classification tasks and adopt the fully-connected neural networks. Furthermore, we treat several important and challenging classes of mathematical optimization problems in wireless communications. Although a few works have explored transfer learning on certain wireless communication problems [33, 34, 23], the techniques used are specific to the application settings or objective characteristics. On the other hand, our approach is more general and readily adapts to different problems or objectives. We illustrate this by experimenting over two application scenarios: the power control utility optimization for device-to-device (D2D) wireless networks, as well as the downlink beamforming and localization problems for multiple-input-single-output (MISO) wireless networks. Specifically, for each of the application scenario, we explore transfer learning between a pair of distinct yet correlated objectives: the min rate and sum rate objectives in the D2D networks; and the beamforming gain and the localization accuracy in the MISO networks. Optimization results suggest that the proposed approach achieves better knowledge transfer and mitigates over-fitting on limited target task data more effectively than the conventional transfer learning method.

The remaining of the paper is organized as follows. Section II formulates the general transfer learning problem and establishes the concept of common information. Section III introduces the three applications to be studied in details: the MNIST classification problem, the D2D network power control problems, and the downlink MISO network beamforming and localization problems. Section IV proposes the novel transfer learning approach for general neural networks, along with the proper selections of the common information for the three applications. The performances of the proposed method over three applications are presented and analyzed in Section V. Finally, conclusions are drawn in Section VI.

For mathematical symbols, we use lower-case letters for scalar variables, lower-case bold-face letters for vector variables, and upper-case bold-face letters for matrix variables. By default, we regard any vector as a single-column matrix. We use the superscript $(\cdot)^{\text{H}}$ to denote the Hermitian transpose of a matrix (or a vector regarded as a single-column matrix). We use $\mathcal{CN}(0,\Sigma)$ to denote the zero-mean circularly-symmetric complex normal distribution, with $\Sigma$ being the covariance matrix. We use $[\cdot]_{k}$ to denote the $k$-th element of a vector. We use $|\cdot|$ to denote the absolute value of a complex number, and $||\cdot||_{2}$ to denote the $L^{2}$ norm of a vector. Lastly, we use the operator $\leftarrow$ to denote the assignment to a specified variable.

Among many variants of transfer learning formulations, we focus on the transfer learning setting where the source task and the target task share the same input distribution but differ in their respective objectives. Let $\mathcal{S}$ denote the source task and $\mathcal{T}$ denote the target task. Consider the optimization problems summarized by the following components:

• •

  Input parameters $\mathbf{p}$ summarizing all environment information essential for optimization, which follow the same distribution in both $\mathcal{S}$ and $\mathcal{T}$;

• •

  Optimization variables for $\mathcal{S}$: $\mathbf{x}_{s}$;

• •

  Objective (or utility) for $\mathcal{S}$: $u_{s}(\mathbf{x}_{s})$;

• •

  Optimization variables for $\mathcal{T}$: $\mathbf{x}_{t}$;

• •

  Objective (or utility) for $\mathcal{T}$: $u_{t}(\mathbf{x}_{t})$.

If $\mathcal{T}$ and $\mathcal{S}$ are supervised learning tasks, $u_{s}(\mathbf{x}_{s})$ and $u_{t}(\mathbf{x}_{t})$ are also dependent on the ground-truth labels, which we omit in our notations as they are not variables to be optimized.

To optimize for $\mathcal{S}$ and $\mathcal{T}$, we utilize neural networks to compute the optimal values of the optimization variables. Specifically:

• •

  Let $\mathcal{F}_{\Theta_{s}}$ with trainable parameters $\Theta_{s}$ denote the neural network mapping for optimizing $\mathcal{S}$:

  $\displaystyle\mathcal{F}_{\Theta_{s}}(\mathbf{p})=\mathbf{x}_{s}\>.$

  (1)

• •

  Let $\mathcal{F}_{\Theta_{t}}$ with trainable parameters $\Theta_{t}$ denote the neural network mapping for optimizing $\mathcal{T}$:

  $\displaystyle\mathcal{F}_{\Theta_{t}}(\mathbf{p})=\mathbf{x}_{t}\>.$

  (2)

Under this transfer learning setting, $\mathcal{S}$ and $\mathcal{T}$ are correlated in the sense that for optimizing $u_{s}(\mathbf{x_{s}})$ and $u_{t}(\mathbf{x_{t}})$, the information extracted from $\mathbf{p}$ should be similar. To mathematically formalize this concept, we introduce the concept of common information, which we denote by $\mathcal{I}(\mathbf{p})$, with the following definition:

Essentially, the common information is the information required by both $\mathcal{S}$ and $\mathcal{T}$. We note that for all ($\mathcal{S}$, $\mathcal{T}$) pairs, there exists at least one choice of the common information, which is the problem input $\mathbf{p}$ itself.

Despite the correlations, $\mathcal{S}$ and $\mathcal{T}$ are still two different tasks. For a neural network model exclusively trained on one task, its learned features are not optimal for the other one: some features learned from optimizing $\mathcal{S}$ could be totally irrelevant or even counter-productive for optimizing $\mathcal{T}$. For successful transfer learning between $\mathcal{S}$ and $\mathcal{T}$, the features need to be general, i.e., it should contain more knowledge than required for optimizing just a single task.

Moreover, due to reasons such as the cost and overhead of data acquisition, the data is assumed to be highly limited for training $\Theta_{t}$ on $\mathcal{T}$. This assumption is particularly relevant for scenarios where the target task $\mathcal{T}$ is adapted on the fly after the model training and deployment. To effectively learn $\mathcal{F}_{\Theta_{t}}$ with limited data, we need to utilize the correlation between $\mathcal{S}$ and $\mathcal{T}$ and transfer over the knowledge already learned in $\mathcal{F}_{\Theta_{s}}$ which is also useful for solving $\mathcal{T}$. In essence, features and representations computed both in $\mathcal{F}_{\Theta_{s}}$ and in $\mathcal{F}_{\Theta_{t}}$ should contain the common information $\mathcal{I}(\mathbf{p})$ as defined in Section II-B.

MNIST handwritten digits classification [32] is one classical problem explored by many machine learning algorithms. Specifically, the MNIST dataset consists of images of handwritings on single digits from 0 to 9. The original classification task requires predicting the correct digit from each handwriting image.

We formulate the following transfer learning problem based on this original MNIST digit classification problem. We modify the classification objectives to obtain a source task and a target task. Specifically, let $\mathcal{S}$ be the task of identifying whether the input image represents the digit 1 or not; and let $\mathcal{T}$ be the task of identifying whether the input image represents the digit 8 or not. We select the input images (which are the problem inputs $\mathbf{p}$) to only include handwritings representing 0, 1, or 8, such that for both $\mathcal{S}$ and $\mathcal{T}$, the positive and negative samples are relatively well balanced with a ratio around 1-to-2.

The reason for the specific number selections is as follows: as described in Section II-C, $\mathcal{S}$ and $\mathcal{T}$ should be similar and at the meantime with significant differences. Both tasks of identifying the number 1 and identifying the number 8 require the common information of the complete pixel patterns such as edges and corners over the entire input image. Nonetheless, the difference is also apparent, as handwritings of the digit 1 should resemble a linear pattern of pixels throughout the image, while handwrittings of the digit 8 should resemble local circular patterns of pixels. As the defined input set only includes handwrittings of 0, 1, and 8, a machine learning model exclusively trained for solving $\mathcal{S}$ only needs to discover and rely on a simple high-level feature: whether the pixels form a global linear pattern or not. Needless to say, the knowledge in this model would not transfer well to solving $\mathcal{T}$, as the learned feature is not informative for distinguishing handwrittings of 8 from 0. Therefore, an effective transfer learning approach is needed for obtaining transferable features from $\mathcal{S}$ to $\mathcal{T}$.

Corresponding to the notions in Section II-A, $\mathbf{x}_{s}$ and $\mathbf{x}_{t}$ are the neural network predictions on the probabilities that the handwriting input $\mathbf{p}$ represents the respective digit specified by the corresponding task. Specifically, with $\mathcal{S}$ and $\mathcal{T}$ both being binary classification tasks, we have $\mathbf{x}_{s}\in[0,1]$ being the probability of $\mathbf{p}$ representing the digit 1; and $\mathbf{x}_{t}\in[0,1]$ being the probability of $\mathbf{p}$ representing the digit 8. Let $\mathbf{p}_{\text{digit}}$ be the actual digit the handwriting $\mathbf{p}$ represents. The source task and target task objectives are respectively:

The definitions in (7) and (8) represent the classification accuracy, i.e., the average values of $u_{s}(\mathbf{x}_{s})$ and $u_{t}(\mathbf{x}_{t})$ over a set of samples are the percentages of correct predictions across that sample set.

For the first wireless communication application, we study the power control problem for D2D wireless networks. Consider a wireless network with $N$ D2D links that transmit independently over the frequency band of bandwidth $w$ with full frequency reuse. Let $\mathbf{G}=\{g_{ij}\}_{i,j\in\{1\dots N\}}$ denote the set of channel gains, with $g_{ij}$ being the channel gain from the $j$-th transmitter to the $i$-th receiver. The power control problem takes the wireless channel state information as the input, so we have:

Let $P_{i}$ denote the maximum transmission power for the $i$-th transmitter. The problem of power control is to find the optimal values of the variables $\mathbf{x}=\{x_{i}\}_{i\in\{1\dots N\}}$, where $x_{i}\in[0,1]$ denotes the percentage of maximum power that the $i$-th transmitter transmits at. Under a specific power control solution $\mathbf{x}$, the $i$-th link has the following achievable rate:

where $\sigma^{2}$ denotes the background noise power level.

For the transfer learning problem, we consider two link rate utility functions that are important under different application scenarios:

• •

  The sum rate optimization as $\mathcal{S}$:

  $\displaystyle u_{s}(\mathbf{x}_{s})=\sum_{i=1}^{N}r_{i}$

  (11)

• •

  The min rate optimization as $\mathcal{T}$:

  $\displaystyle u_{t}(\mathbf{x}_{t})=\min_{i=1\dots N}r_{i}$

  (12)

where $\mathbf{x}_{s}$ and $\mathbf{x}_{t}$ are the power control solutions for the sum-rate optimization $\mathcal{S}$ and the min-rate optimization $\mathcal{T}$ respectively.

Examining (11) and (12), they are correlated in the sense that a set of higher rates over all links leads to a higher objective value. Both objectives would benefit from proper interference mitigation. On the other hand, these two objectives differ significantly in term of fairness among links: (12) ensures complete fairness by optimizing the worst link rate; (11) largely ignores fairness since the optimal sum rate might be achieved through heavily utilizing strong links. With this distinction, conducting transfer learning between (11) and (12) is challenging, as certain features crucial in optimizing (11) could potentially lead to degraded performance in (12).

As a second application in wireless communications, we study the transfer learning from the downlink MISO beamforming task to the localization task. The source task aims to design the optimal downlink beamformers based on the uplink received pilots, while the target task aims to find the locations of the users. These two seemingly unrelated tasks nevertheless share the common input as the estimated channel state information.

Consider a downlink MISO network of $M$ base stations (BS) collectively serving a single user equipement (UE), where each BS is equipped with $K$ antennas while the UE is equipped with a single antenna. We assume that the locations of all BSs are fixed. We denote the UE location by the coordinate $(x_{\text{UE}},y_{\text{UE}},z_{\text{UE}})$, with $x_{\text{UE}}$ and $y_{\text{UE}}$ being unknown and $z_{\text{UE}}$ being fixed, which is usually the case in practical indoor scenarios where the UE is located on the ground or at a certain level of known height. We assume reciprocity of uplink and downlink channels, with the set of channel coefficients denoted by $\mathbf{H}=\{\mathbf{h}_{m}\}_{m\in\{1\dots M\}}$, where $\mathbf{h}_{m}\in\mathbb{C}^{K}$ is the vector of channel coefficients from the $K$ antennas of the $m$-th BS to the UE. We use the Rician fading model for modeling the wireless channels, with the channels from the $m$-th BS to the UE modelled as follows:

with $d_{m}$ being the distance from the UE to the $m$-th BS and $\rho(d_{m})$ being the associated pathloss as a function over this distance, $\mathbf{h}_{m}^{\text{LOS}}\in\mathbb{C}^{K}$ and $\mathbf{h}_{m}^{\text{NLOS}}\in\mathbb{C}^{K}$ being the channel coefficients for the line-of-sight (LOS) path and non-line-of-sight (NLOS) paths respectively, and $\epsilon$ being the Rician factor as the ratio of power between the LOS and NLOS channel components. Specifically, we model the LOS channel component as:

where $\theta_{m}$ and $\phi_{m}$ are the azimuth and elevation angle-of-arrival (AoA) respectively from the UE to the $m$-th BS, and $\mathbf{a}_{m}\in\mathbb{C}^{K}$ is the steering vector, which is a function of $\theta_{m}$ and $\phi_{m}$. Let $\delta$ and $\lambda$ denote the antenna spacing and the signal wavelength respectively, the $k$-th component of $\mathbf{a}_{m}$, which corresponds to the $k$-th antenna of the $m$-th BS, is computed as follows:

with $i_{r}(n)$ and $i_{c}(n)$ being the row and column index of the $n$-th antenna in the antenna array respectively. As for the NLOS paths, we use the Rayleigh fading model to model these path components collectively:

We assume a maximum transmission power level of $P_{\text{BS}}$ for each BS and $P_{\text{UE}}$ for the UE, and a noise level of $\sigma^{2}$ both at the BSs and at the UE. We use an uplink pilot stage for the BS to infer the wireless channel, thereby allowing it to perform downlink beamforming or localization tasks. Specifically, the UE sends $T$ uplink pilot signals, which each BS measures through applying a sequence of sensing vectors over its $K$ antennas. Let $\mathbf{c}=\{c^{t}\}_{t\in\{1\dots T\}}$ be the set of uplink pilot signals sent by UE, and $\mathbf{V}=\{\mathbf{v}_{m}^{t}\}_{m\in\{1\dots M\},\ t\in\{1\dots T\}}$ be the set of sensing vectors employed by all the BSs, where $\mathbf{v}_{m}^{t}\in\mathbb{C}^{K}$ is the sensing vector adopted by the $m$-th BS for receiving the $t$-th pilot. According to the power constraint, we have $|c^{t}|^{2}=P_{\text{UE}},\forall t$. Also, we require $||\mathbf{v}_{m}^{t}||_{2}^{2}=1,\forall m,t$. For the $t$-th uplink pilot the UE transmits, the $m$-th BS receives the measurement $r_{m}^{t}\in\mathbb{C}$:

where $n_{m}^{t}\sim\mathcal{CN}(0,\sigma^{2}\mathbf{I})$ is the noise experienced at the $m$-th BS when receiving the $t$-th uplink pilot transmission through $\mathbf{v}_{m}^{t}$.

We study the transfer learning problem between two tasks: the downlink beamforming task $\mathcal{S}$, and the UE localization task $\mathcal{T}$, both of which are based on the uplink pilot measurements as the input. Specifically, we assume both the UE uplink pilots $\mathbf{c}=\{c^{t}\}_{t\in\{1\dots T\}}$ and the sensing vectors adopted at each BS $\mathbf{V}=\{\mathbf{v}_{m}^{t}\}_{m\in\{1\dots M\},\ t\in\{1\dots T\}}$ are fixed. Correspondingly, we have the problem inputs as:

We now provide detailed descriptions of the source task and the target task as follows.

A neural network consists of consecutive hidden layers of neurons computing non-linear functions (i.e. activations), forming a computation flow. For a regular neural network learning one specific input-to-output mapping, the features computed at each hidden layer follow a general information flow pattern: from the input to the output, the amount of information describing the input gradually reduces layer by layer, while only the information necessary for predicting the output is maintained [35, 36, 37]. While being efficient for learning a single mapping, this information flow pattern may not be desirable in the transfer learning setting. Instead, we desire the features learned by the model to be generalizable and retain sufficient information for being transferable to new target tasks. In other words, the learned model features need to contain the common information among tasks of interest in order to be effective for transfer learning.

For transfer learning on CV or NLP applications, it is relatively clear which features are likely to hold the common information. Specifically, with highly structured inputs, the neural network computation flow learned under regular training is likely to have structures already: the entire flow can be divided into a feature learning stage and an optimization stage. Take convolutional neural networks solving CV tasks as examples, the feature learning stage includes the convolution layers that compute general high-level features which contain the common information for many relevant CV tasks (such as edge patterns, pixel intensities, or color gradients over the input image), followed by the optimization stage consisting of fully connected layers that process these high-level features and compute the task-specific outputs (e.g. classification class scores). To conduct transfer learning over correlated CV tasks, the convolutional layers are shared among the models as the feature learning stage, while the fully connected layers of each model are further trained on a per-task basis [12, 14, 15].

However, for general mathematical optimization problems, the inputs lack clear structures in most cases. Under regular training methods, the resulting neural network computation flow is not clearly divided by stages, with internal features gradually becoming more and more task-specific layer-by-layer. As the result, it is difficult to identify or explicitly encourage the learning of transferable features, or equivalently, to obtain features that incorporate the desired common information.

To tackle the challenges of transfer learning for mathematical optimization problems, we propose a novel transfer learning approach to encourage the learning of transferable features. Specifically, we first identify a proper selection of the common information for the source task and the target task. When training the neural network model on the source task, on top of the regular task-based loss, we introduce in addition a loss term for the common information reconstruction.

Fig. 1 illustrates the proposed transfer learning approach. We adopt the most general fully-connected neural network architecture. Within the neural network, we select a hidden layer as the feature layer where we encourage the transferable features to be computed. Correspondingly, the part of the neural network computation flow from the input layer to the feature layer forms the feature learning stage, while the part of the computation flow from the feature layer to the output layer forms the optimization stage. There is no particular constraint on which hidden layer to select as the feature layer, as long as there is sufficient transformation capacity (by having a sufficient number of hidden layers) both from the input layer to the selected layer and from the selected layer to the output layer.

From the selected feature layer, we add in a reconstruction stage as a separate branch in the computation flow, in parallel to the optimization stage. The reconstruction stage aims to reconstruct the common information $\mathcal{I}(\mathbf{p})$ by using the features in the feature layer. We denote the corresponding reconstruction loss by $\mathcal{L}_{R}$. Together with the original loss associated with optimizing the source task utility $u_{s}(\mathbf{x}_{s})$, which we denote by $\mathcal{L}_{S}$, the loss function $\mathcal{L}$ that we use to train $\mathcal{F}_{\Theta_{s}}$ is:

where $\alpha$ is the relative weighting scalar between the loss terms.

With $\mathcal{F}_{\Theta_{s}}$ trained on $\mathcal{L}$ as in (24), the feature layer computes features that are pertinent to the source task optimization, while also containing knowledge of the common information $\mathcal{I}(\mathbf{p})$. Therefore, these features are highly general and transferable to different tasks that are relevant to the source task. We emphasize that when the choice of $\mathcal{I}(\mathbf{p})$ is generic (e.g., by choosing the problem inputs $\mathbf{p}$ as the common information), the proposed approach is target-task agnostic, since no prior knowledge of the target task is needed throughout the training procedure.

After training the neural network $\mathcal{F}_{\Theta_{s}}$ on the source task $\mathcal{S}$ as described in Section IV-B, transfer learning on the target task $\mathcal{T}$ is straightforward. We first transfer the subset of the neural network parameters $\Theta_{s}$ up to the feature layer as the shared feature learning stage to the neural network parameters $\Theta_{t}$, and leave the remaining parameters in $\Theta_{t}$ unassigned. When training $\mathcal{F}_{\Theta_{t}}$, we freeze all the transferred parameters so they would remain unchanged as the already-optimized feature learning stage. With the proposed transfer learning approach, through the transferred feature learning stage, the features computed in the feature layer of $\mathcal{F}_{\Theta_{t}}$ are already valuable for optimizing $\mathcal{T}$ before starting any target task training.

For the parameters after the feature layer in $\Theta_{t}$, we train these parameters with the regular loss associated with the target task utility $u_{t}(\mathbf{x}_{t})$, which we denote as $\mathcal{L_{T}}$. This further training leads to the task-specific optimization stage in $\Theta_{t}$. With the number of trainable parameters greatly reduced, along with the fact that the features from the feature layer are also optimized for $\mathcal{T}$, only a small amount of additional training data is needed to obtain a well-performing model $\mathcal{F}_{\Theta_{t}}$. In essence, the computational complexity of target-task training is significantly reduced when compared to the conventional transfer learning approach.

We emphasize that the selection of $\mathcal{L_{T}}$ is not required during the source task training phase specified in Section IV-B. Instead, we only need to decide the proper $\mathcal{T}$ after the target task $\mathcal{T}$ occurs. Therefore, the specification on $\mathcal{L_{T}}$ does not affect the proposed approach on being target-task agnostic.

For the actual testing or implementations on either $\mathcal{S}$ or $\mathcal{T}$, only the feature learning stage and the corresponding optimization stage of the neural network model need to be executed. Therefore, the proposed approach does not introduce any additional parameters or computational complexity into the model for actual inference. In fact, our approach only increases the model complexity during the source task training. Since in practical scenarios, as the initial stage of the model development, training the model on $\mathcal{S}$ is usually done offline with sufficient time and data resources, the extra complexity from our approach is likely to be negligible.

According to Definition Definition, selecting the proper common information for a given set of tasks is non-trivial. Fortunately, as mentioned in Section II-B, we can always resort to using the problem inputs $\mathbf{p}$ as the common information when no alternative selection is apparent (which also enables the proposed approach to be target task agnostic, as previously discussed). In the following, we examine the three applications for numerical simulations and propose a proper choice for the common information for each of the applications.