# Learning to Retrieve Iteratively for In-Context Learning

Yunmo Chen, Tongfei Chen, Harsh Jhamtani, Patrick Xia, Richard Shin
Jason Eisner, Benjamin Van Durme
Microsoft
<EMAIL_ADDRESS><EMAIL_ADDRESS>Johns
Hopkins University; performed while interning at Microsoft. Google; performed
while at Microsoft.

###### Abstract

We introduce _iterative retrieval_ , a novel framework that empowers
retrievers to make iterative decisions through _policy optimization_. Finding
an optimal portfolio of retrieved items is a combinatorial optimization
problem, generally considered NP-hard. This approach provides a learned
approximation to such a solution, meeting specific task requirements under a
given family of large language models (LLMs). We propose a training procedure
based on reinforcement learning, incorporating feedback from LLMs. We
instantiate an iterative retriever for composing in-context learning (ICL)
exemplars and apply it to various semantic parsing tasks that demand
synthesized programs as outputs. By adding only 4M additional parameters for
state encoding, we convert an off-the-shelf dense retriever into a stateful
iterative retriever, outperforming previous methods in selecting ICL exemplars
on semantic parsing datasets such as SMCalFlow, TreeDST, and MTOP.
Additionally, the trained iterative retriever generalizes across different
inference LLMs beyond the one used during training.

Learning to Retrieve Iteratively for In-Context Learning

Yunmo Chen††thanks:  Johns Hopkins University; performed while interning at
Microsoft., Tongfei Chen, Harsh Jhamtani, Patrick Xia, Richard Shin††thanks:
Google; performed while at Microsoft. Jason Eisner, Benjamin Van Durme
Microsoft<EMAIL_ADDRESS><EMAIL_ADDRESS>

##  1 Introduction

A significant emergent capability of large language models (LLMs) is _in-
context learning_ (ICL; Brown et al., 2020), which facilitates few-shot
learning. In ICL, a set of _exemplars_ 111 An exemplar is a tuple of input and
output, demonstrating the mapping relationship between the two. is usually
provided to build the mapping relationship between inputs and outputs. These
exemplars can either be hand-crafted and fixed or retrieved from a training
set. However, if retrieving from the dataset, the retrievers used in such
applications are typically off-the-shelf models (e.g., Contriever (Izacard et
al., 2022)) that do not consider interactions among retrieved items when
multiple targets are required, nor the specific characteristics of the
inference LLMs and downstream task requirements. Research (Gao et al., 2021;
Liu et al., 2022; Lu et al., 2022, _i.a._) has shown that ICL is sensitive to
both the exemplars provided and their order within prompts. Off-the-shelf
retrievers, which generally rank items based solely on semantic similarity
(Lee et al., 2019; Reimers and Gurevych, 2019a, i.a.), do not ensure optimal
conditions for either criterion, leading to suboptimal performance in
downstream LLM generation. Hence, there is a need for a retriever capable of
constructing a portfolio of items tailored to achieve optimal generation with
LLMs.

Figure 1: Above: ICL under a single retriever call. Below: ICL under our
proposed iterative retriever.

We propose _iterative retrieval_ to address this problem. Unlike traditional
retrievers that perform a single call to obtain a list of similar items
ordered by their similarities, iterative retrieval involves a sequence of
retrieval calls, each using different query vectors. This makes the retriever
stateful, maintaining an _internal state_. The process can be likened to
navigating the encoding space of exemplars, with each step adjusting direction
based on previously selected exemplars, thus building a trajectory of exemplar
selections.

This approach can be formulated as _Markov decision processes_ (MDPs). At each
step, the _action_ taken by the retriever is a retrieval call that fetches
(potentially multiple) documents from the dataset $\mathcal{D}$.222 The action
space is at least as large as $\mathcal{D}$. The policy is trained to
optimally select exemplars at each step so that the overall trajectory
maximizes the reward, leading to better ICL performance. By leveraging the
LLMs as environments, we create simulators that allow a policy to roll out in
the environment and receive feedback on the effectiveness of the composed
prompts, measured by a reward (metric). Thus, exemplar selection and prompt
composition can be framed as _policy optimization_ aimed at maximizing
rewards, which can be addressed through reinforcement learning.

We situate our study in in-context semantic parsing due to its difficulty,
popularity, and practical value.333 Code generation is considered one of the
most useful but challenging techniques in the era of LLMs. Some semantic
parsing tasks share structural similarity with code generation and program
synthesis. We instantiate an iterative retriever and investigate the
performance of policy learning under this setup. Our contributions include:

  * •

We propose a novel iterative retrieval framework that builds a portfolio of
exemplars for ICL, considering both interactions among retrieved exemplars and
their relationship with LLMs;

  * •

We instantiate this iterative retriever for the in-context semantic parsing
task and train its policy via reinforcement learning, demonstrating superior
performance over strong baselines from prior work, thereby proving its
effectiveness;

  * •

Through a series of analyses, we provide insights into the behaviors of an
iterative retriever initialized with an off-the-shelf retriever.

##  2 Overview of an Iterative Retriever

We consider the problem of _in-context learning_ (ICL): given a dataset
$\mathcal{D}=\\{(x_{i},y_{i})\\}_{i}$ of _exemplars_ , a retriever $R$
retrieves a sequence of exemplars $R(x)$ based on input query $x$ and generate
the answer $y$ based on the distribution $P_{\rm LM}(\cdot|x;R(x))$.

This retriever $R:\mathcal{X}\to\mathcal{D}^{K}$ retrieves an ordered list (of
length $K$) of exemplars for the LM. The goal of the retriever $R$ is to
select a sequence of exemplars $((x_{i},y_{i}))_{1\leq i\leq K}$ such that the
probability of the expected output $y$ is maximized:

$\underset{\begin{subarray}{c}(x_{i},y_{i})\in\mathcal{D}\end{subarray}}{\arg\max}P_{\rm
LM}(y|x;((x_{i},y_{i}))_{1\leq i\leq K}).$ (1)

However, this is a _combinatorial optimization_ problem that is
computationally infeasible to solve exactly. Much of prior work resort to
selecting top-$k$ exemplars based on a scoring function $S$:

$R(x)=\underset{(x^{\prime},y^{\prime})\in\mathcal{D}}{\arg\mathrm{top}_{k}}~{}S(x,(x^{\prime},y^{\prime}))$
(2)

Prior work has differed on the choice of the scoring function $S$: BM25 Roy et
al. (2023), coverage Gupta et al. (2022), etc. However, such method did not
model the interaction between the retrieved exemplars and the language model.
We propose an iterative version, where we create a _retrieval state_ $s$, and
for each step $i$ one exemplar $(x,y)\in\mathcal{D}$ is retrieved. This is an
approximation to the optimization problem in Equation 1.

$\displaystyle(x_{i},y_{i})$ $\displaystyle\leftarrow R_{\rm step}(s_{i})$ (3)
$\displaystyle s_{i+1}$ $\displaystyle\leftarrow\tau(s_{i},(x_{i},y_{i}))$ (4)

After $K$ steps, the retrieved sequence would be $R_{\rm
iter}(x)=((x_{i},y_{i}))_{1\leq i\leq K}$. This formulation of an _iterative
retriever_ naturally fits in the definition of a Markov decision process
(MDP). Here, our decision process comprises of
$(\mathcal{D}^{*},\mathcal{D},\tau,r)$, where

  * •

The state set $\mathcal{D}^{*}$ contains exemplar sequences whose elements are
in $\mathcal{D}$;

  * •

The action set is just $\mathcal{D}$: each action selects one exemplar from
$\mathcal{D}$. In theory, more than 1 exemplar can be selected at each step,
but we proceed with just 1 exemplar for simplicity;

  * •

The transition function
$\tau:\mathcal{D}^{*}\times\mathcal{D}\to\mathcal{D}^{*}$ appends an exemplar
to the existing sequence;

  * •

The reward function $r:\mathcal{D}^{*}\times\mathcal{D}\to\mathbb{R}$ funnels
signal from the LLM back to the retriever. It will be discussed in §4.

By situating our proposed iterative retriever under this RL scenario, we can
utilize all sorts of RL techniques to train this retriever from the
environment, which is the LLM itself. In the next section, we instantiate a
neural iterative retriever and situate it under a common task, namely
_semantic parsing_ , under this ICL framework.

Figure 2: ICL prompt construction for an example in SMCalFlow. _Above:_ ICL
with BM25 as the retriever. _Below:_ An instance of our iterative retriever.
BM25 retrieves examples that overlaps lexically with the query, whereas the
trained iterative retriever is better at retrieving structurally similar
exemplars since it is trained to maximize the probability of the LM generating
the reference parse.

##  3 Instantiating an Iterative Retriever

Figure 3: Samples of $(x,y)$ pairs for semantic parsing under different
datasets used in this paper.

We consider an _instance_ of in-context learning, namely few-shot semantic
parsing. Given a natural language query $x$, a model is expected to output a
semantic representation $y$ of $x$ given a sequence of exemplars (see Figure
3).

We instantiate a _neural iterative retriever_ based on the formulation we
proposed above:

  * •

