Remote Sensing for Water Quality: A Multi-Task, Metadata-Driven Hypernetwork Approach

1 Introduction

Inland water quality has deteriorated globally, primarily due to excessive nutrient loading and global warming, leading to issues such as harmful algal blooms [Hou et al., 2022]. Water quality is a critical factor influencing the health of aquatic ecosystems, the safety of drinking water supplies, and the sustainability of agriculture and industry. Traditional methods for monitoring water quality involve in-situ sampling and laboratory analysis, which are expensive, offer poor spatial and temporal coverage, and are limited by site accessibility. Remote sensing provides large-scale, continuous, and noninvasive observations, motivating research into water quality remote sensing models [Pahlevan et al., 2020]. Effective monitoring aligns with the United Nations Sustainable Development Goals 6 and 14 [United Nations, 2015], aiming for universal access to clean water and promoting the Leave No One Behind principle [United Nations, 2016].

Remote sensing models aim to estimate chlorophyll-a concentration in lakes to use as an indicator of water quality, given multispectral satellite data that represent water colour (Figure 1). These models make pixel-level predictions, based on single-point ground truth samples corresponding to a single satellite pixel. This problem is, therefore, formulated as a regression task with tabular features. Training accurate models is challenging as datasets contain dozens or hundreds of lakes with varied relationships between water colour and chlorophyll-a concentration, due to differences in biological and geological conditions [Neil et al., 2019]. Furthermore, training data is scarce due to expensive ground truth labels, with as few as five data points per lake, motivating dataefficient methods [Graffeuille et al., 2022].

Multi-Task Learning (MTL) approaches can enable learning of lake differences, by modelling lake functions as distinct but related tasks and jointly learning these tasks with knowledge transfer to improve generalisation performance. MTL neural network architectures can be broadly categorised into hard parameter sharing and soft parameter sharing [Ruder, 2017]. Hard parameter sharing, which shares most parameters across tasks, offers parameter efficiency but limited flexibility due to a shared feature representation [Long et al., 2017]. In contrast, soft parameter sharing allows for flexible representations by learning unique weights for each task, constrained by mechanisms like regularization, at the cost of reduced parameter efficiency [Misra et al., 2016]. This highlights a trade-off between parameter-efficient task scaling and learning flexible task networks. As such, MTL approaches struggle to learn problems with many diverse tasks from few samples, such as water quality remote sensing. This necessitates the development of MTL architectures which are both parameter efficient and flexible.

An opportunity to enhance learning for water quality remote sensing is in auxiliary data sources which describe the lakes themselves, such as depth, temperature and weather. These task-level features, which we refer to as metadata, are informative of the environmental processes within lakes, indicating that lakes with similar features are likely to be highly related tasks. For example, a warm, shallow lake is likely to be similar to another lake of similar characteristics, but different from a cold, deep lake [Yang et al., 2022]. When available, metadata may help the model learn betweentask relationships, particularly as lakes have insufficient data to accurately learn task relationships implicitly during training [Zheng et al., 2019; Zamir et al., 2018]. However, existing MTL algorithms cannot leverage task metadata effectively because they are designed to learn from data only at the sample level, whereas metadata is task-level data and is therefore invariant across all samples within a task. This necessitates the development of MTL architectures that can effectively leverage metadata to improve learning.

To address these gaps, we propose Multi-Task Hypernetworks, a novel MTL architecture which learns flexible task networks with substantially fewer parameters per task than other soft parameter sharing approaches, and can effectively leverage task metadata. We achieve this by learning a unique low-dimensional embedding vector of each task, which captures task differences and relationships. Our architecture can generate flexible task networks and enable abstract knowledge transfer using task embeddings as the only task-specific parameters; in our experiments, we use only ten parameters per task. This task-level parameter framework also allows us to naturally inject task metadata into the model.

