# Multigrid solver with super-resolved interpolation

Francisco Holguin
University of Michigan Ann Arbor, Department of Astronomy
Los Alamos National Laboratory &GS Sidharth
Los Alamos National Laboratory &Gavin Portwood
Lawrence Livermore National Laboratory

###### Abstract

The multigrid algorithm is an efficient numerical method for solving a variety
of elliptic partial differential equations (PDEs). The method damps errors at
progressively finer grid scales, resulting in faster convergence compared to
standard iterative methods such as Gauss-Seidel. The prolongation, or coarse-
to-fine interpolation operator within the multigrid algorithm lends itself to
a data-driven treatment with ML super resolution, commonly used to increase
the resolution of images. We (i) propose the novel integration of a super
resolution generative adversarial network (GAN) model with the multigrid
algorithm as the prolongation operator and (ii) show that the GAN-
interpolation improves the convergence properties of the multigrid in
comparison to cubic spline interpolation on a class of multiscale PDEs
typically solved in physics and engineering simulations.

##  1 Introduction

Partial differential equations describe the behavior of many physical
phenomena, from the smallest quantum scales to the largest scales in the
universe. Efficient numerical methods are crucial to simulate these complex
phenomena with current and foreseeable computational resources. The
application of machine learning to accelerate and improve the accuracy of
physical simulations is an increasingly more active area of research,
leveraging the surge of advancements in machine learning architectures over
the last decade (e.g. Wei et al., 2018; Peurifoy et al., 2018; Ranade et al.,
2020; Sanchez-Gonzalez et al., 2020). While replacing conventional numerical
solvers with machine learning models is promising, improving heuristic
operators within existing mathematically-formally derived numerical methods
can allow for interpretable computational gains, easier implementation and
more rapid deployment.

##  2 Proposal

###  2.1 Multigrid method

Multigrid methods are particularly effective for solving multiscale
differential equations relevant to physics and engineering (e.g. Trottenberg
et al., 2001; Sidharth, 2020). Multigrid methods formally scale linearly with
the number of unknowns in a discretized physical system, which make them
attractive as high performance computing increasingly enables higher-fidelity
and larger scale computations. Generally, the favorable scaling of multigrid
methods is due to their ability to iteratively reduce solution errors at
multiple scales by obtaining coarse solutions then utilizing the coarsened
solutions to solve increasingly higher-fidelity solutions. The utilization of
coarse fidelity solutions to accelerate finer fidelity solutions is enabled by
a prolongation, or interpolation, operator.

Prolongation within a multigrid algorithm lends itself to a data-driven
treatment due to the ability for machine learning techniques to perform data-
informed interpolation, as seen frequently seen in super-resolution methods as
used in image analysis (c.f. Farsiu et al., 2004; Park et al., 2003). In this
work, we consider the applications of enhancing multigrid methods by using a
data-driven prolongation operator akin to super-resolution operators used in
image analysis. In particular we consider a common Poisson differential
equation, frequently solved in engineering and physics calculations, written
as

$\nabla^{2}p=\frac{\partial^{2}p}{\partial^{2}x}+\frac{\partial^{2}p}{\partial^{2}y}=f(x,y),$
(1)

where the solution $p$ is discretized on a grid. Figure 1(a) shows a schematic
of the two-level multigrid, where the top horizontal level represents the
highest resolution grid $h_{2}$. Applying the Poisson operator $\nabla^{2}p-f$
to the grid $h_{2}$, using a relaxation method such as Gauss-Seidel, results
in an residual $r_{f}$ relative to the true solution. For these relaxation
methods, high frequency errors (relative to the grid scale) are quickly
reduced, while low spatial frequency errors remain. It can be shown
(Trottenberg et al., 2001) that the residual also satisfies the Poisson
equation, so we recursively solve the Poisson equation for this error. In
order to efficiently suppress low frequency errors, we apply a restriction
operator $\displaystyle{\bm{R}}$ to coarsen the grid, so that previously low
spatial frequency errors on the fine grid become high frequency errors on the
coarse grid. Once the coarsest level is reached, we apply a prolongation
operator $\displaystyle{\bm{P}}$ to the solution in order to correct the
solution on the fine grid. In a uniformly spaced grid, the restriction
operator $\displaystyle{\bm{R}}$, can be a simple as taking only the black
points in red-black ordering. The prolongation operator requires more work, as
it involves an interpolation scheme to create a higher resolution grid. For
example, this interpolation scheme can be a two-dimensional mesh
interpolation, such as the bivariate spline function in the SciPy Python
library (Virtanen et al., 2020).

(a) Schematic of a two-level multigrid cycle. Adapted from Chen & Sun (2001).

(b) Structure of discriminator and generator we implement, modified from Ledig
et al. (2017).

Figure 1: Diagrams for two-level multigrid and GAN interpolation.

###  2.2 GAN and integration within multigrid

Deep learning (LeCun & Hinton, 2015) methods have been able to capture and
reproduce the structure of extremely complex physical systems (e.g. Portwood
et al., 2021; Wang et al., 2020; He et al., 2019). Even without any prior
knowledge of the system, the network can learn the behavior and patterns with
adequate data, relatively quickly. Super resolution refers to the application
of these networks to low resolution images in order to predict a high
resolution version. A GAN (Goodfellow et al., 2014) is particularly effective
as an accurate super resolution model. Here, we use the GAN architecture from
Ledig et al. (2017), following the example in Birla , as shown in Figure 1(b).
Unlike Ledig et al. (2017) who use the pre-trained VGG loss designed for
photographic image recognition, we use mean squared error (MSE) loss function
for the generator.

The GAN takes $n_{s}\times n_{s}$ (here $n_{s}=6$) point input tiles from the
coarse solution for spatially-invariant training that generates more training
samples from limited datasets. Using the windowed inputs, the GAN can operate
on grids with different resolution without any change in input formats. Ghost
cells are implemented by splitting the domain up into overlapping windows and
utilize only a part of the any particular GAN’s predictions for the high
resolution grid.

The GAN prolongation operator works as follows:

  * •

Find the max and min data of the grid $p_{\rm{max}},p_{\rm{min}}$, and
normalize the grid data using a symmetric log function from
$\displaystyle[-\infty,\infty]$ to $\displaystyle[-1,1]$, such that for
positive data, $\rm{log}(p_{\rm{max,min}})=(1,0)$ and for negative data,
$\rm{log}\ p_{\rm{max,min}}=(-1,0)$.

  * •

