The explosive growth of spatio-temporal data emphasizes the needs for automatically discovering spatial-related knowledge (Shekhar et al. 2003). Spatio-temporal data are complex due to inherent data characteristics such as implicit spatial relationships between variables and the data variability across locations (Huang et al. 2018;Zheng et al. 2020;Huang et al. 2020). For example, Figure 1 shows the normalized average corn yield for every county in the Midwestern United States. The yield data exhibit a strong spatial variability due to the variation in weather, soils, and management practices across different counties. Hence, a global model trained over large regions may not perform well for every individual county (Karpatne et al. 2018). In addition, it is often expensive to collect a large number of labeled samples in real-world societal applications, which makes it challenging to train advanced deep neural network models separately for each location. Therefore, the development of effective machine learning techniques for spatial-related tasks with strong spatial variability is urgently needed.
Transfer learning methods have been widely explored for adapting machine learning models across space. For example, previous works use a Long-Short Term Memory (LSTM) structure with the attention mechanism to transfer spatial-related information (Nigam et al. 2019;Sharma, Rai, and Krishnan 2020;Jiang et al. 2022). However, these methods directly learn on global data but do not consider regional discrepancy across space. Recent deep learning-based domain adaptation approaches (Nevavuori, Narra, and Lipping 2019;Elavarasan and Vincent 2020) have demonstrated success on several tasks when trained with sufficiently labeled data; however, their performance can be degraded given limited labeled data in regression tasks.
Few-shot learning has shown promise in reducing the need for large labeled samples. Nevertheless, standard fewshot learning methods often perform worse if data are from heterogeneous distributions, which is a common issue in real spatial datasets. Meta learning addresses this issue through the idea of task-adaptive learning. Specifically, meta learning aims to extract meta knowledge from multiple training tasks, which can then be used to facilitate task-adaptive learning for a single task using a small number of data samples. Meta learning has shown encouraging results in many important societal problems, such as agricultural monitoring and traffic management (Pan et al. 2020;Li, Zhang, and Huang 2020;Tseng et al. 2021). Existing meta-learning methods can be categorized based on how they leverage meta knowledge in new tasks, e.g., the optimization-based methods (Finn, Abbeel, and Levine 2017;Li et al. 2017;Antoniou, Edwards, and Storkey 2019), the feed-forward model-based methods (Mishra et al. 2018;Qiao et al. 2018), and metric-learning-based methods (Sung et al. 2018;Willard et al. 2021). For example, the Model-Agnostic Meta Learning (MAML) algorithm (Finn, Abbeel, and Levine 2017) aims to learn an initial model (i.e., metamodel) that can be quickly adapted to new tasks. However, most existing meta-learning methods have limits in handling a large number of heterogeneous tasks, e.g., modeling data from a large number of locations with non-stationary relationships between input and output variables. This can be a common issue in many societal applications. For example, the variation of weather and soils over space interact with the complex carbon, nitrogen, and water cycles during crop growth, which ultimately leads to a strong variability in crop yield patterns.
In this paper, we develop a task-adaptive meta-learning framework by adapting the predictive model gradually over space via a "spatialized" easy-to-hard task hierarchy. In particular, we first train a standard MAML model by considering each location as a separate task. Then we iteratively split the set of tasks to create new branches of harder tasks. Moreover, we synchronously transform the meta-learning model following the obtained easy-to-hard task hierarchy. Given a new task, we can first identify its difficulty level and then adapt the meta-model from the corresponding layer of the hierarchy to the new task. Our contributions can be summarized as follows:
• We create the first meta-learning method that uses spatialrelated tasks in crop yield prediction, which is critical for ensuring food supply and estimating farmers' insurance and subsidies; • We propose a new meta-learning strategy to learn different difficulty levels of tasks in an easy-to-hard hierarchy that can be quickly adapted to new tasks; • We extend existing meta-learning methods to handle a large number of heterogeneous tasks; • Our evaluation on real crop yield data over large regions show the superiority of our proposed approach over standard machine learning and meta-learning baselines.
Few-shot learning has been widely adopted for addressing real-world small data problems due to its great diversity and feasibility (Thrun and Pratt 1998;Finn, Abbeel, and Levine 2017;Wang et al. 2020). Typically, few-shot learning has gained attention in three fields: (1) metric learning-based methods that learn a similarity space, which helps build the connections between new few-shot examples with existing data (Vinyals et al. 2016;Snell, Swersky, and Zemel 2017;Jiang et al. 2020
Multi-task learning (MTL) aims to learn shared representations jointly from multiple training tasks (Caruana 1997).
It assumes the shared information across different tasks can be leveraged to improve the overall performance in all tasks (Zhang and Yang 2018;Ma, Du, and Matusik 2020). These approaches assume that such shared representations could transfer to other tasks, such as object detection (Zhang et al. 2014;Li et al. 2016), image segmentation (Kendall, Gal, and Cipolla 2018), multi-lingual machine translation (Dong et al. 2015;Zhou et al. 2019) and understanding (Liu et al. 2019;Wu, Zhang, and Ré 2020). However, the spatialrelated tasks cannot be directly learned with the multi-task objectives. This is because location-based data (e.g., crop yields across the United States) have significantly different distributions based on their geographic features. In addition, tasks are relatively independent in real scenario problems thus unable to jointly learn an overall model that benefits every task. A potential solution is to explore task-based feature relations (Xie et al. 2021;Zhao et al. 2020); however, it requires sufficient labels thus has limitations in many real applications.
In this section, we introduce the notation and definition used in the crop yield prediction problem. This is essentially a regression task that involves real spatio-temporal data collected in the Midwestern United States. The inputs are an array of daily-collected time-series features, including weather and soil conditions and plant property. The target variable is the county-level crop yield (for corns) at a yearly scale. Specifically, let x i ∈ R s×t denote input features and y i ∈ R denote crop yield label of the U.S. county . An array of task-specific R 2 is used as the input of Algorithm 2 that returns an easy-hard task splitting threshold (γ). In the first task layer (r = 0), the task T 1 has R 2 greater than the threshold (γ = γ 1 ) thus is used to train an easy (E-1) model, otherwise, the tasks (T 2 to T 5 ) are used to train a hard (H-1) model. The adaptive training will run multiple splitting iterations until the maximum split number is reached (e.g., 3 splits). The splitting thresholds (γ) will update in each epoch training as shown in Table 3. (f) In the adaptive test phase, input tasks (T ′ 1 to T ′ 4 ) are split using the thresholds (γ) updated in the training, and apply the same splitting strategies as training to obtain its best adaptive meta model.
i, where i ∈ {1, 2 . . . n}, s is the spatial dimension (i.e., the number of locations sampled from each county), and t is the temporal dimension. Let T i = (X i , Y i ) denote the crop yield prediction task of county i with the input data
, p ∈ {1, 2, . . . , k}, and q ∈ {1, 2 . . . , l}. Here x
is a task-specific training sample from the Train Support Set (Train S-Set) of size k, and x
is a validation sample from the Train Query Set (Train Q-Set) of size l, which is reserved for evaluating the training task performance.
MAML (Finn, Abbeel, and Levine 2017) aims to solve meta-learning problems by optimizing the adaptability of the meta model F θ . It learns parameter θ over a set of training tasks T train where T train = {T 1 , T 2 , T 3 , ..., T n }, such that the learned meta model F θ is able to quickly solve new tasks T ′ j ∈ T test by slightly fine-tuning F θ with a small amount of task-specific samples (X ′ j , Y ′ j ).
is a sample used to fine-tune the learned meta model from the Test Support Set (Test S-Set), and x ′ (j) q ′ , y ′ (j)
is the evaluation sample for the performance in the Test Query Set (Test Q-Set).
Specifically, we train F θ on the task T i ∈ T train with gradient descent optimization
where L Ti (F θ ) is the task-related training (outer) loss, α is a meta learning rate, and ℓ is the associated (inner) loss (e.g., mean squared loss, MSE). In the adaptation stage, MAML optimizes θ such that the following meta loss is minimized using the task-wise fine-tuned parameter θ i over validation samples of each training task x
(3) The meta parameter θ is then updated by gradient descent θ ← θ -β∇ θ L MAML (F θ ). The learned meta model F θ can be used to fine-tune a new task T ′ j ∈ T test through Eq. 1.
In this section, we introduce the LSTM-Attention network (Xu et al. 2020) for the corp yield prediction, as shown in Figure 2 (b). The inputs are fed to an LSTM layer to learn hidden states and their attentions. The outputs are model predictions learned by a multi-layer perception. In particular, the LSTM module trains on the spatial and temporal input features to learn its weights by computing multiple gates (input gate i, forget gate f , and output gate o) that determine whether the incoming data stream should retain or forget:
(8) where σ denotes the sigmoid function, c is the cell state, W is the weight matrix, b is the bias term, h is the hidden state, and t is the time step. In addition, we use an attention module that contains several dense layers and a softmax layer to learn the attention α for each hidden state h t at time step t
where W att and b att are attention weight and bias, respectively. Afterward, the aggregated attention α and hidden state h over all the time steps are fed to a multi-layer perception that returns predicted corp yield ŷ, as
Task-adaptive MAML
Existing meta-learning approaches require new testing task T ′ ∈ T test to be from the same distribution as the training tasks T train . The adaptation performance to the new task T ′ can often be degraded when the training task distribution p(T train ) is highly heterogeneous due to a large number of training tasks T train . To address this issue, we consider decomposing the training task distribution p(T train ) based on the task difficulty level and have the model be adapted gradually following an easy-to-hard task hierarchy, as shown in Figure 2 (e). In particular, we start with building an initial predictive model using all the task samples in the Train S-Set. Different predictive models can be used in the proposed framework thus we adopt the LSTM-Attention network (introduced from Eq. 4 to 10) to train a meta model on the Train S-Set and optimize its MSE loss using Eq. 1.
The performance of the learned meta model on the validation data of each task T i from the Train Q-Set can serve as a proxy measure for the task difficulty level. When the validation loss for a specific task is higher, it indicates that this task has different patterns compared to the majority of tasks in T train that dominates the training of the initial model. Hence, we can split the current set of training tasks T train into Easy Task (E-r) and Hard Task (H-r), where r = 0, 1, 2, . . . , u indicates the task layer (difficulty level of the easy-to-hard task hierarchy shown in Figure 2 (c)), by using a threshold γ on the validation performance. Specifically, each task T i on the task layer r is categorized as
q ] < γ Easy Task (E-r) others (11) where R 2 is the performance metric measured on the validation data of task T i from the Train Q-Set. We repeat this process for every task and gather an array of R 2 for all the tasks. The threshold γ is selected based on a statistical test over the obtained R 2 array, which will be discussed later.
Training phase to build the hierarchy. In the training phase, we iteratively bi-partition the set of hard tasks obtained from the previous task layer (H-r-1), where r ≥ 1. The underlying intuition is to identify the set of tasks that Algorithm 1: Task-adaptive Meta Learning
Output: Optimized meta model weight θ * , Easy Task E-r, Hard Task H-r in the task layer r Input: Tasks T , meta model F θ Initialization: learning rate α, β, task layer r = 0 while not done do if task layer r ≥ 1 then Current task T ← H-r end while not done do for support tasks T i in T do Adapt meta model θ on task T i by Eq. 1:
cannot be well captured by the current model (H-r Model). Starting from the second task layer, we build a meta initial model to be fine-tuned to the tasks in the current task set (i.e., all the hard tasks from the previous task layer) via only a few gradient descent steps, following the standard MAML method. Again, we use the validation performance (measured by R 2 ) of each task-specific fine-tuned model to split the current task set into Easy Task (E-r) and Hard Task (H-r) sets via the threshold-based method (Eq. 11). Here a higher validation loss for a specific task indicates that the meta model cannot generalize well on the task with a small refinement. To accelerate the training, we initialize the meta model with the predictive model (if r = 1) or the meta model (if r ≥ 2) learned from the previous task layer. The process is summarized in Algorithm 1.
Selection of split threshold. We discuss the selection of the threshold γ for splitting the task set at each layer r into Easy (E-r) and Hard (H-r) tasks. Given the obtained validation performance metrics e i (e.g., R 2 ) for each task i, we aim to identify a subset of tasks that have significantly {e i } values compared to the remaining tasks. Hence, we adopt a statistical test, where the null hypothesis H 0 states that e i for all the tasks follow a single normal distribution while the alternative hypothesis H 1 states that there exists a subset of tasks U , and they follow a different normal distribution from the remaining tasks U ′ . Here U can be either the hard or easy tasks. The optimal set U can be obtained by solving the following optimization problem:
According to the prior work (Xie et al. 2021), this can be solved by minimizing the sum of the variance of U and U ′ . Hence we can select the threshold γ that leads to the smallest value of the sum of Var(U ) and Var(U ′ ). This process is summarized in Algorithm 2.
Testing phase using the hierarchy. In the testing phase, given any new task T ′ ∈ T test , we need to identify its difficulty level so that the learned meta models (E-r Model and H-r Model) on the corresponding task layer r can be adapted to the new task. Specifically, starting from the initial model in the easy-to-hard task hierarchy shown in Figure 2 (f), we adapt the corresponding meta model to the new task and measure its validation performance in the Test Q-Set. The obtained R 2 will be compared against the threshold γ on the task layer r. Then we move to the next task layer based on the comparison outcome. This process is repeated until reaching a leaf node of the easy-to-hard task hierarchy. Then we will adapt the final selected model (either the E-r Model where r = 1, 2, . . . , u or the H-r Model w.r.t. r = u) to the new task T ′ i for testing.
The crop yield data are provided by USDA -National Agricultural Statistics Service (NASS) across 630 counties in the United States, and each county has 300 sampled datapoints. Each sample collects 19 daily features including weather and soil conditions, and plant properties, such as temperature, sand content, silt content, and crop yields in 21 consecutive years from 2000 to 2020 (Liu et al. 2021). Inspired by prior work (Jia et al. 2021a,b), we generate the simulated data using a physics-based Ecosys model (Zhou et al. 2021) over Illinois, Indiana and Iowa for pretraining the predictive model, which yields 10K county-level simulation with the same 19 daily features as NASS in 18 consecutive years. We define the meta-learning tasks as county-level crop yield prediction thus we have 630 tasks in total. We construct the meta-learning dataset using the NASS data of which we sample 80% counties in the Midwest U.S.
We implement a diverse set of baselines and meta-learningbased models for model comparison.
We implement the proposed task-adaptive meta-learning framework based on the learn2learn backbone (Arnold et al. 2020) with Pytorchfoot_1 . We pretrain the predictive model with the simulated dataset that achieves 0.9898 R 2 . Afterward, the pretrained predictive model is used to initialize the initial meta model weights in the first task layer of the easy-to-hard hierarchy. We use Adam optimizer with 0.001 learning rate.
In the training phase, we use 32 tasks (out of 630 tasks) as a batch to learn the meta model. In the test phase, we quickly fine-tune the learned meta model for every separate task on its support set and report the performance on its query set. We use the Mean Squared Error (MSE) as the loss function and R 2 as the evaluation metric. We set hyper-parameter a and b in Algorithm 2 as 0.35 and 0.65, respectively, which are experimental values that would guide the model to learn a good threshold (γ) through a subset of the tasks that are between these two bounds. The default adaption number is 1, the (inner) epoch number is 1, and the maximum splitting number is 3 if not specified. We run 30 (outer) epochs with an Nvidia Titan X GPU and report the best performance.
In Table 1, we compared the R 2 performance of different methods for the crop (corn) yield predictions. We evaluate the performance under three different sets of counties:
(1) all the counties, (2) low-yield counties that exclude the top 1/3 high-yield counties, and (3) high-yield counties that exclude the bottom 1/3 low-yield counties. As observed, our proposed adaptive-MAMLs outperform baselines and state-of-the-art meta-learning frameworks. In particular, the Adaptive-MAML-B achieves the best performance in all three sets. In the low-yield counties, the adaptive-MAMLs strongly dominate the baselines and state-of-the-art MAMLs because the low-yield counties are usually along the state boundaries (as shown in Figure 1) and are more difficult to model with existing machine learning methods. Hence, adaptive training on hard tasks will help learn more on these poorly-performed counties. In terms of high-yield counties, they usually have lower data heterogeneity due to the technical improvement in crop cultivation thus the performance is relatively higher than the low-yield. In addition, Table 2 shows the MAML performance with 2 adaptations on trained meta models. As exhibited, the Adaptive-MAML-B improves the performance on the low-yield counties, which
Figure 3 shows county-level predicted performance in Indiana. The Adaptive-MAML-B has mostly higher R 2 (light green in the figures) than the other models. In the baseline models, a region of the difficult tasks (counties) in the upper left boundary shows poor performance; however, our proposed adaptive-MAML greatly reduces the poorlyperformed area, which indicates that our model is able to adaptively learn difficult tasks regardless of the spatialrelated data heterogeneity. This is because our proposed framework learns a set of meta models that can be quickly adapted to different difficulty levels of tasks in the easy-tohard hierarchy.
Figure 4 shows predicted R 2 performance in 16 consecutive years from 2005 to 2020. The Adaptive-MAML-A and Adaptive-MAML-B perform better than the other models in most of the year, but they perform relatively worse in the year 2012. This is because the national-wide low crop yield occurs due to extreme weather and nature conditions. As for the years 2019 and 2020, Adaptive-MAMLs have relatively poor performance because they are only designed to learn spatial heterogeneity and thus may have limitations on capturing temporal data variability, which will be future work.
In this section, we discuss the parameter sensitivity in adaptive-MAMLs. As shown in Figure 2 (e) and (f), the splitting threshold γ is updated each time either starting a 3. As exhibited, the threshold γ becomes higher as the number of epochs increases. During the training process, the adpative-MAML model is gradually updated to better fit training samples, and the γ becomes higher as the model is setting a higher standard for hard tasks. In addition, the threshold γ decreases with the splits going deeper of the easy-to-hard hierarchy, which indicates the tasks are more difficult in bottom task layers (e.g., r = 3), thus the γ becomes lower.
In addition, we test different settings of inner epochs for Adaptive-MAML-A and the maximum split for Adaptive-MAML-B. The results (Figure 6) show that the performance decreases given the increase of the inner epoch training. This shows that over-training under a premature task hierarchy can degrade the model performance. As for the maximum split, both 1, 2, and 3 are acceptable numbers. The small numbers mean fewer easy and hard models will be trained, which is helpful when the data has low heterogeneity.
In Figure 5, we study examples to show the effectiveness of adaptive-MAMLs. The areas (shown as counties on the map) marked with the red and blue rectangles in (a) exhibit low R 2 before adaptive learning. However, the Algorithm 2 splits current training tasks in (a) into (b) hard tasks and (c) easy tasks, and the Algorithm 1 iteratively trains on the hard tasks until a maximum splitting number is reached. As exhibited in (d), the poorly-performed areas (hard tasks) were greatly improved while well-performed areas (easy tasks) maintained excellence compared to (a).
Standard meta-learning methods, e.g., MAML, can have degraded performance given a large number of heterogeneous tasks because spatial data variability is one of the most common issues in many spatial datasets. To bridge the gap, we proposed a task-adaptive MAML to learn spatial-related tasks with an easy-to-hard hierarchy that helps adapt the meta model to new tasks. Extensive experiments show that our methods are superior to a diverse set of baselines and state-of-the-art models on real crop yield data in the Midwest of the United States. The proposed framework demonstrates meta-learning generalizability on a substantial number of spatial-sensitive meta models. In future work, we plan to develop robust MAMLs that adaptively learn both spatial and temporal data heterogeneity and ultimately promote models' feasibility to a wide range of societal problems.
The Thirty-Seventh AAAI Conference on Artificial Intelligence
The code is available at https://github.com/ZhexiongLiu/Task-Adaptive-Meta-Learning.