Our approach accomplishes this by utilising hypernetworks [Ha et al., 2017] for soft parameter knowledge transfer across tasks. Notably, this represents the first attempt within the research area to explore this methodology. Hypernetworks are neural networks that generate the weights of another neural network. Instead of directly learning the weights of a "target" network, which models the task of interest, hypernetworks produce them dynamically based on some input. Our proposed Multi-Task Hypernetwork uses a hypernetwork to generate a different set of weights for multiple target networks which each model a different task. Network weights for all target networks are generated by the same hypernetwork, but use task-specific embeddings as input to the hypernetwork to generate entirely distinct weights (Figure 2). The use of a shared deep hypernetwork enables abstract knowledge transfer between tasks, while the task embeddings capture task differences and relationships. This is a conceptual departure from previous soft parameter sharing methods, which share knowledge at the parameter level. Additionally, by inputting metadata to the shared hypernetwork with the task embedding, the target network weights become a function of their metadata. This allows our architecture to learn task relationships and task functions explicitly from task-level metadata, as well as implicitly by joint task optimisation.

We show that our architecture can improve the accuracy of water quality estimation and other tabular MTL datasets compared to existing approaches, both with and without available metadata. Code and datasets are available at https://github.com/OGraffeuille/Multi-Task-Hypernetworks.

2 RelatedWork

2.1 Multi-Task Learning Architectures

Hard sharing MTL architectures [Caruana, 1997; Lopes et al., 2023] have been adapted to share knowledge in task specific layers using a tensor normal distribution which learns task similarities [Long et al., 2017], to learn task differences in hard sharing layers by learning task masks [Pascal et al., 2021], to reduce generalisation error with an adversarial loss to encourage similar task latent distributions [Shui et al., 2019], and to use architecture search [Guo et al., 2020; Sun et al., 2020]. However, hard sharing approaches have limited flexibility. Some soft parameter sharing architectures [Liu et al., 2019; Sun et al., 2021] instead transfer knowledge across tasks by linearly combining intermediate representations of each task [Misra et al., 2016; Ruder et al., 2019], or using multiple expert networks [Ma et al., 2018]. However, these approaches are parameter-inefficient as they require learning of entire network weights for each task. Transferring knowledge by factorising model weights across tasks can improve parameter efficiency but has limited expressiveness [Yang and Hospedales, 2017]. By instead transferring knowledge through a deep hypernetwork weight generator, our approach can learn abstract task relationships and is substantially more parameter efficient than these architectures as it only learns a small embedding vector for each task.

Multi-task learning with metadata. Leveraging natural language task descriptions has been explored in natural language processing [Ye and Ren, 2021; You et al., 2016] and reinforcement learning [Sodhani et al., 2021]. One prior study leverages metadata for MTL, by clustering tasks on their metadata and then modelling task groups with hard parameter sharing [Zheng et al., 2019]. However, it is incompatible with continuous metadata and therefore not applicable to the problem of water quality remote sensing. Further, this unsupervised metadata approach cannot learn metadata feature importances, interactions, or relationships with task functions, and is limited by its hard parameter sharing architecture.

Task embeddings. In a setting with many available historical models, task embeddings are used to select a model for a new MTL problem [Zhang et al., 2018]. Task embeddings are also used to condition task-specific decoders [Sun et al., 2021]; while this approach is limited to dense image prediction tasks, by using a hypernetwork our work is able to condition general network architectures on task embeddings.

2.2 Hypernetworks

Another field of neural architecture research is hypernetworks, originally introduced as a method to compress model weights [Ha et al., 2017], but since applied to diverse domains [Littwin and Wolf, 2019; Shamsian et al., 2021]. Transformer-based architectures use hypernetworks to condition task-specific prompts and adapters [Ye and Ren, 2021; ¨Ust¨un et al., 2022; Liu et al., 2022].

Hypernetworks can inject model-level information into deep learning systems, such as our motivation of exploiting task-level metadata. In continual learning, frozen taskspecific embeddings are a memory-efficient way to store previous task models [von Oswald et al., 2020]. In MTL, hypernetworks can incorporate user-defined task preferences and compute requirements in a hard parameter sharing architecture search [Raychaudhuri et al., 2022], or learn the pareto front of multi-objective task optimisation [Navon et al., 2020]. These works generate weights for an entire hard parameter sharing MTL network conditioned on a modellevel input of interest. In contrast, our work generates the weights for a single task network conditioned on task features, using a hypernetwork as a method for soft parameter sharing between tasks by generating multiple target networks.

3 Multi-Task Hypernetwork

