# Evolution Transformer: In-Context Evolutionary Optimization

Robert Tjarko Lange  TU Berlin, Google DeepMindGermany, Japan
<EMAIL_ADDRESS>,  Yingtao Tian  Google DeepMindJapan
<EMAIL_ADDRESS>and  Yujin Tang  Google DeepMindJapan<EMAIL_ADDRESS>
(2024)

###### Abstract.

Evolutionary optimization algorithms are often derived from loose biological
analogies and struggle to leverage information obtained during the sequential
course of optimization. An alternative promising approach is to leverage data
and directly discover powerful optimization principles via meta-optimization.
In this work, we follow such a paradigm and introduce Evolution Transformer, a
causal Transformer architecture, which can flexibly characterize a family of
Evolution Strategies. Given a trajectory of evaluations and search
distribution statistics, Evolution Transformer outputs a performance-improving
update to the search distribution. The architecture imposes a set of suitable
inductive biases, i.e. the invariance of the distribution update to the order
of population members within a generation and equivariance to the order of the
search dimensions. We train the model weights using Evolutionary Algorithm
Distillation, a technique for supervised optimization of sequence models using
teacher algorithm trajectories. The resulting model exhibits strong in-context
optimization performance and shows strong generalization capabilities to
otherwise challenging neuroevolution tasks. We analyze the resulting
properties of the Evolution Transformer and propose a technique to fully self-
referentially train the Evolution Transformer, starting from a random
initialization and bootstrapping its own learning progress. We provide an open
source implementation under https://github.com/RobertTLange/evosax.

evolution strategies, machine learning

††copyright: rightsretained††journalyear: 2024††copyright:
rightsretained††ccs: Computing methodologies Evolutionary Robotics Figure 1.
Left. Evolution Transformer overview. The model encodes three different types
of information (solution space, fitness & search distribution). The resulting
per-search dimension embedding is processed by an update network using
parameter sharing. The Evolution Transformer is trained by generating
optimization trajectories from teacher BBO algorithms and minimizing the KL
divergence between student and teacher search distribution updates. Right.
Aggregated results on 8 Brax tasks and 5 individual runs (see SI, Figure 10).
We report the interquartile mean and optimality gap after normalizing by the
task-specific OpenAI-ES (Salimans et al., 2017) performance. Evolution
Transformer robustly outperforms Evolutionary Optimization baselines.

##  1\. Introduction

Evolutionary optimization (EO; e.g. Salimans et al., 2017; Such et al., 2017)
provides a scalable and simple-to-use tool for black-box optimization (BBO).
Instead of having to rely on access to well-behaved gradients, such methods
only require the evaluation of the function (e.g. neural network loss,
training run, or physics simulator). This makes them especially well suited
for otherwise challenging applications such as meta-learning or neural
architecture search. Most of the specific algorithms are derived from loose
biological metaphors and are oftentimes not very adaptive to the problem at
hand. Therefore, recent work has tried to address these limitations by aiming
to meta-optimize attention-based modules of EO algorithms (Lange et al.,
2023a, b).

Here, we extend this line of work and introduce the first fully Transformer-
based Evolution Strategy (ES), which we call Evolution Transformer (EvoTF;
Figure 1, Left Top). It leverages both self- and cross-attention to implement
two crucial inductive biases desirable for BBO algorithm: Invariance of the
distribution update with respect to the within-generation ordering of the
solutions and equivariance with respect to the ordering of the individual
search dimensions. We show that this Transformer can be successfully trained
via supervised learning using optimization trajectories obtained by teacher
algorithms. We name this method Evolutionary Algorithm Distillation (EAD;
Figure 1, Left Bottom). After training, the resulting model can imitate
various teacher BBO algorithms and performs in-context evolutionary
optimization (Figure 1, Right). Furthermore, the resulting Evolution Strategy
generalizes to previously unseen optimization problems, the number of search
dimensions, and population members. We propose two avenues to improve upon the
teacher’s performance by introducing additions to the supervised algorithm
distillation training setup: Finally, we introduce Self-Referential
Evolutionary Algorithm Distillation (SR-EAD), which can be used to self-train
Evolution Transformer models: Given a parametrization, we can generate
multiple perturbed trajectories, filter such, and use them to bootstrap the
observed performance improvements. Hence, it provides a promising direction
for the open-ended discovery of in-context evolutionary optimizers. Our
contributions are summarized as follows:

  1. (1)

We introduce Evolution Transformer, an architecture inducing a population-
order invariant and dimension-order equivariant search update (Section 3).
After supervised training with Evolutionary Algorithm Distillation, the
Transformer has distilled various evolutionary optimizers and performs in-
context learning on unseen tasks (Section 4).

  2. (2)

We analyze the Transformer BBO algorithm and show empirically that it captures
desirable properties such as scale-invariance and perturbation strength
adaptation (Section 5).

  3. (3)

We compare supervised EAD with the meta-evolution of Evolution Transformer
network parameters. We find that meta-evolution is feasible but tends to
overfit the meta-training tasks and requires more accelerator memory (Section
6).

  4. (4)

We introduce Self-Referential Algorithm Distillation, which alleviates the
need for teacher algorithm specification. By perturbing the parameters of the
Evolution Transformer, we construct a set of diverse self-generated
trajectories. After performance filtering, the sequences can be used to self-
referentially train the model from scratch (Section 7).

Figure 2. Evolution Transformer. We construct features resembling information
from solution evaluations and the search distribution. They are processed by
self-attention and Perceiver modules to obtain four separate embeddings. The
stacked per-dimension embeddings are processed by standard Transformer encoder
blocks. An MLP outputs the distribution update predictions. The model is
invariant to the order of the population members and equivariant to the search
dimension order.

##  2\. Related Work & Background

Black-Box Optimization. Our work develops efficient BBO algorithms with
flexible Transformer-based parametrization. Given a function
$f(\mathbf{x}):\mathbb{R}^{D}\to\mathbb{R}$ with unknown functional form, i.e.
we cannot compute its derivative (or it is not well behaved), BBO seeks to
find its global optimum leveraging only function evaluations:

$\min_{\mathbf{x}}f(\mathbf{x}),\text{s.t.}\
\mathbf{x}_{d}\in[l_{d},u_{d}]\subset[-\infty,\infty],\forall d=1,...,D.$

Evolutionary Optimization Methods. EO provides a set of BBO algorithms
inspired by the principles of biological evolution. Roughly speaking, they can
be grouped into Evolution Strategies (Rechenberg, 1978) and Genetic Algorithms
(GA), which differ in the way they perform the selection and mutation of
solution candidates. Here, we focus on diagonal Gaussian ES. Given a
population size $N$ and the search distribution summary statistics
$\bm{\mu}\in\mathbb{R}^{D},\Sigma=\bm{\sigma}\mathbf{1}_{D\times
D}\in\mathbb{R}^{D\times D}$, ES sample a population of candidate solutions
$X=[x_{1},...,x_{N}]\in\mathbb{R}^{N\times D}$ at each generation. Afterwards,
the performance (or fitness) of each candidate is evaluated on the task of
interest and one obtains fitness scores
$F=[f_{1},...,f_{N}]\in\mathbb{R}^{N}$. The search distribution is then
updated to increase the likelihood of sampling well-performing solutions,
$\bm{\mu}^{\prime},\bm{\sigma}^{\prime}\leftarrow\texttt{UPDATE}(\bm{\mu},\bm{\sigma},X,F,H)$,
where $H\in\mathbb{R}^{D\times D_{H}}$ denotes a set of ES summary statistics
(e.g. momentum-like terms). There exist various types of ES including
estimation-of-distribution (CMA-ES, Hansen, 2006), natural (SNES, Schaul et
al., 2011; Wierstra et al., 2014) and finite-difference-based ES (OpenAI-ES,
Salimans et al., 2017). Evolution Transformer provides a flexible
parametrization of the ES-update
$\texttt{UPDATE}_{\theta}(\bm{\mu},\bm{\sigma},X,F,H)$ and optimizes the set
of Transformer weights $\theta$ to distill teacher BBO algorithm updates using
gradient-descent.
Meta-Learned Evolutionary Optimizers. Formally, ES-update rules are inherently
set of operations, i.e. the order of the population members within a
generation should not affect the performed distribution change. Self-attention
provides a natural inductive bias for such an operation. Previously, Lange et
al. (2023b, a) constructed ES and GA algorithms, which used self- and cross-
attention to process the information within a single generation. The
associated parameters were then meta-evolved on a small task distribution of
BBO problems. Unlike Evolution Transformer, the meta-evolved EO was not a
sequence model, i.e. does not leverage a causal Transformer with positional
encoding.
Self-Referential Higher-Order Evolution. Schmidhuber (1987) first articulated
the vision of self-refining Genetic Algorithms. Later on Metz et al. (2021)
showed that randomly initialized black-box gradient-descent optimizers can
optimize themselves using large-scale population-based training (Jaderberg et
al., 2017). Kirsch and Schmidhuber (2022) further formulated a rule to
automatically allocate resources for self-referential meta-learning. Finally,
(Lu et al., 2023b) provides theoretical results for infinite-order meta-
evolution.
Autoregressive Models for BBO. Chen et al. (2017) proposed to learn RNN-based
BBO algorithms using privileged access to gradient computations of the fitness
(Gaussian Process sampled) functions at training time. Furthermore, Dery et
al. (2022); Chen et al. (2022) used large pre-collected datasets to train a T5
Encoder-Decoder architecture (Raffel et al., 2020) for BBO, called OptFormer.
Krishnamoorthy et al. (2022) later on investigated the usage of generating
offline data for training autoregressive BBO models. Lange et al. (2024b) show
that language models can be used as recombination operators for ES. In
contrast, our work introduces an architecture specialized to flexibly
represent ES algorithms and proposes several complementary enhancements to the
supervised training pipeline to improve upon teacher algorithms. Finally,
OptFormer is intended to be applied to hyperparameter optimization problems
and not to train large numbers of neural network weights, unlike Evolution
Transformer.
Meta-Learned In-Context Learning & AD for RL. Kirsch et al. (2022); Lu et al.
(2023a); Team (2023) trained large sequence models using supervised and
reinforcement learning without access to teacher algorithms. Instead, they
only leveraged large programmatically-generated or augmented task spaces to
induce in-context learning across long timescales. Laskin et al. (2022), on
the other hand, introduced the concept of Transformer-based AD in the context
of Reinforcement Learning. They showed that it is possible to perform AD for
actor-critic-based teacher algorithms. Afterwards, a Transformer policy shows
in-context learning capabilities when evaluated on unseen environments. Our
work extends these ideas to zero-order optimization and introduces a tailored
architecture and AD enhancements in order to increase performance.
Self-Attention & Perceiver Cross-Attention. The Transformer (Vaswani et al.,
2017) stacks blocks of multi-head self-attention (MHSA) operations,
feedforward MLP transformation, dropout and layer normalization. At its core
self-attention is a set operation which projects an input matrix
$X\in\mathbb{R}^{N\times D}$ onto $D_{K}$-dimensional vectors
$Q,K,V\in\mathbb{R}^{N\times D_{K}}$ called queries, keys and values,
respectively:

$\displaystyle\text{Attention}(X)$
$\displaystyle=\text{softmax}(QK^{T}/\sqrt{D_{K}})V$
$\displaystyle=\text{softmax}(XW_{Q}(XW_{K})^{T}/\sqrt{D_{K}})XW_{V}.$

Permuting the rows of $X$ will apply the same permutation to
$\text{Attention}(X)$ (e.g., Lee et al., 2019; Tang and Ha, 2021; Kossen et
al., 2021). MHSA has quadratic complexity with respect to the sequence length
$N$. The Perceiver module (Jaegle et al., 2021) was introduced to partially
alleviate this problem. Instead, it performs a cross-attention operation:

$\displaystyle\text{Perceiver}(X,Z)=\text{softmax}(ZW_{Q}(XW_{K})^{T}/\sqrt{D_{K}})XW_{V},$

where $Z\in\mathbb{R}^{N_{Z}\times D_{Z}}$ denotes a set of learned latent
vectors. Hence, the dimensionality of
$\text{Perceiver}(X,Z)\in\mathbb{R}^{N_{Z}\times D_{K}}$ is fixed and
independent of the sequence length.

##  3\. Evolution Transformer: Population-Order In- & Dimension- Order
Equivariant Search Updates

We design a Transformer-based architecture that can flexibly represent an ES
for different population sizes and the number of search space dimensions. To
enable this endeavor, we leverage self-attention and Perceiver cross-attention
to induce population-order invariance and batch over the individual search
dimensions to obtain a dimension-order equivariant neural network ES. At a
high level the architecture (Figure 2) constructs a set of embedding features
from evaluated solution candidates, their fitness scores, and search
distribution statistics per generation. Such features consist of common
Evolutionary Optimization ingredients such as finite difference gradients,
normalized fitness scores, and evolution paths/momentum statistics (see SI for
detailed input feature specification). After combining these features, we
perform causal self-attention over the different generations. Finally, we
output the search distribution updates per search dimension.
Solution Perceiver. Given solution candidates $X\in\mathbb{R}^{N\times D}$
sampled within a generation, we pre-compute the set of normalized features for
each dimension, $\tilde{X}\in\mathbb{R}^{N\times D\times D_{X}}$. The
information is processed by a Perceiver module applied per individual search
dimension $j$. The weights are shared across the dimensions. Hence, we perform
the following operations:

$\displaystyle\bm{h}_{d}^{S}=\text{softmax}\left(\frac{ZW_{Q}(\tilde{X}[:,d]W_{K})^{T}}{\sqrt{D_{K}}}\right)\tilde{X}[:,d]W_{V}\in\mathbb{R}^{N_{Z}\times
D_{K}},$

for $d=1,\dots,D$. Note that this information compresses the population
information into a set of $D$ tensors with a fixed shape. Thereby, the
dimensionality of the resulting representation is independent of the number of
population members. This together with per-dimension updates allows for
general applicability to different BBO settings.
Fitness Perceiver. The fitness scores $F\in\mathbb{R}^{N}$ are similarly first
processed into a set of features $\tilde{F}\in\mathbb{R}^{N\times D_{F}}$ and
afterwards compressed into a set of Perceiver latent vectors,
$\bm{h}^{F}=\mathbb{R}^{N_{Z}\times D_{K}}$.
Distribution Attention & Cross-Dimension Perceiver. We further process a set
of features capturing the search distribution dynamics (e.g. momentum and
evolution paths with different time scales for both $\bm{\mu}$ and
$\bm{\sigma}$). We perform both MHSA within a single dimension and Perceiver
compression across dimensions.
MHSA Transformer Blocks Across Time. The four embedding steps can be computed
in parallel over time, making the implementation efficient on hardware
accelerators. Afterwards, we repeat and concatenate the different embeddings
to construct a single stacked per-dimension representation, $\bm{h}$. The
aggregated representation is further processed by a set of standard MHSA
Transformer blocks with positional encoding over the timesteps/generations.
Distribution Update Module. The final representation is projected per
dimension to output the proposed change in the search distribution and we
perform a multiplicative update inspired by SNES (Schaul et al., 2011):

$\displaystyle\bm{\mu}_{\text{EvoTF}}^{\prime}$
$\displaystyle=\bm{\mu}+\eta_{\mu}\times\bm{\sigma}\times\texttt{EvoTransformer}_{\mu}(X,F,\bm{h})$
$\displaystyle\bm{\sigma}_{\text{EvoTF}}^{\prime}$
$\displaystyle=\bm{\sigma}\times\exp(\eta_{\sigma}\times\texttt{EvoTransformer}_{\sigma}(X,F,\bm{h})),$

with tune-able learning rates $\eta_{\mu},\eta_{\sigma}$. During supervised
algorithm distillation training both of them are set to 1.
Evolutionary Algorithm Distillation. We train Evolution Transformer to distill
the search procedure of various teacher algorithms. More specifically, we
collect batches of teacher trajectories
$\\{X_{g},F_{g},\mu_{g},\sigma_{g}\\}_{g=1}^{G}$ on a set of synthetic BBO
tasks. We minimize the Kullback-Leibler (KL) divergence between the teacher’s
and proposed updated distribution at each generation with
$S=\bm{\sigma}^{2}\mathbf{I}_{D}$:

$\displaystyle\text{KL}((\bm{\mu}_{\text{EvoTF}},\bm{S}_{\text{EvoTF}})||(\bm{\mu}_{\text{T}},\bm{S}_{\text{T}}))=1/2(\text{tr}(\bm{S}_{\text{T}}^{-1}\bm{S}_{\text{EvoTF}})+$
$\displaystyle\log\frac{|\bm{S}_{T}|}{|\bm{S}_{\text{EvoTF}}|}+(\bm{\mu}_{\text{T}}-\bm{\mu}_{\text{EvoTF}})^{T}\bm{S}_{\text{T}}^{-1}(\bm{\mu}_{\text{T}}-\bm{\mu}_{\text{EvoTF}})-D).$

Deploying EvoTF-based Evolution Strategies. After successful teacher
distillation, we can use the trained Transformer in ’inference mode’ and use
it for optimization on different BBO tasks, e.g. with previously unseen
numbers of search dimensions or population members. Note that during training
the Evolution Transformer is never actively acting on the problems. Instead,
it only has to predict the on-policy behavior of the teacher algorithm. Hence,
on new evaluation problems, it has to effectively act ’off-policy’ since the
previous $\bm{\mu},\bm{\sigma}$ updates were generated by the Evolution
Transformer strategy itself and not by a teacher algorithm.
General Implementation Details. We use JAX (Bradbury et al., 2018) to
implement the Evolution Transformer architecture and dimension-batched search
updating. We further leverage evosax (Lange, 2022a) to parallelize the
generation of teacher ES algorithm sequences on hardware accelerators. The
architecture uses standard Transformer encoder blocks and during training, we
leverage a causal (lower diagonal) mask for autoregressively training
Evolution Transformer (Radford et al., 2018) to predict the teacher
algorithm’s search distribution updates. Training is conducted using a single
to four A100 NVIDIA GPU. For inference, we use a sliding window context of the
most recent $K$ generations. This is largely due to memory constraints arising
for $D>1000$ search dimensions. All Evolution Transformer ES evaluation
experiments are conducted using a single V100S NVIDIA GPU. The code is
published under https://github.com/<anonymized_repository>.

##  4\. Supervised Evolutionary Algorithm Distillation clones various teacher
BBO algorithms

Figure 3. Evolutionary Algorithm Distillation allows EvoTF to distill teacher
algorithms. Top. KL distillation loss with different Transformer modules.
Middle. Evaluation on a 14x14 MNIST CNN Classification task throughout
distillation. Bottom. ’S+F’ uses only the Solution and Fitness Perceiver.
’S+F+D’ also uses the Distribution Attention and ’S+F+D+CD’ uses all network
modules. Evaluation on a Pendulum MLP Control task throughout distillation.
Results are averaged across 3 independent runs.

Experiment Outline. We investigate whether the proposed Evolution Transformer
architecture is capable of distilling different individual BBO and ES teacher
algorithms. More specifically, we focus on 4 common classes of teacher BBO
algorithms: Finite-difference ES including OpenAI-ES (Salimans et al., 2017),
estimation-of-distribution ES (Sep-CMA-ES, Ros and Hansen, 2008), Gaussian
Hill Climbing and Natural Evolution Strategies (Schaul et al., 2011; Wierstra
et al., 2014). For each individual algorithm, we generate Evolution
Transformer AD training trajectories by executing it on a set of synthetic
tasks (see below) with different search settings. Afterwards, we train the
Evolution Transformer to auto-regressively predict the search distribution
updates of the teacher algorithm.
Implementation Details. The teacher optimization trajectories are sampled from
a subset of BBOB (Hansen et al., 2010) benchmark functions containing
different fitness landscape characteristics. For all trajectories, we use a
fixed set of search dimensions ($D=5$) and population size ($N=10$). We
randomize the initial mean of the search algorithm as well as its search
variance. Each trajectory consists of 32 generations and we collect the
solution, fitness, and distribution features. At each update step, we sample a
rollout batch of 32 trajectories online before constructing the KL loss. We
use pre-attention layer normalization, no dropout, and the Adam optimizer with
a cosine warmup learning rate schedule. The Evolution Transformer modules use
a single attention/Perceiver block and in total contain approximately 300k
parameters. Throughout supervised learning, we evaluate the performance of the
current network weight checkpoint with a fixed context window of $K=5$ on a
set of holdout tasks including both unseen BBOB and small neuroevolution
tasks.
Results Interpretation. In the first row of Figure 3 we show that this
procedure successfully distills the considered source teacher algorithms. We
note that the teacher algorithms vastly differ in their types of search update
equations. Hence, the Evolution Transformer architecture design is flexible
enough to represent all of these. Furthermore, the successful distillation of
a teacher BBO algorithm requires only a handful BBOB tasks (Figure 11). The
performance increase does not significantly increase with the number of
considered BBOB problems. Throughout the supervised AD training procedure, we
evaluate the Evolution Transformer ES on various downstream tasks. In this
case, the trained Evolution Transformer model is used to perform BBO without
any explicit teacher guidance. Instead, it has to perform optimization using
frozen network weights and only make use of the in-context information
provided by the ongoing BBO trajectory. We find that the Evolution Transformer
ES is capable of generalizing to completely unseen neuroevolution tasks such
as the ant control task (Figure 1). This indicates that the supervised AD
procedure did not lead to overfitting, but instead has led to successful in-
context evolutionary optimization. Finally, we perform ablation studies on the
modules used in the Evolution Transformer architecture (Figure 3). We find
that adding both the fitness and distribution information helps the model in
distilling the different algorithms. The cross-dimensional Perceiver is not
required. Note, that especially the distribution features are required to
distill SNES, Sep-CMA-ES, and OpenAI-ES.

##  5\. Analysis: Evolution Transformer Captures Desirable Evolution Strategy
Properties