The state of the MDP, i.e. the sequence of exemplars, is modeled by a fixed-
length vector $\mathbf{s}\in\mathbb{R}^{d}$. The initial state
$\mathbf{s}_{0}$ is a parameter.

  * •

At each step 1 exemplar is retrieved. We define a _policy distribution_ that
picks one exemplar from the training set $\mathcal{D}$, similar to Lu et al.
(2023):

$\pi((x_{i},y_{i})|s_{i})\propto\exp(\mathbf{Q}(\mathbf{s}_{i})\cdot\mathbf{F}_{\rm
enc}(x_{i})/\beta)$ (5)

where $\mathbf{Q}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ maps a state vector
$\mathbf{s}_{i}$ to a query vector $\mathbf{q}_{i}$, $\mathbf{F}_{\rm
enc}:V^{*}\to\mathbb{R}^{d}$ is a text embedder that maps a text sequence into
a vector, and $\beta$ is a temperature hyperparameter. In our experiments,
$\mathbf{F}_{\rm enc}$ is initialized with the weights of Contriever Izacard
et al. (2022), a general-purpose text embedder trained for retrieval.

Under this policy, if we take greedy decoding, the retrieval step would just
be

$\displaystyle(x_{i},y_{i})$ $\displaystyle\leftarrow R_{\rm
step}(\mathbf{s}_{i})=\underset{(x^{\prime},y^{\prime})\in\mathcal{D}}{\arg\max}~{}\pi((x_{i},y_{i})|s_{i})$
$\displaystyle=\underset{(x^{\prime},y^{\prime})\in\mathcal{D}}{\arg\max}~{}\mathbf{Q}(\mathbf{s}_{i})\cdot\mathbf{F}_{\rm
enc}(x_{i}).$ (6)

This is a _maximum inner product search_ (MIPS) problem, and thus can be
solved with a vector index such as FAISS (Douze et al., 2024).

  * •

State transition is modeled by a gated recurrent unit (GRU; Chung et al.,
2014) update:

$\mathbf{s}_{i+1}\leftarrow\mathrm{GRU}(\mathbf{s}_{i},\mathbf{F}_{\rm
enc}(x_{i}))$ (7)

where the encoded vector of the retrieved exemplar $x_{i}$ is passed to the
GRU to update the state.444 Using a Transformer decoder here results in more
unstable training as we discovered in our experiments. See §6.1.

Note that the only additional parameters we included in this neural iterative
retriever is the state transition model, where we instantiate as a GRU.

This is different from a regular retriever, where a single retrieval call to
the training set
$R(x)=\arg\max_{(x,y)\in\mathcal{D}}\mathbf{q}\cdot\mathbf{F}_{\rm enc}(x)$ is
made. The iterative retriever navigates the encoding space of exemplars,
adjusting the query vector $\mathbf{q}^{\prime}$ at each step based on
previously selected exemplars $s$, thus steering the search process to find
new candidates. Figure 2 demonstrates the process of such an iterative
retriever. This stateful design allows for optimized retrieval results through
iterative interactions, incorporating signals from both external sources
(LLMs) and internal states (previously retrieved items tracked via state
transitions).

##  4 Training

#### Environment Simulator

To construct feedback (or reward) from the underlying LLMs, we treat LLMs as
environments where actions are performed and evaluated. We design an iterative
prompting schedule within this LLM environment to simulate the process of
iterative retrieval and corresponding ICL prompt execution. At each step $i$,
the current sequence of chosen exemplars, $s_{i}$, is turned into an LLM
prompt using a predefined template,555 Refer to Appendix A.3 for the template
used in this work. then used for LLM generation. This schedule effectively
simulates the real-world scenario of prompting LLMs, allowing us to observe
various execution dynamics, such as generated hypotheses and their
probabilities.

#### Reward Design

Technically, if the final task metric were available, it can be used directly
as the reward to optimize for. However, such a reward is often too coarse to
reflect differences in partially correct results. For example, if the metric
is exact match, which is common in semantic parsing tasks, the reward would
simply be the Kronecker delta $\delta(y^{*},y)$, yielding 1 only if the
prediction $y$ exactly matches the reference $y^{*}$, and 0 otherwise.

Given that the LLM simulator provides access to the probabilities of generated
sequences,666 This is generally accessible in many LLM inference
implementations such as vLLM Kwon et al. (2023). For OpenAI-style APIs, this
can be accessed using the “echo” parameter. we employ a more general reward
design that is not task-specific. Our reward leverages the LM completion
probability of the reference sequence $P_{\rm LM}(y^{*}|x)$ Shin et al.
(2021); Shi et al. (2023), which captures subtle changes in the likelihood of
the LM generating the target sequence with respect to changes in the input
$x$. In ICL, more exemplars typically result in better performance before
reaching saturation. Inspired by Zhang et al. (2022), We further refine the
reward to reflect the _increase in the likelihood_ of the reference $y^{*}$
given the prompt. This is a proxy value that measure _how much this exemplar
contribute_ to generating the reference parse. This design encourages the
model to select exemplars that most significantly contribute to the final
result given the existing exemplar sequence $s_{i}$:

$\displaystyle r(s_{i},x_{i})$ $\displaystyle=P_{\rm LM}(y^{*}\mid
x;s_{i},(x_{i},y_{i}))$ $\displaystyle-P_{\rm LM}(y^{*}\mid x;s_{i}).$ (8)

#### Policy Optimization

We employ _proximal policy optimization_ (PPO; Schulman et al., 2017) to train
an iterative retriever for its stability and efficiency.777 We experimented
with various other RL algorithms (including policy gradient (Sutton et al.,
1999) and advantage actor critic (A2C; Mnih et al., 2016)) and found that PPO
is the most stable one for our scenario. One core idea of PPO is to define a
clipping term that controls the policy optimization process, so that variance
is reduced. Given a trajectory $(x_{1},\cdots,x_{T})$, we have

$\displaystyle\mathcal{L}_{i}^{\text{clip}}(\theta)=\hat{\mathbb{E}}_{i}\big{[}\min(\rho_{i},\mathrm{clip}_{\varepsilon}(\rho_{i}))\cdot\hat{A}_{i})\big{]},$
(9)

where $\rho_{i}=\frac{\pi_{\theta}(x_{i}|s_{i})}{\pi_{\theta_{\rm
old}}(x_{i}|s_{i})}$ is a probability ratio between action $x_{i}$888 $x_{i}$
describes that the action of an iterative retriever is to retrieve an exemplar
from a candidate set, hence $x_{i}\in\mathcal{D}$. performed against the
current policy $\pi_{\theta}$ and old policy $\pi_{\theta_{\rm old}}$ at state
$s_{i}$, $\mathrm{clip}_{\varepsilon}(\rho)$ clips $\rho$ to be within
$(1-\varepsilon,1+\varepsilon)$ and $\hat{A}$ is the advantage.

Advantage $\hat{A}_{i}$ at step $i$ describes how much better it is to take a
specific action $x_{i}$ at state $s_{i}$, over randomly selecting an action
according to $\pi(x_{i}|s_{i})$. To compute it, besides the neural model
defined in §3, we follow common practice in reinforcement from human feedback
(RLHF; Huang et al., 2024) to add a single linear layer to serve as a state-
value function $V(s)=\mathbf{v}\cdot\mathbf{s}$ that maps states to values.
Generalized advantage estimate (GAE; Schulman et al., 2016) is then used for
variance-reduced advantage estimation atop the learned state-value function:

$\displaystyle\hat{A}_{i}$
$\displaystyle=\delta_{i}+(\gamma\lambda)\delta_{i+1}+\cdots+(\gamma\lambda)^{T-i+1}\delta_{T-1}$
(10) $\displaystyle\delta_{i}$ $\displaystyle=r_{i}+\gamma
V(s_{i+1})-V(s_{i})$ (11)

where $r_{i}$ is the reward obtained at step $i$, $\gamma$ is the discount
factor, $\lambda$ downweighs rewards corresponding to delayed effects.
Following Schulman et al. (2017) on PPO in Actor-Critic style, we minimize the
value function error term by a squared-error loss, with an additional entropy
bonus term $-H$:

$L_{i}^{\rm
PPO}=\mathbb{E}_{i}\big{[}\mathcal{L}_{i}^{\text{clip}}(\theta)+c_{1}\hat{A}_{i}^{2}(\theta)-c_{2}H_{\pi_{\theta}}(s_{t})\big{]}$
(12)

where $c_{1}$, $c_{2}$ are coefficients.

#### Sampling & Collecting Experience

In a single retrieval step, the retriever selects an exemplar from a candidate
set, with the policy $\pi_{\theta}(x_{i}|s_{i})$ defining a probability
distribution over candidates $x_{i}\in\mathcal{D}$. In this RL simulation, it
is crucial to sample different actions at each step to enable the model to
explore various trajectories and benefit from those that yield higher rewards.
However, constructing the entire distribution and sampling from it at each
step is computationally infeasible, especially when the number of candidates
exceeds 100K. Furthermore, these distributions often exhibit a long-tailed
nature, where many candidates have low scores, suggesting that a significant
portion of candidates may be less similar and potentially less useful for ICL.

Figure 4: Stratified sampling employed in our approach. Our sampling method
retains the top $k/N_{s}$ samples and split the rest into $(N_{s}-1)$ strata
to perform stratified sampling. The resulting $k$ samples are renormalized to
construct action distribution.