Problem formulation. We define the multi-task learning problem as follows. Consider a set of T related tasks {Tt}Tt=1 = {(Dt, ft)}Tt=1 . Each task Tt has a deep learning function {ft : X ↦→ Y}Tt=1 , and a set of training data {Dt}Tt=1 which contains nt training instances Dt = {xti , yt i }nt i=1 where xti ∈ X, yt i ∈ Y. The aim of MTL is to jointly learn each task function {ft}Tt=1 from the training data of all tasks {Dt}Tt=1 .

Metadata. Metadata is data that describes and gives information about other data [Zheng et al., 2019]. In the context of MTL, we define metadata as task-level data, that is, data describing the tasks themselves. For example, lake attributes in a problem where each lake is associated with a different task, or mechanical information about robot arms when modelling each arm's movements as its own task. Metadata may come from auxiliary datasets or from domain knowledge. We note that for some MTL applications, metadata is unavailable or undefined. Consider the learning of semantic segmentation and surface normal of an image as a multi-task problem [Misra et al., 2016]; no relationship between these two tasks exists which can be expressed with task-level features.

Problem formulation with metadata. When metadata is available, we modify the multi-task learning problem formulation as follows. Each task Tt additionally has an associated metadata feature vector mt of size dm, such that Tt = (Dt,mt, ft). The aim of MTL with metadata is therefore to jointly learn the task functions {ft}Tt=1 from the metadata {mt}Tt=1 as well as from the task training data {Dt}Tt=1 . When metadata is unavailable, we use the original MTL problem formulation. If this metadata feature vector is informative with regards to task functions, leveraging this additional source of knowledge may improve the performance of MTL models. However, current approaches are not designed to effectively leverage metadata. A naive approach to learn from metadata with existing techniques would be to append it as supplementary features to training data {ˆxti }nt i=1 = {(xti ,mt)}nt i=1 . This may not be an effective way to exploit this data as the appended features are constant for all data within a task.

3.1 Method Overview

The foundation of our MTL architecture is a hypernetwork h, which is shared by all tasks. The network weights θt of target function ft are not trained directly. Instead, they are generated by h according to θt = h(et), where et is a de dimensional embedding vector associated with task t.

The hypernetwork generates the weights for each task's target network. Task-specific embeddings provide the flexibility for the hypernetwork to generate different weights for each target network. The learnt embeddings are a low-dimensional representations of target functions, allowing our model to efficiently learn task similarities and relationships.

Multi-Task Hypernetworks are therefore parameterised by the weights θh that define h, and T task-specific embedding vectors {et}Tt=1 . Most model parameters are shared between all tasks since |θh| ≫ Tde. Despite this, our approach is best classified as soft parameter sharing, as the target networks are generated with flexible weights at all network layers. To perform inference on task t, we compute the weights of ft with h then predict x with ft. This is performed for all tasks in parallel. During training, loss gradients are backpropagated through the target networks, then pass through the hypernetwork weights and task embeddings, such that the entire architecture can be trained directly with any gradient optimizer. Task embeddings are trained identically to other parameters.

When metadata mt is available, Multi-Task Hypernetworks can naturally leverage it by appending it to the task embeddings as input to the hypernetwork θt = h(et,mt). Unlike task embeddings, metadata are treated as constants and are frozen during training. Considering that task-specific embeddings allow Multi-Task Hypernetworks to learn task differences and relationships, by appending static metadata to the embeddings, the metadata acts as pre-learnt embedding priors from an auxiliary source. Including informative metadata may therefore improve the model's ability to learn task relationships. Further, under the hypernetwork framework, the weights that define a task's target network are a function of that task's metadata. As our model is trained with metadata over multiple tasks, the hypernetwork will directly learn the relationship between the metadata and the task functions. Our architecture is therefore able to leverage metadata in learning both task relationships and task functions explicitly.

3.2 Hypernetwork Architecture

The hypernetwork maps from a low-dimensional input task embedding (and optionally task metadata) to the weight matrices that define a task network. Our hypernetwork architecture, displayed in Figure 3, achieves this with two components: a feature extractor and a weight generator.

To model the water quality remote sensing problem, the hypernetwork architecture described in this section generates linear target networks. Our approach can model different target network architectures by modifying the hypernetwork weight generator to instead generate different weight tensors, such as the convolutional layer generation in the original hypernetwork implementation [Ha et al., 2017]. We note that hypernetworks are used for efficient weight generation of various target network architectures [Ha et al., 2017; Lin et al., 2020; Ye and Ren, 2021], indicating that our approach would be viable in diverse settings.