Figure 4. Self-Attention and Perceiver maps for Evolution Transformer (EAD-
trained on SNES) with a single attention block at a single generation. Top.
Separable Sphere problem. Bottom. Non-separable Rosenbrock problem. All
problems are 3-dimensional and use 5 population members. The fitness attention
assigns higher credit to the best-performing population members. The
distribution attention indicates that the EvoTF correctly infers whether the
fitness landscape is separable.

EvoTF Properties on Characteristic Problems. After having demonstrated that
EAD can successfully distill teacher BBO into the Evolution Transformer, we
next analyze the properties of the resulting Evolution Strategy. We focus on
three core properties desirable for an Evolution Strategy (Hansen, 2006):

  1. (1)

Unbiasedness/Stationarity: Given a random fitness function, e.g.
$f(\bm{x}\sim\mathcal{N}(0,1)$, we do not want to observe a drift in the mean
statistic, i.e. $\mathbb{E}(\bm{\mu}^{\prime}|\bm{\mu})=\bm{\mu}$.

  2. (2)

Translation Invariance: Given an offset to the fitness function, the
performance of the ES is not going to change, e.g. $f_{b}(\bm{x})=f(\bm{x}-b)$
we desire $\bm{x}_{b}^{\star}=\bm{x}^{\star}-b$.

  3. (3)

Scale Self-Adaptation: Given a linear fitness function,
$f(\bm{x})=\sum_{d=1}^{D}\bm{x}_{d}$, we desire that the perturbation strength
$\bm{\sigma}$ increases over the course of the evolutionary optimization
trajectory.

We test these properties for an Evolution Transformer trained to distill the
SNES teacher algorithm. After supervised EAD training, we roll out the
resulting EvoTF Evolution Strategy and track the relevant statistics. In
Figure 5 we show that the model correctly incorporates all the desired
characteristics. The mean on a random fitness function does not show any clear
bias and remains at its initialization. Furthermore, it correctly finds the
optimal $\bm{x}_{b}^{\star}$ for different translation offset levels. Finally,
we find that the scale quickly increases for the linear fitness function,
which allows for quick adaptation and fitness improvements.

Evolution Transformer Attention Maps. Next, we visualize the attention maps
for an Evolution Transformer trained to distill the SNES teacher (Schaul et
al., 2011). We consider two different 3-dimensional BBO tasks: The separable
Sphere task and the Rosenbrock function moderate condition number. For both
cases, we consider 3 search dimensions and 5 population members. For both
problems, the fitness feature-based attention scores across the population
members are highest for the best-performing members within the considered
generation. Interestingly, we observe that the distribution self-attention
maps correctly identify the separability of the fitness landscape. For the
Sphere task, it attributes the majority of its attention to its diagonal,
indicating that each dimension-specific update mostly considers the dimension-
specific evaluation information. For the non-separable Rosenbrock task, on the
other hand, the attention is split across the different dimensions, hinting at
a transfer of evaluation information. The multi-head self-attention over the
generation sequence, on the other hand, attributes attention across all
previous generations. This indicates that the Transformer integrates
information accumulated through multiple generations, providing evidence for
in-context evolutionary optimization. This observation is consistent for both
considered BBO problems. Finally, the solution attention does not appear
visually interpretable. This is not too surprising, given the non-linear
transformations performed by the Perceiver module with the solution features.

Figure 5. Evolution Transformer (SNES) properties on characteristic problems.
EvoTF correctly implements the core properties of unbiasedness on the random
fitness (top row), translation invariance on a Sphere Task (middle row), and
scale self-adaptation on the linear fitness (bottom row). All tasks consider 3
search dimensions and 5 population members The search mean is initialized
between $[-3,3]$.

##  6\. Meta-Evolution of Evolution Transformer Weights Can Overfit The Meta-
Training Task Distribution

Figure 6. Meta-Evolution of Evolution Transformer weights. We train the neural
network weights using meta-black-box optimization (Lange et al., 2023b, a) on
a set of 5 BBOB problems. We compare evolving a EvoTF parametrization from
scratch with fine-tuning an EAD-pretrained SNES-EvoTF initialization. While
meta-evolution quickly improves performance on BBOB tasks, it tends to not
generalize to neuroevolution tasks. Results are averaged across 3 independent
runs.

Experiment Outline. Is supervised EAD more efficient than meta-learning the
weights of the Evolution Transformer directly via evolutionary optimization?
Can one potentially obtain performance improvements by fine-tuning a
previously distilled Evolution Transformer checkpoint via meta-evolution? To
answer these questions we turn to the meta-black-box optimization (MetaBBO,
Figure 7) paradigm introduced by Lange et al. (2023b, a). More specifically,
we use a standard diagonal ES to evolve the EvoTF weights on the same set of
BBOB tasks previously used for EAD. We again evaluate the performance of the
resulting EvoTF checkpoint throughout meta-evolution training.

Figure 7. Evolution Transformer Meta-Black-Box Optimization (MetaBBO)
procedure. We sample neural network parameters for EvoTF from a meta-
optimizer. Afterwards, all EvoTransformer candidates are evaluated by running
BBO on a set of tasks. We compute an aggregated meta-fitness score across the
tasks and update the meta-optimizer.

Implementation Details. We use Sep-CMA-ES (Ros and Hansen, 2008) as the meta-
ES and sample a population of 256 different EvoTF parametrizations, which we
evaluate on a set of 64 BBOB tasks. Afterwards, we construct normalized meta-
fitness scores for all meta-population members and update the meta-ES. This
iterate meta-evolution procedure is executed for 1000 meta-generations. We
consider two settings: 1) Meta-evolving EvoTF parameters starting from a
random initialization and 2) starting from a previously distilled SNES
teacher. Again, the entire meta-evaluation procedure is conducted using
hardware acceleration and the vectorization and parallelization capabilities
of JAX (Bradbury et al., 2018; Lange, 2022a). The training is done using 4
A100 NVIDIA GPUs. Due to memory constraints, we consider a smaller EvoTF model
with ca. 200k parameters and do not use the cross-dimension Perceiver module.
Results Interpretation. In Figure 6 we show that meta-evolving an EvoTF
parametrization from a randomly initialized meta-search mean is indeed
feasible and results in an ES capable of generalizing to unseen BBOB tasks and
partially to neuroevolution tasks (e.g. CartPole). Only 500 meta-generations
are required to obtain a competitive EvoTF-based ES. We also investigated
whether meta-evolution can act as a fine-tuning strategy after performing
supervised Evolutionary Algorithm Distillation. To do so we pre-train an EvoTF
to distill SNES, afterwards we initialize the meta-search mean to the
previously obtained checkpoint and continue training via meta-evolution. While
such a fine-tuning exercise quickly improves performance on BBOB tasks, it
dramatically decreases the performance on the neural network weight evolution
evaluation tasks. This indicates a degree of overfitting to the meta-train
task distribution and highlights the need for diverse meta-training coverage.
This opens up question with regards to generating a curriculum of tasks, that
automatically increase the difficulty of meta-training tasks according to the
capabilities of the current EvoTF.

##  7\. Self-Referential Evolutionary Algorithm Distillation is feasible but
can be unstable

Figure 8. Self-Referential Evolutionary Algorithm Distillation of Evolution
Transformer weights. By iterating the generation of offspring EvoTF
parametrizations, the generation and filtering of BBO trajectories, and self-
distillation of such filtered trajectories, we obtain a self-referentially
trained evolution strategy. Results are averaged across 3 independent runs.

Experiment Outline. We wondered whether it is possible to completely alleviate
the need for a teacher algorithm (in EAD) or a meta-black-box optimization
algorithm (e.g. Sep-CMA-ES). More specifically, we propose to leverage random
perturbations of an EvoTF checkpoint to generate different so-called ’self-
referential offspring’ (Figure 9). Afterwards, these offspring generate
diverse optimization trajectories on a fixed set of tasks. We then filter the
collected trajectories by their performance. Only the best trajectories are
used to train the EvoTF with algorithm distillation. This procedure iterates
between perturbation, trajectory generation, performance filtering and
distillation. Thereby, at each iteration of the procedure, the EvoTF from the
previous iteration can bootstrap its behavior based on previously observed
performance improvements of its perturbed versions.
Implementation Details. As a proof of concept, we use simple scalar Gaussian
perturbations to the current EvoTF weight checkpoint to generate $N=64$ EvoTF
offspring. We compare two different perturbation strengths used to generate
offspring, $\sigma_{0}\in\\{0.004,0.005\\}$. We very slowly (exponentially)
decay the strength throughout training. At each iteration, we collect the
performance of the offspring on the previously used subset of 5 BBOB tasks.
All of the offspring are evaluated on the same randomly sampled tasks and we
select the best BBO trajectory observed by all offspring to be added to the
distillation batch. We proceed by performing a single gradient descent
distillation update and iterating the procedure. We conduct training on 4 A100
NVIDIA GPUs leveraging JAX-parallelization of the trajectory generation step.

Figure 9. Self-Referential Evolutionary Algorithm Distillation Procedure for
Evolution Transformer without the need for a teacher algorithm or meta-
optimizer.

Results Interpretation. In Figure 8 we show that is indeed feasible to perform
self-referential learning without the need for an explicit improvement
algorithm in the form of a meta-optimizer or teacher algorithm. The
performance on the BBOB tasks quickly improves above random. Furthermore, on a
subset of neuroevolution tasks such as CartPole and MNIST classification, the
SR-EAD-based EvoTF can perform the optimization of MLP weights. On the
Pendulum control task, on the other hand, self-referential training appears
unstable. After an initial improvement, the performance drops again. This
observation appears robust to the choice of perturbation strength.

##  8\. Conclusion

Summary. We introduced the Evolution Transformer, a sequence model fully based
on the Transformer architecture that characterizes a family of ES. Using
Evolutionary Algorithm Distillation, we cloned various BBO algorithms, which
can afterwards be deployed on unseen optimization tasks. It provides a novel
point of view for data-driven evolutionary optimization. While using meta-
evolution to directly optimize the Transformer parameters is feasible, it
results in overfitting to the considered meta-training problems. Finally, we
introduced Self-Referential Evolutionary Algorithm Distillation, which
alleviates the need for a teacher or meta-optimization algorithm. Instead, it
uses random parameter perturbations to bootstrap observed self-improvements.
We hypothesize that this procedure can facilitate the open-ended discovery of
novel in-context evolutionary optimization algorithms.
Limitations. While our Transformer-based model can be trained efficiently on
modern hardware accelerators, the deployment comes with large memory
requirements, especially for a large number of search dimensions or population
members. Therefore, we had to employ a sliding context window. This ultimately
limits the power of in-context learning. While we provide empirical insights
into the behavior of the Evolution Transformer, we still lack a full
mechanistic understanding of such a ’black-box BBO’. Finally, self-referential
training can be unstable and appears to jump between local optima.
Stabilization will require a better theoretical understanding of the learning
dynamics induced by mixing evolutionary perturbation/data generation and
gradient descent-based self-distillation.
Future Work. We are interested in the open-ended discovery of novel algorithms
leveraging SR-EAD. While we experimented solely with evolutionary optimization
in the inner loop, the framework is more general and could be applied to any
sequential algorithm with a suitable filtering mechanism. Furthermore, recent
developments in state space models may provide an opportunity to address the
context length limitations. Finally, further benchmarking on a more diverse
task sets will be required (Lange et al., 2024a).