To address these challenges and reduce the computational cost of sampling
trajectories while managing the trade-offs between exploration and
exploitation, we propose a stratified sampling ($N_{s}$ strata) method to
construct a modified policy $\tilde{\pi}$ that contains $k$ candidates.

To start, we construct a buffer with top-$B$ exemplars retrieved with Equation
5. Retain the top $k/N_{s}$ samples in the policy. Split the rest into
$(N_{s}-1)$ strata, and sample $k/N_{s}$ from each. Combine all these selected
exemplars and renormalize these scores with softmax (with temperature
$\beta_{\rm renorm}$). This method enables the model to focus on more
promising candidates while still allowing for exploration (see Figure 4 for an
illustration).

During training, experience replay (Lin, 1992) is employed to improve training
efficiency. To collect experience, we run inference with the current policy
fixed on several training examples to generate trajectories. At each step,
information such as policy, reward, and value is recorded. These trajectories
are stored in a replay buffer, then shuffled and split into mini-batches for
policy optimization. This approach allows the same experiences to be replayed
multiple times, reducing the number of required simulation runs.

##  5 Experimental Setup

#### Datasets

We validate our pilot iterative retriever for ICL on a set of semantic parsing
datasets, namely SMCalFlow (Andreas et al., 2020), TreeDST (Cheng et al.,
2020), and MTOP (English portion only; Li et al., 2021), following the
BenchClamp benchmark Roy et al. (2023). Samples of representations are shown
in Figure 3. For statistics, see Appendix A.1.

#### Baselines

We compare our iterative retriever (henceforth denoted as IterR) with a range
of off-the-shelf retrievers, including BM25 (Robertson and Zaragoza, 2009) and
a dense encoder, Contriever (Izacard et al., 2022). Additionally, we benchmark
against two strong baselines from prior work on improving exemplar selection:
EPR (Rubin et al., 2022) and CEIL (Ye et al., 2023). EPR is an efficient
exemplar retrieval method for in-context learning (ICL) that leverages a
scoring LM to label positive and negative training examples, then using this
dataset for a contrastively learned dense retriever. CEIL uses _determinantal
point processes_ (DPPs) to model the interaction between the given input and
in-context exemplars.

For the EPR baseline, we replace the base dense retrieval encoder with
Contriever instead of S-BERT (Reimers and Gurevych, 2019b) for fair
comparison. Following Ye et al. (2023), we use the trained EPR model as
initialization for CEIL. Similarly, the same EPR checkpoint is used to
initialize the text encoder in IterR. Note that in IterR, we freeze the
weights of the EPR encoder and only train the GRU-based state transition
function, policy network, and value network, resulting in 4M more parameters
compared to the original Contriever ($110{\rm M}\to 114{\rm M}$).

For retrievers without iterative capabilities, we adapt them by taking only
the top-$k$ retrieved items and keeping their original ranks. For EPR, CEIL,
and IterR, we selected the best performing model checkpoints on the validation
set. All generation is run with 10 exemplars; i.e. $k=10$.

#### Generation with LLMs

The inference LLM is essential for executing input prompts to generate
responses. In our experiments, we use Llama-2-7b to build the environment
simulator and train the policy using its signals. With the learned policy, we
investigate both intra-family and inter-family generalization by replacing the
inference LLMs. For models within the same Llama-2 family, we explore various
model sizes and finetuned versions, including CodeLlama-70b-Instruct (Rozière
et al., 2023), a model further fine-tuned for code generation. For inter-
family experiments, we choose Mistral-7b (Jiang et al., 2023). For decoding
configurations, we consistently use beam search with beam size 3 and sampling
temperature 0.5.

#### Hyperparameters

Please refer to Appendix A.2.

#### Evaluation Metrics

We follow prior work in evaluating semantic parsing (Roy et al., 2023), where
_exact match_ at $k$ (EM@$k$) is used. Exact match results for top-$k$ decoded
hypotheses reflects beam search decoding used in LLMs, where multiple parsing
results are generated simultaneously.

However, EM is a stringent metric, penalizing even minor mismatches. For
instance, a parse with a substructure reordered differently (a && b) from the
reference (b && a) is still correct but would score zero under EM. This is
problematic in semantic parsing, where target parses are compositional, making
it important to assess the correctness of substructures. Since SMCalFlow and
TreeDST involve deeply nested structures, we also adopt _SMatch_ (Cai and
Knight, 2013), following Chen et al. (2023), to evaluate performance on
substructures. SMatch is designed to evaluate AMRs Langkilde and Knight
(1998). Generated code can be transformed to AMRs by treating each function’s
return value as an _entity_ and each argument to a function as a _value_ ,
where the parameter name is the _relation_. See Appendix B for details.

##  6 Results & Analyses

We evaluate the performance of different retrievers by comparing their
downstream ICL performance on semantic parsing (Table 1). IterR outperforms
all baselines across three datasets on all metrics.

The gain in EM is intuitive since it aligns with the training objective, which
involves the probability of generating target parses. The improvement in
SMatch indicates that IterR optimizes retrieval results to improve
compositionality to some extent, even with a simple objective.999 While a more
dedicated reward design, such as incorporating various linearizations of
target structures, might further enhance IterR’s performance. This work
focuses on demonstrating the framework’s effectiveness rather than dedicatedly
optimizing for a specific task design.

Retriever | Exact Match | SMatch
---|---|---
| @1 | @2 | @3 | P | R | F
SMCalFlow
BM25 | 39.8 | 43.7 | 44.0 | 66.6 | 64.2 | 65.3
Contriever | 44.0 | 48.5 | 48.9 | 68.4 | 66.8 | 67.6
EPR | 48.5 | 52.0 | 52.3 | 73.3 | 76.7 | 75.0
CEIL | 51.1 | 54.2 | 55.8 | 74.9 | 75.2 | 75.1
IterR | 54.1 | 58.4 | 58.5 | 76.6 | 78.1 | 77.3
TreeDST
BM25 | 50.8 | 56.1 | 56.6 | 81.8 | 81.8 | 81.8
Contriever | 54.7 | 60.4 | 61.0 | 83.3 | 82.5 | 82.9
EPR | 54.0 | 58.2 | 58.8 | 84.7 | 83.4 | 84.0
CEIL | 56.2 | 58.3 | 61.6 | 81.3 | 84.4 | 84.9
IterR | 58.2 | 63.4 | 63.8 | 85.5 | 85.8 | 85.7
MTOP
BM25 | 57.4 | 63.2 | 63.9 | - | - | -
Contriever | 59.3 | 64.2 | 64.7 | - | - | -
EPR | 62.3 | 68.8 | 69.2 | - | - | -
CEIL | 63.6 | 69.4 | 69.8 | - | - | -
IterR | 63.9 | 70.9 | 71.0 | - | - | -
Table 1: Comparison of our approach, IterR against baselines. “EM@$k$” denotes
exact match at top-$k$; “P”, “R” and “F” denote precision, recall, and $\rm
F_{1}$ score respectively. Experiment results are run with 10 exemplars in the
prompt, averaged over 3 inference runs, and significance tests using paired
$t$-test confirm that the improvements over Contriever, EPR, and CEIL are
statistically significant ($p<0.05$). Figure 5: Performance comparisons on
using various LLMs for inference (top row: SMCalFlow; mid: TreeDST; bottom:
MTOP). Our IterR used in these experiments are trained with Llama-2-7b but
performs retrieval of ICL exemplars used on other LLMs.

#### Generalization across Inference LLMs

IterR benefits from interactive training with an underlying LLM. While
training incurs costs, these can be minimized by training only once, ideally
using a smaller LM. Hence in this section we investigate the generalization
capabilities of IterR trained with a smaller LM $A$, but used for generation
under a larger LM $B$.

In the following experiments, IterR is trained with Llama-2-7b as the
environment, but used for (a) _intra-family LMs_ : variants within the Llama-2
model family; and (b) _inter-family LMs_ : Mistral (Jiang et al., 2023) from a
different model family. We follow the setups described in §6, substituting
only the LLM. As shown in Figure 5, IterR significantly outperforms ($>1\%$
gain) baselines for $75\%$ of the settings and is comparable to a prior strong
baseline (within $1\%$ in absolute performance) for $15\%$ of settings,
demonstrating its generalization within and beyond its own model family.

In intra-family generalization, performance metrics improve with larger model
sizes, and IterR consistently outperforms all baselines. This improvement is
most evident with larger models such as Llama-2-70b and
CodeLlama-70b-Instruct. For inter-family generalization, IterR maintains its
advantage across datasets, though this is less pronounced than within the same
model family. This is expected, as the signal from LLM simulator is more
representative for models sharing the same pre-training procedure. Notably,
with Mistral, Contriever performs worse than BM25 on MTOP, but IterR still
shows improvement. This suggests that IterR, comprising a frozen EPR and
additional GRU layers, can learn task-specific abilities not present in the
vanilla EPR.

#### ICL & Number of Exemplars

Figure 6: Performance comparisons across the various numbers of exemplars used
for ICL.

