# Dataset Distillation for Offline Reinforcement Learning

Jonathan Light∗<EMAIL_ADDRESS>
Department of Computer Science
Rensselaer Polytechnic Institute
Troy, NY 12180, USA Yuanzhe Liu∗<EMAIL_ADDRESS>
Department of Computer Science
Rensselaer Polytechnic Institute
Troy, NY 12180, USA Ziniu Hu<EMAIL_ADDRESS>
Computing + Mathematical Sciences Department
California Institute of Technology
Los Angeles, CA 91125, USA

###### Abstract

Offline reinforcement learning often requires a quality dataset that we can
train a policy on. However, in many situations, it is not possible to get such
a dataset, nor is it easy to train a policy to perform well in the actual
environment given the offline data. We propose using data distillation to
train and distill a better dataset which can then be used for training a
better policy model. We show that our method is able to synthesize a dataset
where a model trained on it achieves similar performance to a model trained on
the full dataset or a model trained using percentile behavioral cloning. Our
project site is available here. We also provide our implementation at this
GitHub repository.

**footnotetext: These authors contributed equally to this work

##  1 Introduction

A significant challenge in reinforcement learning (RL) is that the data
generation process is coupled with the training process, and data generation
requires frequent online interaction with the environment, which is not
possible in many settings. Offline RL aims to solve this problem by decoupling
the two and training the agent on a given a static, fixed dataset (Levine et
al. (2020),Prudencio et al. (2023)). However, offline RL relies on a dataset
generated by a good expert policy. We often do not have access to data
generated by good policies, only mediocre ones. Offline training also means we
face the distributional shift problem, where the policy trained on the dataset
is produces a different data distribution than the one in the dataset.

Instead of taking the usual offline RL approach of finding a better way to
train a model given the offline dataset, we take an alternate approach of
asking, is there a way to distill a better offline dataset to train on? We
believe that this approaches offers several advantages over finding a better
training method. First of all, it is easier to _interpret_ a distilled dataset
vs a better trained model. Secondly, distillation tends to lead to better
generalization capabilities since we learn the key features of the input space
(Stanton et al. (2021), Sachdeva and McAuley (2023)). Thirdly, a distilled
dataset is much smaller than the original offline dataset, which improves
sample efficiency.

We propose using a method from data distillation Wang et al. (2018) known as
gradient matching to train a smaller synthetic dataset on the offline dataset.
We evaluated the students trained using our procedure on the Procgen
environment Cobbe et al. (2020), which consists of procedurally generated
games. Specifically, the student is only given access to offline expert policy
data on some of the procedurally generated maps, and must generalize the
knowledge they learn on those maps to other unseen, out-of-distribution
settings. We show that students trained using synthetic data are able to
perform similarly or better than students trained on the original offline
policy dataset or students trained using percentile behavorial cloning both in
distribution and out of distribution, despite the fact that they are _trained
on a smaller dataset._

Why does training on a smaller dataset help in RL settings specifically? RL is
a learning paradigm that is natural very prone to randomness and over-fitting
due to the fact that the agent also has control of the data generation
process. The insight that we have here is that a smaller and well controlled
dataset can reduce randomness and overfitting. Just like how humans learn more
effectively when read a well written book instead of reading many low quality
articles, reinforcement learning agents can also learn a better, more
generalizable policy by training on a high quality dataset.

To summarize, our main contributions are as follows:

  * •

We propose a new method Synthetic that synthesizes a new dataset given a
offline dataset of trajectories generated by an expert policy using dataset
distillation.

  * •

We show that a RL-model trained on a dataset synthesized using our method is
able to perform similarly or better in the environment than directly training
on the expert data or other techniques such as percentile behavioral cloning

  * •

We demonstrate how we are able to achieve similar performance with a far
smaller dataset when training the RL-model on our synthesized dataset

##  2 Methodology

We describe our general problem setting, the baseline methods of tackling the
problem, and our method here.

###  2.1 Offline reinforcement learning problem setting

In our reinforcement learning setting, the environment is modelled as a Markov
decision process (MDP)
$\mathcal{M}=\langle\mathcal{S},\mathcal{A},T,R,s_{0},s_{-1}\rangle$ with
state space $\mathcal{S}$, action space $\mathcal{A}$, $T(s_{t+1}|s_{t},a)$ is
the probabilistic transition function, $R(s_{t},a)\in\mathbb{R}$ is the
immediate transition reward, $s_{0}$ is the start state, and $s_{-1}$ is the
end state since we are in an episodic setting. A policy function
$\pi\mathrel{\mathop{\mathchar
58\relax}}\mathcal{S}\rightarrow\Delta(\mathcal{A})$ is a mapping from a state
to a probability distribution over the action space. We will assume that both
the $\mathcal{S}$ and $\mathcal{A}$ are discrete in our setting. When a policy
can be parameterized by some parameter $\theta$, we denote the policy as
$\pi_{\theta}$. Since we are using deep reinforcement learning, $\pi_{\theta}$
will be a neural network with weights $\theta$.

The goal of parametrized reinforcement learning is to learn the optimal
$\theta^{*}$ that maximizes the cumulative of an episode. We define the
cumulative reward in terms of the trajectory distribution induced by the
policy $\pi_{\theta}$. A trajectory $\tau$ is a sequence of states and actions
that starts with $s_{0}$ and ends with $s_{-1}$, i.e.
$\tau=((s_{0},a_{0}),(s_{1},a_{1}),...,(s_{-1},a_{-1}))$. Then the trajectory
distribution $p_{\pi}$ induced by policy $\pi$ and environment $\mathcal{M}$
is given by

$p_{\pi}(\tau)=\prod_{(s_{t},a_{t})\in\tau}\pi({a_{t}|s_{t}})T(s_{t+1}|s_{t},a_{t})$

The expected cumulative reward is then

$J(\pi)=\mathbb{E}_{\tau\sim p_{\pi}}\left[\sum_{(s,a)\in\tau}R(s,a)\right]$

and our goal is to find

$\theta^{*}=\operatorname*{arg\,max}_{\theta}J(\pi_{\theta}),\quad\max_{\theta}J(\pi_{\theta})\approx\max_{\pi}J(\pi)$

In offline reinforcement learning, the agent is not allowed to interact with
the environment and collect data through that. Instead, we are given a static
data set $\mathcal{D}=\\{(s_{t}^{i},a_{t}^{i},s_{t+1}^{i},r_{t}^{i})\\}$ of
transitions to learn the best policy $\pi_{\theta}$ from, where
$(s_{t}^{i},a_{t}^{i})\in\tau^{i}$ are part of the trajectory of episode $i$.
Sometimes the dataset also includes the future return
$G_{t}^{i}=\sum_{(s_{k}^{i},a_{k}^{i})\in\tau^{i},k\geq t}R(s,a)$ at both the
current state and for the whole episode, so
$\mathcal{D}=\\{(s_{t}^{i},a_{t}^{i},s_{t+1}^{i},r_{t}^{i},G_{t}^{i},G_{0}^{i})\\}$.
We assume that the data is generated using some policy $\pi_{\beta}$, i.e.
$\tau^{i}\sim p_{\pi_{\beta}}$.

###  2.2 Behavioral cloning

A common way to approach offline RL is to simply attempt to train a policy
$\pi_{\theta}$ to imitate $\pi_{\beta}$. This is done in a supervised learning
fashion by training $\pi_{\theta}$ to predict $a_{t}^{i}$ given $s_{t}^{i}$
and minimizing the loss