The hypernetwork takes as input a task embedding (and optionally task metadata) of dimension de (or de + dm). The feature extractor is a feedforward neural network with lh linear layers of dh neurons, allowing our model to learn abstract representations of task embeddings and metadata. Similarly to features in a typical neural network, metadata is normalised to unit mean and variance.

The weight generator consists of linear layers mapping from the resulting dh dimensional latent space to weight matrices which each define a layer in the target network. Let layer n of the target network be a linear layer with dnin inputs and dnout outputs. This layer has a weight matrix of size dnin ×dnout and a bias of size dnout , for a total of (dnin +1)×dnout parameters. For simplicity, let dn∗ in = dnin + 1, such that layer n has dn∗ in × dnout parameters. To generate these parameters, the hypernetwork requires a linear mapping from dh to dn∗ in × dnout . This can be achieved with a hypernetwork weight matrix W of size dh × dn∗ in × dnout and a bias b of size dn∗ in × dnout . However, W contains dh times as many parameters as the target network layer, which is generally prohibitively large.

Weight compression. To reduce the number of parameters in the hypernetwork weight generator, we implement a simple tensor compression technique for hypernetworks generating linear layers. This technique is analogous to the convolution chunking technique in the original hypernetwork implementation [Ha et al., 2017]. Instead of learning W directly, we learn three smaller weight matrices, W A,W B ,W C , of size dh × dn∗ in , dh × dnout , dn∗ in ×dnout respectively. These matrices expand into W as follows: Wijk = W AijW BikW Cjk . This weight compression substantially reduces the total number of model parameters while remaining trainable via standard backpropagation. These three matrices each learn interaction terms between two of the three dimensions of W: task features, target network layer input features and target network layer output features.

Complexity. The total number of parameters in our MultiTask Hypernetwork is given by:

where lt is the number of layers in the target network. Note that a single target network has ∑lt i=1 di∗ indiout total parameters, or half of the first term representing the hypernetwork weight generator parameters. Despite having more parameters than a single target network, the hypernetwork weights are shared for all tasks, such that our approach requires only de additional parameters per task. This is substantially fewer than other soft parameter sharing MTL architectures, and in practice, our method has fewer parameters when modelling more than a few tasks. An experimental analysis of computational complexity validates that our method is suitable for training large target networks, but is omitted due to space constraints.

4 Experiments

We evaluate the effectiveness of our approach by comparing its performance to baseline methods, in the typical MTL setting without metadata, and with metadata. We include metadata in baseline methods naively as described in Section 3.

Experimental setup. Experiments were repeated across 50 seeds for each combination of hyperparameters. Data was partitioned into training, validation, and test sets in a 60/20/20 ratio. Performance on the validation set was used for hyperparameter selection and early stopping. Hyperparameter selection was performed independently with and without metadata. Regarding general hyperparameters, all methods used the learning rate η ∈ {10p |p = -3,-3.5,-4,-4.5} and random batches of size 64. All models used a network architecture of three fully connected layers of 32 neurons with ReLU activations. We empirically found that networks of this size had sufficient representation capacity to effectively model the test datasets. Regarding Multi-Task Hypernetwork hyperparameters, we tuned the number of layers in the hypernetwork feature extractor lh ∈ {0, 1, 2} where each layer had 32 neurons with ReLU activations. Task embeddings used de = 10 parameters and were initialised with a uniform distribution and variance 1, using identical initialisations across tasks. All experiments are implemented in PyTorch with a GeForce RTX 3080 GPU.

Water Quality dataset. This dataset was produced in collaboration with local council. We combined the GLORIA dataset for hyperspectral remote sensing reflectance measurements and co-located water quality attributes [Lehmann et al., 2023] with the LakesATLAS dataset for global lake attribute data [Lehner et al., 2022]. Lakes with at least five data points were included, resulting in 88 diverse lakes spanning glacial, riverine, volcanic origins, and artificial reservoirs, with mean chlorophyll-a concentrations ranging 0.7 to 290 mg/m3 .