We investigated how the performance of IterR changes with the number of
exemplars ($\\{1,\cdots,10\\}$) used for ICL on the SMCalFlow dataset (Figure
6). IterR consistently outperforms baseline models across various metrics and
numbers of exemplars, with one exception for the EM@3 metric when using 6
exemplars. This aligns with our training objective, where actions that boost
performance at each step receive higher advantages. IterR achieves comparable
performance with fewer exemplars.

CEIL shows a similar trend in EM, but its SMatch performance lags
significantly, indicating poorer quality in near-miss predictions compared to
IterR. Practically, this means our method allows for a trade-off between
performance and cost, enabling effective ICL with fewer exemplars and reducing
the number of tokens processed by LLMs.

###  6.1 Ablation Study

Variant | EM@1 | SMatch-F
---|---|---
Contriever | 44.0 | 67.6
IterR | 54.1 | 77.3
$-$ EPR intialization | 45.1 | 68.8
$-$ GRU; $+$ Transformer decoder | 50.1 | 75.1
$-$ Stratified sampling | 52.3 | 75.7
Table 2: Results on ablation study. $-$ EPR intialization indicates the model
is trained from Contriever instead of a EPR finetuned checkpoint. $+$
Transformer decoder replaces GRU with a Transformer decoder. $-$ Stratified
sampling replaces the stratified sampling described in Figure 4 with sampling
directly from the buffer.

We further conduct ablation study on components of an iterative retriever,
focusing on the SMCalFlow dataset and use Llama-2-7b while changing the
configuration of the iterative retriever. Results are reported in Table 2.

#### EPR Initialization

Although we follow prior work in using EPR as initialization for
$\mathbf{F}_{\rm enc}$, our iterative retriever is agnostic to the choice of
base encoders for similarity search. Even without EPR initialization, our
training procedure still improves performance against Contriever ($\approx
1\%$ gain under Contriever, but $\approx 6\%$ gain under EPR). We see that
IterR benefits more when using EPR initialization, significantly outperforming
the baselines. We hypothesize that this advantage stems from two sources: (1)
EPR is fine-tuned on the target dataset, making it more domain-specific; (2)
EPR restructures the action space, subsequently enhancing sample efficiency in
RL training.

#### State Transition with Transformer Decoder

In §3, we parameterize the state transition function in the iterative
retriever with a GRU. To explore alternatives, we conducted an ablation
experiment by replacing the GRU with a more powerful Transformer decoder,
configured with 3 layers, 1024 hidden dimensions, with learnable positional
encodings. Despite the increased expressiveness of the Transformer decoder, we
observed a performance drop. During training, employing the warmup technique
Xiong et al. (2020) led to a trivial solution where the policy learned to
predict a nearly fixed trajectory across test examples. Disabling the warmup
stabilized the training but did not improve performance. Developing a
stabilized approach to train the Transformer decoder as a state encoder is
beyond the scope of this work, as our focus is on demonstrating the overall
framework of iterative retrieval rather than optimizing a specific model for
the state transition function. Notably, even with the less powerful GRU, our
iterative retriever successfully learns a policy that retrieves a more
optimized sequence of ICL exemplars.

#### Effectiveness of Stratified Sampling

To collect diverse experience from policy rollouts, we introduce a stratified
sampling method (described in §4) that balances the trade-off between
exploration and exploitation. We found that sampling from the raw policy in
Equation 5 results in a significant performance drop. Additionally,
qualitative examination of several such distributions revealed a preference
for exploitation over exploration, as similar items at the top of the
retrieved list all had higher probabilities.

##  7 Additional Related Work

#### LLMs as Environment in RL

Lu et al. (2023) used policy gradient to learn a dense retriever for ICL
exemplar retrieval, but the state does not contain previously selected
examples, and thus is not iterative and unable to model exemplar order. Zhang
et al. (2022) used $Q$-learning RL for ICL exemplar reordering, with a similar
reward design like ours. However, the proposed method does not extend to
exemplar _retrieval_ , since the policy space is too large to be handled by
$Q$-learning.

#### Few-shot Semantic Parsing

Few-shot semantic parsing using LLMs has shown impressive capabilities in
understanding new examples with minimal training data (Shin et al., 2021; Shin
and Van Durme, 2022). However, these parsers often struggle with
generalization and fail to parse unobserved local structures due to their
limited access to information encoded through exemplars (Bogin et al., 2022).
To this end, recent research has explored various approaches to improving
exemplar selection. EPR (Rubin et al., 2022) used a proxy LM to score outputs
from an unsupervised retriever, enabling better training of a dense retriever.
Oren et al. (2021), Gupta et al. (2022), and Levy et al. (2023) emphasize
learning to select exemplars based on particular criteria, such as diversity
measures and coverage of local structures, to enhance compositional
generalization. While these approaches have shown performance improvements in
semantic parsing tasks, these are highly based on heuristics constructed from
researcher’s experience. Our approach could be seen as an _automated_ version
(through RL) of seeking information useful for semantic parsing.

##  8 Conclusion

We proposed _iterative retrievers_ that iteratively builds a prompt to perform
in-context learning. Such retrievers are framed as Markov decision processes
and trained via policy optimization from LLM feedback, where the policy
directs which exemplar to append to the existing exemplar sequence.
Experiments on semantic parsing demonstrated performance gain of iterative
retrievers over various datasets and state-of-the-art baselines, showing that
they are able to construct prompts that improves in-context learning and
downstream LLM generation.

## Limitations

In our instantiation of the iterative retriever, at each step a single
exemplar is retrieved. One could envision multiple exemplars being retrieved
at each step, thus making the RL trajectory shorter. This could make RL
training easier and inference faster.

Our reward design depends on a particular linearization of the target
structure. A more structured reward function may exhibit better training
behavior and lead to better performance.

The encoder for queries in the iterative retriever is frozen in our current
setup. A trainable query encoder that receives feedback from LLMs may be
desired, but we left that for future work.

While we believe that semantic parsing / code generation is one of the most
useful but challenging task for LLMs, as such is a representative task for ICL
research, we have not tested the effectiveness of iterative retrievers under
other LLM tasks.

## Acknowledgements

This work has been supported by the U.S. National Science Foundation under
grant 2204926. Any opinions, findings, and conclusions or recommendations
expressed in this article are those of the authors and do not necessarily
reflect the views of the National Science Foundation.