$\mathcal{L}_{\textbf{BC}}(\theta|\mathcal{D})=\sum_{(s_{t}^{i},a_{t}^{i},...)\in\mathcal{D}}w(s_{t}^{i},a_{t}^{i},...)||\pi_{\theta}(s_{t}^{i})-a_{t}^{i}||$
(1)

where $||$ is some norm implying some notion of distance and $w$ is some
weighting on the data-points in $\mathcal{D}$. Usually we take the uniform
weighting $w(*)=1$. Since we usually use stochastic gradient descent (SGD) to
minimize the loss, in this case we can simply use
$p(*)=\frac{w(*)}{\sum_{(s_{t}^{i},a_{t}^{i},...)\in\mathcal{D}}w(s_{t}^{i},a_{t}^{i},...)}$
as the probability of selecting the sample.

Since $\pi_{\beta}$ might not be optimal, we can attempt to train a better
policy $\pi_{\theta}$ by filtering out observations in $\mathcal{D}$ that lead
to poor outcomes. In other words, we let
$w(s_{t}^{i},a_{t}^{i},...,G_{0}^{i})=\mathbb{I}_{G_{0}^{i}\geq b}$ where $b$
is some threshold on good vs bad outcomes. We call the method of choosing $b$
such that we only end up with x% of the initial dataset, and training a policy
$\pi_{\theta}$ using BC on it as BCx%, or percentile behavior cloning. In
other words, the goal of percentile behavior cloning is to filter out a better
training dataset.

###  2.3 Synthetic dataset

Instead of filtering out bad samples in order to create a good training
dataset, our method _directly learns a good training sample_ instead. This is
done through dataset distillation, a technique used in supervised learning
Wang et al. (2018). We described how we ‘train’ a synthesized dataset
$\mathcal{D}_{\phi}$ here, parameterized by $\phi$, given the offline dataset
$\mathcal{D}_{\text{real}}$. Our method aims to reduce the gradient matching
loss of $\phi$ with respect to some random initialization of the model weights
$\theta$ according to some distribution $\theta\sim p_{\theta}$. Given some
model parameters $\theta$, we first get the gradient of $\theta$ with respect
to the BC loss we defined in 1 on both the real dataset,
$\nabla_{\theta}\mathcal{L}_{\text{BC}}(\theta|\mathcal{D}_{\text{real}})$,
and our synthetic one,
$\nabla_{\theta}\mathcal{L}_{\text{BC}}(\theta|\mathcal{D}_{\phi})$. Then we
define the gradient matching loss as

$\mathcal{L}_{\text{grad match}}(\phi|\theta_{i})=\mathbb{E}_{\theta\sim
p_{\theta}}\left[||\nabla_{\theta}\mathcal{L}_{\text{BC}}(\theta|\mathcal{D}_{\text{real}})-\nabla_{\theta}\mathcal{L}_{\text{BC}}(\theta|\mathcal{D}_{\phi})||\right]$
(2)

We then use SGD to minimize this loss. This method helps guarantee that the
synthesized dataset $\mathcal{D}_{\text{syn}}$ will produce a gradient similar
to that of $\mathcal{D}$ when a model is trained on it.

Figure 1: Overview of our dataset distillation process. On the left we train
the dataset by taking the matching gradient loss between the real offline
dataset and our synthetic dataset. On the right we then use the trained
synthetic dataset to train a RL model, which we then evaluate on the real
environment.

##  3 Experimental Setup

We provide details on how we implemented our experiments below, including what
environments we tested our method on, the architecture for our models, and how
we trained the models.

###  3.1 Environment

We used the Procgen environment developed by OpenAI, a suite of 16
procedurally generated environments that allow the creation of adaptive
environments with the use of different seeds Cobbe et al. (2020). The inherent
mechanics of these environments serve as an ideal platform for evaluating a
student model’s ability to learn and adapt to variations introduced by
different seeds. Furthermore, the diverse environments help our study as the
agent is trained for a variety of challenges that may not occur during
standard training procedures. This way we are allowed to scrutinize the
agents’ performance across different scenarios that could arise in practical
implementations. Some sample games from the Procgen environment are shown in
Figure 2.

Figure 2: Screenshots of games in Procgen Benchmark Cobbe et al. (2020)

In Procgen environments, the state space, consists of 64x64 pixel images (RGB
array with three channels and values $0$ to $255$). The action space is
discrete in nature and usually includes movements (up, down, left, right) and
interactions like collecting items or opening doors. For our experiments, we
consider three procedurally generated games: Bigfish, Starpilot, and Jumper.
In Starpilot for example, the player must navigate a space ship to avoid being
hit by bullets and shoot down enemies in an arcade game fashion. Enemies and
obstacles are procedurally generated, so each ‘map’ is different. The key
characteristic of the Procgen benchmark is that, given a different seed, the
player encounters a different ‘map’ though the game rules are unchanged. _A
key challenge for AI agents lies in how well they can adapt to seeds that they
have not seen before_.

###  3.2 Model architectures and training

####  3.2.1 Model training

There are two model architectures that we consider – the expert policy model
and the student policy model. The expert model is the neural network that we
use as an expert policy, which is then used to produce the offline dataset
$\mathcal{D}$. The student model is the neural network which we train on the
offline data to produce a policy $\pi_{\theta}$. Recall that the goal of
offline RL is to optimize the parameters $\theta$ in order to produce a good
policy, as described in section 2.1.

The expert model is an agent with convolutional architecture found in IMPALA
Espeholt et al. (2018), following the convention in Procgen paper Cobbe et al.
(2020). We trained the expert policy on the environment using proximal policy
optimization (PPO) Schulman et al. (2017). We used the PFRL package and
followed one community-created pytorch implementation on GitHub Fujita et al.
(2021); Lerrytang . We trained the expert model for 25 million steps on 200
seeds until it achieved a satisfying level of performance on the environment.
Hence, we trained three expert models in total for each of the three
environments. This IMPALA network has three convolutional blocks. The first
convolutional block has the output channel 16, the next two blocks has the
output channel 32. This setup has a total number of
$9712+41632+41632+524544+3855+257=621632$ trainable parameters. Table 1 shows
more details regarding expert model.

For the student (the model that tries to mimic the expert), we use CNN
(convolutional neural network) as our base model. The CNN model has 4
convolutional modules followed by a fully connected layer to the logits. The
convolutional layers utilize a $3\times 3$ kernel with $3$ output channels and
coupled with a ReLU activation and a average $2$ dimensional pooling layer
with pool kernal size $2$ and stride $2$. For first $2$ layers, the dimension
of output channel is $4$ times larger than that of the input channel. For the
last $2$ layers, the dimension of output channel is $4$ times smaller than
that of the input channel. So the output channel of the last layer is the same
as the input channel of the first layer. The output of the convolutional
layers is then passed to a fully connected layer, which maps to the logits.
This leaner setup has a total number of $336+5232+336+327+735=6966$ trainable
parameters, reducing the size of the fully connect layer. Student model has
much fewer trainable parameters comparing to that of the expert model. Table 1
contains hyperparameters of training student models.

For each database collection method, we train 10 students and take the average
of reward mean and reward standard deviation. We use Adam optimizer with
learning rate 5e-3 Kingma and Ba (2017). For behavioral cloning students, we
train them with 1000 steps and batch size 256. For Synthetic students, we
train them with only 100 steps and batch size 15.