Divide the normalized coarse grid into overlapping $n_{s}\times n_{s}$ windows
(i.e. assuming a $n_{s}^{2}$ kernel with stride = 2). Apply GAN interpolation
to produce a set of high resolution $n_{l}^{2}$ window kernels. Here
$n_{l}=24$.

  * •

Construct the fine grid by assigning the central $(n_{s}+2)^{2}$ fine window
values to the fine grid. Any non-overlapping areas (i.e the edges) are
assigned to the window covering that area.

  * •

Transform and scale back to the original pressure data.

###  2.3 Data and training of the GAN

The pressure-Poisson formulation of the incompressible Navier-Stokes equations
is a physical system with the form of Eq. 1. The pressure-Poisson equation is
used to ensure mass conservation in incompressible flows and is derived from
taking the divergence of the fluid momentum equation

$\nabla^{2}p=\nabla\cdot(\nu\nabla^{2}u-(u\cdot\nabla)u)=f(x,y)$ (2)

Here $\nu$ is the fluid viscosity, $u$ is the fluid velocity, and $p$ is the
fluid pressure. Evolving the system in time involves calculating the source
term $f(x,y)$ from the velocity field $u$, solving the Poisson equation for
pressure (e.g. using a multigrid method), and then using that pressure to
update the velocity field. We produce a training set of 200 pressure fields on
a grid size $192^{2}$ by evolving a biperiodic velocity field from a broadband
initial condition using the incompressible Navier-Stokes.

We construct the training set of 1000 pressure grids as follows:

  * •

Select one of the 200 pressure fields randomly and restrict (downsample) the
grid by a random power $l$ of 2 on each grid length, such that the smallest
grid size is $(n_{s}=6)^{2}$. The result is a random grid of size
$(192/2^{l})^{2}$.

  * •

From the grid above, pick a random $6^{2}$ window. The total size of the space
of such windows is $17,000$.

  * •

Transform and rescale the window data range from
$\displaystyle[-\infty,\infty]$ to $\displaystyle[-1,1]$ in the same manner as
described in the GAN prolongation method in 2.2, using the global min and max
of all the pressure grids. We used $|p_{\rm{min}}|=10^{-10}$ and
$|p_{\rm{max}}|=10^{-3}$.

##  3 Proposal implementation

Figure 2: Comparison of the first 15 iterations of solving Eq. 1 with the
multigrid solver starting from an initially random grid. Top: interpolation
with traditional mesh interpolation. Bottom: interpolation with GAN super
resolution.  Figure 3: Averaged results of multigrid algorithm solving $100$
different grids. Norm of the difference of grid between iterations as a
function of iteration, shown for two choices of prolongation operator, two-
dimensional spline and GAN super resolution interpolation. The multigrid
parameters used are $N_{\rm{smooth,pre}}$ = 10, $N_{\rm{smooth}}$ = 20,
$N_{\rm{step}}$ = 4, and $r_{\rm{min}}$ = 12. Top: The GAN operator is used
for interpolation in every iteration. Bottom: The GAN and spline operators are
alternated every other iteration. Figure 4: Convergence in $p$ spectral
density with increasing multigrid iterations.

We solve Eq. 2 using a test set source grid. Our implementation of the
multigrid algorithm is characterized by four parameter choices:
$N_{\rm{smooth,pre}}$ = number of smoothing operations on the initial grid;
$N_{\rm{smooth}}$ = number of smoothing operations at the beginning of each
iteration; $N_{\rm{step}}$ = number of times the the grid resolution is
reduced by a factor of two on each dimension; $r_{\rm{min}}$ = coarsest
resolution. Pre-smoothing is often necessary in order to ensure convergence.
To study the efficacy of a GAN prolongation operator under controlled
settings, we implement a two-level V-cycle multigrid algorithm (fine scale
$192^{2}$ grid to a coarse scale $12^{2}$ grid, corresponding to
$N_{\rm{step}}$ = 5 and $r_{\rm{min}}=12$). A two-level multigrid method with
a coarsening factor of 16 is expected to perform poorly, limited by
prolongation and interpolation, and therefore serves as an ideal test bed for
evaluating the gains made with a super-resolution GAN. We sample $100$
pressure and source term grid pairs from our training set, following the
method described in Section 2.3.

Figure 3 shows a proof-of-concept test of our multigrid solver on $100$ grids.
We plot the mean profiles and $1-\sigma$ bounds of the norm difference between
the pressure grids at the current and previous iteration. The top row of the
figure shows the result of using the GAN operator at every iteration compared
to the spline operator. The multigrid solver converges in 100 iterations
(compared to $10^{4}$ iterations for a simple iterative solver).
Statistically, we find that using the super resolution prolongation operator
can result in faster convergence compared to bivariate spline operator; the
large $1-\sigma$ lines on the ratio between the two convergence lines shows
that the performance is distributed. For some fields, super resolution
exhibits faster convergence, but for other fields, the spline interpolation
can fare better. The bottom row shows the result of alternating the GAN and
spline operators compared with using only the spline operator. There is faster
convergence using the alternating operators for all of the fields tested. This
improvement is likely due to the fact that the GAN has been trained on real
pressure field solutions and not random fields. Therefore, as soon as the
pressure field evolves from random to one with physical properties, the
advantages of the GAN prolongation operator set in. In practice, the initial
conditions will not necessarily be random, but will be physically realistic,
such as in a time-stepping physical simulation.

These tests demonstrate strongly that there is potential for super-resolution
to improve the efficiency of the multigrid algorithm by data-driven
prolongation operators. Figure 2 shows a sequence of the first 15 iterations
of a particular field, for the solver shown in the top row of Figure 3. While
both solvers converge to machine zero, with the conventional spline-based
interpolation, the early iterations are unable to capture mid to high-
frequency information. With the GAN-based interpolation however, the solver
feeds in higher frequency information in the early iterations. This is
confirmed via the pressure field spectra shown in Figure 4, which show compare
the spectra for the first 15 iterations.

##  4 Conclusion