## References

  * Andreas et al. (2020) Jacob Andreas, John Bufe, David Burkett, Charles Chen, Josh Clausman, Jean Crawford, Kate Crim, Jordan DeLoach, Leah Dorner, Jason Eisner, Hao Fang, Alan Guo, David Hall, Kristin Hayes, Kellie Hill, Diana Ho, Wendy Iwaszuk, Smriti Jha, Dan Klein, Jayant Krishnamurthy, Theo Lanman, Percy Liang, Christopher H. Lin, Ilya Lintsbakh, Andy McGovern, Aleksandr Nisnevich, Adam Pauls, Dmitrij Petters, Brent Read, Dan Roth, Subhro Roy, Jesse Rusak, Beth Short, Div Slomin, Ben Snyder, Stephon Striplin, Yu Su, Zachary Tellman, Sam Thomson, Andrei Vorobev, Izabela Witoszko, Jason Wolfe, Abby Wray, Yuchen Zhang, and Alexander Zotov. 2020.  Task-oriented dialogue as dataflow synthesis.  _Transactions of the Association for Computational Linguistics_ , 8:556–571.
  * Bogin et al. (2022) Ben Bogin, Shivanshu Gupta, and Jonathan Berant. 2022.  Unobserved local structures make compositional generalization hard.  In _Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing_ , pages 2731–2747, Abu Dhabi, United Arab Emirates. Association for Computational Linguistics.
  * Brown et al. (2020) Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020.  Language models are few-shot learners.  In _Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual_.
  * Cai and Knight (2013) Shu Cai and Kevin Knight. 2013.  Smatch: an evaluation metric for semantic feature structures.  In _Proceedings of the 51st Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers)_ , pages 748–752, Sofia, Bulgaria. Association for Computational Linguistics.
  * Chen et al. (2023) Yunmo Chen, William Gantt, Tongfei Chen, Aaron White, and Benjamin Van Durme. 2023.  A unified view of evaluation metrics for structured prediction.  In _Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing_ , pages 12868–12882, Singapore. Association for Computational Linguistics.
  * Cheng et al. (2020) Jianpeng Cheng, Devang Agrawal, Héctor Martínez Alonso, Shruti Bhargava, Joris Driesen, Federico Flego, Dain Kaplan, Dimitri Kartsaklis, Lin Li, Dhivya Piraviperumal, Jason D. Williams, Hong Yu, Diarmuid Ó Séaghdha, and Anders Johannsen. 2020.  Conversational semantic parsing for dialog state tracking.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 8107–8117, Online. Association for Computational Linguistics.
  * Chung et al. (2014) Junyoung Chung, Çaglar Gülçehre, KyungHyun Cho, and Yoshua Bengio. 2014.  Empirical evaluation of gated recurrent neural networks on sequence modeling.  _CoRR_ , abs/1412.3555.
  * Douze et al. (2024) Matthijs Douze, Alexandr Guzhva, Chengqi Deng, Jeff Johnson, Gergely Szilvasy, Pierre-Emmanuel Mazaré, Maria Lomeli, Lucas Hosseini, and Hervé Jégou. 2024.  The faiss library.
  * Gao et al. (2021) Tianyu Gao, Adam Fisch, and Danqi Chen. 2021.  Making pre-trained language models better few-shot learners.  In _Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_ , pages 3816–3830, Online. Association for Computational Linguistics.
  * Gupta et al. (2022) Shivanshu Gupta, Sameer Singh, and Matt Gardner. 2022.  Structurally diverse sampling for sample-efficient training and comprehensive evaluation.  In _Findings of the Association for Computational Linguistics: EMNLP 2022_ , pages 4966–4979, Abu Dhabi, United Arab Emirates. Association for Computational Linguistics.
  * Huang et al. (2024) Shengyi Huang, Michael Noukhovitch, Arian Hosseini, Kashif Rasul, Weixun Wang, and Lewis Tunstall. 2024.  The N+ implementation details of RLHF with PPO: A case study on tl;dr summarization.  _CoRR_ , abs/2403.17031.
  * Izacard et al. (2022) Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave. 2022.  Unsupervised dense information retrieval with contrastive learning.  _Trans. Mach. Learn. Res._ , 2022.
  * Jiang et al. (2023) Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de Las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. 2023.  Mistral 7b.  _CoRR_ , abs/2310.06825.
  * Kwon et al. (2023) Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. 2023.  Efficient memory management for large language model serving with pagedattention.  In _Proceedings of the 29th Symposium on Operating Systems Principles, SOSP 2023, Koblenz, Germany, October 23-26, 2023_ , pages 611–626. ACM.
  * Langkilde and Knight (1998) Irene Langkilde and Kevin Knight. 1998.  Generation that exploits corpus-based statistical knowledge.  In _36th Annual Meeting of the Association for Computational Linguistics and 17th International Conference on Computational Linguistics, Volume 1_ , pages 704–710, Montreal, Quebec, Canada. Association for Computational Linguistics.
  * Lee et al. (2019) Kenton Lee, Ming-Wei Chang, and Kristina Toutanova. 2019.  Latent retrieval for weakly supervised open domain question answering.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pages 6086–6096, Florence, Italy. Association for Computational Linguistics.
  * Levy et al. (2023) Itay Levy, Ben Bogin, and Jonathan Berant. 2023.  Diverse demonstrations improve in-context compositional generalization.  In _Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 1401–1422, Toronto, Canada. Association for Computational Linguistics.
  * Li et al. (2021) Haoran Li, Abhinav Arora, Shuohui Chen, Anchit Gupta, Sonal Gupta, and Yashar Mehdad. 2021.  MTOP: A comprehensive multilingual task-oriented semantic parsing benchmark.  In _Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume_ , pages 2950–2962, Online. Association for Computational Linguistics.
  * Lin (1992) Long Ji Lin. 1992.  Self-improving reactive agents based on reinforcement learning, planning and teaching.  _Mach. Learn._ , 8:293–321.
  * Liu et al. (2022) Jiachang Liu, Dinghan Shen, Yizhe Zhang, Bill Dolan, Lawrence Carin, and Weizhu Chen. 2022.  What makes good in-context examples for GPT-3? In _Proceedings of Deep Learning Inside Out (DeeLIO 2022): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures_ , pages 100–114, Dublin, Ireland and Online. Association for Computational Linguistics.
  * Lu et al. (2023) Pan Lu, Liang Qiu, Kai-Wei Chang, Ying Nian Wu, Song-Chun Zhu, Tanmay Rajpurohit, Peter Clark, and Ashwin Kalyan. 2023.  Dynamic prompt learning via policy gradient for semi-structured mathematical reasoning.  In _The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023_. OpenReview.net.
  * Lu et al. (2022) Yao Lu, Max Bartolo, Alastair Moore, Sebastian Riedel, and Pontus Stenetorp. 2022.  Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 8086–8098, Dublin, Ireland. Association for Computational Linguistics.
  * Mnih et al. (2016) Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. 2016.  Asynchronous methods for deep reinforcement learning.  _CoRR_ , abs/1602.01783.
  * Oren et al. (2021) Inbar Oren, Jonathan Herzig, and Jonathan Berant. 2021.  Finding needles in a haystack: Sampling structurally-diverse training sets from synthetic data for compositional generalization.  In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_ , pages 10793–10809, Online and Punta Cana, Dominican Republic. Association for Computational Linguistics.
  * Reimers and Gurevych (2019a) Nils Reimers and Iryna Gurevych. 2019a.  Sentence-bert: Sentence embeddings using siamese bert-networks.  In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing, EMNLP-IJCNLP 2019, Hong Kong, China, November 3-7, 2019_ , pages 3980–3990. Association for Computational Linguistics.
  * Reimers and Gurevych (2019b) Nils Reimers and Iryna Gurevych. 2019b.  Sentence-BERT: Sentence embeddings using Siamese BERT-networks.  In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_ , pages 3982–3992, Hong Kong, China. Association for Computational Linguistics.
  * Robertson and Zaragoza (2009) Stephen E. Robertson and Hugo Zaragoza. 2009.  The probabilistic relevance framework: BM25 and beyond.  _Found. Trends Inf. Retr._ , 3(4):333–389.
  * Roy et al. (2023) Subhro Roy, Samuel Thomson, Tongfei Chen, Richard Shin, Adam Pauls, Jason Eisner, and Benjamin Van Durme. 2023.  Benchclamp: A benchmark for evaluating language models on syntactic and semantic parsing.  In _Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023_.
  * Rozière et al. (2023) Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Tal Remez, Jérémy Rapin, Artyom Kozhevnikov, Ivan Evtimov, Joanna Bitton, Manish Bhatt, Cristian Canton-Ferrer, Aaron Grattafiori, Wenhan Xiong, Alexandre Défossez, Jade Copet, Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usunier, Thomas Scialom, and Gabriel Synnaeve. 2023.  Code llama: Open foundation models for code.  _CoRR_ , abs/2308.12950.
  * Rubin et al. (2022) Ohad Rubin, Jonathan Herzig, and Jonathan Berant. 2022.  Learning to retrieve prompts for in-context learning.  In _Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pages 2655–2671, Seattle, United States. Association for Computational Linguistics.
  * Schulman et al. (2016) John Schulman, Philipp Moritz, Sergey Levine, Michael I. Jordan, and Pieter Abbeel. 2016.  High-dimensional continuous control using generalized advantage estimation.  In _4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings_.
  * Schulman et al. (2017) John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. 2017.  Proximal policy optimization algorithms.  _CoRR_ , abs/1707.06347.
  * Shi et al. (2023) Weijia Shi, Sewon Min, Michihiro Yasunaga, Minjoon Seo, Rich James, Mike Lewis, Luke Zettlemoyer, and Wen-tau Yih. 2023.  REPLUG: retrieval-augmented black-box language models.  _CoRR_ , abs/2301.12652.
  * Shin et al. (2021) Richard Shin, Christopher Lin, Sam Thomson, Charles Chen, Subhro Roy, Emmanouil Antonios Platanios, Adam Pauls, Dan Klein, Jason Eisner, and Benjamin Van Durme. 2021.  Constrained language models yield few-shot semantic parsers.  In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_ , pages 7699–7715, Online and Punta Cana, Dominican Republic. Association for Computational Linguistics.
  * Shin and Van Durme (2022) Richard Shin and Benjamin Van Durme. 2022.  Few-shot semantic parsing with language models trained on code.  In _Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pages 5417–5425, Seattle, United States. Association for Computational Linguistics.
  * Sutton et al. (1999) Richard S. Sutton, David A. McAllester, Satinder Singh, and Yishay Mansour. 1999.  Policy gradient methods for reinforcement learning with function approximation.  In _Advances in Neural Information Processing Systems 12, [NIPS Conference, Denver, Colorado, USA, November 29 - December 4, 1999]_ , pages 1057–1063. The MIT Press.
  * Xiong et al. (2020) Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie-Yan Liu. 2020.  On layer normalization in the transformer architecture.  In _Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event_ , volume 119 of _Proceedings of Machine Learning Research_ , pages 10524–10533. PMLR.
  * Ye et al. (2023) Jiacheng Ye, Zhiyong Wu, Jiangtao Feng, Tao Yu, and Lingpeng Kong. 2023.  Compositional exemplars for in-context learning.  In _International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA_ , volume 202 of _Proceedings of Machine Learning Research_ , pages 39818–39833. PMLR.
  * Zhang et al. (2022) Yiming Zhang, Shi Feng, and Chenhao Tan. 2022.  Active example selection for in-context learning.  In _Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing_ , pages 9134–9148, Abu Dhabi, United Arab Emirates. Association for Computational Linguistics.

##  Appendix A Experiment Details

###  A.1 Dataset Statistics

Dataset | Train | Dev | Test
---|---|---|---
SMCalFlow | 108,753 | 12,271 (500 used) | 13,496
TreeDST | 121,652 | 22,910 (500 used) | 22,841
MTOP | 15,667 | 2,235 (500 used) | 4,386
Table 3: Dataset statistics.

###  A.2 Hyperparameters