Other datasets. To analyse the general effectiveness of our approach, we include three additional multi-task tabular regression datasets with metadata and limited data instances (Table 1). Cubic is a synthetic dataset which we design to have a simple and complete relationship between task metadata and the task functions. Each task is a third order 1D polynomial defined over x ∈ [-1, 1] with random coefficients, where the goal is to estimate y given x. Polynomial coefficients are task metadata. The Algorithms dataset [Brazdil et al., 1994] aims to estimate the performance of traditional ML algorithms on datasets, given dataset summary statistics, where each algorithm is a task. Metadata is not from auxiliary sources but manually constructed from domain knowledge by naively categorising algorithms across binary labels e.g. "is tree", "is kernel". The Robot Arm dataset [Duka, 2014] aims to estimate robot arm joint angles given the end arm position. Each task is highly nonlinear and represents an arm with different joint lengths, which are the task metadata.

Baseline methods. We include Hard sharing [Caruana, 1997], MRN (Multilinear Relational Networks) [Long et al., 2017], MaxRoam (Maximum Roaming Networks) [Pascal et al., 2021], Cross-stitch networks [Misra et al., 2016], Sluice networks [Ruder et al., 2019], DMTRL (Deep Multi-Task Representation Learning) [Yang and Hospedales, 2017]. Regarding traditional approaches, STL-naive is a single task learning network with no task information. STL is a single task learning network which is given task information to learn task differences, as metadata if available or as onehot task embeddings otherwise. We exclude methods which use domain-specific components rendering them incompatible with tabular problems, from computer vision [Liu et al., 2022; Sun et al., 2021; Liu et al., 2019; Bhattacharjee et al., 2022; Lopes et al., 2023] and natural language processing [Tay et al., 2021; ¨Ust¨un et al., 2022; Ye and Ren, 2021].

4.1 Overall Performance

Results are displayed in Table 2 as mean ± standard error. To estimate the overall difference in performance between the methods for all datasets, we use a Friedman test with Nemenyi post-hoc test on the results, displayed in Figure 4.

Without metadata. On the Water Quality dataset, MultiTask Hypernetworks achieve RMSE 20% lower than all other methods. STL achieved the second lowest RMSE, indicating that other MTL algorithms struggle to model problems with many diverse tasks, and highlighting the benefits of model flexibility and parameter efficiency for this application. In this setting, our approach also achieves the lowest RMSE for Algorithms and Robot Arm. The Friedman-Nemenyi test shows that overall, our method outperforms all baseline methods other than Sluice networks on these datasets, indicating that our architecture can outperform many existing MTL architectures in this setting.

With metadata. Multi-Task Hypernetworks again outperform other methods by 20% on Water Quality. Overall, practitioners using our algorithm and leveraging lake metadata can substantially increase the accuracy of their chlorophyll-a estimations compared to existing modelling approaches. Our approach also achieves the lowest RMSE for all datasets. The Friedman-Nemenyi test shows that overall, our method substantially outperforms all baseline methods. Interestingly, despite not being designed for multi-task problems, STL overall outperforms all the multi-task learning methods except sluice networks. This indicates that current MTL architectures are not able to leverage metadata effectively.

Effects of Metadata

To analyse the impact of metadata on model performance, we consider the completeness, noisiness and complexity of the metadata. The metadata of the Water Quality dataset are insufficient to entirely represent lake task functions, as lakes have many bio-geo-optical factors which cannot be entirely represented by 11 features, so we describe this metadata as incomplete. The metadata is sensor data, and hence contains noise. The metadata is also complicated, as these metadata features have abstract interactions. We observe that this metadata improves the performance of Multi-Task Hypernetworks and STL, but not other methods. This indicates that our approach can extract knowledge even from complex metadata, whereas other MTL approaches cannot.

Regarding the Cubic dataset, the metadata is complete as the metadata features define the task functions, is generated without noise, and is simple since features represent coefficients which are linearly related to the task function. We observe that this straightforward metadata improves the performance of all methods except MaxRoam, but that Multi-Task Hypernetworks outperform other methods by 25%. Robot Arm metadata is also complete and generated noiselessly, but is complicated, as the relationship between the robot arm lengths and the task functions is highly nonlinear. Similar to Water Quality, we observe that this complicated metadata only improves the performance of Multi-Task Hypernetworks and STL, further indicating that our approach can extract knowledge from complex metadata while other MTL approaches cannot. Algorithms metadata is manually labelled and therefore noiseless. It is simple, but incomplete as algorithm outcomes cannot be entirely represented by high-level categorical descriptions. Despite this, metadata improves the performance of Multi-Task Hypernetworks, Cross-stitch and Sluice networks. This demonstrates the utility of metadata in enhancing model performance, even when naively labelled.