## References

  * (1)
  * Bradbury et al. (2018) James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. 2018.  JAX: composable transformations of Python+NumPy programs.  (2018).  http://github.com/google/jax
  * Chen et al. (2017) Yutian Chen, Matthew W Hoffman, Sergio Gómez Colmenarejo, Misha Denil, Timothy P Lillicrap, Matt Botvinick, and Nando Freitas. 2017.  Learning to learn without gradient descent by gradient descent. In International Conference on Machine Learning. PMLR, 748–756.
  * Chen et al. (2022) Yutian Chen, Xingyou Song, Chansoo Lee, Zi Wang, Richard Zhang, David Dohan, Kazuya Kawakami, Greg Kochanski, Arnaud Doucet, Marc’aurelio Ranzato, et al. 2022\.  Towards learning universal hyperparameter optimizers with transformers.  Advances in Neural Information Processing Systems 35 (2022), 32053–32068.
  * Dery et al. (2022) Lucio M Dery, Abram L Friesen, Nando De Freitas, Marc’Aurelio Ranzato, and Yutian Chen. 2022.  Multi-step Planning for Automated Hyperparameter Optimization with OptFormer.  arXiv preprint arXiv:2210.04971 (2022).
  * Freeman et al. (2021) C Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. 2021.  Brax–A Differentiable Physics Engine for Large Scale Rigid Body Simulation.  arXiv preprint arXiv:2106.13281 (2021).
  * Hansen (2006) Nikolaus Hansen. 2006.  The CMA evolution strategy: a comparing review.  Towards a new evolutionary computation: Advances in the estimation of distribution algorithms (2006), 75–102.
  * Hansen et al. (2010) Nikolaus Hansen, Anne Auger, Steffen Finck, and Raymond Ros. 2010.  Real-parameter black-box optimization benchmarking 2010: Experimental setup.  Ph.D. Dissertation. INRIA.
  * Harris et al. (2020) Charles R Harris, K Jarrod Millman, Stéfan J Van Der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J Smith, et al. 2020\.  Array programming with NumPy.  Nature 585, 7825 (2020), 357–362.
  * Hunter (2007) John D Hunter. 2007.  Matplotlib: A 2D graphics environment.  IEEE Annals of the History of Computing 9, 03 (2007), 90–95.
  * Jaderberg et al. (2017) Max Jaderberg, Valentin Dalibard, Simon Osindero, Wojciech M Czarnecki, Jeff Donahue, Ali Razavi, Oriol Vinyals, Tim Green, Iain Dunning, Karen Simonyan, et al. 2017\.  Population based training of neural networks.  arXiv preprint arXiv:1711.09846 (2017).
  * Jaegle et al. (2021) Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira. 2021.  Perceiver: General perception with iterative attention. In International conference on machine learning. PMLR, 4651–4664.
  * Kirsch et al. (2022) Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. 2022.  General-purpose in-context learning by meta-learning transformers.  arXiv preprint arXiv:2212.04458 (2022).
  * Kirsch and Schmidhuber (2022) Louis Kirsch and Jürgen Schmidhuber. 2022.  Eliminating meta optimization through self-referential meta learning.  arXiv preprint arXiv:2212.14392 (2022).
  * Kossen et al. (2021) Jannik Kossen, Neil Band, Clare Lyle, Aidan N Gomez, Thomas Rainforth, and Yarin Gal. 2021.  Self-attention between datapoints: Going beyond individual input-output pairs in deep learning.  Advances in Neural Information Processing Systems 34 (2021), 28742–28756.
  * Krishnamoorthy et al. (2022) Siddarth Krishnamoorthy, Satvik Mehul Mashkaria, and Aditya Grover. 2022.  Generative pretraining for black-box optimization.  arXiv preprint arXiv:2206.10786 (2022).
  * Lange et al. (2023a) Robert Lange, Tom Schaul, Yutian Chen, Chris Lu, Tom Zahavy, Valentin Dalibard, and Sebastian Flennerhag. 2023a.  Discovering Attention-Based Genetic Algorithms via Meta-Black-Box Optimization. In Proceedings of the Genetic and Evolutionary Computation Conference. 929–937.
  * Lange et al. (2023b) Robert Lange, Tom Schaul, Yutian Chen, Tom Zahavy, Valentin Dalibard, Chris Lu, Satinder Singh, and Sebastian Flennerhag. 2023b.  Discovering evolution strategies via meta-black-box optimization. In Proceedings of the Companion Conference on Genetic and Evolutionary Computation. 29–30.
  * Lange et al. (2024a) Robert Lange, Yujin Tang, and Yingtao Tian. 2024a.  NeuroEvoBench: Benchmarking Evolutionary Optimizers for Deep Learning Applications.  Advances in Neural Information Processing Systems 36 (2024).
  * Lange (2021) Robert Tjarko Lange. 2021.  MLE-Infrastructure: A Set of Lightweight Tools for Distributed Machine Learning Experimentation.  (2021).  http://github.com/mle-infrastructure
  * Lange (2022a) Robert Tjarko Lange. 2022a.  evosax: JAX-based Evolution Strategies.  arXiv preprint arXiv:2212.04180 (2022).
  * Lange (2022b) Robert Tjarko Lange. 2022b.  gymnax: A JAX-based Reinforcement Learning Environment Library.  (2022).  http://github.com/RobertTLange/gymnax
  * Lange et al. (2024b) Robert Tjarko Lange, Yingtao Tian, and Yujin Tang. 2024b.  Large Language Models As Evolution Strategies.  arXiv preprint arXiv:2402.18381 (2024).
  * Laskin et al. (2022) Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, et al. 2022\.  In-context reinforcement learning with algorithm distillation.  arXiv preprint arXiv:2210.14215 (2022).
  * Lee et al. (2019) Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. 2019.  Set transformer: A framework for attention-based permutation-invariant neural networks. In International conference on machine learning. PMLR, 3744–3753.
  * Lu et al. (2023a) Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahani. 2023a.  Structured state space models for in-context reinforcement learning.  arXiv preprint arXiv:2303.03982 (2023).
  * Lu et al. (2023b) Chris Lu, Sebastian Towers, and Jakob Foerster. 2023b.  Arbitrary Order Meta-Learning with Simple Population-Based Evolution.  arXiv preprint arXiv:2303.09478 (2023).
  * Metz et al. (2021) Luke Metz, C Daniel Freeman, Niru Maheswaranathan, and Jascha Sohl-Dickstein. 2021.  Training learned optimizers with randomly initialized learned optimizers.  arXiv preprint arXiv:2101.07367 (2021).
  * Radford et al. (2018) Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. 2018\.  Improving language understanding by generative pre-training.  (2018).
  * Raffel et al. (2020) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. 2020.  Exploring the limits of transfer learning with a unified text-to-text transformer.  The Journal of Machine Learning Research 21, 1 (2020), 5485–5551.
  * Rechenberg (1978) Ingo Rechenberg. 1978.  Evolutionsstrategien.  In Simulationsmethoden in der Medizin und Biologie. Springer, 83–114.
  * Ros and Hansen (2008) Raymond Ros and Nikolaus Hansen. 2008.  A simple modification in CMA-ES achieving linear time and space complexity. In International conference on parallel problem solving from nature. Springer, 296–305.
  * Salimans et al. (2017) Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Ilya Sutskever. 2017.  Evolution strategies as a scalable alternative to reinforcement learning.  arXiv preprint arXiv:1703.03864 (2017).
  * Schaul et al. (2011) Tom Schaul, Tobias Glasmachers, and Jürgen Schmidhuber. 2011.  High dimensions and heavy tails for natural evolution strategies. In Proceedings of the 13th annual conference on Genetic and evolutionary computation. 845–852.
  * Schmidhuber (1987) Jürgen Schmidhuber. 1987.  Evolutionary principles in self-referential learning, or on learning how to learn: the meta-meta-… hook.  Ph.D. Dissertation. Technische Universität München.
  * Such et al. (2017) Felipe Petroski Such, Vashisht Madhavan, Edoardo Conti, Joel Lehman, Kenneth O Stanley, and Jeff Clune. 2017.  Deep neuroevolution: Genetic algorithms are a competitive alternative for training deep neural networks for reinforcement learning.  arXiv preprint (2017).
  * Tang and Ha (2021) Yujin Tang and David Ha. 2021.  The sensory neuron as a transformer: Permutation-invariant neural networks for reinforcement learning.  Advances in Neural Information Processing Systems 34 (2021), 22574–22587.
  * Tang et al. (2022) Yujin Tang, Yingtao Tian, and David Ha. 2022.  EvoJAX: Hardware-Accelerated Neuroevolution.  arXiv preprint arXiv:2202.05008 (2022).
  * Team (2023) Adaptive Agents Team. 2023.  Human-timescale adaptation in an open-ended task space.  arXiv preprint arXiv:2301.07608 (2023).
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  Advances in neural information processing systems 30 (2017).
  * Waskom (2021) Michael L Waskom. 2021.  Seaborn: statistical data visualization.  Journal of Open Source Software 6, 60 (2021), 3021.
  * Wierstra et al. (2014) Daan Wierstra, Tom Schaul, Tobias Glasmachers, Yi Sun, Jan Peters, and Jürgen Schmidhuber. 2014.  Natural evolution strategies.  The Journal of Machine Learning Research 15, 1 (2014).

.

Figure 10. Full Learning Curve on Brax Tasks using a population size of
$N=128$ and a Tanh Control Network. A SNES-distilled EvoTF performs very
competitively. Results are averaged across 5 independent runs.

##  Appendix A Additional Results

###  A.1. Brax Learning Curves

We use the Brax (Freeman et al., 2021) set of continuous control tasks and the
EvoJAX (Tang et al., 2022) evaluation wrapper. We use the standard 4-layer
tanh activation network to control various robots. The population size is
$N=128$ and we tune all baselines and EvoTF hyperparameters using a small grid
search over $\sigma_{0}$ and the learning rates. See Figure 10.

###  A.2. Impact of Task Distribution for EAD

We compare the impact of the task distribution used to generate distillation
trajectories. See Figure 11.

Function |  Reference |  Property |  Tasks
---|---|---|---
Sphere |  Hansen et al. (p. 5, 2010) |  Separable (Indep.) |  Small
Rosenbrock |  Hansen et al. (p. 40, 2010) |  Moderate Condition |  Medium
Discus |  Hansen et al. (p. 55, 2010) |  High Condition |  Medium
Rastrigin |  Hansen et al. (p. 75, 2010) |  Multi-Modal (Local) |  Medium
Schwefel |  Hansen et al. (p. 100, 2010) |  Multi-Modal (Global) |  Medium
BuecheRastrigin |  Hansen et al. (p. 20, 2010) |  Separable (Indep.) |  Large
AttractiveSector |  Hansen et al. (p. 30, 2010) |  Moderate Condition |  Large
Weierstrass |  Hansen et al. (p. 80, 2010) |  Multi-Modal (Global) |  Large
SchaffersF7 |  Hansen et al. (p. 85, 2010) |  Multi-Modal (Global) |  Large
GriewankRosen |  Hansen et al. (p. 95, 2010) |  Multi-Modal (Global) |  Large
Table 1. Meta-BBO BBOB-Based Task Families.