In this research, we consider novel integration of data-driven interpolation
into the robust multigrid differential equation solution technique as applied
to a common differential system utilized in many physics and engineering
domains. We have demonstrated the ability for super-resolution techniques,
GANs in particular, to serve as a prolongation operator which can improve
solution time when compared to high order spline interpolation techniques. As
we continue our investigations for this approach, we propose that the
integration of data-driven methods into formally-derived numerical solution
procedures is a promising approach for computational acceleration of
simulations.

## References

  * (1) D. Birla.  URL https://github.com/deepak112/Keras-SRGAN.
  * Chen & Sun (2001) Shu-hua Chen and Wen-yih Sun.  Application of the multigrid method and a flexible hybrid coordinate in a nonhydrostatic model.  _Monthly weather review_ , 129(11):2660–2676, 2001.
  * Farsiu et al. (2004) Sina Farsiu, Dirk Robinson, Michael Elad, and Peyman Milanfar.  Advances and challenges in super-resolution.  _International Journal of Imaging Systems and Technology_ , 14(2):47–57, 2004.
  * Goodfellow et al. (2014) I. J. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Ward-Farley, S. Ozair, A. Courville, and Y. Bengio.  _Adv. Neural Procs. Sys._ , 27:2672 – 2680, 2014.
  * He et al. (2019) Siyu He, Yin Li, Yu Feng, Shirley Ho, Siamak Ravanbakhsh, Wei Chen, and Barnabás Póczos.  Learning to predict the cosmological structure formation.  _Proceedings of the National Academy of Sciences_ , 116(28):13825–13832, 2019.
  * LeCun & Hinton (2015) Y. LeCun and G. Hinton.  Deep learning.  _Nature_ , 521:436 – 444, 2015.
  * Ledig et al. (2017) Christian Ledig, Lucas Theis, Ferenc Huszár, Jose Caballero, Andrew Cunningham, Alejandro Acosta, Andrew Aitken, Alykhan Tejani, Johannes Totz, Zehan Wang, et al.  Photo-realistic single image super-resolution using a generative adversarial network.  pp. 4681–4690, 2017.
  * Park et al. (2003) Sung Cheol Park, Min Kyu Park, and Moon Gi Kang.  Super-resolution image reconstruction: a technical overview.  _IEEE signal processing magazine_ , 20(3):21–36, 2003.
  * Peurifoy et al. (2018) John Peurifoy, Yichen Shen, Li Jing, Yi Yang, Fidel Cano-Renteria, Brendan G DeLacy, John D Joannopoulos, Max Tegmark, and Marin Soljačić.  Nanophotonic particle simulation and inverse design using artificial neural networks.  _Science advances_ , 4(6):eaar4206, 2018.
  * Portwood et al. (2021) Gavin D Portwood, Balasubramanya T Nadiga, Juan A Saenz, and Daniel Livescu.  Interpreting neural network models of residual scalar flux.  _Journal of Fluid Mechanics_ , 907:A23, 2021.
  * Ranade et al. (2020) Rishikesh Ranade, Chris Hill, and Jay Pathak.  Discretizationnet: A machine-learning based solver for navier-stokes equations using finite volume discretization.  _arXiv preprint arXiv:2005.08357_ , 2020.
  * Sanchez-Gonzalez et al. (2020) Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter Battaglia.  Learning to simulate complex physics with graph networks.  In _International Conference on Machine Learning_ , pp. 8459–8468. PMLR, 2020.
  * Sidharth (2020) G S Sidharth.  A multiscale subgrid decomposition.  _AIAA Scitech 2020 Forum_ , pp. 0820, 2020.
  * Trottenberg et al. (2001) U. Trottenberg, C. W. Oosterlee, and A. Schuller.  _Multigrid_.  2001\.
  * Virtanen et al. (2020) Pauli Virtanen, Ralf Gommers, Travis E Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, et al.  Scipy 1.0: fundamental algorithms for scientific computing in python.  _Nature methods_ , 17(3):261–272, 2020.
  * Wang et al. (2020) Rui Wang, Karthik Kashinath, Mustafa Mustafa, Adrian Albert, and Rose Yu.  Towards physics-informed deep learning for turbulent flow prediction.  In _Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_, pp. 1457–1466, 2020.
  * Wei et al. (2018) Qianshi Wei, Ying Jiang, and Jeff ZY Chen.  Machine-learning solver for modified diffusion equations.  _Physical Review E_ , 98(5):053304, 2018.


# MemoryPrompt: A Light Wrapper to Improve Context Tracking in Pre-trained
Language Models

###### Abstract

Transformer-based language models (LMs) track contextual information through
large, hard-coded input windows. We introduce MemoryPrompt, a leaner approach
in which the LM is complemented by a small auxiliary recurrent network that
passes information to the LM by prefixing its regular input with a sequence of
vectors, akin to soft prompts, without requiring LM finetuning. Tested on a
task designed to probe a LM’s ability to keep track of multiple fact updates,
a MemoryPrompt-augmented LM outperforms much larger LMs that have access to
the full input history. We also test MemoryPrompt on a long-distance dialogue
dataset, where its performance is comparable to that of a model conditioned on
the entire conversation history. In both experiments we also observe that,
unlike full-finetuning approaches, MemoryPrompt does not suffer from
catastrophic forgetting when adapted to new tasks, thus not disrupting the
generalist capabilities of the underlying LM.

Keywords:  memory-augmented language model, prompting

MemoryPrompt: A Light Wrapper to Improve Context Tracking in Pre-trained
Language Models