| Expert | BC Student |  Synthetic Student
---|---|---|---
Model Params | 621632 | 6966 | 6966
Optimizer | Adam | Adam | Adam
Learning Rate | 1e-5 | 5e-3 | 5e-3
Batch Size | 8 | 256 | 15
Steps | 25M | 1000 | 100
Table 1: Hyperparameters: Expert V.S. BC Student V.S. Synthetic Student

####  3.2.2 Data construction

To construct offline RL dataset, we ran 100 episodes of our trained experts on
all three environments, and store the data as mentioned in Section 2.1. To
construct the synthetic data, given the synthetic data size, we sample data
randomly from the offline RL data generated by expert, and then use gradient
matching loss to udpate the synthetic data as illustrated in Figure 1. We use
our experts to the RL Model to obtain gradients and compute the loss. Here, we
use SGD optimizer with learning rate $0.1$ and momentum $0.5$, training with
1000 epochs.

##  4 Results

We compare our method to (1) the expert policy that we trained in an online
fashion by interacting with the environment until we achieved good
performance, which was then used to generate the offline dataset $\mathcal{D}$
and (2) a student trained on datasets with different levels of filtering using
percentile behavioral cloning as described in section 2.2. In other words, we
want to benchmark how well our method performs at _generating a quality
dataset that can be used for training an offline RL model_.

The student model was trained with the same number of stochastic gradient
descent steps and same batch size for all baseline methods. We show the in-
distribution performance, i.e. the performance on seeds contained in the
offline dataset $\mathcal{D}$, of the various methods in table 2 and figure 4.
We see here that Synthetic outperforms all percentile behavior cloning methods
in both Jumper and Bigfish environments. Synthetic does not outperform
percentile BC on Starpilot. We observed that the Starpilot expert mainly takes
one action during game – the ‘shoot’ action – compared to Bigfish and Jumper
where the expert takes a more even distribution of actions. Hence the dataset
is hard to distill because of the imbalanced samples of different actions.

Since we can procedurally generate out of distribution (OOD) scenarios in
Procgen, we also tested the OOD performance of the various dataset generation
methods, as shown in table 3 and figure 4. We see here that similar, Synthetic
outperforms percentile behavior cloning methods in Jumper environment.
Synthetic also matches all percentile BC performances on Bigfish. Synthetic
does not outperform percentile BC on Starpilot since the dataset of Starpilot
is imbalanced.

Our student model trained by Synthetic only uses $150$ data samples, as shown
in table 4. Given much smaller dataset size (less than half of BC10%) and much
fewer training steps as mentioned in table 1, Synthetic achieves competitive
results compared to behavioral cloning in different setups. Synthetic also
generalizes well out of distribution as the OOD performances matches the ID
results, and often times in both Starpilot and Jumper outperforms the ID
results.

ID Performance
---
Environment | Expert | BC 10% | BC 25 % | BC 40 % | BC 100 % | Synthetic
Bigfish | 14.27 $\pm$ 15.53 | 0.90 $\pm$ 1.44 | 0.93 $\pm$ 1.47 | 1.01 $\pm$ 1.62 | 1.00 $\pm$ 1.67 | 1.03 $\pm$ 1.99
Starpilot | 28.88 $\pm$ 19.41 | 1.73 $\pm$ 2.20 | 2.10 $\pm$ 2.66 | 2.17 $\pm$ 2.58 | 1.85 $\pm$ 2.22 | 1.5 $\pm$ 1.96
Jumper | 8.79 $\pm$ 3.26 | 1.79 $\pm$ 3.54 | 2.15 $\pm$ 4.12 | 1.95 $\pm$ 3.86 | 2.32 $\pm$ 4.13 | 2.76 $\pm$ 4.40
Table 2: Average in distribution performance of student trained on various
data collection methods. OOD Performance
---
Environment | Expert | BC 10% | BC 25 % | BC 40 % | BC 100 % | Synthetic
Bigfish | 6.03 $\pm$ 9.84 | 0.93 $\pm$ 1.38 | 0.85 $\pm$ 1.19 | 0.83 $\pm$ 1.28 | 0.87 $\pm$ 1.32 | 0.83 $\pm$ 1.04
Starpilot | 23.34 $\pm$ 18.30 | 1.83 $\pm$ 2.39 | 1.95 $\pm$ 2.39 | 2.12 $\pm$ 2.35 | 1.82 $\pm$ 2.20 | 1.54 $\pm$ 1.93
Jumper | 5.61 $\pm$ 4.96 | 1.81 $\pm$ 3.51 | 1.8 $\pm$ 3.73 | 1.82 $\pm$ 3.70 | 2.50 $\pm$ 4.26 | 2.86 $\pm$ 4.43
Table 3: Average out of distribution performance of student trained on various
data collection methods. Dataset Size
---
Environment | BC 10% | BC 25 % | BC 40 % | BC 100 % | Synthetic
Bigfish | 2027 | 5014 | 7336 | 10450 | 150
Starpilot | 1116 | 2796 | 4192 | 6830 | 150
Jumper | 392 | 919 | 1337 | 4837 | 150
Table 4: Dataset size used on various data collection methods with respect to
different environments.

Figure 3: In distribution performance of various data collection methods

Figure 4: Out of distribution performance of various data collection methods

##  5 Related work

###  5.1 Deep Reinforcement Learning

Deep reinforcement learning has seen incredible success recently in tackling
wide-ranging problems, from chess and Go to Atari games and robotics Silver et
al. (2016). We have also seen great improvements in the architecture used to
design such agents, from PPO Schulman et al. (2017) to decision transformers
Chen et al. (2021). As deep neural networks have shown their effectiveness in
various tasks, researchers in reinforcement learning have increasingly turned
their attention to them. Many complex reinforcement learning situations
require these versatile neural networks for tasks such as encoding the states
of the agents, learning complex policies, and assessing their performance.
Arulkumaran et al. (2017) give a good overview of various ways deep neural
networks were incorporated into reinforcement learning settings.

###  5.2 Knowledge Distillation

As using large deep neural networks started bringing remarkable success in
multiple real-world scenarios related to large-scale data, it became important
to deploy deep models within mobile devices and embedded systems. Bucilua et
al. (2006) first addressed this issue and proposed compression of large models
for transferring the information from large models to train a small model such
that accuracy is not hampered. Hinton et al. (2015) popularized the term
’knowledge distillation’ as the process of learning a small model from a large
model (teacher) to a small student model. In recent times, there have been
many extensions to knowledge distillation where the focus is on compressing
deep neural networks. The lightweight student models have paved the way for
integrating knowledge distillation in various applications like adversarial
attacks Papernot et al. (2016), security and privacy of data Wang et al.
(2019), data augmentation Lee et al. (2020) etc. KD has been a key instrument
in the study of natural language processing (NLP) Devlin et al. (2018). Sun et
al. (2019) and Tang et al. (2019) have used some lightweight variations of
BERT (called BERT model compression) through knowledge distillation. Jiao et
al. (2019) proposed a TinyBERT, a two-stage transformer knowledge
distillation, to make the framework even lighter.

###  5.3 Policy distillation

There have also been attempts to distill the policy of a expert (teacher)
network down to a student network directly Rusu et al. (2015), known as policy
distillation. Policy distillation is a specialized application of knowledge
distillation where it adapts the principles of KD in the context of
Reinforcement Learning. It is used to transfer knowledge from one policy to
another in deep RL. Czarnecki et al. (2019)identified three techniques for
distillation in DRL, making comparisons of their motivations and strengths
through theoretical and empirical analysis, including expected entropy
regularized distillation, which ensures convergence while learning quickly.
Policy distillation can also be used to extract an RL agent’s policy to train
a more efficient and smaller network that performs expertly Rusu et al.
(2015).