Throughout all main experiments, we use the ’Medium’ set of BBOB tasks.

Figure 11. Impact of task distribution for EAD performance. We compare using 1
(small), 5 (medium) and 10 (large) BBOB functions with random offsets and mean
intialization. Successful distillation does not require many optimization
tasks. Results are averaged across 5 independent runs.

##  Appendix B Evolution Transformer Features & Model Hyperparameters

Parameter | Value | Parameter | Value
---|---|---|---
Number Transformer Blocks | 1 | Number of MHSA Heads | 2
Number of Perceiver Latents | 4 | Perceiver Latent Dim | 32
Embedding Dimension | 64 | Batchsize | 32
Sequence Length | 32 | Learning rate | 0.0015
Gradient Clip Norm | 1 | Train steps | 5000
Table 2. Evolution Transformer Hyperparameters.

The features used as inputs include:

  1. (1)

Solution: Normalized difference from mean, squared normalized difference from
mean, difference from best seen solution, normalized difference from best seen
solution, etc.

  2. (2)

Fitness: Improvement boolean, z-score, centered rank, normalized in range,
SNES weights, etc.

  3. (3)

Distribution: Finite-difference gradient, SNES gradient, evolution paths,
momentum paths

##  Appendix C Meta-Evolution Hyperparameters

Parameter | Value | Parameter | Value
---|---|---|---
Population Size | 256 | # Generations | 1000
MetaBBO | Sep-CMA-ES | Perturbation Strength | 0.005
Task Batch Size | 64 |  |
Table 3. Meta-evolution hyperparameters.

##  Appendix D Self-Referential Evolutionary Algorithm Distillation (SR-EAD)
Hyperparameters

Parameter | Value | Parameter | Value
---|---|---|---
Population Size | 64 | Perturbation Strength | 0.004
Exponential Decay | 0.99999 | Filtering | Best
Table 4. SR-EAD hyperparameters.

##  Appendix E Software Dependencies

The codebase will be open-sourced under Apache 2.0 license and publicly
available under https://github.com/<anonymous>. All training loops and ES are
implemented in JAX (Bradbury et al., 2018). All visualizations were done using
Matplotlib (Hunter, 2007) and Seaborn (Waskom, 2021, BSD-3-Clause License).
Finally, the numerical analysis was supported by NumPy (Harris et al., 2020,
BSD-3-Clause License). Furthermore, we used the following libraries: Evosax:
Lange (2022a), Gymnax: Lange (2022b), Evojax: Tang et al. (2022), Brax:
Freeman et al. (2021).

##  Appendix F Compute Requirements & Experiment Organization

The experiments were organized using the MLE-Infrastructure (Lange, 2021, MIT
license) training management system.

Simulations were conducted on a high-performance cluster using between 3 and 5
independent runs (random seeds). We mainly rely on individual V100S and up to
4 A100 NVIDIA GPUs. The experiments take between 2 and 5 hours wall clock
time.


# Multi-Modal Image Captioning for the Visually Impaired

Hiba Ahsan Nikita Bhalla11footnotemark: 1 Daivat Bhatt11footnotemark: 1
Kaivankumar Shah11footnotemark: 1
<EMAIL_ADDRESS>
University of Massachusetts Amherst  Equal contribution

###### Abstract

One of the ways blind people understand their surroundings is by clicking
images and relying on descriptions generated by image captioning systems.
Current work on captioning images for the visually impaired do not use the
textual data present in the image when generating captions. This problem is
critical as many visual scenes contain text. Moreover, up to 21% of the
questions asked by blind people about the images they click pertain to the
text present in them Bigham et al. (2010). In this work, we propose altering
AoANet, a state-of-the-art image captioning model, to leverage the text
detected in the image as an input feature. In addition, we use a pointer-
generator mechanism to copy the detected text to the caption when tokens need
to be reproduced accurately. Our model outperforms AoANet on the benchmark
dataset VizWiz, giving a 35% and 16.2% performance improvement on CIDEr and
SPICE scores, respectively.

.

##  1 Introduction

Image Captioning as a service has helped people with visual impairments to
learn about images they take and to make sense of images they encounter in
digital environments. Applications such as TapTapSee (2012) allow the visually
impaired to take photos of their surroundings and upload them to get
descriptions of the photos. Such applications leverage a human-in-the-loop
approach to generate descriptions. In order to bypass the dependency on a
human, there is a need to automate the image captioning process.
Unfortunately, the current state-of-the-art (SOTA) image captioning models are
built using large, publicly available, crowd-sourced datasets which have been
collected and created in a contrived setting. Thus, these models perform
poorly on images clicked by blind people largely because the images clicked by
blind people differ dramatically from the images present in the datasets. To
encourage solving this problem, Gurari et al. (2020) released the VizWiz
dataset, a dataset comprising of images taken by the blind. Current work on
captioning images for the blind do not use the text detected in the image when
generating captions (Figures 1(a) and 1(b) show two images from the VizWiz
dataset that contain text). The problem is critical as many visual scenes
contain text and up to 21% of the questions asked by blind people about the
images clicked by them pertain to the text present in them. This makes it more
important to improvise the models to focus on objects as well as the text in
the images.

With the availability of large labelled corpora, image captioning and reading
scene text (OCR) have seen a steady increase in performance. However,
traditional image captioning models focus only on the visual objects when
generating captions and fail to recognize and reason about the text in the
scene. This calls for incorporating OCR tokens into the caption generation
process. The task is challenging since unlike conventional vocabulary tokens
which depend on the text before them and therefore can be inferred, OCR tokens
often cannot be predicted from the context and therefore represent independent
entities. Predicting a token from vocabulary and selecting an OCR token from
the scene are two rather different tasks which have to be seamlessly combined
to tackle this task.

In this work, we build a model to caption images for the blind by leveraging
the text detected in the images in addition to visual features. We alter
AoANet, a SOTA image captioning model to consume embeddings of tokens detected
in the image using Optical Character Recognition (OCR). In many cases, OCR
tokens such as entity names or dates need to be reproduced exactly as they are
in the caption. To aid this copying process, we employ a pointer-generator
mechanism. Our contributions are 1) We build an image captioning model for the
blind that specifically leverages text detected in the image. 2) We use a
pointer-generator mechanism when generating captions to copy the detected text
when needed.

(a) Model: a bottle of water is on top of a table
Ground Truth: a clear plastic bottle of Springbourne brand spring water

(b) Model: A piece of paper with text on it
Ground Truth: In a yellow paper written as 7259 and totally as 7694

##  2 Related Work

Automated image captioning has seen a significant amount of recent work. The
task is typically handled using an encoder-decoder framework; image-related
features are fed to the encoder and the decoder generates the caption Aneja et
al. (2018); Yao et al. (2018); Cornia et al. (2018). Language modeling based
approaches have also been explored for image captioning Kiros et al. (2014);
Devlin et al. (2015). Apart from the architecture, image captioning approaches
are also diverse in terms of the features used. Visual-based image captioning
models exploit features generated from images. Multi-modal image captioning
approaches exploit other modes of features in addition to image-based features
such as candidate captions and text detected in images Wang et al. (2020); Hu
et al. (2020).

The task we address deals with captioning images specifically for the blind.
This is different from traditional image captioning due to the authenticity of
the dataset compared to popular, synthetic ones such as MS-COCO Chen et al.
(2015) and Flickr30k Plummer et al. (2015) . The task is relatively less
explored. Previous works have solved the problem using human-in-the-loop
approaches Aira (2017); BeSpecular (2016); TapTapSee (2012) as well as
automated ones Microsoft ; Facebook . A particular challenge in this area has
been the lack of an authentic dataset of photos taken by the blind. To address
the issue, Gurari et al. (2020) created VizWiz-Captions, a dataset that
consists of descriptions of images taken by people who are blind. In addition,
they analyzed how the SOTA image captioning algorithms performed on this
dataset. Concurrent to our work, Dognin et al. (2020) created a multi-modal
transformer that consumes ResNext based visual features, object detection-
based textual features and OCR-based textual features. Our work differs from
this approach in the following ways: we use AoANet as our captioning model and
do not account for rotation invariance during OCR detection. We use BERT to
generate embeddings of the OCR tokens instead of fastText. Since we use
bottom-up image feature vectors extracted using a pre-trained Faster-RCNN, we
do not use object detection-based textual features. Similarly, since the
Faster-RCNN is initialized with ResNet-101 pre-trained for classification, we
do not explicitly use classification-based features such as those generated by
ResNext.

We explored copy mechanism in our work to aid copying over OCR tokens from the
image to the caption. Copy mechanism has been typically employed in textual
sequence-to-sequence learning for tasks such as summarization See et al.
(2017); Gu et al. (2016). It has also been used in image captioning to aid
learning novel objects Yao et al. (2017); Li et al. (2019). Also, Sidorov et
al. (2020) introduced an M4C model that recognizes text, relates it to its
visual context, and decides what part of the text to copy or paraphrase,
requiring spatial, semantic, and visual reasoning between multiple text tokens
and visual entities such as objects.

##  3 Dataset

The Vizwiz Captions dataset Gurari et al. (2020) consists of over $39,000$
images originating from people who are blind that are each paired with five
captions. The dataset consists of $23,431$ training images, $7,750$ validation
images and $8,000$ test images. The average length of a caption in the train
set and the validation set was 11. We refer readers to the VizWiz Dataset
Browser Bhattacharya and Gurari (2019) as well as the original paper by Gurari
et al. (2020) for more details about the dataset.

##  4 Approach

We employ AoANet as our baseline model. AoANet extends the conventional
attention mechanism to account for the relevance of the attention results with
respect to the query. An attention module $f_{att}(Q,K,V)$ operates on queries
$Q$, keys $K$ and values $V$. It measures the similarities between $Q$ and $K$
and using the similarity scores to compute a weighted average over $V$.

$\displaystyle a_{i,j}=f_{sim}(q_{i},k_{j})$
$\displaystyle,\alpha=\frac{e^{a_{i,j}}}{\sum_{j}e^{a_{i,j}}}$ (1)
$\displaystyle\hat{v}_{i}$ $\displaystyle=\sum_{j}\alpha_{i,j}v_{i,j}$ (2)
$\displaystyle f_{sim}(q_{i},k_{j})$
$\displaystyle=\text{softmax}(\frac{q_{i}k_{j}^{T}}{\sqrt[]{D}})v_{i}$ (3)

where $q_{i}\in Q$ is the $i^{th}$ query, $k_{j}\in K$ and $v_{j}\in V$ are
the $j^{th}$ key/value pair, $f_{sim}$ is the similarity function, $D$ is the
dimension of $q_{i}$ and $\hat{v}_{i}$ is the attended vector for query
$q_{i}$.