Name Search Bounds Encoder $\\{$ GTR-T5, Contriever, SBert $\\}$ Learning rate
$\\{5\times 10^{-6},1\times 10^{-5},{\bf 3\times 10^{-5}},5\times 10^{-5}\\}$
LR scheduler $\\{$ reduce-on-plateau, cosine-annealing $\\}$ State transition
$\\{$ GRU, LSTM, Transformer Decoder$\\}$ $\beta_{\rm renorm}$
$\\{0.5,1.0,{\bf 5.0},10.0\\}$ $c_{1}$ $\\{0.1,0.3,{\bf 0.5},0.7\\}$ $c_{2}$
$\\{0,0.005,{\bf 0.01},0.05,0.1,0.15\\}$ $\gamma$ 0.99 $\lambda$ 0.95 Action
buffer size 768 PPO ratio cutoff 1.2 PPO batch size 128 Replay buffer size
2048 Avg. training time 24 hrs GPU used 4 Nvidia V100 32 GB # of parameters*
$\sim$114M (w/ 110M frozen)

Table 4: Hyperparameters and other reproducibility information for IterR.
$\beta_{\rm renorm}$ is the temperature used to create a renormalized action
distribution. $c_{1}$ and $c_{2}$ are coefficients used in the PPO loss.
$\gamma$ and $\lambda$ are discount factors used in GAE.

###  A.3 Prompt Template

The prompt template used across all our experiments is shown in Table 5.

Let’s translate what a human user says
---
into what a computer might say.
Human: $x_{1}$
Computer: $y_{1}$
$\cdots$
Human: $x_{N}$
Computer: $y_{N}$
Human: $x$
Computer:
Table 5: Prompt template used in our experiments. This template will be
instantiated as prompts when filled with retrieved exemplars
$R(x)=((x_{1},y_{1}),\cdots,(x_{N},y_{N}))$ and the test example $x$.

##  Appendix B SMatch Evaluation

For evaluation of semantic parse or code generation on partial results, we
utilize SMatch Cai and Knight (2013). Generated code can be transformed to
AMRs by treating each function’s return value as an _entity_ and each argument
to a function as a _value_ , where the parameter name is the _relation_. An
example is given below.

Consider the following parse in SMCalFlow, expressed in Lisp:

⬇

(Yield

:output (Event.start

:obj (FindNumNextEvent

:constraint (Event.subject_?

:obj (?~= "staff␣meeting"))

:number 1L)))

This will be transformed into the following AMR:

Figure 7: Example AMR based on the previous parse.

This AMR can be easily converted to the following triples.

⬇

instance($0, Yield)

output($0, $1)

instance($1, Event.start)

obj($1, $2)

instance($2, FindNumNextEvent)

constraint($2, $3)

instance($3, Event.subject_?)

obj($3, $4)

instance($4, ?~=)

ARG0($4, "staff␣meeting")

number($2, 1L)


# Improvements to short-term weather prediction with recurrent-convolutional
networks ††thanks: This project benefited from parallel development in the
fellowship “Seamless Artificially Intelligent Thunderstorm Nowcasts” from the
European Organisation for the Exploitation of Meteorological Satellites
(EUMETSAT).

1st Jussi Leinonen  Federal Office of Meteorology and Climatology MeteoSwiss
Locarno-Monti, Switzerland
<EMAIL_ADDRESS>ORCID 0000-0002-6560-6316

###### Abstract

The Weather4cast 2021 competition gave the participants a task of predicting
the time evolution of two-dimensional fields of satellite-based meteorological
data. This paper describes the author’s efforts, after initial success in the
first stage of the competition, to improve the model further in the second
stage. The improvements consisted of a shallower model variant that is
competitive against the deeper version, adoption of the AdaBelief optimizer,
improved handling of one of the predicted variables where the training set was
found not to represent the validation set well, and ensembling multiple models
to improve the results further. The largest quantitative improvements to the
competition metrics can be attributed to the increased amount of training data
available in the second stage of the competition, followed by the effects of
model ensembling. Qualitative results show that the model can predict the time
evolution of the fields, including the motion of the fields over time,
starting with sharp predictions for the immediate future and blurring of the
outputs in later frames to account for the increased uncertainty.

###### Index Terms:

weather forecasting, neural networks, gated recurrent units, optimizers, model
ensembling

##  I Introduction

Neural networks have recently been investigated for many areas of weather
forecasting, including nowcasting, data assimilation, acceleration of the
evaluation of model physics, and the postprocessing of model outputs (e.g. [1,
2, 3, 4]), with ambitious plans in place to further expand their use in the
future [5]. As a complex spatiotemporal problem, weather forecasting benefits
from the ability of convolutional networks to learn spatial relationships and
the ability of recurrent networks to learn the rules of temporal evolution.
Therefore, improvements to the architecture, training and application of these
networks can benefit several different applications within the area of weather
forecasting.

The Weather4cast 2021111https://www.iarai.ac.at/weather4cast/ Stage 1
challenge[6] was organized to encourage the development of neural networks for
the weather forecasting field. The challenge gave the competitors a task of
developing machine learning models that predict the future evolution of
satellite-derived weather data fields given several time steps of the past
state of those fields. The four variables considered were temperature (either
cloud top temperature when cloudy, or the surface temperature otherwise),
convective rainfall rate (crr_intensity), probability of tropopause folding
(asii_turb_trop_prob), and cloud mask (cma). The challenge was divided into
two parts: in the “Core” competition, the networks were evaluated in regions
where they were also trained, while in the “Transfer Leaning” competition, the
evaluation was performed in regions where training data were not available,
emphasizing the ability of the model to generalize outside the training data.
The competition entry by the present author [7] used a recurrent-convolutional
neural network with modified Convolutional Gated Recurrent Unit (ConvGRU, [8])
layers organized in a multiscale encoder-forecaster architecture. This
architecture proved to be an effective solution to the problem posed in the
challenge and achieved first place in both the Core and Transfer Learning
competitions.

After the conclusion of Stage 1, Weather4cast 2021 continued with the IEEE Big
Data Cup competition. This competition was identical in structure to the
original except for the addition of new data regions for both training and
evaluation. Therefore, improvements could be made by making optimal use of the
new data and by improving upon the models used in Stage 1. This paper presents
the improvements made by the author in the IEEE Big Data Cup competition that
resulted in better model performance. The code that can be used to replicate
the results can be found in https://github.com/jleinonen/weather4cast-bigdata;
the weights for the pre-trained models are provided at [9].

##  II Models and training

The predictions posted by the author to the IEEE Big Data Cup competition
utilized the same basic model design used in the Stage 1 competition and
described in [7]. The IEEE Big Data Cup competition dataset included more
training data in the form of new regions. These data were utilized by
retraining the models with the new dataset; furthermore, the static data
provided by the competition organizers, consisting of surface elevation,
latitude and longitude, were used for each of the regions. No external data
sources were used, nor were any calculations performed to generate additional
input data, except for data augmentation with mirroring and rotation.
Modifications were also made to the models and the training procedure to
further improve the performance. These included a shallower variant of the
model, a change to a different optimizer, conditional modeling to address
representativeness issues with crr_intensity, and model ensembling to reduce
the loss. The sections below describe the improvements in more detail. Other
than these improvements, the models and the training procedure were the same
as in [7], and the reader is referred to that paper for the details.

###  II-A Shallower model variant

In the Stage 1 competition, two variants of the recurrent-convolutional model
were investigated: one that used standard ConvGRU gates and another (ResGRU)
that utilized residual units in place of the convolutions. In this work, the
model selection was expanded with a variant where the deepest layer of the
original ResGRU model was removed. This variant was found to outperform the
deeper network in some cases, and had the further advantage of using much
fewer weights ($4.6$ million, compared to $18.6$ million in the full-size
model).

###  II-B AdaBelief optimizer

While the original worked used the widely adopted Adam optimizer, in this work
AdaBelief [10], a modification of the algorithm used in Adam [11], was
adopted. According to its authors, AdaBelief improves the generalization
ability of the solutions it finds compared to those found by Adam. Further
information can be found in the repository at https://github.com/juntang-
zhuang/Adabelief-Optimizer.

The hyperparameters of AdaBelief and comparisons to Adam were examined using
the temperature variable. After some experimentation with the parameters,
AdaBelief with the default hyperparameters was adopted, with the $\epsilon$
parameter set to $10^{-14}$ (the default of the TensorFlow implementation), as
non-default parameters either degraded the results or failed to provide a
significant benefit. The runs with variables other than temperature were
performed with AdaBelief using these hyperparameters, as the time remaining in
the contest did not allow for parameter tuning for each variable separately.

###  II-C Adjustments to rain rate prediction

Predicting crr_intensity involves particular challenges that do not apply to
the other variables because often it is not raining at all, while much of the
rainfall is delivered by only a few events. In Stage 1, it was already noted
that a different model performed best for crr_intensity compared to the other
variables. In further investigation of the statistics of this variable, it was
also found that there were issues with the representativeness of the training
dataset: The validation set had an average crr_intensity of $7.26\times
10^{-4}$ on the $[0,1]$ normalized scale, while for the training set it was
considerably lower, $6.01\times 10^{-4}$. The equivalent discrepancies for the
other variables were at least an order of magnitude smaller.