Model parameters. In our experiments, our architecture has fewer network parameters than other soft sharing approaches of equal capacity, with an order of magnitude fewer parameters than Cross-stitch and Sluice networks, and fewer than DMTRL despite this method being designed to minimise network parameters.

4.2 Ablation Study

We carry out an ablation study to provide insight into the contributions of the different components in our architecture with and without metadata, displayed in Table 3.

Weight matrix compression reduces overparameterisation. Using full-sized weight matrices in the hypernetwork weight generator without the matrix compression technique described in Section 3.2 substantially reduces performance for all datasets except Cubic without metadata. Without compression, our architecture has many more parameters and is more flexible, indicating that the weight matrix compression reduces overparameterisation and increases generalisability.

Feature extractor learns metadata representations. Removing the feature extractor layers from the hypernetwork decreases performance for Water Quality and Robot Arm with metadata. This may be because these datasets have complicated metadata, and so benefit from a non-linear feature extractor to learn abstract metadata representations. Datasets with simple metadata, or with learnt task embeddings but no metadata, may not benefit from this deep learning component. An optimal Multi-Task Hypernetwork architecture for each dataset can be found by tuning the depth of the feature extractor lh, as is done in experiments in Section 4.1.

Task embeddings provide flexibility. Using only metadata as input to the hypernetwork without task-specific embeddings decreases performance in all datasets except Robot Arm. This may be because trainable embeddings give the hypernetwork degrees of freedom between tasks, to learn task relationships and differences more flexibly than with only static metadata. As Robot Arm has complete metadata which captures all task differences, this dataset may not require the additional flexibility from learnable task embeddings. We experimentally demonstrate that task embeddings learn "meaningful" task representations which are predictive of task-level knowledge, but omit this due to space constraints.

Metadata is informative. Randomising metadata by independently shuffling metadata features between tasks decreases performance in all datasets, indicating that our architecture can extract knowledge from informative metadata.

Dataset size affects metadata performance gain. The relationship between dataset size and the benefits of leveraging metadata is examined through sensitivity analysis on Cubic and Robot Arm datasets, which can be produced in various sizes. We vary the number of training instances per task (Figure 5) and the number of tasks in each dataset (Figure 6). Increasing the number of training instances per task decreases the gain in performance from leveraging metadata. Intuitively, as the task size increases, the model is better able to learn task functions and relationships implicitly from joint optimisation without metadata. This indicates that metadata is most useful for small datasets. Increasing the number of tasks increases the gain in performance from leveraging metadata for Robot Arm and for Cubic with fewer than ten tasks. Intuitively, more tasks provide the hypernetwork with more metadata samples to learn the relationship between metadata and task functions. Cubic, with simpler metadata, may need fewer tasks to learn this relationship, while the complex Robot Arm metadata benefits from more tasks.

5 Conclusion

We propose Multi-Task Hypernetworks, a novel architecture for multi-task learning, which learns flexible task functions with fewer parameters than existing soft parameter sharing approaches. We formalise task-level data, called metadata, into the MTL problem. Uniquely, our architecture can leverage metadata to learn task functions and relationships explicitly. We reveal the potential for metadata to enhance performance on MTL problems, particularly with few data samples, but show that existing methods cannot effectively leverage it. We show that our architecture is effective for multi-task learning and effectively extracts knowledge from informative metadata. Our approach substantially improves the accuracy of remote sensing water quality estimation, a challenging problem which contains many diverse tasks with few samples, enabling more accurate global water monitoring without additional expensive data collection. A limitation of our work is the current lack of understanding on evaluating whether metadata will improve learning. Future research could explore the potential of transfer learning to lakes with no training data, generating task networks from metadata alone.

Acknowledgments

This research is partially funded by the New Zealand MBIE TAIAO data science programme. We thank Henry Gouk for his insight.