Nathanaël Carraz Rakotonirina1, Marco Baroni1,2
---
1Universitat Pompeu Fabra, Barcelona, Spain
2ICREA, Barcelona, Spain
{nathanael.rakotonirina<EMAIL_ADDRESS>

Abstract content

Figure 1: Unfolded graph of MemoryPrompt at training time. The input is
divided into segments and, for each segment, the augmented system produces
both the LM output and the memory vectors (blue) which are concatenated to the
embeddings of the next segment.

## 1\. Introduction

Transformer-based language models (LMs) can only track user-provided
contextual information if it fits into their context window. The brute-force
solution of using a huge window suffers obvious problems of scale. While there
is promising work in the direction of making long contexts more
computationally efficient (Dai et al., 2019; Beltagy et al., 2020; Chen et
al., 2023), we introduce MemoryPrompt, a cheaper and complementary solution
that augments a pre-trained LM with a smaller auxiliary recurrent network,
trained to carry context-relevant information through time. Inspired by work
on “soft prompting” (Lester et al., 2021; Liu et al., 2021; Zhong et al.,
2021), this information is passed to the LM at each time step as a continuous
token prefixed to its regular input. Our method does not architecturally alter
or finetune the underlying pre-trained LM, and it can thus leverage any
available pre-implemented model. Memory augmentation is moreover unobtrusive,
in the sense that it does not greatly affect the underlying LM behaviour in
standard next-token prediction tasks, and thus the augmented model can
successfully manage contextually-driven updates, while preserving the
knowledge encoded in the underlying LM.111The code and data to reproduce our
analysis is available at https://github.com/ncarraz/MemoryPrompt.

## 2\. Related work

Methods to enhance a sequence processing network with an external
differentiable memory have been explored since the comeback of neural networks
during the last decade (e.g., Joulin and Mikolov, 2015; Sukhbaatar et al.,
2015; Graves et al., 2016).

After the advent of the Transformer (Vaswani et al., 2017), much work on long-
range memory has focused on how to make its attention mechanisms more
efficient, in order to handle a larger span (e.g., Dai et al., 2019; Beltagy
et al., 2020). Similarly to us, Fan et al. (2020) and Hutchins et al. (2022)
introduce a recurrence to allow the Transformer to carry information through
time steps, but they do it by modifying the core Transformer architecture.

The idea of using a differentiable external memory has also made a comeback in
the context of Transformer-based language models. A recent example is the
Memformer architecture of Wu et al. (2022). Like in our approach, Memformer
has an external memory component that interacts with a Transformer via reading
and writing operations. However, this approach demands architectural changes
to the Transformer, special reading/writing operations and end-to-end training
of the memory and the underlying Transformer, preventing its use with pre-
trained LMs.

The closest approach to ours is the recently introduced Recurrent Memory
Transformer (RMT) model of Bulatov et al. (2022) (see also Bulatov et al.,
2023). They divide the input into segments, and add the same real-valued
memory vectors both at the beginning and at the end of the segment. Then, the
activations of the memory vectors added at the end of the segment form the
memory vectors of the next segment in a recurrent manner. The key difference
between RMT and our approach is that we use a lightweight module to produce
the memory vectors and keep our base model frozen, updating only the
parameters of the module. As we will show below, this is crucial to prevent
catastrophic forgetting of the knowledge that was originally encoded in the
base model.

Our idea of passing information to the Transformer in the form of continuous
tokens fed to its standard input comes from the literature on soft prompting
(e.g., Lester et al., 2021; Li and Liang, 2021; Liu et al., 2021; Zhong et
al., 2021), where sequences of vectors living in the target LM’s embedding
space are prefixed to task-specific inputs, to implicitly adapt the model to
the task without fine-tuning. We extend continuous prompts to a setup where,
instead of a fixed prefix, a dynamic one must be generated at each step, in
order to carry constantly updated information.

## 3\. The MemoryPrompt model

We augment an autoregressive language model with a recurrent memory module,
allowing it to extend its effective context-length and to keep track of
information updated through time.

The input is divided into segments that are sequentially processed by the
model. For each segment, the contextual representation of the last token is
fed as input to the memory module which is composed of an MLP followed by an
LSTM (Hochreiter and Schmidhuber, 1997). The output of the memory module is a
series of memory vectors $\mathbf{P}\in\mathbb{R}^{m\times e}$, where $e$ is
the word embedding space and $m$ is the number of vectors.

The augmented system is trained end-to-end with backpropagation through time
(see Figure 1), but only the parameters of the memory module are updated: the
LM components are kept frozen. Algorithm 1 illustrates the forward process of
MemoryPrompt in detail.

Hyperparameters : number of memory vectors $m$

number of blocks $B$

1 Initialize the hidden state $\mathbf{h}$ and cell state $\mathbf{c}$

2 Segment the input into $B$ blocks of $n$ tokens
$\mathbf{X}=\\{\mathbf{X_{1}},...,\mathbf{X_{B}}\\}$

3 Embed the input $\mathbf{X}$ to get
$\mathbf{E}=\\{\mathbf{E_{1}},...,\mathbf{E_{B}}\\}$ where
$\mathbf{E_{i}}\in\mathbb{R}^{n\times e}$

4 Get the activations before the last linear layer
$\mathbf{A}=\textit{LM}_{\phi}(\mathbf{E_{1}})$ where
$\mathbf{A}\in\mathbb{R}^{n\times e}$

5 Compute the memory vector from the activations of the last token
$\mathbf{p},\mathbf{h},\mathbf{c}=\textit{LSTM}_{\theta}(\textit{MLP}_{\theta}(\mathbf{A}[-1,:]),\mathbf{h},\mathbf{c})$
where $\mathbf{p}$ is a vector of length ${m\times e}$

6 Reshape $\mathbf{p}$ into the matrix $\mathbf{P}\in\mathbb{R}^{m\times e}$

7 Get the probabilities over the vocabulary using the last linear layer
$\mathbf{O}=\textit{Softmax}(\mathbf{A}\mathbf{W}_{\phi})$ where
$\mathbf{W}_{\phi}\in\mathbb{R}^{e\times V}$

8 for  _$b=2$ to $B$_ do

9 Concatenate the memory vectors to the embedding before feeding to the model
$\mathbf{A}=\textit{LM}_{\phi}([\mathbf{P};\mathbf{E_{b}}])$ where
$[\mathbf{P};\mathbf{E_{b}}]\in\mathbb{R}^{(m+n)\times e}$

10 Compute the next memory vector, hidden state and cell state
$\mathbf{p},\mathbf{h},\mathbf{c}=\textit{LSTM}_{\theta}(\textit{MLP}_{\theta}(\mathbf{A}[-1,:]),\mathbf{h},\mathbf{c})$

11 Reshape $\mathbf{p}$ into the matrix $\mathbf{P}$

12 Get the probabilities over the vocabulary
$\mathbf{O}=\textit{Softmax}(\mathbf{A}\mathbf{W}_{\phi})$

13

14 end for

Algorithm 1 Forward flow of MemoryPrompt for a single input

## 4\. Experimental setup

### 4.1. Datasets

#### Fact updating

We simulate a scenario in which the model is exposed to realistic fact updates
that it needs to track. We use sequences of facts gathered from the version of
TREx (Elsahar et al., 2018) curated by Elazar et al. (2021). Each fact is
represented as a triple $\langle subject,relation,object\rangle$—for example,
$\langle\textit{Antim Peak},\textit{continent},\textit{Antarctica}\rangle$. As
this example shows, not all facts can be plausibly updated. We thus identified
3 mutable relations, for which an update would be credible, namely employer
(Paul Allen/Microsoft), position held (Otto Suhr/mayor) and work location
(Lucio Fontana/Milan). These were updated by randomly selecting other objects
from the same TREx relation pool (e.g., Paul Allen’s employer might be
randomly updated to Apple or BBC). We used the remaining 36 relations to
generate stable facts that are not updated. Each fact is instantiated as a
natural-language statement using a template from Elazar et al. (2021) (e.g.,
the template for the employer relation is [subject] works for [object]).

Figure 2: Sequence from the short-fd dataset (see Table 1). Non-highlighted
text contains stable facts. The pivot (in blue) is the mutable fact to track,
and the final answer is the most up-to-date object of the pivot. Distractors
(in orange) are mutable facts distinct form the pivot, which might belong to
the same relation and might be updated (here, the object associated with Guido
Pepoli is updated from cardinal to bishop).

A sequence in the dataset is composed of multiple statements. A statement can
either introduce a new fact or update a previously introduced mutable fact.
The model is asked to predict the most up-to-date object of a single mutable
fact called the pivot fact. The other mutable facts in the sequence, that can
also be updated, are called distractors. Note that a statement can pertain to
a pivot fact, to another mutable fact (a distractor), or to a stable fact.
Nothing distinguishes pivots from distractors, so a model must be able to
track updates to both fact types.

An example sequence is shown in Figure 2. We generate different versions of
the fact-updating dataset as shown in Table 1. Each dataset contains 26,892
sequences for training, 150 for validation and 346 for testing. The number of
examples in the validation and test sets is limited by the number of mutable
facts in TREx. For training, we could generate an arbitrarily large number of
sequences by randomly selecting objects. However, performing this augmentation
on validation or test data does not affect the final performance; hence the
smaller size of the latter splits.

Note that, while the datasets are automatically generated, they contain
plausible statements and updates expressed in natural language, and one could
easily imagine concrete scenarios in which similar fact streams arise (e.g.,
processing a stream of messages from a news agency).

#### Multi-Session Chat (MSC)

We also test MemoryPrompt in the less-constrained scenario of long-distance
language modeling on the MSC dataset (Xu et al., 2021). This is a long-term
conversation dataset which consists of multi-session crowd-worker chats, where
the speakers might refer to their whole shared history. The dataset contains
35,880 training dialogues, 5,000 validation examples and 5,008 test cases.

Dataset |  | no.
---
statements
| no.
---
distractors
| no.
---
updates
short no distractors (nd) | $[10\mathrel{{.}\,{.}}\nobreak 30]$ | $0$ | $[0\mathrel{{.}\,{.}}\nobreak 4]$
short few distractors (fd) | $[10\mathrel{{.}\,{.}}\nobreak 30]$ | $[3\mathrel{{.}\,{.}}\nobreak 7]$ | $[0\mathrel{{.}\,{.}}\nobreak 4]$
long no distractors (nd) | $[150\mathrel{{.}\,{.}}\nobreak 190]$ | $0$ | $[0\mathrel{{.}\,{.}}\nobreak 9]$
long few distractors (fd) | $[150\mathrel{{.}\,{.}}\nobreak 190]$ | $[3\mathrel{{.}\,{.}}\nobreak 10]$ | $[0\mathrel{{.}\,{.}}\nobreak 9]$
long many distractors (md) | $[150\mathrel{{.}\,{.}}\nobreak 190]$ | $[25\mathrel{{.}\,{.}}\nobreak 50]$ | $[0\mathrel{{.}\,{.}}\nobreak 9]$
long many updates (mu) | $[150\mathrel{{.}\,{.}}\nobreak 190]$ | $[3\mathrel{{.}\,{.}}\nobreak 10]$ | $[0\mathrel{{.}\,{.}}\nobreak 49]$
Table 1: Different fact-updating dataset configurations. Datasets are
generated using three parameters whose values are specified by ranges: the
number of statements, the number of distractors and the number of updates to
the pivot fact. When generating a sequence, the value of each parameter is
randomly selected within the corresponding range. The positions of the pivot
and the distractors are also randomly selected.

### 4.2. Models

The MemoryPrompt memory module consists of a 1-layer 1024-dimensional MLP
followed by a 1-layer LSTM. We use 5 memory vectors. We use LMs from the OPT
family (Zhang et al., 2022). For reference, our MemoryPrompt configuration
adds about 10% new parameters to OPT-350M. Since the number of parameters of
the memory module only depends on the embedding dimension of the base model,
this ratio decreases as the base model gets larger. Due to computational
limitations, and because we are interested in external memories as an
alternative to scaling up models, we only applied MemoryPrompt and RMT to
OPT-125M and OPT-350M.

We compare the following setups. The LM-full-context setting involves an OPT
LM that takes a whole fact-updating dataset sequence or MSC conversation
history as input. In the LM + MemoryPrompt setup, the input sequence is
segmented into blocks that are separately processed by the model. Each block
contains 5 facts for the short fact-updating datasets, 10 for the long ones,
and a turn for MSC. Only the parameters of the memory module are updated. LM-
finetuned + RMT is our re-implementation of RMT, the memory-based approach of
Bulatov et al. (2022), which adds the same memory vectors both at the
beginning and at the end of a segment when, as in our case, it is applied to
decoder-only models. The activations of the memory vectors added at the end of
the segment form the vectors for the next segment. The LM is finetuned to
produce those vectors without resorting to a separate module. The input and
the number of memory vectors are the same as for MemoryPrompt.

### 4.3. Training setup

We train the models using the AdamW optimizer (Loshchilov and Hutter, 2017)
with a learning rate of 7e-5 (7e-6 for the MSC dataset) and weight decay of
0.1, without hyperparameter tuning. We follow a linear learning rate schedule
for MSC, warming up from 0 to the maximum learning rate over the first 12,000
steps. We clip gradient norms at 1.0 to avoid exploding gradients for MSC.
Moreover, we apply L2 regularization to the memory vectors, as we informally
observed this to mitigate catastrophic forgetting.

In the fact-updating experiments, during training and evaluation we provide
the models with demonstrations using facts from the same relation as the
pivot. MemoryPrompt is robust to the number of demonstrations. In our
experiments, we use 4 examples. Moreover, following Bulatov et al. (2023), we
use curriculum learning when training on longer sequences for better
performance and faster convergence. The model is first trained on the short
fact-updating dataset until convergence. Then, we use this as the starting
point for a long dataset.

Making RMT converge on the fact-updating dataset was not trivial. Unlike
MemoryPrompt, which directly converged on short fact-updating datasets, RMT
required a more elaborate curriculum learning process. We first had to train
it on very short sequences whose length was gradually increased upon
convergence. Moreover, RMT is more prone to exploding and vanishing gradients
(Hochreiter, 1998; Bengio et al., 1994). Training RMT is even more difficult
when the model is larger. We failed to make OPT-350M-finetuned + RMT converge
on any of the datasets due to this problem.

## 5\. Results

Model short long Perplexity WikiText-103 Forgetting rate TREx nd fd nd fd md
mu random pivot object 45.55 44.26 30.88 29.85 31.02 9.61 - - OPT-125M-full-
context 47.39 36.41 34.97 31.21 19.36 21.96 27.68 0.0 OPT-350M-full-context
46.24 39.30 23.41 22.83 19.07 9.82 22.04 0.0 OPT-1.3B-full-context 44.79 39.30
42.48 44.21 37.57 19.94 14.64 0.0 OPT-2.7B-full-context 42.19 37.86 34.97
39.30 40.75 17.91 12.47 0.0 OPT-IML-1.3B-full-context 50.86 44.79 41.32 43.35
33.81 18.49 14.68 0.0 OPT-IML-MAX-1.3B-full-context 49.71 47.68 40.75 45.08
36.70 18.49 14.70 0.0 OPT-125M-finetuned + RMT 93.33 55.60 89.65 45.49 23.12
72.42 11455.11 97.4 OPT-125M + MemoryPrompt 89.99 58.84 92.14 51.38 25.83
80.69 27.80 13.0 OPT-350M + MemoryPrompt 86.35 53.0 87.05 46.82 22.19 74.33
22.10 10.0

Table 2: Accuracy of the different models on the fact-updating datasets,333The dataset names nd, fd, md and mu come from Table 1. WikiText-103 perplexity and forgetting rate (%) on TREx facts, which measure catastrophic forgetting. The non-trivial random pivot object baseline picks an object appearing in any pivot statement at random. Model |  | Perplexity
---
MSC
| Perplexity
---
WikiText-103
| Forgetting rate
---
TREx
OPT-125M (conditioned on last turn) | 23.26 | 27.68 | 0.0
OPT-125M (conditioned on history) | 20.79 | 27.68 | 0.0
OPT-125M-finetuned + RMT | 17.82 | 202.83 | 80.0
OPT-125M + MemoryPrompt | 20.60 | 28.37 | 13.0
Table 3: Perplexity on MSC. We also report perplexity on WikiText-103 and
forgetting rate (%) on TREx facts, which measure catastrophic forgetting.
Figure 3: Accuracy of OPT-1.3B full-context (top) and OPT-125M + MemoryPrompt
(bottom) as a function of the number of updates on the long many-updates (mu)
fact-updating dataset.

### 5.1. Fact updating

Memory-augmented models outperform their full-context counterparts on all the
fact-updating datasets (see Table 3). Indeed, MemoryPrompt applied to
OPT-125M, the smallest model in the OPT family, outperforms the largest full-
context models we were able to run, including instrucion-tuned (IML) variants,
often dramatically.

The average input length to the full-context models in the long setups is
1,698 tokens, whereas it is only 92 tokens for MemoryPrompt, showing that, for
this task, memory vectors are better and much more efficient representations
compared to full textual context.

RMT shows good performance overall, but it lags in general behind
MemoryPrompt. Importantly, as we show in Section 5.3 below, the good
performance of RMT comes at the cost of having lost much of the knowledge
stored in the underlying model, making this approach unviable for scenarios in
which we want a LM to both effectively track contextual updates and maintain
its core general knowledge.

Surprisingly, OPT-125M outperforms the larger OPT-350M model in both the full-
context and memory-augmented scenarios. This might be related to the recent
observation by Voita et al. (2023) that OPT-350M is an “outlier” model in the
OPT family. In any case, the MemoryPrompt wrapper clearly benefits OPT-350M as
well.

As shown in Figure 3, unlike for the full-context baselines, MemoryPrompt’s
performance is essentially stable across the number of updates within a
dataset. On the other hand, in Table 3 we observe a clear slump for all
memory-augmented models on the long md dataset, the setting with the largest
number of facts to keep track of (recall that models can’t tell distractors
apart from the pivot, so they must track updates of both). This suggests that
MemoryPrompt is robust to the number of updates of specific facts, but it
struggles when there are many different facts to update (a problem that RMT
also displays). If MemoryPrompt, as suggested by our qualitative analysis in
Section 5.4 below, is learning to specialize specific memory vectors to
specific facts, it figures that, with just 5 memory vectors and too many facts
to keep track of, its performance drops. In other words, the problem is not
with the updating and retrieval mechanism, but with the amount of different
facts to store, clearly an issue to be pursued in future work.

### 5.2. MSC

MemoryPrompt’s performance on MSC is as good as conditioning on the whole
conversation history (see Table 3).444We also evaluate MemoryPrompt on the IRC
Disentanglement (Kummerfeld et al., 2019) dialog dataset and obtain similar
results as with MSC. However, we could not train RMT on it due to
computational limitations This implies that the memory vectors can effectively
track long previous interactions without requiring a large context window,
even when the relevant information in the previous interactions is not as
clearly encoded as in the fact updating datasets. We note that the entire
history contains 1,412 tokens on average vs. 18 tokens of context for
MemoryPrompt. RMT achieves lower perplexity, but, as it requires full-model
finetuning, this comes at the cost of catastrophic forgetting, as we will
discuss next.

### 5.3. (No) catastrophic forgetting

Our intention with MemoryPrompt is to provide a light way for the model to
track contextual information without affecting the model base performance. In
a realistic scenario, a generalist LM should be able to update information
during a conversation, but still retain the knowledge it acquired during pre-
training. We thus want to ensure that adding the MemoryPrompt vectors to the
LM input is not affecting the model on its base task (next token prediction).
To test this, we collect 4 prefixes generated at the end of sequences from
each fact-updating dataset (24 prefixes in total). For MSC, we collect
prefixes generated at the end of conversations. These prefixes are then
combined with the input of the base LM (as we normally do when using
MemoryPrompt). We apply the same approach to evaluating RMT for catastrophic
forgetting.

We evaluate MemoryPrompt and RMT on two tasks probing how their memory vectors
change the model’s original behaviour. First, we compute their perplexity on a
standard corpus that was not used to train the memory component. In
particular, we use WikiText-103 (Stephen et al., 2017). We also compute a
forgetting rate score, that is, the proportion of (non-updated) TREx facts for
which a memory-augmented LM predicts a different object than the same model
would without any modification (so, trivially, this measure is 0% for the
unchanged full-context models). We compute forgetting rate on facts that were
not used in the fact updating datasets.

As the Perplexity columns of tables 3 and 3 show, MemoryPrompt’s WikiText-103
perplexities are almost identical to those of the equivalent non-augmented
models. The situation is very different for RMT, whose WikiText-103
perplexities suggest that the augmented model has become incapable to perform
the standard LM token prediction task. Concerning forgetting rate (last column
of tables 3 and 3), we confirm again that RMT is dramatically affected by the
memory-augmentation process, making new predictions for a large majority of
facts that were not updated. MemoryPrompt’s forgetting rate is also non-
negligible, but much lower than that of RMT.

Overall, we confirm that a LM augmented with MemoryPrompt’s vectors is not
greatly affected in its default performance, so that it can both take
contextual updates into account, but continue using the rich general knowledge
it acquired during pretraining.

Figure 4: Cosine similarity between one of 5 memory vectors and the embeddings
of the objects in a sequence. Each figure represents a sequence, with objects
of different facts on the x-axis. The pivot objects are in bold. The
subject/relation of the pivots are Louis Charles Delescluze/work location
(left) and Rao Remala/employer (right).

### 5.4. Memory vector analysis

Memory vectors lie in the same space as the model token embeddings. We can
thus get insights on what they are recording by directly measuring their
cosine similarity to that of tokens of interests, such as those representing
updated objects. Figure 4 shows representative cosine similarity profiles
across several updating and non-updating statements for the same memory vector
and all the object embeddings (examples taken from the short fact-updating
dataset). As the figure shows, this memory vector tends to become more similar
to the pivot objects that occur in update statements, and often retains this
high-similarity across a number of statements involving other, irrelevant
objects. For example, on the left panel, the memory vector increases its
similarity to Ghent, Boston, Ghent and Hamburg, respectively, as it encounters
fact updates involving these objects. Irrelevant facts are ignored, even when
they involve cities as objects, which could in principle confuse the model.

This analysis suggests that the MemoryPrompt vectors possess a certain degree
of interpretability, although we typically failed to find more than a single
memory vector per dataset that possessed this degree of transparency, and we
must leave further vector decoding work to future research.

## 6\. Discussion

We introduced MemoryPrompt, a simple approach that allows a Transformer-based
LM to carry and update contextual information across long spans without
requiring a long attention context window. The main intuition of our approach
is that relevant information can be passed across processing steps by letting
the LM read a set of “memory” vectors as a prefix to its standard input, akin
to the idea of soft prompting. By piggybacking on the LM standard interface,
we can add our memory module to a pre-trained LM without touching its
architecture and without fine-tuning the model.

On a dataset designed to test the ability of models to track multiple fact
updates, a smaller LM augmented with MemoryPrompt can greatly outperform a
much larger LM that gets the whole context as its input. On the long-span
dialogue MSC dataset, MemoryPrompt provides comparable performance to that of
a model that can access the whole dialogue history from its orders-of-
magnitudes larger context window.

We compared MemoryPrompt to our re-implementation of RMT, a state-of-the-art
memory-augmented model that serves a similar purpose as MemoryPrompt. While
RMT achieves comparable performance in the tasks we considered, it does so at
the cost of catastrophic forgetting, that is, erasing the general token-
prediction capabilities of the underlying LM. MemoryPrompt is only slightly
affected by this problem.

MemoryPrompt still needs to be tested on a more varied set of challenges and
applied to larger LMs (our experiments were constrained by computational
limitations). Moreover, we have seen that MemoryPrompt is not robust to the
number of mutable facts in a sequence, as demonstrated by its performance on
the many-distractors dataset (see Table 3). In this setup, OPT-125M +
MemoryPrompt outperforms its full-context counterpart, but it lags behind the
larger full-context models.

Moving forward, we believe that the light-touch memory tracking abilities of
MemoryPrompt have a natural application in adapting pre-trained LMs to
specific users, and our main goal for the future is to extend and test the
model in a realistic setup in which a LM needs to adapt to a persistent user,
by storing user-specific knowledge in its memory. This, in turn, raises
interesting questions concerning the nature of memories. Can the memory
system, for example, learn which types of facts are user-dependent and highly
mutable, and should be constantly tracked and updated? Can our method be
extended to track different streams of information coming from different
interlocutors? We leave these and other questions to future work.

## 7\. Ethics Statement

As we are using pre-trained language models without re-training them, we are
neither amplifying nor alleviating the ethical issues that come with them. We
are also using datasets derived from existing resources, and we are not aware
of specific ethical issues pertaining to these resources.

Current state-of-the-art language models are extremely large, and they can
only be trained and run by a few companies that have access to enormous
computational resources. We see efforts, such as ours, to improve the
performance of smaller language models as positively contributing to the need
to democratize access to high-performance AI systems.

## 8\. Acknowledgements

We thank Marco Del Tredici and the COLT group for many interesting discussions
about the topic of memory in language models. We thank Corentin Kervadec for
helpful feedback and suggestions, as well as for help in constructing the
fact-updating datasets. Our work was funded by the European Research Council
(ERC) under the European Union’s Horizon 2020 research and innovation
programme (grant agreement No. 101019291). This paper reflects the authors’
view only, and the funding agency is not responsible for any use that may be
made of the information it contains.

## 9\. Bibliographical References

##

  * Beltagy et al. (2020) Iz Beltagy, Matthew Peters, and Arman Cohan. 2020.  Longformer: The long-document transformer.  https://arxiv.org/abs/2004.05150.
  * Bengio et al. (1994) Yoshua Bengio, Patrice Simard, and Paolo Frasconi. 1994.  Learning long-term dependencies with gradient descent is difficult.  _IEEE transactions on neural networks_ , 5(2):157–166.
  * Bulatov et al. (2023) Aydar Bulatov, Yuri Kuratov, and Mikhail Burtsev. 2023.  Scaling Transformer to 1M tokens and beyond with RMT.  https://arxiv.org/abs/2304.11062.
  * Bulatov et al. (2022) Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev. 2022.  Recurrent memory transformer.  In _Proceedings of NeurIPS_ , pages 11079–11091, New Orleans, LA.
  * Chen et al. (2023) Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. 2023.  Extending context window of large language models via positional interpolation.  _arXiv preprint arXiv:2306.15595_.
  * Dai et al. (2019) Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. 2019.  Transformer-XL: Attentive language models beyond a fixed-length context.  In _Proceedings of ACL_ , pages 2978–2988, Florence, Italy.
  * Elazar et al. (2021) Yanai Elazar, Nora Kassner, Shauli Ravfogel, Abhilasha Ravichander, Eduard Hovy, Hinrich Schütze, and Yoav Goldberg. 2021.  Measuring and improving consistency in pretrained language models.  _Transactions of the Association for Computational Linguistics_ , 9:1012–1031.
  * Elsahar et al. (2018) Hady Elsahar, Pavlos Vougiouklis, Arslen Remaci, Christophe Gravier, Jonathon Hare, Frederique Laforest, and Elena Simperl. 2018.  T-rex: A large scale alignment of natural language with knowledge base triples.  In _Proceedings of the Eleventh International Conference on Language Resources and Evaluation (LREC 2018)_.
  * Fan et al. (2020) Angela Fan, Thibaut Lavril, Edouard Grave, Armand Joulin, and Sainbayar Sukhbaatar. 2020.  Addressing some limitations of transformers with feedback memory.  https://arxiv.org/abs/2002.09402.
  * Graves et al. (2016) Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska-Barwinska, Sergio Gomez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, Adrià Puigdomènech Badia, Karl Moritz Hermann, Yori Zwols, Georg Ostrovski, Adam Cain, Helen King, Christopher Summerfield, Phil Blunsom, Koray Kavukcuoglu, and Demis Hassabis. 2016.  Hybrid computing using a neural network with dynamic external memory.  _Nature_ , 538(7626):471–476.
  * Hochreiter (1998) Sepp Hochreiter. 1998.  The vanishing gradient problem during learning recurrent neural nets and problem solutions.  _International Journal of Uncertainty, Fuzziness and Knowledge-Based Systems_ , 6(02):107–116.
  * Hochreiter and Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber. 1997.  Long short-term memory.  _Neural Computation_ , 9(8):1735–1780.
  * Hutchins et al. (2022) DeLesley Hutchins, Imanol Schlag, Yuhuai Wu, Ethan Dyer, and Behnam Neyshabur. 2022\.  Block-recurrent transformers.  _arXiv preprint arXiv:2203.07852_.
  * Joulin and Mikolov (2015) Armand Joulin and Tomas Mikolov. 2015.  Inferring algorithmic patterns with stack-augmented recurrent nets.  In _Proceedings of NIPS_ , Montreal, Canada.  Published online: https://papers.nips.cc/book/advances-in-neural-information-processing-systems-28-2015.
  * Kummerfeld et al. (2019) Jonathan K. Kummerfeld, Sai R. Gouravajhala, Joseph J. Peper, Vignesh Athreya, Chulaka Gunasekara, Jatin Ganhotra, Siva Sankalp Patel, Lazaros C Polymenakos, and Walter Lasecki. 2019.  A large-scale corpus for conversation disentanglement.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pages 3846–3856, Florence, Italy. Association for Computational Linguistics.
  * Lester et al. (2021) Brian Lester, Rami Al-Rfou, and Noah Constant. 2021.  The power of scale for parameter-efficient prompt tuning.  In _Proceedings of EMNLP_ , pages 3045–3059, Punta Cana, Dominican Republic.
  * Li and Liang (2021) Xiang Li and Percy Liang. 2021.  Prefix-tuning: Optimizing continuous prompts for generation.  In _Proceedings of ACL_ , pages 4582–4597, Online.
  * Liu et al. (2021) Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig. 2021.  Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing.  https://arxiv.org/abs/2107.13586.
  * Loshchilov and Hutter (2017) Ilya Loshchilov and Frank Hutter. 2017.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_.
  * Stephen et al. (2017) Merity Stephen, Xiong Caiming, Bradbury James, and Richard Socher. 2017.  Pointer sentinel mixture models.  _Proceedings of ICLR_.
  * Sukhbaatar et al. (2015) Sainbayar Sukhbaatar, Arthur Szlam, Jason Weston, and Rob Fergus. 2015.  End-to-end memory networks.  http://arxiv.org/abs/1503.08895.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  In _Proceedings of NIPS_ , pages 5998–6008, Long Beach, CA.
  * Voita et al. (2023) Elena Voita, Javier Ferrando, and Christoforos Nalmpantis. 2023.  Neurons in large language models: Dead, n-gram, positional.  https://arxiv.org/abs/2309.04827.
  * Wu et al. (2022) Qingyang Wu, Zhenzhong Lan, Kun Qian, Jing Gu, Alborz Geramifard, and Zhou Yu. 2022\.  Memformer: A memory-augmented transformer for sequence modeling.  In _Findings of ACL_ , pages 308–318, Online.
  * Xu et al. (2021) Jing Xu, Arthur Szlam, and Jason Weston. 2021.  Beyond goldfish memory: Long-term open-domain conversation.  _arXiv preprint arXiv:2107.07567_.
  * Zhang et al. (2022) Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. 2022.  Opt: Open pre-trained transformer language models.  _arXiv preprint arXiv:2205.01068_.
  * Zhong et al. (2021) Zexuan Zhong, Dan Friedman, and Danqi Chen. 2021.  Factual probing is [MASK]: Learning vs. learning to recall.  In _Proceedings of NAACL_ , pages 5017–5033, Online.