The problem of different distributions in the training and validation datasets
was approached by building a conditional model containing two copies of the
same model. The results from one of these copies would be used if the maximum
rain rate in the input time series was less than a threshold value of $0.026$
on the $[0,1]$ normalized scale; otherwise the other copy would be used. The
threshold is was chosen such that it divides the training dataset
approximately evenly over the two models. The conditional approach helps
adjust for the representativeness problem because, as the validation set
contains more raining cases than the training set, the model used for the
higher rain rates would be used more often to compensate for this.

Two such conditional models were trained for crr_intensity. In the first, the
conditional model contained two copies of the ConvGRU architecture used for
crr_intensity in [7]; before training, the weights of both copies were
initialized to a model trained for the entire dataset. In the second, slightly
differently constructed ConvGRU models were used to avoid certain inconvenient
architectural features of the older model, and the entire conditional model
was trained from random initialization.

###  II-D Model ensembling

The experimentation with different types of models and hyperparameters
produced multiple relatively good models. This presented an opportunity to
further improve the results by ensembling the models. As most of the
experimentation and tuning was done with the temperature variable, this
variable had more ensemble members available. For temperature, the ensemble
consisted of variants of the shallow model (see Sect. II-A); the deep model
weights were not used as they became incompatible with the final model
architecture due to changes during the development process. For
asii_turb_trop_prob and cma, the ensemble consisted of a mixture of deep and
shallow models, while the crr_intensity ensemble used the two different
conditional models described in Sect. II-C. The final ensembles had five
members for temperature, three for asii_turb_trop_prob and two each for
crr_intensity and cma.

Three ensembling strategies were attempted to produce predictions for the
final leaderboard. In the first, members were equally weighted with the
weights summing to $1$. In the second, optimal weights $\mathbf{w}$ were
derived using ridge regression [12], also called Tikhonov regularization [13],
which gives weights that optimize mean square error as

$\mathbf{w}=\left(\mathbf{X}^{\mathrm{T}}\mathbf{X}+\lambda\mathbf{I}\right)^{-1}\mathbf{X}^{\mathrm{T}}\mathbf{y}$
(1)

where $\mathbf{X}$ is a dimension $(n,p)$ matrix with $p$ models and $n$
predictions per model, $\mathbf{y}$ is a $(n,1)$ vector with the correct
values, and $\lambda$ is a regularization parameter. For all variables,
$\lambda$ was set to $10^{-4}$ times the mean of the diagonal of
$\mathbf{X}^{\mathrm{T}}\mathbf{X}$. The matrices
$\mathbf{X}^{\mathrm{T}}\mathbf{X}$ and $\mathbf{X}^{\mathrm{T}}\mathbf{y}$
were computed using the validation dataset due to suspicion that slight
overfitting of the models to the training set might weaken the derivation of
the weights. This also motivated a third strategy as a variant of the second,
as it was noted that the weights for crr_intensity summed to a total of
$1.0526$, even as the sum of weights for the other three variables deviated
from $1$ by less than $0.004$. This likely reflects the representativeness
issue that was already noted in Sect. II-C. While the nonzero weighting
considerably improved the metric for the validation set, it was not clear if
the systematic difference would also be valid for the held-out dataset.
Therefore, a solution was also created where the weights for crr_intensity
were derived as optimal weights with the sum constrained to $1$. These can be
derived using the method of Lagrange multipliers [13], which gives

$\begin{bmatrix}\mathbf{w}\\\
\mu\end{bmatrix}=\begin{bmatrix}\mathbf{X}^{\mathrm{T}}\mathbf{X}&-\frac{1}{2}\mathbf{q}\\\
\mathbf{q}^{\mathrm{T}}&0\\\
\end{bmatrix}^{-1}\begin{bmatrix}\mathbf{X}^{\mathrm{T}}\mathbf{y}\\\
1\end{bmatrix}$ (2)

where $\mu$ is a constraint parameter that can be discarded after solving the
equation, and $\mathbf{q}=[1,1,...,1]^{\mathrm{T}}$ is a $(p,1)$ vector with
all elements set to $1$.

The ensemble weights are shown in Table I. The constrained optimal weights
produced by Eq. 2 are also shown for all variables for reference purposes.
However, no submission was attempted with these other than for crr_intensity,
as for the other variables the weights given by Eq. 1 were already close to
$1$ and the negative weight given by Eq. 2 for one of the temperature ensemble
member was deemed suspicious.

TABLE I: Ensemble weights | temperature | crr_intensity | asii_turb_ | cma
---|---|---|---|---
|  |  | trop_prob |
Ridge |  |  |  |
(1) | $0.1455$ | $0.5206$ | $0.4344$ | $0.4864$
(2) | $0.2666$ | $0.5320$ | $0.2722$ | $0.5165$
(3) | $0.0904$ |  | $0.2941$ |
(4) | $0.2487$ |  |  |
(5) | $0.2457$ |  |  |
Constrained |  |  |  |
(1) | $0.2094$ | $0.5122$ | $0.4592$ | $0.4846$
(2) | $0.3224$ | $0.4878$ | $0.2613$ | $0.5154$
(3) | $-0.0664$ |  | $0.2795$ |
(4) | $0.2660$ |  |  |
(5) | $0.2687$ |  |  |

##  III Results

In the course of model development, the various model candidates were
evaluated using the validation dataset after training and the results were
recorded. A summary of these is shown in Table II. The first row shows the
results using the best models from Stage 1 without retraining for the new
dataset. Models D1 and D2 are different training runs of the deeper variant,
models S1–S5 use the shallow variant, while C1 and C2 are the conditional
models used for crr_intensity. The results for the equally weighted ensemble
and the ensemble with weights obtained using ridge regression are also shown.
It should be noted that models were trained for each variable separately; for
example, the S1 model for temperature and the S1 model for cma only share a
common architecture, not common weights. The different versions of the shallow
model for temperature are as follows: S1 was trained from random
initialization using the default settings of the AdaBelief optimizer, S2 was
trained by initializing the weights with those of S1 and resetting the
learning rate to $10^{-3}$ in the beginning of training, S3 was trained like
S2 but setting the $\epsilon$ parameter to $10^{-7}$ (another value suggested
by the authors), S4 was trained like S1 but enabling weight decay in
AdaBelief, and S5 was trained like S1 but using the Adam optimizer instead.

TABLE II: Summary of the validation set metrics of the different variables and models | temperature | crr_intensity |  asii_turb_ trop_prob | cma
---|---|---|---|---
Stage 1 | $0.004539$ | $0.00009126$ | $0.002165$ | $0.1512$
Deep |  |  |  |
D1 | $0.004370$ | $0.00009107$ | $0.002033^{\mathrm{1}}$ | $0.1440^{\mathrm{1}}$
D2 | $0.004335$ | $0.00009122$ | $0.002021^{\mathrm{2}}$ |
Shallow |  |  |  |
S1 | $0.004345^{\mathrm{1}}$ |  | $0.002065^{\mathrm{3}}$ | $0.1438^{\mathrm{2}}$
S2 | $0.004349^{\mathrm{2}}$ |  |  |
S3 | $0.004344^{\mathrm{3}}$ |  |  |
S4 | $0.004355^{\mathrm{4}}$ |  |  |
S5 | $0.004378^{\mathrm{5}}$ |  |  |
Cond. |  |  |  |
C1 |  | $0.00008960^{\mathrm{1}}$ |  |
C2 |  | $0.00008968^{\mathrm{2}}$ |  |
Ensemble |  |  |  |
Equal | $0.004285$ | $0.00008872$ | $0.001989$ | $0.1423$
Ridge | $0.004282$ | $0.00008588$ | $0.001987$ | $0.1424$
1-5Included in the ensembles (Table I)

Multiple general trends can be identified from the results of Table II. First,
it can be seen that simply retraining the model with the new data and the
AdaBelief optimizer improved the results considerably (up to $4.8\%$ for cma)
over the models that produces the best results in Stage 1, except for
crr_intensity where the improvement was more marginal ($0.2\%$). Between the
retraining and the new optimizer, the addition of new training data seems to
be the more important factor, as the S5 model trained with Adam was only
slightly worse than the other models trained with AdaBelief. Second, the
adoption of the conditional model for crr_intensity provided an improvement
($1.8\%$) for that variable that was more significant than the retraining.
Third, the ensembling further improves the scores for all variables. The
ensemble of equally weighted models provided benefits of $1.0\%$–$1.6\%$ over
the best individual model, depending on the variables. Ridge regression
weighting provided results that were very similar to the equal-weights
ensemble for all variables except for crr_intensity, where the improvement was
large ($3.2\%$) compared to the other benefits from ensembling. This is likely
because the ensemble weights for crr_intensity, whose sum differs considerably
from $1$, correct for the bias between the training and validation sets notes
in Sect. II-D. However, it was not clear a priori if this bias would also
apply to the test and held-out datasets.

Figs. 1–4 show the results of prediction with the ensemble models using the
ridge regression weights. Like the predictions in [7], these show that the
model is able to detect and predict the motion of the input fields.
Predictions start sharp and become blurrier with longer lead times, reflecting
the increasing uncertainty of the prediction.