While our work also teaches a student a policy based on some teacher policy,
we take more indirect, offline approach where the student is only allowed to
see offline data generated by the expert (teacher).

###  5.4 Dataset Distillation

Dataset distillation is a dataset reduction method that synthesizes a smaller.
In the original work, this is by feeding a randomly initialized model with
samples from the real data and samples from the synthetic dataset and taking
the gradient of the model with respect to these two data samples Wang et al.
(2018). The difference between the two gradients is taken as the loss, and the
values of the data in the synthetic dataset are updated using SGD (while
keeping the model weights fixed) Since then, a wide variety of different
distillation methods have been proposed (Yu et al. (2023), Lei and Tao
(2023)). In one such work, instead of matching the gradients for a single
sample, the sum of the gradients (total change in model parameters) after
training on a series of samples is matched instead (Cazenavette et al.
(2022)). Despite recent interest in this technique, to the best of the
author’s knowledge, there have not been any applications of dataset
distillation to reinforcement learning yet.

###  5.5 Task Generalization

The goal of task generalization is to transfer what is learned from doing one
task to other tasks, broadening the capabilities of the model. In the ideal
scenario, the learned model should be able to apply its knowledge to changing
tasks by using the core knowledge learned. (Taylor and Stone (2007)) suggests
a new transfer method called ”Rule Transfer” which aims to learn the rules of
a source task and apply them to other target tasks. (Taylor et al. (2008))
aims to learn mappings between the transitions from the source to the target
task. In the problem suggested in (Oh et al. (2017)), agents are required to
learn to execute sequences of instructions after mastering subtask-solving
skills. The problem gives out a good basis for generalizing to unseen tasks.
In (Lehnert et al. (2020)), authors suggest that using reward predictions
gives the agents better generalization capabilities.

###  5.6 Other Works

In the work (Yuan et al. (2021)), the authors propose a Reinforcement Learning
based method for Knowledge Distillation for scenarios where multiple teachers
are available. Their work is focused on NLPs and uses large pre-trained models
like BERT and RoBERTa, where the framework dynamically assigns weights to
different teacher models on each training instance, in order to maximize the
performance of the distilled student model. In (Xie et al. (2021)), the
authors present a novel framework (DRL-Rec) for knowledge distillation between
RL-based models in list-wise recommendation, where they introduce one module
by which the teacher decides on which lesson to teach to the student and
another confidence-guided distillation through which it is determined how much
the student should learn from each lesson.

##  6 Conclusion and limitations

There are several limitations to our study. The first one is that, limited by
the computation resources and time we had, we only tested on three
environments in Procgen. However, we believe that the experiments on these
environments demonstrate the potential of our method, and we look forward to
future work on other environments. We also focus on imitation policy learning
in our work since our emphasis is on the dataset distillation, not the policy
learning method. However, it is also possible to use other RL methods such as
q-learning or actor-critic to train policies on the synthetic dataset. We
mainly benchmark against percentile behavior cloning, since that is that is
the closest method in the existing literature that ‘filters’ for a better
quality training dataset.

In conclusion, we proposed and tested a method that synthesizes a better
quality training dataset for offline reinforcement learning. The performance
of our method suggests that the quality of the dataset a key component to
training a better model, and a smaller but higher quality dataset can lead to
similar or better performance compared to a larger one. We believe that these
methods can be highly effective in settings with low amounts of data or noisy
data, and where data cannot be collected online. There is still much to
explore in this research space.

## References

  * Arulkumaran et al. (2017) Kai Arulkumaran, Marc Peter Deisenroth, Miles Brundage, and Anil Anthony Bharath.  Deep reinforcement learning: A brief survey.  _IEEE Signal Processing Magazine_ , 34(6):26–38, 2017.
  * Bucilua et al. (2006) C Bucilua, R Caruana, and A Niculescu-Mizil.  Model compression, in proceedings of the 12 th acm sigkdd international conference on knowledge discovery and data mining.  _New York, NY, USA_ , 4, 2006.
  * Cazenavette et al. (2022) George Cazenavette, Tongzhou Wang, Antonio Torralba, Alexei A Efros, and Jun-Yan Zhu.  Dataset distillation by matching training trajectories.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 4750–4759, 2022.
  * Chen et al. (2021) Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch.  Decision transformer: Reinforcement learning via sequence modeling, 2021.
  * Cobbe et al. (2020) Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman.  Leveraging procedural generation to benchmark reinforcement learning.  In Hal Daumé III and Aarti Singh, editors, _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pages 2048–2056. PMLR, 13–18 Jul 2020.  URL https://proceedings.mlr.press/v119/cobbe20a.html.
  * Czarnecki et al. (2019) Wojciech M Czarnecki, Razvan Pascanu, Simon Osindero, Siddhant Jayakumar, Grzegorz Swirszcz, and Max Jaderberg.  Distilling policy distillation.  In _The 22nd international conference on artificial intelligence and statistics_ , pages 1331–1340. PMLR, 2019.
  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Espeholt et al. (2018) Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymir Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, Iain Dunning, Shane Legg, and Koray Kavukcuoglu.  Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures, 2018.
  * Fujita et al. (2021) Yasuhiro Fujita, Prabhat Nagarajan, Toshiki Kataoka, and Takahiro Ishikawa.  Chainerrl: A deep reinforcement learning library.  _Journal of Machine Learning Research_ , 22(77):1–14, 2021.  URL http://jmlr.org/papers/v22/20-376.html.
  * Hinton et al. (2015) Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network, 2015.
  * Jiao et al. (2019) Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu.  Tinybert: Distilling bert for natural language understanding.  _arXiv preprint arXiv:1909.10351_ , 2019.
  * Kingma and Ba (2017) Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization, 2017.
  * Lee et al. (2020) Hankook Lee, Sung Ju Hwang, and Jinwoo Shin.  Self-supervised label augmentation via input transformations.  In _International Conference on Machine Learning_ , pages 5714–5724. PMLR, 2020.
  * Lehnert et al. (2020) Lucas Lehnert, Michael L Littman, and Michael J Frank.  Reward-predictive representations generalize across tasks in reinforcement learning.  _PLoS computational biology_ , 16(10):e1008317, 2020.
  * Lei and Tao (2023) Shiye Lei and Dacheng Tao.  A comprehensive survey of dataset distillation.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 2023.
  * (16) Lerrytang.  Lerrytang/train-procgen-pfrl: Pytorch code to train and evaluate procgen tasks.  URL https://github.com/lerrytang/train-procgen-pfrl/tree/main.
  * Levine et al. (2020) Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu.  Offline reinforcement learning: Tutorial, review, and perspectives on open problems.  _arXiv preprint arXiv:2005.01643_ , 2020.
  * Oh et al. (2017) Junhyuk Oh, Satinder Singh, Honglak Lee, and Pushmeet Kohli.  Zero-shot task generalization with multi-task deep reinforcement learning.  In _International Conference on Machine Learning_ , pages 2661–2670. PMLR, 2017.
  * Papernot et al. (2016) Nicolas Papernot, Patrick McDaniel, Xi Wu, Somesh Jha, and Ananthram Swami.  Distillation as a defense to adversarial perturbations against deep neural networks.  In _2016 IEEE symposium on security and privacy (SP)_ , pages 582–597. IEEE, 2016.
  * Prudencio et al. (2023) Rafael Figueiredo Prudencio, Marcos ROA Maximo, and Esther Luna Colombini.  A survey on offline reinforcement learning: Taxonomy, review, and open problems.  _IEEE Transactions on Neural Networks and Learning Systems_ , 2023.
  * Rusu et al. (2015) Andrei A Rusu, Sergio Gomez Colmenarejo, Caglar Gulcehre, Guillaume Desjardins, James Kirkpatrick, Razvan Pascanu, Volodymyr Mnih, Koray Kavukcuoglu, and Raia Hadsell.  Policy distillation.  _arXiv preprint arXiv:1511.06295_ , 2015.
  * Sachdeva and McAuley (2023) Noveen Sachdeva and Julian McAuley.  Data distillation: A survey, 2023.
  * Schulman et al. (2017) John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms, 2017.
  * Silver et al. (2016) David Silver, Aja Huang, Christopher Maddison, Arthur Guez, Laurent Sifre, George Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis.  Mastering the game of go with deep neural networks and tree search.  _Nature_ , 529:484–489, 01 2016.  doi: 10.1038/nature16961.
  * Stanton et al. (2021) Samuel Stanton, Pavel Izmailov, Polina Kirichenko, Alexander A Alemi, and Andrew G Wilson.  Does knowledge distillation really work?  _Advances in Neural Information Processing Systems_ , 34:6906–6919, 2021.
  * Sun et al. (2019) Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu.  Patient knowledge distillation for bert model compression.  _arXiv preprint arXiv:1908.09355_ , 2019.
  * Tang et al. (2019) Raphael Tang, Yao Lu, Linqing Liu, Lili Mou, Olga Vechtomova, and Jimmy Lin.  Distilling task-specific knowledge from bert into simple neural networks.  _arXiv preprint arXiv:1903.12136_ , 2019.
  * Taylor and Stone (2007) Matthew E Taylor and Peter Stone.  Cross-domain transfer for reinforcement learning.  In _Proceedings of the 24th international conference on Machine learning_ , pages 879–886, 2007.
  * Taylor et al. (2008) Matthew E Taylor, Gregory Kuhlmann, and Peter Stone.  Autonomous transfer for reinforcement learning.  In _AAMAS (1)_ , pages 283–290, 2008.
  * Wang et al. (2019) Ji Wang, Weidong Bao, Lichao Sun, Xiaomin Zhu, Bokai Cao, and S Yu Philip.  Private model compression via knowledge distillation.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 33, pages 1190–1197, 2019.
  * Wang et al. (2018) Tongzhou Wang, Jun-Yan Zhu, Antonio Torralba, and Alexei A Efros.  Dataset distillation.  _arXiv preprint arXiv:1811.10959_ , 2018.
  * Xie et al. (2021) Ruobing Xie, Shaoliang Zhang, Rui Wang, Feng Xia, and Leyu Lin.  Explore, filter and distill: Distilled reinforcement learning in recommendation.  In _Proceedings of the 30th ACM International Conference on Information & Knowledge Management_, pages 4243–4252, 2021.
  * Yu et al. (2023) Ruonan Yu, Songhua Liu, and Xinchao Wang.  Dataset distillation: A comprehensive review.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 2023.
  * Yuan et al. (2021) Fei Yuan, Linjun Shou, Jian Pei, Wutao Lin, Ming Gong, Yan Fu, and Daxin Jiang.  Reinforced multi-teacher selection for knowledge distillation.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 35, pages 14284–14291, 2021.