The AoANet model introduces a module AoA which measures the relevance between
the attention result and the query. The AoA module generates an "information
vector", $i$, and an "attention gate", $g$, both of which are obtained via
separate linear transformations, conditioned on the attention result and the
query:

$\displaystyle i$ $\displaystyle=W^{i}_{q}q+W^{i}_{v}\hat{v}+b^{i}$ (4)
$\displaystyle g$ $\displaystyle=\sigma(W^{g}_{q}q+W^{g}_{v}\hat{v}+b^{g})$
(5)

where $W^{i}_{q},W^{i}_{v},b^{i},W^{g}_{q},W^{g}_{v},b^{g}$ are parameters.
AoA module then adds another attention by applying the attention gate to the
information vector to obtain the attended information $\hat{i}$.

$\displaystyle\hat{i}=g\odot i$ (6)

The AoA module can thus be formulated as:

$AoA(f_{att},Q,K,V)=\sigma(W^{g}_{q}Q+W^{g}_{v}f_{att}(Q,\\\
K,V)+b^{g})\odot(W^{i}_{q}Q+W^{i}_{v}f_{att}(Q,K,V)+b^{i})$ (7)

The AoA module is applied to both the encoder and decoder. The model is
trained by minimizing the cross-entropy loss:

$\displaystyle
L(\theta)=-\sum^{T}_{t=1}\log(p_{\theta}(y^{*}_{t}|y^{*}_{1:t-1}))$ (8)

where $y^{*}_{1:T}$ is the ground truth sequence. We refer readers to the
original work Huang et al. (2019) for more details. We altered AoANet using
two approaches described next.

###  4.1 Extending Feature Set with OCR Token Embeddings

Our first extension to the model was to increase the vocabulary by
incorporating OCR tokens. We use an off-the-shelf text detector available -
Google Cloud Platform’s vision API Google . After extracting OCR tokens for
each image using the API, we use a standard stopwords
list111https://gist.github.com/sebleier/554280 as part of necessary pre-
processing. We use this API to detect text in an image and then generate an
embedding for each OCR token that we detect using a pre-trained base, uncased
BERT Devlin et al. (2019) model. The image and text features are fed together
into the AoANet model. We expect the BERT embeddings to help the model direct
its attention towards the textual component of the image. Although we also
experiment with a pointer-generator mechanism explained in Section 4.2, we
wanted to leverage the model’s inbuilt attention mechanism that currently
performs as a state of the art model and guide it towards using these OCR
tokens.

Once the OCR tokens were detected, we conducted two different experiments with
varying sizes of thresholds. We first put a count threshold of 5 i.e. we only
add words to the vocabulary which occur 5 or more times. With this threshold,
the total words added were $4,555$. We then put a count threshold of 2. With
such a low threshold, we expect a lot of noise to be present in the OCR tokens
vocabulary - half-detected text, words in a different language, or words that
do not make sense. With this threshold, the total words added were $19,781$. A
quantitative analysis of the OCR tokens detected and their frequency is shown
in Figure 2.

Figure 2: Word counts and frequency bins for threshold = 2 and threshold = 5

###  4.2 Copying OCR Tokens via Pointing

In sequence-to-sequence learning, there is often a need to copy certain
segments from the input sequence to the output sequence as they are. This can
be useful when sub-sequences such as entity names or dates are involved.
Instead of heavily relying on meaning, creating an explicit channel to aid
copying of such sub-sequences has been shown to be effective Gu et al. (2016).

In this approach, in addition to augmenting the input feature set with OCR
token embeddings, we employ the pointer-generator mechanism See et al. (2017)
to copy OCR tokens to the caption when needed. The decoder then becomes a
hybrid that is able to copy OCR tokens via pointing as well as generate words
from the fixed vocabulary. A soft-switch is used to choose between the two
modes. The switching is dictated by generation probability, $p_{gen}$,
calculated at each time-step, $t$, as follows:

$\displaystyle
p_{gen}=\sigma(w_{h}^{T}c_{t}+w_{s}^{T}h_{t}+w_{x}^{T}x_{t}+b_{ptr})$ (9)

where $\sigma$ is the sigmoid function and $w_{h},w_{s},w_{x}\text{ and
}b_{ptr}$ are learnable parameters. $c_{t}$ is the context vector, $h_{t}$ is
the decoder hidden state and $x_{t}$ is the input embedding at time $t$ in the
decoder. At each step, $p_{gen}$ determines whether a word has to be generated
using the fixed vocabulary or to copy an OCR token using the attention
distribution at time $t$. Let extended vocabulary denote a union of the fixed
vocabulary and the OCR words. The probability distribution over the extended
vocabulary is given as:

$\displaystyle P(w)=p_{gen}P_{vocab}(w)+(1-p_{gen})\sum_{i:w_{i}=w}a_{i}^{t}$
(10)

$P_{vocab}$ is the probability of $w$ using the fixed vocabulary and $a$ is
the attention distribution. If $w$ does not appear in the fixed vocabulary,
then $P_{vocab}$ is zero. If $w$ is not an OCR word, then
$\sum_{i:w_{i}=w}a_{i}^{t}$ is zero.

##  5 Experiments

In our experiments, we alter AoANet as per the approaches described in Section
4 and compare these with the baseline model. AoANet-E refers to AoANet altered
as per the approach described in Section 4.1. To observe the impact of the
number of OCR words added to the extended vocabulary, we train two Extended
variants: (1) E5: Only OCR words with frequency greater than or equal to 5.
(2) E2: Only OCR words that occur with frequency greater than or equal to 2.
AoANet-P refers to AoANet altered as per the approach described in Section
4.2. The extended vocabulary consists of OCR words that occur with frequency
greater than or equal to 2.

We use the code222https://github.com/husthuaan/AoANet released by the authors
of AoANet to train the model. We cloned the repository and made changes to
extend the feature set and the vocabulary using OCR tokens as well as to
incorporate the copy mechanism during decoding
333https://github.com/hiba008/AlteredAoA. We train our models on a Google
Cloud VM instance with 1 Tesla K80 GPU. Like the original work, we use a
Faster-RCNN Ren et al. (2015) model pre-trained on ImageNet Deng et al. (2009)
and Visual Genome Krishna et al. (2017) to extract bottom-up feature vectors
of images. The OCR token embeddings are extracted using a pre-trained base,
uncased BERT model. The AoANet models are trained using the Adam optimizer and
a learning rate of $2e-5$ annealed by 0.8 every 3 epochs as recommended in
Huang et al. (2019). The baseline AoANet is trained for 10 epochs while
AoANet-E and AoANet-P are trained for 15 epochs.

Model | Validation Scores | Test Scores
---|---|---
BLEU-4 | ROUGE_L | SPICE | CIDEr | BLEU-4 | ROUGE_L | SPICE | CIDEr
AoANet | 21.4 | 43.8 | 11.1 | 40.0 | 19.5 | 43.1 | 12.2 | 40.8
AoANet-E5 | 21.4 | 43.6 | 10.8 | 41.4 | 19.8 | 42.9 | 11.9 | 40.2
AoANet-E2 | 24.3 | 46.1 | 12.9 | 54.1 | 22.3 | 45.0 | 14.1 | 53.8
AoANet-P | 21.6 | 43.6 | 11.5 | 46.1 | 19.9 | 42.7 | 12.8 | 45.4
Table 1: Validation and Test scores for AoANet, AoANet-E5 (extended vocabulary
variant with OCR frequency threshold as 5), AoANet-E2 (extended vocabulary
variant with OCR frequency threshold as 2) and AoANet-P (pointer-generator
variant).

##  6 Results

We show quantitative metrics for each of the models that we experimented with
in Table 1. We show qualitative results where we compare captions generated by
different models in Table 2. Note that none of the models were pre-trained on
the MS-COCO dataset as Gurari et al. (2020) have done as part of their
experimenting process.

We compare different models and find that merely extending the vocabulary
helps to improve model performance on the dataset. We see that the AoAnet-E5
matches the validation scores for AoANet but we see an improvement in the
CIDEr score. Moreover, we see a massive improvement in validation and test
CIDEr scores for AoANet-E2. Similarly, we see a gain in the other metrics too.
This goes to show that the BERT embeddings generated for each OCR token for
the images do provide an important context to the task of generating captions.
Moreover, we see the AoANet-P scores, where we use pointer-generator to copy
OCR tokens after extending the vocabulary also perform better than our
baseline AoANet model. This goes to show that an OCR copy mechanism is an
essential task in generating image captions. Intuitively, it makes sense
because we would expect to humans to use these words while generating lengthy
captions ourselves.

We feel that top-k sampling is a worthwhile direction of thought especially
when we would like some variety in the captions. Beam-search is prone to
preferring shorter captions, as the probability values for longer captions
accumulates smaller values as discussed by Holtzman et al. (2019).

##  7 Error Analysis

Although there have been concerns about the robustness of the GCP API towards
noise Hosseini et al. (2017), we focused our attention on the model’s
captioning performance and on the pointer-generator mechanism. We agree that
the API’s performance might hinder the quality of the captions generated but
we expected it to not have a large enough impact.

We first look at how the Extended variants compare with the baseline. We
observe that adding text-based features to the feature set imparts useful
information to the model. In 3(a), AoANet perceives the card as a box of food.
Addition of text features enables AoANet-E5 to perceive it as a box with black
text. While not entirely correct, it is an improvement over the baseline. The
alteration also encourages it to be more specific. When the model is unable to
find the token that entails specificity, it resorts to producing UNK.
Extending the vocabulary to accommodate more OCR words helps address this
problem. In image 3(b), baseline AoANet is unable to recognize that the bottle
is a supplements bottle. AoANet-E5 attempts to be specific but since ’dietary’
and ’supplement’ are not present in the extended vocabulary, it outputs UNK.
AoANet-E2 outputs a much better caption. We see a similar pattern in 3(c).

We now look at how the Pointer variant performs compared to the baseline and
the Extended variant. Incorporating copy mechanism helps the Pointer variant
in copying over OCR tokens to the caption. AoANet-P is able to copy over
‘oats’ and ‘almonds’ in 3(d) and the token ‘rewards’ in 3(e). But the model is
prone to copying tokens multiple times as seen in images 3(b) and 3(f). This
is referred to as repetition which is a common problem in sequence-to-sequence
models Tu et al. (2016) as well as in pointer generator networks. Coverage
mechanism Tu et al. (2016); See et al. (2017) is used to handle this and we
wish to explore this in the future.

##  8 Conclusion

In this work, we propose a pointer-generator based image captioning model that
deals specifically with images taken by people with visual disabilities. Our
alteration of AoANet shows significant improvement on the VizWiz dataset
compared to the baseline. As stated in Section 7, we would like to explore
coverage mechanism in the future. Dognin et al. (2020) recently discussed
their winning entry to the VizWiz Grand Challenge. In addition, Sidorov et al.
(2020) introduced a model that has shown to gain significant performance
improvement by using OCR tokens. We intend to compare our model with these and
improve our work based on the observations made.

##  9 Acknowledgements