Figure 1: An example of a model prediction for the temperature variable using
the ridge regression model ensemble. The left side shows the past temperature,
while the right side show the real temperature in the future (top row) and the
corresponding model prediction (bottom row). $T$ is the index of the frame in
the sequence; $T=0$ refers to the last input frame and $T=1$ is the first
prediction. Scaled temperature normalized to the range $(0,1)$, covering the
data range in the whole dataset, is shown.

Figure 2: As Fig. 1, but for crr_intensity. A different case is shown because
the case in Fig. 1 contains negligible precipitation.

Figure 3: As Fig. 1, but for asii_turb_trop_prob.

Figure 4: As Fig. 1, but for cma. The white lines show the contour at $0.5$,
which is the decision threshold for a binary cloud mask.

The results achieved on the competition leaderboards are shown in Table III.
The first two columns show the test leaderboard scores for the Core and
Transfer Learning competitions, respectively, while the third and fourth
column show the equivalent scores for the final leaderboards which used the
held-out dataset and determined the results of the competition. The retraining
improved the scores considerably, as the best single models outperform the
Stage 1 models by $3.0\%$ for the Core competition and $5.4\%$ for the
Transfer Learning competition. The conditional model for crr_intensity
improved the scores slightly, and the ensembling provided a benefit comparable
to the improvements seen in Table II. Using optimized weights from Eq. 1
improved the scores on the final leaderboards, but only marginally.

TABLE III: Results on the leaderboards | Test | Test | Final | Final
---|---|---|---|---
| Core | Transfer | Core | Transfer
Stage 1 | $0.4757$ | $0.4552$ |  |
Best Single | $0.4615$ | $0.4308$ |  |
Conditional rain | $0.4605$ | $0.4289$ |  |
Equal Ensemble | $\mathbf{0.4530}$ | $\mathbf{0.4229}$ | $0.4734$ | $0.4328$
Ridge Ensemble |  |  | $\mathbf{0.4729}$ | $\mathbf{0.4323}$

##  IV Conclusions

The results of the efforts to improve the model offer various insights into
the training of multi-scale recurrent-convolutional models for image time
series prediction tasks and particularly weather forecasting:

  1. 1.

The models appear to be able to learn better from a larger amount of training
data; this was most clearly evidenced by the fact that the scores in the
Transfer Learning competition increased significantly, indeed more than those
in the Core competition, after retraining with additional data. This indicates
that the additional data improved the ability of the model to generalize.
While the larger effect in the Transfer Learning competition might also have
resulted from the new Core dataset being more representative of the Transfer
Learning regions, it is not clear if this is the case.

  2. 2.

The model introduced in [7] performs well also in a variant that has the
deepest layers removed. This model attained similar results to the original,
deeper variant, sometimes outperforming it and sometimes being slightly worse.
The shallower variant has much fewer weights than the deep variant, although
its computational and memory requirements are not much smaller as these are
determined mostly by the top levels of the network.

  3. 3.

The AdaBelief optimizer seems to somewhat outperform the popular Adam
optimizer at this task, giving better results with the validation dataset. The
results obtained in this study support the claim by the authors of AdaBelief
that results obtained with it generalize better to datasets outside the
training data.

  4. 4.

For datasets where the training data do not represent well the validation
and/or testing data, such as crr_intensity in this competition, using
conditional models (Sect. II-C) can improve the results. However, the solution
adopted here — choosing the model based on the maximum rain rate in the input
sequence — was rather ad hoc, and better can likely be found. Indeed, better
solutions are likely already documented in the literature, but the time
constraints of the competition prohibited an exhaustive search for the best
method.

  5. 5.

Ensembling provided an improvement on the order of $1\%$ to the scores. While
such an advantage can potentially be decisive in a competition environment, it
depends on the use case whether such an improvement is worth the added
training in real-world applications. However, researchers who test different
models and hyperparameters during the development process may benefit from
saving the different variants and finally ensembling them, as this provides
the ensemble members without additional resources spent on training.

Among potential further improvements, the author recommends investigating the
role of regularization, normalization and dropout, which were only examined in
a very limited fashion in this study, as they did not produce obvious benefits
in the model investigated here, but were very successfully used by other
challenge participants (e.g. [14]) to limit overfitting. Another interesting
technique to investigate would be transformer networks, which have proven very
effective for modeling sequences; another competitor used transformers, and
while the model presented here outperformed the transformer model in the
competition, it may be possible to find a way to combine the advantages of the
two approaches. Additional input variables could also further improve the
results; while the author did not find substantial benefits from using inputs
beyond the four target variables (possibly because the experimentation was
performed with the temperature variable), the other competitors’ models did
benefit from them. Furthermore, the prediction of the cma variable, which is
essentially a binary classification problem, could benefit from an alternative
loss function. These were attempted in this study but not studied thoroughly,
so further research is needed.

## Acknowledgment

The author thanks U. Hamann and A. Rigazzi for discussions regarding the model
and training.

## References

  * [1] S. Ravuri, K. Lenc, M. Willson, D. Kangin, R. Lam, P. Mirowski, M. Fitzsimons, M. Athanassiadou, S. Kashem, S. Madge, R. Prudden, A. Mandhane, A. Clark, A. Brock, K. Simonyan, R. Hadsell, N. Robinson, E. Clancy, A. Arribas, and S. Mohamed, “Skilful precipitation nowcasting using deep generative models of radar,” _Nature_ , vol. 597, pp. 672–677, 2021. doi:10.1038/s41586-021-03854-z
  * [2] A. J. Geer, “Learning earth system models from observations: machine learning or data assimilation?” _Phil. Trans. R. Soc. A._ , vol. 379, no. 2194, p. 20200089, 2021. doi:10.1098/rsta.2020.0089
  * [3] A. Seifert and S. Rasp, “Potential and limitations of machine learning for modeling warm-rain cloud microphysical processes,” _J. Adv. Model. Earth Syst._ , vol. 12, no. 12, p. e2020MS002301, 2020. doi:10.1029/2020MS002301
  * [4] S. E. Haupt, W. Chapman, S. V. Adams, C. Kirkwood, J. S. Hosking, N. H. Robinson, S. Lerch, and A. C. Subramanian, “Towards implementing artificial intelligence post-processing in weather and climate: proposed actions from the Oxford 2019 workshop,” _Philos. Trans. R. Soc. London, Ser. A_ , vol. 379, no. 2194, p. 20200091, 2021. doi:10.1098/rsta.2020.0091
  * [5] P. Dueben, U. Modigliani, A. Geer, S. Siemen, F. Pappenberger, P. Bauer, A. Brown, M. Palkovic, B. Raoult, N. Wedi, and V. Baousis, “Machine learning at ECMWF: A roadmap for the next 10 years,” European Centre for Medium-Range Weather Forecasts, Tech. Rep. 878, 2021. doi:10.21957/ge7ckgm
  * [6] A. Gruca, P. Herruzo, P. Rípodas, A. Kucik, C. Briese, M. K. Kopp, S. Hochreiter, P. Ghamisi, and D. P. Kreil, “CDCEO’21 - first workshop on complex data challenges in Earth observation,” in _Proceedings of the 30th ACM International Conference on Information & Knowledge Management_, ser. CIKM ’21. New York, NY, USA: Association for Computing Machinery, 2021, pp. 4878–4879. doi:10.1145/3459637.3482044
  * [7] J. Leinonen, “Spatiotemporal weather data predictions with shortcut recurrent-convolutional networks: A solution for the Weather4cast challenge,” in _1st workshop on Complex Data Challenges in Earth Observation (CDCEO) 2021_ , 2021.
  * [8] N. Ballas, L. Yao, C. Pal, and A. C. Courville, “Delving deeper into convolutional networks for learning video representations,” 2016. [Online]. Available: https://arxiv.org/abs/1511.06432
  * [9] J. Leinonen, “Model weights for a Weather4cast 2021 Challenge IEEE Big Data Cup Stage solution,” 2021. doi:10.5281/zenodo.5646287
  * [10] J. Zhuang, T. Tang, Y. Ding, S. C. Tatikonda, N. Dvornek, X. Papademetris, and J. Duncan, “AdaBelief optimizer: Adapting stepsizes by the belief in observed gradients,” in _Advances in Neural Information Processing Systems_ , H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, Eds., vol. 33. Curran Associates, Inc., 2020, pp. 18 795–18 806. [Online]. Available: https://proceedings.neurips.cc/paper/2020/file/d9d4f495e875a2e075a1a4a6e1b9770f-Paper.pdf
  * [11] D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” in _3rd International Conference for Learning Representations, San Diego, California, USA_ , 2014. [Online]. Available: https://arxiv.org/abs/1412.6980
  * [12] A. E. Hoerl and R. W. Kennard, “Ridge regression: Biased estimation for nonorthogonal problems,” _Technometrics_ , vol. 12, no. 1, pp. 55–67, 1970\. doi:10.2307/1267351
  * [13] S. Boyd and L. Vandenberghe, _Convex Optimization_. Cambridge, UK: Cambridge University Press, 2004.
  * [14] P. H. Kwok and Q. Qi, “A variational U-Net for weather forecasting,” in _1st workshop on Complex Data Challenges in Earth Observation (CDCEO) 2021_ , 2021.