# Insights into LLM Long-Context Failures:
When Transformers Know but Don’t Tell

Taiming Lu✰ Muhan Gao✰ Kuai Yu✰
Adam Byerly Daniel Khashabi
Johns Hopkins University

###### Abstract

Large Language Models (LLMs) exhibit positional bias, struggling to utilize
information from the middle or end of long contexts. Our study explores LLMs’
long-context reasoning by probing their hidden representations. We find that
while LLMs encode the position of target information, they often fail to
leverage this in generating accurate responses. This reveals a disconnect
between information retrieval and utilization, a ‘know but don’t tell’
phenomenon. We further analyze the relationship between extraction time and
final accuracy, offering insights into the underlying mechanics of transformer
models. The code is accessible here: https://github.com/TaiMingLu/know-dont-
tell.

Insights into LLM Long-Context Failures:
When Transformers Know but Don’t Tell

Taiming Lu✰ Muhan Gao✰ Kuai Yu✰ Adam Byerly Daniel Khashabi Johns Hopkins
University

00footnotetext: ✰ indicates equal contributions. 00footnotetext:  Emails:
{tlu37, mgao38<EMAIL_ADDRESS>
##  1 Introduction

The advent of Large Language Models (LLMs), optimized with advanced
transformer architectures, has delivered marked improvement in language
processing capabilities. These models excel at simultaneously processing
extended contexts Ding et al. (2024); Chen et al. (2023), significantly
benefiting various downstream tasks like long-text question answering,
summarization, and inference Wang et al. (2024); Zhang et al. (2024a); Shaham
et al. (2022, 2023).

Despite their advanced capabilities, LLMs often struggle to utilize long
inputs fully. This tendency, known as positional bias, leads LLMs to
disproportionately prioritize information at the beginning or end of the input
sequence Wang et al. (2023) while crucial details in the middle are frequently
overlooked Liu et al. (2023b). Numerous strategies have been proposed to
address these biases Tang et al. (2024); Li et al. (2023); Zhang et al.
(2024b), yet the underlying causes and potential solutions remain unclear.
This underscores the need for a deeper investigation into how LLMs handle
long-context integration. To fully assess the capabilities of LLMs in handling
extended contexts, it is not enough to merely evaluate their final
performance: some important information is hidden in models’ representations.

Figure 1: Following prompts by Liu et al. (2023b), for each transformer layer,
we train a probing classifier to probe the model’s ability to identify useful
information. The peak accuracy among layers indicates the effectiveness of the
model’s processing of context.

In this work, we present a probing analysis of LLMs long-context
generalization. Specifically, we build probes based on the internal
representation of LLMs for various layers and positions to measure the
accuracy of reconstructing the position they correspond to (see Figure 1). A
necessary condition for effective long-context processing by LLMs is their
ability to encode positional information in their intermediate
representations.

We conduct experiments on two tasks from Liu et al. (2023b) and three recent
open-source models. Our findings reveal a gap between the accuracy of LLMs’
predictions and the probes on their representations. Notably, while LLMs can
accurately identify the position of crucial information within the context,
they often fail to utilize this information effectively in their responses,
leading to what we term the ‘know but don’t tell’ phenomenon. To our
knowledge, this is the first work to use probing analysis to highlight this
observation. We hope that our work on distinguishing “knowing” and “telling”
motivates future work on tackling the long-context challenges of LLMs.

In summary, our contributions are as follows: (1) Probing analysis: We
introduce a novel framework to investigate the long-context reasoning
capabilities of LLMs. This framework allows us to measure how accurately LLMs
encode positional information across various layers and positions within their
intermediate representations. (2) Empirical evaluation: We conduct
comprehensive experiments using tasks from Liu et al. (2023b) and three recent
open-source models. Our empirical evaluation provides new insights into the
positional biases of LLMs and their impact on model performance. (3) ‘Know but
Don’t Tell’ phenomenon: Our analysis reveals a critical gap between LLMs’
ability to encode and utilize positional information. We identify the "know
but don’t tell" phenomenon, where LLMs accurately identify the position of
crucial information but fail to leverage this knowledge in generating accurate
responses.

We believe these contributions provide a significant step towards
understanding and improving the long-context processing capabilities of LLMs.
By distinguishing between the encoding and utilization of positional
information, our work lays the foundation for future advancements in LLM
performance and reliability.

##  2 Related Work

#### Positional bias.

LLMs exhibit a positional bias, where their performance is influenced by the
location of crucial context information Zhao et al. (2021). One prominent
example is the “lost in the middle” phenomenon, where comprehension declines
for information in the center of a long context Liu et al. (2023b).
Additionally, recency bias is observed, particularly in few-shot learning
scenarios, where models tend to favor information near the end of the prompt
Zhao et al. (2021). Such biases could stem from the positioning of key data in
pre-training sets, which often places important elements near critical points
Peysakhovich and Lerer (2023). Our work delves into this phenomenon by
examining the underlying mechanisms within the transformer layers of LLMs.

#### Probing.

Probing classifiers are extensively used to elucidate the inner workings of
LLMs Alain and Bengio (2016); Azaria and Mitchell (2023); Jin et al. (2024);
Ju et al. (2024); Templeton et al. (2024). Various works train probes on model
representations to assess how well they encode various linguistic features,
such as phrase-level, syntactic, and semantic information Liu et al. (2023a);
Marks and Tegmark (2023); Li et al. (2024). The efficacy of a classifier in a
given task indicates the degree to which that layer successfully captures
pertinent information. In our study, we employ probing as a proxy to determine
whether the LLMs accurately identify and represent crucial parts of the
context.

##  3 Experimental Setup

We design a layer-wise probing task to examine if the model successfully
identifies the target information from the given prompt. We expect that higher
probing accuracy signifies a stronger connection between the model’s hidden
representations and its internal knowledge of the target information.

Figure 2: Accuracy of directly generating answers by LLMs (blue line) vs
maximum probing accuracy across layers by our probing classifiers (red line).
In both tasks, our probing classifiers outperform the model’s generated
answers, across all the gold positions. This indicates a discrepancy between
_knowing_ the context and _using_ it.

#### Datasets and prompts.

We follow the datasets and prompts used by Liu et al. (2023b). Our datasets
include: (1) Key-Value pairs retrieval (kv-pairs) where the context contains a
collection of keys and their corresponding values (128-bit randomly generated
UUIDs). The goal of this task is to identify a value given its key. Each
prompt for this task contains 100 kv-pairs and a target key. (2) Multi-
document question answering (MDQA) where the context contains multiple sets of
evidence paragraphs. The goal of this task is to, given a question, identify
the relevant document and produce an answer. Each prompt for this task
contains 30 documents, and a target question. Given a set of key-value
pairs/documents, with only one containing target information, LLM is prompted
to output the value/answer given the key/question. Further details on prompt
construction can be found in §A.

#### Probing classifiers.

For each input prompt, we collect the last token embedding of _each layer_. We
then train separate linear classifiers for each layer would receive the
embedding (of the last token) as input and the gold kv-pair/document ID
(position among all pairs/documents) as the target output. The classifier
minimizes the following objective:

$J(\theta)=-\frac{1}{N}\sum_{i=1}^{N}\sum_{j=1}^{C}y_{ij}\log\left(\frac{e^{\theta_{j}^{\top}x_{i}}}{\sum_{k=1}^{C}e^{\theta_{k}^{\top}x_{i}}}\right),$

where $N$ is the number of data points, $C$ is the number of different gold
IDs, $x_{i}$ represents the input embedding, $y_{ij}$ is the one-hot encoded
label for the $i$-th data point and $j$-th ID. Ultimately, this recipe gives
one probing classifier for embeddings of _each layer_. Using these models, we
show results per layer and across layers.

#### Models and hyperparameters.

We employ LLaMa3-8B-Instruct AI@Meta (2024) to our probing analysis. Results
for two other models are in shown §C, where we at the same conclusion. To
reduce uncertainty caused by random initialization, each classifier is trained
ten times. In the results, we report the mean and std (error bars) of
independent ten experiments.

Figure 3: The figures show the probing accuracy for each layer for the two
tasks: kv-pairs (left) and MDQA (right). Different colors indicate the
position of target information in the input context. In both tasks, mid-
context information requires more layers to be extracted.

#### Metrics.

Following Liu et al. (2023b), we use accuracy to quantify the success of our
models. Specifically, we quantify accuracy for two types of targets: (a)
Generation accuracy quantifies how well LLMs generate correct value in kv-pair
retrieval or generate correct answer string in MDQA. (b) Probing accuracy
quantifies how accurately classifiers can predict the gold kv-pair or document
ID, indicating whether the layers sufficiently encode information from the
input context.

##  4 LLMs Know but Don’t Tell

###  4.1 Experiment: maximum probing accuracy across all LLM layers

We focus on the peak accuracy across all transformer layers as a proxy to
determine if the model ever correctly identifies the useful information within
the prompt during the forward pass. Specifically, we select the probing
classifier with the highest accuracy across all layers. In Figure 2, we show
this peak layer probing accuracy. For comparison, we also show the accuracy of
LLMs in generating the answer (independent of our probing classifiers).

#### LLMs know but don’t tell.

Our results indicate that the model’s hidden representations indeed contain
information about the location of the target information. Specifically, in kv-
pairs setup (Fig. 2; left) there is always a layer such that its probe can
near-perfectly identify the location of the correct key-value pair associated
with the prompt. This is true, even for instances where the LLM does not
return the correct answer or abstains from producing any answer. This suggests
a disconnect between the model’s ability to locate the information and
generate a response based on that information.

A similar trend is also observed for MDQA (Fig. 2; right) where the peak
probing accuracy is consistently higher than the direct answer accuracy,
indicating the same disconnect from document grounding to response generation.
These findings highlight that while the model can recognize and encode the
location of relevant information within its layers, this knowledge does not
always translate into an accurate generation answer.

###  4.2 Experiment: probing across per layers

To understand the flow of information across LLMs’ layers, we shift our focus
to probing classifiers’ accuracy across LLM layers. Figure 3 visualizes
probing classifier accuracy per layer. For comparison, in both kv-pairs/MDQA
setups, we show this accuracy for three positions: target information at the
start, middle, or end of the input.

#### Mid-context information requires more layers to be located.

Our results reveal that LLM locates target information gradually at early
layers. Specifically, in the kv-pair setup (Fig. 3; left), probing accuracy
consistently increases until it reaches perfect accuracy at layer 13. Notably,
when the target kv-pair is at the middle position of the input prompt, LLM
requires more layers to locate the target information.

The general trends of the MDQA scenario (Figure 3; right) are similar in
principle, but with nuanced differences. The patterns vary significantly with
the position of target information. Classifiers perform best when the target
document is at the start of the input context, with near-perfect prediction
since early layers , and maintain it in subsequent layers. However, it takes
more layers for the probing classifier to achieve peak accuracy for the middle
and tail gold context. Interestingly, when the target document is in the
middle, classifier accuracy decreases after the peak. As MDQA task requires a
higher level of reasoning, the model is shifting from locating documents to
generating language output.

###  4.3 Experiment: the number of layers taken for locating target
information

Our probing experiments (§4.2) reveal that the model’s encoding of target
information position initially improves but then degrades as layer depth
increases. This motivates, the investigation of the relationship between the
number of layers taken by the model to locate target information from the
prompt and the LLM’s generation accuracy of generating the target information.

Figure 4:  The LLM layer that achieves the peak probing accuracy ($x$-axis)
vs. the accuracy of LLM in generating the correct answer ($y$-axis). We
observe that a _later_ peak correlates with _lower_ accuracy in the language
model’s final output. This implies that the earlier an LLM encodes information
from a specific index, the higher the accuracy of the final output for that
position.

We run additional 35, 40, 45, and 50 multi-document probing tasks. In Fig. 4,
for all IDs that achieve probing accuracy greater than 60% (to suppress the
outliers), on the $x$-axis we show the layer at which probing achieves the
peak accuracy. On the $y$-axis we show the the LLM’s generation accuracy (no
probes involved).

#### Early-layer information localization leads to higher accuracy in LLM
output.

As Fig. 4 shows, there is a statistically significant (two-sided t-test)
negative correlation between the layer with peak accuracy of locating target
information and its final output accuracy ($p<5e-5$). This negative
correlation implies that the earlier the model identifies the target document
within its layers, the more likely it is to generate an accurate final answer.

##  5 Conclusion

Our study investigates LLMs positional bias, indicating that LLM could capture
context information, but do not tell the correct answer. The experiment
results demonstrate that the input context is embedded in the model’s hidden
representation, but such information is not decoded into anticipated output.

##  6 Limitation

Knowledge of the gold document’s location and the ability to cite from it are
distinct but connected; the model might know the location but still, fail to
integrate it into a coherent and accurate answer. This comparison does not
fully capture the nuanced interactions between the model’s internal attention
mechanisms and output generation capabilities. While these limitations are
acknowledged, they do not detract from the core contributions of our work. Our
findings provide valuable insights into the positional effects on model
performance and highlight the importance of document sequence in information
retrieval tasks. By identifying specific areas where the model struggles, we
lay the groundwork for future improvements and optimizations in model design
and training.

##  7 Ethical Considerations

Currently the consequences of misinterpretations or errors in long context
processing can be significant, in fields like healthcare, legal, and public
services. In other cases, LLM long-context failure results in harmful, biased,
and misleading generations Anil et al. (2024). Our research considers the
potential negative impacts of these errors and actively works to uncover the
mechanisms that could minimize such risk.

## References

  * AI@Meta (2024) AI@Meta. 2024.  Llama 3 model card.
  * Alain and Bengio (2016) Guillaume Alain and Yoshua Bengio. 2016.  Understanding intermediate layers using linear classifier probes.  _arXiv preprint arXiv:1610.01644_.
  * Anil et al. (2024) Cem Anil, Esin Durmus, Mrinank Sharma, Joe Benton, Sandipan Kundu, Joshua Batson, Nina Rimsky, Meg Tong, Jesse Mu, Daniel Ford, et al. 2024.  Many-shot jailbreaking.  _Anthropic, April_.
  * Azaria and Mitchell (2023) Amos Azaria and Tom Mitchell. 2023.  The internal state of an llm knows when its lying.  _arXiv preprint arXiv:2304.13734_.
  * Chen et al. (2023) Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. 2023.  Extending context window of large language models via positional interpolation.  _Preprint_ , arXiv:2306.15595.
  * Ding et al. (2024) Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. 2024.  Longrope: Extending llm context window beyond 2 million tokens.  _Preprint_ , arXiv:2402.13753.
  * Jiang et al. (2023) Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. 2023.  Mistral 7b.  _Preprint_ , arXiv:2310.06825.
  * Jin et al. (2024) Mingyu Jin, Qinkai Yu, Jingyuan Huang, Qingcheng Zeng, Zhenting Wang, Wenyue Hua, Haiyan Zhao, Kai Mei, Yanda Meng, Kaize Ding, et al. 2024.  Exploring concept depth: How large language models acquire knowledge at different layers?  _arXiv preprint arXiv:2404.07066_.
  * Ju et al. (2024) Tianjie Ju, Weiwei Sun, Wei Du, Xinwei Yuan, Zhaochun Ren, and Gongshen Liu. 2024.  How large language models encode context knowledge? a layer-wise probing study.  _arXiv preprint arXiv:2402.16061_.
  * Li et al. (2024) Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. 2024.  Inference-time intervention: Eliciting truthful answers from a language model.  _Advances in Neural Information Processing Systems_ , 36.
  * Li et al. (2023) Zongjie Li, Chaozheng Wang, Pingchuan Ma, Daoyuan Wu, Shuai Wang, Cuiyun Gao, and Yang Liu. 2023.  Split and merge: Aligning position biases in large language model based evaluators.  _Preprint_ , arXiv:2310.01432.
  * Liu et al. (2023a) Kevin Liu, Stephen Casper, Dylan Hadfield-Menell, and Jacob Andreas. 2023a.  Cognitive dissonance: Why do language model outputs disagree with internal representations of truthfulness?  _arXiv preprint arXiv:2312.03729_.
  * Liu et al. (2023b) Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. 2023b.  Lost in the middle: How language models use long contexts.  _Preprint_ , arXiv:2307.03172.
  * Marks and Tegmark (2023) Samuel Marks and Max Tegmark. 2023.  The geometry of truth: Emergent linear structure in large language model representations of true/false datasets.  _arXiv preprint arXiv:2310.06824_.
  * Peysakhovich and Lerer (2023) Alexander Peysakhovich and Adam Lerer. 2023.  Attention sorting combats recency bias in long context language models.  _Preprint_ , arXiv:2310.01427.
  * Shaham et al. (2023) Uri Shaham, Maor Ivgi, Avia Efrat, Jonathan Berant, and Omer Levy. 2023.  ZeroSCROLLS: A zero-shot benchmark for long text understanding.  In _Findings of the Association for Computational Linguistics: EMNLP 2023_ , pages 7977–7989, Singapore. Association for Computational Linguistics.
  * Shaham et al. (2022) Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. 2022.  SCROLLS: Standardized CompaRison over long language sequences.  In _Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing_ , pages 12007–12021, Abu Dhabi, United Arab Emirates. Association for Computational Linguistics.
  * Tang et al. (2024) Raphael Tang, Xinyu Zhang, Xueguang Ma, Jimmy Lin, and Ferhan Ture. 2024.  Found in the middle: Permutation self-consistency improves listwise ranking in large language models.  _Preprint_ , arXiv:2310.07712.
  * Team et al. (2024) Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, Pouya Tafti, Léonard Hussenot, Pier Giuseppe Sessa, Aakanksha Chowdhery, Adam Roberts, Aditya Barua, Alex Botev, Alex Castro-Ros, Ambrose Slone, Amélie Héliou, Andrea Tacchetti, Anna Bulanova, Antonia Paterson, Beth Tsai, Bobak Shahriari, Charline Le Lan, Christopher A. Choquette-Choo, Clément Crepy, Daniel Cer, Daphne Ippolito, David Reid, Elena Buchatskaya, Eric Ni, Eric Noland, Geng Yan, George Tucker, George-Christian Muraru, Grigory Rozhdestvenskiy, Henryk Michalewski, Ian Tenney, Ivan Grishchenko, Jacob Austin, James Keeling, Jane Labanowski, Jean-Baptiste Lespiau, Jeff Stanway, Jenny Brennan, Jeremy Chen, Johan Ferret, Justin Chiu, Justin Mao-Jones, Katherine Lee, Kathy Yu, Katie Millican, Lars Lowe Sjoesund, Lisa Lee, Lucas Dixon, Machel Reid, Maciej Mikuła, Mateo Wirth, Michael Sharman, Nikolai Chinaev, Nithum Thain, Olivier Bachem, Oscar Chang, Oscar Wahltinez, Paige Bailey, Paul Michel, Petko Yotov, Rahma Chaabouni, Ramona Comanescu, Reena Jana, Rohan Anil, Ross McIlroy, Ruibo Liu, Ryan Mullins, Samuel L Smith, Sebastian Borgeaud, Sertan Girgin, Sholto Douglas, Shree Pandya, Siamak Shakeri, Soham De, Ted Klimenko, Tom Hennigan, Vlad Feinberg, Wojciech Stokowiec, Yu hui Chen, Zafarali Ahmed, Zhitao Gong, Tris Warkentin, Ludovic Peran, Minh Giang, Clément Farabet, Oriol Vinyals, Jeff Dean, Koray Kavukcuoglu, Demis Hassabis, Zoubin Ghahramani, Douglas Eck, Joelle Barral, Fernando Pereira, Eli Collins, Armand Joulin, Noah Fiedel, Evan Senter, Alek Andreev, and Kathleen Kenealy. 2024.  Gemma: Open models based on gemini research and technology.  _Preprint_ , arXiv:2403.08295.
  * Templeton et al. (2024) Adly Templeton, Tom Conerly, Jonathan Marcus, Jack Lindsey, and Trenton Bricke. 2024.  Scaling monosemanticity: Extracting interpretable features from claude 3 sonnet.  https://transformer-circuits.pub/2024/scaling-monosemanticity/index.html.
  * Wang et al. (2024) Cunxiang Wang, Ruoxi Ning, Boqi Pan, Tonghui Wu, Qipeng Guo, Cheng Deng, Guangsheng Bao, Qian Wang, and Yue Zhang. 2024.  Novelqa: A benchmark for long-range novel question answering.  _Preprint_ , arXiv:2403.12766.
  * Wang et al. (2023) Yiwei Wang, Yujun Cai, Muhao Chen, Yuxuan Liang, and Bryan Hooi. 2023.  Primacy effect of chatgpt.  _arXiv preprint arXiv:2310.13206_.
  * Zhang et al. (2024a) Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang, Zhiyuan Liu, and Maosong Sun. 2024a.  $\infty$bench: Extending long context evaluation beyond 100k tokens.  _Preprint_ , arXiv:2402.13718.
  * Zhang et al. (2024b) Zhenyu Zhang, Runjin Chen, Shiwei Liu, Zhewei Yao, Olatunji Ruwase, Beidi Chen, Xiaoxia Wu, and Zhangyang Wang. 2024b.  Found in the middle: How language models use long contexts better via plug-and-play positional encoding.  _arXiv preprint arXiv:2403.04797_.
  * Zhao et al. (2021) Tony Z. Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. 2021.  Calibrate before use: Improving few-shot performance of language models.  _Preprint_ , arXiv:2102.09690.

##  Appendix A Prompting Details

Following setup by Liu et al. (2023b), we construct key-value pairs retrieval
and multi-document question answering prompting dataset.

#### Key-Value pairs retrieval (kv-pairs)

We generate $n$ pairs of 128-bit randomly generated UUID.

Example Key-Value pair "7f666c61-573f-4212-a0a9-6f90d487cd4a" :
"2a1d0ba0-cfe4-4df5-987a-6ee1be2c6ac0"

The $n$ kv-pairs are composed into one single JSON object. To test at ID $k$,
we choose one pair as gold, insert it at ID $k$, and then construct as a
prompt in the format:

Extract the value corresponding to the specified key in the JSON object below.
JSON data: { "key1: "value1", “key2": "value2", … "keyk": "valuek", … "keyn":
"valuen", } Key: "keyk" Corresponding value:

#### Multi-document question answering (MDQA)

In the $n$ document setting, we randomly select one question answer pair from
the dataset by Liu et al. (2023b). Subsequently we retrieve the document
containing this answer and mark it as gold.

Example retrieval Question: who got the first nobel prize in physics Answer:
Wilhelm Conrad Röntgen Document: (Title: List of Nobel laureates in Physics)
The first Nobel Prize in Physics was awarded in 1901 to Wilhelm Conrad
Röntgen, of Germany, who received…

We then sample $n-1$ distractors, relevant documents that do not contain the
answer. To test at ID $k$, we randomly shuffle the distractors and then insert
the gold document at ID $k$. Example prompt with gold document at ID $k$ is
like:

Write a high-quality answer for the given question using only the provided
search results (some of which might be irrelevant). Document [1](Title: Asian
Americans in science and technology) Prize in physics for discovery of the
subatomic… … Document [$k$](Title: List of Nobel laureates in Physics) The
first Nobel Prize in Physics was awarded in 1901… … Document [$n$] (Title:
Scientist) and pursued through a unique method, was essentially in place.
Ramón y Cajal won … Question: who got the first nobel prize in physics Answer:

##  Appendix B Probing Setup

In the experiment described in §3, we employ linear classifiers as our probing
method.

For any given task, we choose $\\{1,0.1n,0.2n,\dots,1.0n\\}$-th position as
gold ID. Following the prompt format in §A, we generate prompts with all
chosen IDs, for $10,000$ iterations, resulting in a set of 110,000 prompts.

Each prompt is fed into language model, and the embedding from each layer’s
last token is collected. For each layer, separately, we have $110,000$
embeddings corresponding to 11 IDs and train a classifier for ten times, with
embedding as input and ID as output. We calculate their mean accuracy and
standard deviation.

##  Appendix C Experiments Results on Mistral-7B-Instruct-v0.3 Jiang et al.
(2023) and Gemma-7b-it Team et al. (2024)

We conduct same experiment procedure on additional two models, which produce
the same pattern. The experiment is running on one A100 GPU.

Figure 5: In the left figure, Mistral’s probing result in 100 kv-pairs
retrieval task resembles with Llama3’s in most ways, but there’re still a few
differences. First, the trends of locating information from the head and end
are more similar compared to those in Llama3. Also the model takes more layers
than Llama3 to well locate information in middle context. The right figure
highlights the significant discrepancy between the probing peak accuracy and
generation accuracy, indicating a severe ’know don’t tell’ phenomenon. Figure
6: Mistral’s layer-wise probing classifier accuracy performs in the same
pattern as Llama3. Middle context takes more layers to encode and results in a
lower peak accuracy. In right figure, both its peak accuracy and generation
accuracy shows a U-shape curve, with still probing consistently outperforms
generation. Figure 7: In the left figure, Gemma shows a similar pattern like
other models in 100 kv-pairs retrieval task. However, information from all
positions is located more slowly than other models: Information in the head or
end position of input context is located after 5 layers, while information in
the middle of input takes 15 layers to be located. In the right figure, it
shows a significant gap between generation accuracy and probing peak accuracy,
which is similar with Mistral. This highlights a significant ’know don’t tell’
phenomenon. Figure 8: For Gemma MDQA task, although head and end context
follows same observation, it middle context shows a sudden decrease in
accuracy, indicating a sudden information lost but soon retrieve it back. The
right figure follows the same pattern, where generation accuracy is
consistently low, disconnecting from the high U-shape probing accuracy curve.