The authors would like to thank Mohit Iyyer and Kalpesh Krishna at the
University of Massachusetts, Amherst for their invaluable guidance and
suggestions. The authors would also like to thank the University of
Massachusetts, Amherst for providing the necessary resources throughout the
course of this project.

Image |  Captions
---|---
(a)  |  AoANet: the back of a box of food that is yellow
AoANet-E5: the back of a yellow box with black text
AoANet-E2: the back of a card with a barcode on it
AoANet-P: the back of a UNK UNK card
GT1: The back of an EBT card that is placed on a black surface.
GT2: The back of a California EBT debit card.
GT3: A yellow EBT card on a dark fabric surface.
GT4: The backside of a beige EBT card with a magnetic strip.
GT5: back of yellow Quest card with black text on it and a white empty
signature box
(b)  |  AoANet:a person is holding a bottle of seasoning
AoANet-E5: a person is holding a bottle of UNK
AoANet-E2: a person is holding a bottle of dietary supplement
AoANet-P: a person is holding a bottle of super tablets tablets tablets
tablets tablets tablets
GT1: A bottle of Nature’s Blend Vitamin D3 2000 IU with 100 tablets.
GT2: bottle of Nature’s Blend brand vitamin D3 tablets, 100 count, 2000 IU per
tab
GT3: A hand is holding a container of vitamin D.
GT4: Someone is holding a black bottle with a yellow lid.
GT5: A person’s hand holds a bottle of Vitamin D3 tablets.
(c)  |  AoANet: a a green bottle with a green and white label
AoANet-E5: a green bottle of UNK UNK UNK UNK
AoANet-E2: a bottle of body lotion is on a table
AoANet-P: a bottle of vanilla lotion is sitting on a table
GT1: A container of vanilla bean body lotion is on a white table.
GT2: A bottle of body lotion sits on top of a white table
GT3: a plastic bottle of vanilla bean body lotion from bath and body works
GT4: A bottle of body lotion that says Noel on it sitting on a table with a
phone behind it and other items around it.
GT5: A body lotion bottle is on top of table with several papers behind it and
a set of keys in the background.
(d)  |  AoANet: a box of frozen dinner is on top of a table
AoANet-E5: a box of UNK ’s UNK brand UNK UNK
AoANet-E2: a box of granola granola granola granola bars
AoANet-P: a box of oats ’s almond almond bars
GT1: A box of nature valley roasted almond crunchy bars is on a table.
GT2: A box of granola bars sitting on a floral cloth near a wooden object.
GT3: A granola bar box sits on a table cloth with other items.
GT4: Green box with roasted almond granola bar place tablecloth with flower
prints.
GT5: A package of granola bars is lying on top of a table.
(e)  |  AoANet: a hand holding a box of chocolate ’s brand
AoANet-E5: a person is holding a package of food
AoANet-E2: a hand holding a card with a number on it
AoANet-P: a person is holding a box of rewards card
GT1: Appears to be a picture of a reward card
GT2: A plastic card that says speedy rewards membership card.
GT3: A Speedy Rewards membership card with a large gold star displayed on it.
GT4: a human hold some cards like credit cards and reward cards
GT5: Rewards membership card from the Speedway chain of stores.
(f)  |  AoANet: a bottle of water is on top of a table
AoANet-E5: a bottle of water is on top of a table
AoANet-E2: a bottle of vanilla vanilla coffee mate creamer
AoANet-P: a bottle of vanilla vanilla vanilla vanilla vanilla
GT1: A bottle of coffee creamer has a plastic flip top cap that can also be
twisted off.
GT2: A blue bottle of coffee creamer is sitting on a counter top next to a
black cup.
GT3: A container of Coffee Mate French Vanilla showing part of the front and
part of the back.
GT4: A bottle of French vanilla coffee creamer sits in front of a mug on the
table.
GT5: A bottle of creamer is on top of a table.
Table 2:  Examples of captions generated by AoANet, AoANet-E5 (extended
vocabulary variant with OCR frequency threshold as 5), AoANet-E2 (extended
vocabulary variant with OCR frequency threshold as 2) and AoANet-P (pointer-
generator variant) for validation set images along with their respective
ground truth captions.

## References

  * Aira (2017) Aira. 2017.  Aira: Connecting you to real people instantly to simplify daily life.
  * Aneja et al. (2018) Jyoti Aneja, Aditya Deshpande, and Alexander G Schwing. 2018.  Convolutional image captioning.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 5561–5570.
  * BeSpecular (2016) BeSpecular. 2016.  Bespecular: Let blind people see through your eyes.
  * Bhattacharya and Gurari (2019) Nilavra Bhattacharya and Danna Gurari. 2019.  Vizwiz dataset browser: A tool for visualizing machine learning datasets.  _arXiv preprint arXiv:1912.09336_.
  * Bigham et al. (2010) Jeffrey P. Bigham, Chandrika Jayant, Hanjie Ji, Greg Little, Andrew Miller, Robert C. Miller, Robin Miller, Aubrey Tatarowicz, Brandyn White, Samual White, and Tom Yeh. 2010.  Vizwiz: Nearly real-time answers to visual questions.  In _Proceedings of the 23nd Annual ACM Symposium on User Interface Software and Technology_ , UIST ’10, page 333–342, New York, NY, USA. Association for Computing Machinery.
  * Chen et al. (2015) Xinlei Chen, Hao Fang, Tsung-Yi Lin, Ramakrishna Vedantam, Saurabh Gupta, Piotr Dollár, and C. Lawrence Zitnick. 2015.  Microsoft COCO captions: Data collection and evaluation server.  _CoRR_ , abs/1504.00325.
  * Cornia et al. (2018) Marcella Cornia, Lorenzo Baraldi, Giuseppe Serra, and Rita Cucchiara. 2018.  Paying more attention to saliency: Image captioning with saliency and context attention.  _ACM Transactions on Multimedia Computing, Communications, and Applications (TOMM)_ , 14(2):1–21.
  * Deng et al. (2009) Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. 2009.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE conference on computer vision and pattern recognition_ , pages 248–255. Ieee.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019.  BERT: Pre-training of deep bidirectional transformers for language understanding.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 4171–4186, Minneapolis, Minnesota. Association for Computational Linguistics.
  * Devlin et al. (2015) Jacob Devlin, Hao Cheng, Hao Fang, Saurabh Gupta, Li Deng, Xiaodong He, Geoffrey Zweig, and Margaret Mitchell. 2015.  Language models for image captioning: The quirks and what works.  _arXiv preprint arXiv:1505.01809_.
  * Dognin et al. (2020) Pierre Dognin, Igor Melnyk, Youssef Mroueh, Inkit Padhi, Mattia Rigotti, Jarret Ross, Yair Schiff, Richard A Young, and Brian Belgodere. 2020.  Image captioning as an assistive technology: Lessons learned from vizwiz 2020 challenge.  _arXiv preprint arXiv:2012.11696_.
  * (12) Facebook.  How does automatic alt text work on Facebook? — Facebook Help Center.
  * (13) Google.  Google cloud vision.  (Accessed on 11/30/2020).
  * Gu et al. (2016) Jiatao Gu, Zhengdong Lu, Hang Li, and Victor OK Li. 2016.  Incorporating copying mechanism in sequence-to-sequence learning.  _arXiv preprint arXiv:1603.06393_.
  * Gurari et al. (2020) Danna Gurari, Yinan Zhao, Meng Zhang, and Nilavra Bhattacharya. 2020.  Captioning images taken by people who are blind.
  * Holtzman et al. (2019) Ari Holtzman, Jan Buys, Maxwell Forbes, and Yejin Choi. 2019.  The curious case of neural text degeneration.  _CoRR_ , abs/1904.09751.
  * Hosseini et al. (2017) H. Hosseini, B. Xiao, and R. Poovendran. 2017.  Google’s cloud vision api is not robust to noise.  In _2017 16th IEEE International Conference on Machine Learning and Applications (ICMLA)_ , pages 101–105.
  * Hu et al. (2020) Ronghang Hu, Amanpreet Singh, Trevor Darrell, and Marcus Rohrbach. 2020.  Iterative answer prediction with pointer-augmented multimodal transformers for textvqa.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 9992–10002.
  * Huang et al. (2019) Lun Huang, Wenmin Wang, Jie Chen, and Xiao-Yong Wei. 2019.  Attention on attention for image captioning.  In _Proceedings of the IEEE International Conference on Computer Vision_ , pages 4634–4643.
  * Kiros et al. (2014) Ryan Kiros, Ruslan Salakhutdinov, and Richard S Zemel. 2014.  Unifying visual-semantic embeddings with multimodal neural language models.  _arXiv preprint arXiv:1411.2539_.
  * Krishna et al. (2017) Ranjay Krishna, Yuke Zhu, Oliver Groth, Justin Johnson, Kenji Hata, Joshua Kravitz, Stephanie Chen, Yannis Kalantidis, Li-Jia Li, David A Shamma, et al. 2017\.  Visual genome: Connecting language and vision using crowdsourced dense image annotations.  _International journal of computer vision_ , 123(1):32–73.
  * Li et al. (2019) Yehao Li, Ting Yao, Yingwei Pan, Hongyang Chao, and Tao Mei. 2019.  Pointing novel objects in image captioning.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pages 12497–12506.
  * (23) Microsoft.  Add alternative text to a shape, picture, chart, smartart graphic, or other object.  (Accessed on 11/30/2020).
  * Plummer et al. (2015) Bryan A Plummer, Liwei Wang, Chris M Cervantes, Juan C Caicedo, Julia Hockenmaier, and Svetlana Lazebnik. 2015.  Flickr30k entities: Collecting region-to-phrase correspondences for richer image-to-sentence models.  In _Proceedings of the IEEE international conference on computer vision_ , pages 2641–2649.
  * Ren et al. (2015) Shaoqing Ren, Kaiming He, Ross Girshick, and Jian Sun. 2015.  Faster r-cnn: Towards real-time object detection with region proposal networks.  In _Advances in neural information processing systems_ , pages 91–99.
  * See et al. (2017) Abigail See, Peter J Liu, and Christopher D Manning. 2017.  Get to the point: Summarization with pointer-generator networks.  _arXiv preprint arXiv:1704.04368_.
  * Sidorov et al. (2020) O. Sidorov, Ronghang Hu, Marcus Rohrbach, and Amanpreet Singh. 2020.  Textcaps: a dataset for image captioning with reading comprehension.  _ArXiv_ , abs/2003.12462.
  * TapTapSee (2012) TapTapSee. 2012.  Taptapsee - blind and visually impaired assistive technology - powered by the cloudsight.ai image recognition api.
  * Tu et al. (2016) Zhaopeng Tu, Zhengdong Lu, Yang Liu, Xiaohua Liu, and Hang Li. 2016.  Modeling coverage for neural machine translation.  _arXiv preprint arXiv:1601.04811_.
  * Wang et al. (2020) Jing Wang, Jinhui Tang, and Jiebo Luo. 2020.  Multimodal attention with image text spatial relationship for ocr-based image captioning.  In _Proceedings of the 28th ACM International Conference on Multimedia_ , pages 4337–4345.
  * Yao et al. (2017) Ting Yao, Yingwei Pan, Yehao Li, and Tao Mei. 2017.  Incorporating copying mechanism in image captioning for learning novel objects.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 6580–6588.
  * Yao et al. (2018) Ting Yao, Yingwei Pan, Yehao Li, and Tao Mei. 2018.  Exploring visual relationship for image captioning.  In _Proceedings of the European conference on computer vision (ECCV)_ , pages 684–699.


# On the Limits of Regression Adjustment111Extended abstract accepted to
Conference on Digital Experimentation (CODE) 2023, Cambridge, MA, USA.

Daniel Ting  Meta
{dting<EMAIL_ADDRESS>Kenneth Hung  Meta
{dting<EMAIL_ADDRESS>
##  1 Introduction

Regression adjustment [5], sometimes known as Controlled-experiment Using Pre-
Experiment Data (CUPED) [1], is an important technique in internet
experimentation. It decreases the variance of effect size estimates, often
cutting confidence interval widths in half or more while never making them
worse. It does so by carefully regressing the goal metric against pre-
experiment features to reduce the variance.

The tremendous gains of regression adjustment begs the question: How much
better can we do by engineering better features from pre-experiment data, for
example by using machine learning techniques [4, 3] or synthetic controls [7]?
Could we even reduce the variance in our effect sizes arbitrarily close to
zero with the right predictors?

Unfortunately, our answer is negative. A simple form of regression adjustment,
which uses just the pre-experiment values of the goal metric, captures most of
the benefit. Specifically, under a mild assumption that observations closer in
time are easier to predict that ones further away in time, we upper bound the
potential gains of more sophisticated feature engineering, with respect to the
gains of this simple form of regression adjustment. The maximum reduction in
variance is $50\%$ in Theorem 1, or equivalently, the confidence interval
width can be reduced by at most an additional $29\%$.

This result allows us to assess the potential value of investing in more
advanced versions of regression adjustment. Furthermore, this upper bound
suggests a fundamental boundary for how much analysis-based variance reduction
that adjusts random imbalance from sampling can achieve. Many methods, such as
inverse propensity score weighting or synthetic control, fall under this
framework and suffer the same limitations. Breaking this barrier may require
injecting domain knowledge on the causal mechanism which enables us to adjust
for appropriate post-treatment variables [2], or changing the experimental
design directly.

##  2 Main Result

When $\rho\in[0,1]$ is the correlation between the post-experiment values and
the pre-experiment values in the goal metric, the regression adjustment
estimator reduces the variance by a factor of $1-\rho^{2}$, i.e. when
$\rho\approx 1$ and the post-experiment values can be nearly perfectly
predicted by the pre-experiment values, then the variance of the treatment
effect estimator goes to $0$. These pre-experiment covariates can be further
refined via feature engineering. For example, a machine learning model can
generate predictions for the post-experiment metric, which we can use as
covariates [4, 3].

While any covariates that are known to be unaffected by the treatment
assignment is valid [1], practitioners often restrict themselves to “safe”
covariates that are pre-treatment [6]. Under a reasonable assumption (1) that
we cannot predict the goal metric in the post-experiment period any better
than in the pre-experiment period, we show that this “safe” form of advanced
regression adjustment is limited in further variance reduction. Specifically,

$\operatorname{var}(\hat{\delta}_{\text{advanced
r.a.}})\geq\frac{1}{1+\rho}\operatorname{var}(\hat{\delta}_{\text{basic
r.a.}}).$

In particular, when basic regression adjustment is very effective
($\rho\approx 1$), advanced regression adjustment via feature engineering can
only reduce it further by $50\%$; when basic regression adjustment is not
effective ($\rho\approx 0$), advanced regression adjustment will not be
effective either.

##  3 Mathematical setup

Suppose $Y_{\mathrm{pre}}$ and $Y_{\mathrm{post}}$ are the pre-experiment and
post-experiment values in the goal metric of a randomly sampled user. Then the
basic regression adjustment reduces the variance of the estimated average
treatment effect by

$\frac{\operatorname{var}(\hat{\delta}_{\text{basic
r.a.}})}{\operatorname{var}(\hat{\delta}_{\text{original}})}=1-\rho^{2}=1-\operatorname{cor}(Y_{\mathrm{pre}},Y_{\mathrm{post}})^{2}.$

More advanced regression adjustment methods attempt to design a covariate $X$
that is as correlated with $Y_{\mathrm{post}}$ as possible. Typical advanced
regression adjustment methods include:

  * •

Using multiple pre-experiment covariates. We can consider the resulting linear
combination of the multiple pre-experiment covariates as $X$.

  * •

Producing predictions of $Y_{\mathrm{post}}$ based on external datasets. While
the predictions can be used directly in a diff-in-diff way, [3] recommends
regressing $Y_{\mathrm{post}}$ on these predictions to maintain prediction
unbiasedness, a property that gives more robustness to model misspecification.
In this case, these predictions play the role of $X$.

  * •

Producing predictions of $Y_{\mathrm{post}}$ using $Y_{\mathrm{post}}$. Using
post-experiment data of specifically goal metric is technically not allowed in
regression adjustment. [4] circumvents this by cross-fitting to avoid
overfitting. We are now effectively in the previous scenario — the model is
always based on a different split, i.e. data external to this split.

Suppose $X$ is the best possible covariate for regression adjustment, and
denote the correlation matrix of $(X,Y_{\mathrm{pre}},Y_{\mathrm{post}})$ as
follows.

$\begin{pmatrix}1&\sigma&\tau\\\ \sigma&1&\rho\\\ \tau&\rho&1\end{pmatrix}.$
(1)

Our assumption in Section 2 can be stated more formally as 1.

######  Assumption 1.

The correlation between $X$ and $Y_{\mathrm{post}}$ is at most the correlation
between $X$ and $Y_{\mathrm{pre}}$. In other words, $\tau\leq\sigma$.

The variance reduction from advanced regression adjustment is $1-\tau^{2}$, so
the relative variance reduction compared to basic regression adjustment is
given by $(1-\tau^{2})/(1-\rho^{2})$. Our main result gives a lower bound to
this ratio.

######  Theorem 1.

Suppose $X$ is the best possible covariate for regression adjustment. Under 1,
we have

$\frac{1-\tau^{2}}{1-\rho^{2}}\geq\frac{1}{1+\rho}.$

###### Proof.

Since $X$ is the best possible covariate for regression adjustment, the
maximum correlation between $Y_{\mathrm{post}}$ and any linear combination in
form of $aX+bY_{\mathrm{pre}}$ must be attained at $(a,b)=(1,0)$. Hence the
optimization problem

$\text{maximize
}\operatorname{cov}(aX+bY_{\mathrm{pre}},Y_{\mathrm{post}})\text{ subject to
}\operatorname{var}(aX+bY_{\mathrm{pre}})=1$

is solved at $(a,b)=(1,0)$. With (1) and adding in $\lambda$ as a Lagrange
multiplier, we can rewrite the optimization problem as

$\text{maximize }(a\tau+b\rho)-\lambda(a^{2}+b^{2}-2ab\sigma-1).$

In particular, the partial derivatives with respect to $a$ and $b$ should be
$0$ at $(a,b)=(1,0)$, so

$\tau-2\lambda=0\text{ and
}\rho-2\sigma\lambda=0\Longrightarrow\rho=\sigma\tau.$

Finally, we have $\tau^{2}\leq\sigma\tau=\rho$ and hence

$\frac{1-\tau^{2}}{1-\rho^{2}}\geq\frac{1-\rho}{1-\rho^{2}}=\frac{1}{1+\rho}.\qed$

##  4 Discussion of 1

While highly intuitive, 1 itself is not generally testable, in the sense that
even if it holds for some $X$, there may be other covariates $X^{\prime}$ that
we can engineer where it no longer holds.

There are some heuristics why 1 plausibly holds. For example, if we believe in
some loose form of stationarity, then our ability to predict
$Y_{\mathrm{post}}$ after $t$ days of treatment should be similar to
predicting $Y_{\mathrm{pre}}$ using data $t$ days before the experiment
begins, which in turns less than our ability to predict $Y_{\mathrm{pre}}$
using all available pre-experiment data.

We also note some cases where the assumption may reasonably fail. Seasonality
can be one reason the assumption fails. Suppose our outcome is post-treatment
time period is winter and our pre-treatment period is fall. If our goal metric
measures home heating costs, then it is reasonable to believe that last
winter’s home heating costs ($X$) could be a better predictor of next winter’s
heating costs ($Y_{\mathrm{post}}$) than it is of the fall’s
($Y_{\mathrm{pre}}$).

Another reason is if the post-treatment time period is significantly longer
than the pre-treatment period. In this case, the post-treatment period has
smoothed over more daily fluctuations and may be easier to predict.

## References

  * [1] Alex Deng, Ya Xu, Ron Kohavi and Toby Walker  “Improving the Sensitivity of Online Controlled Experiments by Utilizing Pre-Experiment Data”  In _Proceedings of the Sixth ACM International Conference on Web Search and Data Mining_ , WSDM ’13  Rome, Italy: Association for Computing Machinery, 2013, pp. 123–132  DOI: 10.1145/2433396.2433413
  * [2] Alex Deng, Michelle Du, Anna Matlin and Qing Zhang  “Variance Reduction Using In-Experiment Data: Efficient and Targeted Online Measurement for Sparse and Delayed Outcomes”  In _Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining_ , KDD ’23  Long Beach, CA, USA: Association for Computing Machinery, 2023, pp. 3937–3946  DOI: 10.1145/3580305.3599928
  * [3] Kevin Guo and Guillaume Basse  “The Generalized Oaxaca-Blinder Estimator”  In _Journal of the American Statistical Association_ 118.541 Taylor & Francis, 2023, pp. 524–536  DOI: 10.1080/01621459.2021.1941053
  * [4] Yongyi Guo et al.  “Machine Learning for Variance Reduction in Online Experiments”  In _Advances in Neural Information Processing Systems_ 34 Curran Associates, Inc., 2021, pp. 8637–8648  URL: https://proceedings.neurips.cc/paper_files/paper/2021/file/488b084119a1c7a4950f00706ec7ea16-Paper.pdf
  * [5] Winston Lin  “Agnostic notes on regression adjustments to experimental data: Reexamining Freedman’s critique”  In _The Annals of Applied Statistics_ 7.1 Institute of Mathematical Statistics, 2013, pp. 295–318  DOI: 10.1214/12-AOAS583
  * [6] Huizhi Xie and Juliette Aurisset  “Improving the Sensitivity of Online Controlled Experiments: Case Studies at Netflix”  In _Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining_ , KDD ’16  San Francisco, California, USA: Association for Computing Machinery, 2016, pp. 645–654  DOI: 10.1145/2939672.2939733
  * [7] Congshan Zhang, Dominic Coey, Matt Goldman and Brian Karrer  “Regression Adjustment with Synthetic Controls in Online Experiments”, 2021  URL: https://research.facebook.com/publications/regression-adjustment-with-synthetic-controls-in-online-experiments/