# How Transferable are the Representations Learned by Deep Q Agents?

Jacob Tyo & Zachary Lipton
Machine Learning Department
Carnegie Mellon University
Pittsburgh, PA 15213, USA
<EMAIL_ADDRESS><EMAIL_ADDRESS>

###### Abstract

In this paper, we consider the source of Deep Reinforcement Learning (DRL)’s
sample complexity, asking how much derives from the requirement of learning
useful representations of environment states and how much is due to the sample
complexity of learning a policy. While for DRL agents, the distinction between
representation and policy may not be clear, we seek new insight through a set
of transfer learning experiments. In each experiment, we retain some fraction
of layers trained on either the same game or a related game, comparing the
benefits of transfer learning to learning a policy from scratch.
Interestingly, we find that benefits due to transfer are highly variable in
general and non-symmetric across pairs of tasks. Our experiments suggest that
perhaps transfer from simpler environments can boost performance on more
complex downstream tasks and that the requirements of learning a useful
representation can range from negligible to the majority of the sample
complexity, based on the environment. Furthermore, we find that fine-tuning
generally outperforms training with the transferred layers frozen, confirming
an insight first noted in the classification setting.

##  1 Introduction

Deep Reinforcement Learning (DRL) agents learn policies by selecting actions
directly from raw perceptual data. Despite numerous breakthroughs (Mnih et
al., 2015; Vinyals et al., 2019; Silver et al., 2016), DRL’s prohibitive
sample complexity limits its real world application. The sample complexity of
DRL may derive from many sources, including the requirement of learning useful
state representations and the requirement of learning good policies given a
suitable representation. This paper provides several simple experiments to
provide new insights into this breakdown.

For a DRL agent, where precisely the “representation” ends and the “policy”
begins is not clear. In our experiments, we consider multiple interpretations
of this divide by partitioning the network at various layers. To evaluate the
extent to which representation learning contributes to the sample complexity
of DRL, we execute a series of transfer learning experiments aimed to
determine how quickly an agent can learn given pre-learned representations
(from either the same or a different game). Our experiments proceed in the
following manner: 1) Train a parent network until best reported performance is
achieved. 2) Transplant the first $k$ layers into a child network, re-
initializing the remaining $l-k$ layers randomly. 3) Train the child network,
either fine-tuning the transplanted layers or keeping them frozen, following
the methodology of Yosinski et al. (2014).

In particular, our experiments address the transferability of DQN
representations among the three Atari games _Berzerk_ , _Krull_ , and _River
Raid_. We find the benefits of transferring representations from pretrained
networks to be remarkably variable in general and non-symmetric across pairs
of tasks. While we have not evaluated enough tasks to draw definitive
conclusions, our preliminary results suggest that transfer from simple
environments may improve performance on more complex tasks more than the
reverse transfer. Furthermore, our results suggest that the contribution of
learning a useful representation to the overall sample complexity of the
problem can range from negligible to the majority, based on the destination
environment. Lastly, fine-tuning the transferred layers outperforms training
with those layers frozen in general.

##  2 Background and Experimental Setup

Reinforcement Learning (RL) methods address the problem where an _agent_ must
learn to act in an unknown _environment_ to maximize a _reward_ signal.
Initially, the environment provides a state $s_{0}$ to the agent, and then the
agent selects an action $a_{0}$ based on the provided state. Thereafter, at
each time step the environment provides an agent with a state $s_{t}$ and a
reward $r_{t}$ each influenced by the agents previous action. The agent must
then select subsequent actions $a_{t}$, and so on, until the episode
terminates. Formally, this interaction is described by a Markov Decision
Process (MDP)
${\mathcal{M}}=\langle{\mathcal{S}},{\mathcal{A}},{\mathcal{T}},r,\gamma\rangle$,
where ${\mathcal{S}}$ is the set of states, ${\mathcal{A}}$ is the set of
actions,
${\mathcal{T}}(s,a,s^{\prime})={\mathbb{P}}(s_{t+1}=s^{\prime}|s_{t}=s,a_{t}=a)$
is the transition function, $r(s,a)=\mathbb{E}[r_{t+1}|s_{t}=s,a_{t}=a]$ is
the reward function, and $\gamma\in[0,1]$ is a discount factor.

We focus on Q-learning, an off-policy value-based RL algorithm. The value of
each state is represented by
$v^{\pi}(s)=\mathbb{E}_{\pi}[G_{t}(s_{t})|s_{t}=s]$, and the value of each
state-action pair is represented by
$q^{\pi}(s,a)=\mathbb{E}_{\pi}[G_{t}(s_{t})|s_{t}=s,a_{t}=a]$ where
$G_{t}(s)=\sum_{i=t}^{T}\gamma^{i-t}r_{t}(s)$. In Q-learning, the agent
estimates the state-action value function by predicting the expected
discounted return.

Many interesting problems, including Atari games, have a large state and
action space, making tabular estimates of the Q-function intractable. In these
cases, the Q-function can be approximated. DRL denotes methods that
approximate either the value function (or, in other algorithms, the policy
directly) by deep neural networks.

In this paper, we build on a DRL Q-learning implementation called _Rainbow_
(Hessel et al., 2018): a 5-layer convolutional neural network based on Mnih et
al. (2015) that incorporates double Q-learning (Van Hasselt et al., 2016),
prioritized replay (Schaul et al., 2015), dueling networks (Wang et al.,
2015), multi-step learning (Sutton, 1988), distributional RL (Bellemare et
al., 2017), and Noisy Nets (Fortunato et al., 2017). Furthermore, Rainbow
maintains two separate Q-networks: one with parameters $\theta$, and a second
with parameters $\theta_{\text{target}}$ that are updated from $\theta$ every
fixed number of iterations. In order to capture the game dynamics, a state is
represented by a sequence (four in our case) of history frames.

We tested the transferability of features learned by Rainbow agents on Atari
games (i.e. environments). To make transfer learning experiments feasible
within our resources, we selected environments according to the speed that a
Rainbow agents can reach high performance, requiring that the cardinality of
the state and action spaces of each environment be equivalent and that two
among the three environments were qualitatively similar (same genre of game).
We are interested in game similarity to test a hypothesis that more similar
games have more similar representations, and therefore agent transfer should
be more effective between them. Environments with the same state and action
space cardinality is required as it made agent transfer possible without
modification. Pulling from the results of Hessel et al. (2018), we selected
Berzerk, Krull, and River Raid from the Arcade Learning Environment (Bellemare
et al., 2013).

Drawing inspiration from Yosinski et al. (2014), our experiments proceeded as
follows:

  1. 1.

For all environments, train a parent network (a Rainbow agent) until best-
reported performance is reached.

  2. 2.

For every permutation of environment pairs, transplant the first $k$ layers of
the parent network into a child network (also a Rainbow agent), then
reinitialize the remaining $l-k$ layers randomly (where $l$ is the total
number of layers in the network and $k\in\\{2,4\\}$).

  3. 3.

Train the child network, either fine-tuning or freezing the transplanted
portion of the network (we explore both settings).

With $3$ runs per pair of environments, setting of $k$, and each choice among
{freezing, fine-tuning}, we ran a total of $111$ trials taking over $35$ days
on the available resources.

##  3 Experimental Results and Discussion

In this section, we analyze the performance and learned representations of the
child agents, where one child exists for every environment pair, $k$-value,
and choice among freezing/fine-tuning. For brevity, we will denote child
agents with 4 layers transferred from a parent trained on environment1 then
frozen as child4-frozen-environment1. We will refer to the parents as
baselines with respect to their performance on the environment that they were
trained.

We evaluate each network separated at two places, transferring $2$, and
transferring $4$ (out of $5$) layers and compare freezing vs fine-tuning over
the transplanted layers. For each respective experiment, we refer to the
output of the last layer of the transplanted portion of the network as the
“representation.”

Figure 1: Child agent (blue, red, and green) performance over iterations of
training when the transplanted layers are frozen. Compare against the parent
network trained from scratch (black).

Figures 1a, 1b, and 1c demonstrate the performance of all environment pairs
when $4$ layers are transplanted and then frozen during subsequent training.
Because only the output layer is not frozen, these plots show the performance
of a linear policy. trained on the transplanted representation. Levine et al.
(2017) deem representations that can be used with a linear policy as a “good”
representation, thus our analysis starts with these corresponding experiments.

Figure 1a presents the performance of the agents on Berzerk for runs with $4$
transplanted layers, all of which are frozen during the subsequent training.
The child4-frozen-Berzerk agent (blue) does not out perform the baseline
(either by final performance or training speed), which we find surprising,
since four of five layers of this agent were transferred from a high-
performing parent trained on the same environment, leaving only a linear
policy to be relearned. We might deduce that for this game, the entire
difficulty of DRL can be attributed simply to the “RL” since starting off with
the representations upon which a strong linear policy can be learnt seems to
confer no benefit. Notably, the agents transferred from foreign games
(child4-frozen-Krull and child4-frozen-RiverRaid) have worse final performance
than the baseline. As a linear policy could not be learnt from a known good
representation, we would not expect any other representation to perform
better. Figure 1d shows these trends continuing when $2$ layers are
transplanted. The difference in performance between the children and the
baseline is negligible, and therefore indicating that the difficulty does lie
within the RL problem. Even with more model complexity allowed to the policy
and different representation, no benefits are seen.

However, as shown in Figure 1b, the transfer from Krull-to-Krull results in
faster training and higher final performance than the baseline, perhaps
suggesting that for this game, representation learning is a more significant
part of the challenge. Note that with only two layers transferred, the
representations pretrained on both Krull and Berzerk, outperform the baseline
as shown in Figure 1e. This reflects our intuition that these games are
similar, while the representations pretrained on River Raid confer no benefit.

Figure 2: Child agent (blue, red, and green) performance over iterations of
training when the transplanted layers are fine-tuned. Compare against the
parent network trained from scratch (black).

Transferring and freezing $4$ of $5$ layers from the parent to child agent on
Berzerk and Krull showed negligible and positive changes in performance,
respectively. However, River Raid paints a different picture. As shown in
Figure 1c, no agent is able to learn a useful policy. This is especially
surprising because we know that if nothing else, the transfer from the parent
trained on River Raid itself can replicate the baseline results by learning a
linear policy. Similarly, we see in Figure 1f that the agent with transplanted
layers trained on the _same_ environment performs the worst. Again this is
surprising as, theoretically, the baseline could be replicated. Interestingly,
this may indicate that a representation learned from a more simple environment
may be helpful in solving more complex environments. The converse does not
appear to hold from our experiments, but a deeper analysis on more tasks would
be required to confirm these intuitions.

In general, we see that when transplanted layers are frozen, there is a trade-
off between negative transfer and training speed. As the number of unfrozen
layers increases, the effect of negative transfer decreases but the training
speed slows slightly. Furthermore, Figure 1 shows that transfer between
environments is not symmetric. In other words, if a representation $x$ is
learned on environment $X$ and is shown to perform well on environment $Y$,
then this _does not_ imply that a representation $y$ learned on environment
$Y$ will be effective on $X$. This relation continues to hold, even when fine-
tuning over the transplanted layers. However, when fine-tuning over the
transplanted layers we do not see the same level of negative transfer.

Figures 2a, 2b, and 2c show the performance of all environment pairs when four
layers are transplanted from the parents to the children and then frozen,
whereas Figures 2d, 2e, and 2f show the performance when only two layers are
transplanted and frozen. Intuitively, these experiments which allow the
transplanted representations to change during training, examine a more
flexible notion of their utility (as initializations only). For these
experiments, in all environments except for River Raid, negative transfer is
no longer an issue. In River Raid, negative transfer is seen in $2$ of $6$
children, vs in $6$ of $6$ children when the transferred layers are frozen.

The performance of all agents on Berzerk with $4$ and $2$ layers transplanted
and fine-tuned respectively, shown in Figures 2a and 2d, are similar to the
performance seen for the same game under frozen transplanted layers. There is
little difference between each of the agents, regardless of the environment on
which the parent was trained. This further supports the earlier conclusions
that the difficulty of this game is due entirely to learning the policy as no
transferred representation improves performance.

Transplanting and fine-tuning layers results in a large decrease in training
time, and a large increase in final performance on the Krull environment. This
holds for both transplanting and fine-tuning over $4$ layers as shown in
Figure 2b, and over $2$ layers as shown in Figure 2e. As expected, the
transplanted layers from the parent trained on Krull performed the best when
$4$ layers were transplanted. But interestingly when only $2$ layers were
transplanted, the transplanted layers from the parent trained on Berzerk were
more effective. Overall, this supports our earlier conclusion of the
importance of representations for Krull—in all cases, transfer is preferred to
random initialization.

In Figures 1 and 2, we see the general trend that the transfer of pretrained
layers has negligible effect on Berzerk, positive effects on Krull, and
negative effects on River Raid. Figures 2c and 2f show the performance of all
agents with $4$ and $2$ layers transplanted and fine-tune respectively.
Interestingly, the trend observed in Krull, where the best performing child
was the one with features transferred from a parent trained on the same
environment is the best performer when $4$ layers are transferred, yet it is
not the best performer when only $2$ layers are transferred, is also observed
with River Raid. However, the baseline trains faster and reaches better final
performance than all of the children on River Raid.

##  4 Conclusions

This paper presents an empirical evaluation of $111$ transfer learning
experiments on agents trained on the Atari 2600 games Berzerk, Krull, and
River Raid. We compare the effect of transplanting initial layers of a
pretrained network into a child network while either freezing or fine-tuning
over the transplanted layers. Surprisingly, the benefits of transferring
portions of pretrained networks are highly variable and non-symmetric across
tasks. Furthermore, the requirements of learning a useful representation can
range from nothing to the majority of the sample complexity based on the
destination environment. We present analyses for why each task transfer occurs
as shown, and give intuition for understanding representations and policies in
DQNs. We show that, in general, fine tuning is better than freezing portions
of networks, as performance gains can still be expected with less likelihood
of negative transfer.

Zahavy et al. (2016) have shown that DQNs find hierarchical abstractions
automatically. Our work suggests that the similarity of the high level task
abstraction may be a good metric to determine the transferabitliy of DQN
agents on. Future work includes heavier analysis on these experiments to
determine how agents pretrained on different environments “focus” differently,
how their representations differ, and how to numerically quantify the
contribution of representation learning to the overall sample complexity.
Furthermore, this methodology can pass insight to questions about the benefit
of unsupervised reinforcement learning in pre-training, e.g. techniques based
on intrinsic motivation (Chentanez et al., 2005). To the extent that intrinsic
motivation serves to learn representations suitable for fine-tuning to a given
reward signal, quantifying just how much representation learning is the
bottleneck to learning in the first place can provide insight in assessing its
potential.

## References

  * Asadi & Huber (2007) Mehran Asadi and Manfred Huber.  Effective control knowledge transfer through learning skill and representation hierarchies.  In _IJCAI_ , volume 7, pp. 2054–2059, 2007.
  * Bellemare et al. (2013) Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling.  The arcade learning environment: An evaluation platform for general agents.  _Journal of Artificial Intelligence Research_ , 47:253–279, 2013.
  * Bellemare et al. (2017) Marc G Bellemare, Will Dabney, and Rémi Munos.  A distributional perspective on reinforcement learning.  In _Proceedings of the 34th International Conference on Machine Learning-Volume 70_ , pp. 449–458. JMLR. org, 2017.
  * Calandriello et al. (2014) Daniele Calandriello, Alessandro Lazaric, and Marcello Restelli.  Sparse multi-task reinforcement learning.  In _Advances in Neural Information Processing Systems_ , pp. 819–827, 2014.
  * Chentanez et al. (2005) Nuttapong Chentanez, Andrew G Barto, and Satinder P Singh.  Intrinsically motivated reinforcement learning.  In _Advances in neural information processing systems_ , pp. 1281–1288, 2005.
  * Fortunato et al. (2017) Meire Fortunato, Mohammad Gheshlaghi Azar, Bilal Piot, Jacob Menick, Ian Osband, Alex Graves, Vlad Mnih, Remi Munos, Demis Hassabis, Olivier Pietquin, et al.  Noisy networks for exploration.  _arXiv preprint arXiv:1706.10295_ , 2017.
  * Foster & Dayan (2002) David Foster and Peter Dayan.  Structure in the space of value functions.  _Machine Learning_ , 49(2-3):325–346, 2002.
  * Hessel et al. (2018) Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver.  Rainbow: Combining improvements in deep reinforcement learning.  In _Thirty-Second AAAI Conference on Artificial Intelligence_ , 2018\.
  * Lazaric (2008) Alessandro Lazaric.  _Knowledge transfer in reinforcement learning_.  PhD thesis, PhD thesis, Politecnico di Milano, 2008.
  * Levine et al. (2017) Nir Levine, Tom Zahavy, Daniel J Mankowitz, Aviv Tamar, and Shie Mannor.  Shallow updates for deep reinforcement learning.  In _Advances in Neural Information Processing Systems_ , pp. 3135–3145, 2017.
  * Mnih et al. (2015) Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al.  Human-level control through deep reinforcement learning.  _Nature_ , 518(7540):529, 2015.
  * Rusu et al. (2016) Andrei A Rusu, Neil C Rabinowitz, Guillaume Desjardins, Hubert Soyer, James Kirkpatrick, Koray Kavukcuoglu, Razvan Pascanu, and Raia Hadsell.  Progressive neural networks.  _arXiv preprint arXiv:1606.04671_ , 2016.
  * Schaul et al. (2015) Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver.  Prioritized experience replay.  _arXiv preprint arXiv:1511.05952_ , 2015.
  * Silver et al. (2016) David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al.  Mastering the game of go with deep neural networks and tree search.  _nature_ , 529(7587):484, 2016.
  * Singh (1992) Satinder Pal Singh.  Transfer of learning by composing solutions of elemental sequential tasks.  _Machine Learning_ , 8(3-4):323–339, 1992.
  * Sutton (1988) Richard S Sutton.  Learning to predict by the methods of temporal differences.  _Machine learning_ , 3(1):9–44, 1988.
  * Taylor & Stone (2009) Matthew E Taylor and Peter Stone.  Transfer learning for reinforcement learning domains: A survey.  _Journal of Machine Learning Research_ , 10(Jul):1633–1685, 2009.
  * Van Hasselt et al. (2016) Hado Van Hasselt, Arthur Guez, and David Silver.  Deep reinforcement learning with double q-learning.  In _Thirtieth AAAI Conference on Artificial Intelligence_ , 2016.
  * Vinyals et al. (2019) Oriol Vinyals, Igor Babuschkin, Junyoung Chung, Michael Mathieu, Max Jaderberg, Wojciech M. Czarnecki, Andrew Dudzik, Aja Huang, Petko Georgiev, Richard Powell, Timo Ewalds, Dan Horgan, Manuel Kroiss, Ivo Danihelka, John Agapiou, Junhyuk Oh, Valentin Dalibard, David Choi, Laurent Sifre, Yury Sulsky, Sasha Vezhnevets, James Molloy, Trevor Cai, David Budden, Tom Paine, Caglar Gulcehre, Ziyu Wang, Tobias Pfaff, Toby Pohlen, Yuhuai Wu, Dani Yogatama, Julia Cohen, Katrina McKinney, Oliver Smith, Tom Schaul, Timothy Lillicrap, Chris Apps, Koray Kavukcuoglu, Demis Hassabis, and David Silver.  AlphaStar: Mastering the Real-Time Strategy Game StarCraft II.  https://deepmind.com/blog/alphastar-mastering-real-time-strategy-game-starcraft-ii/, 2019\.
  * Walsh et al. (2006) Thomas J Walsh, Lihong Li, and Michael L Littman.  Transferring state abstractions between mdps.  In _ICML Workshop on Structural Knowledge Transfer for Machine Learning_ , 2006.
  * Wang et al. (2015) Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Van Hasselt, Marc Lanctot, and Nando De Freitas.  Dueling network architectures for deep reinforcement learning.  _arXiv preprint arXiv:1511.06581_ , 2015.
  * Yosinski et al. (2014) Jason Yosinski, Jeff Clune, Yoshua Bengio, and Hod Lipson.  How transferable are features in deep neural networks?  In _Advances in neural information processing systems_ , pp. 3320–3328, 2014.
  * Zahavy et al. (2016) Tom Zahavy, Nir Ben-Zrihem, and Shie Mannor.  Graying the black box: Understanding dqns.  In _International Conference on Machine Learning_ , pp. 1899–1908, 2016.

##  Appendix A Related Work

Hessel et al. (2018) present Rainbow, the DQN used in this paper,
incorporating Double Q-learning (Van Hasselt et al., 2016), prioritized replay
(Schaul et al., 2015), dueling networks (Wang et al., 2015), multi-step
learning (Sutton, 1988), distributional reinforcement learning (Bellemare et
al., 2017), and noisy networks Fortunato et al. (2017).

Zahavy et al. (2016) analyze Deep Q Networks (DQNs) by observing the
activation’s of the model’s last layer and saliency maps. They show that DQNs
learn temporal abstractions, such as hierarchical state aggregation and
options, automatically. Levine et al. (2017) show that the last layer in a
deep architecture can be seen as a linear representation, and thus can be
learned using standard shallow reinforcement learning algorithms. They then
show that this hybrid approach improves performance on the Atari benchmark.

Yosinski et al. (2014) present a large scale study of feature transferabitlity
in deep neural networks. They show that transferability is negatively affected
by the specialization of higher layer neurons to their original task.
Furthermore, optimization difficulties can arise when co-adapted neurons are
split during transfer, and freezing vs fine-tuning over transferred layers are
compared. The authors show that transferring features, even if from a very
different task, can improve generalization performance even after substantial
fine-tuning on a new task. Lastly, a relation between the effectiveness of
transfer and the distance between tasks is presented, but even in the worst
case is shown to be better than random.

Taylor & Stone (2009) provides a survey of transfer learning techniques in
reinforcement learning. Here we will focus on the tasks that allow variation
in the reward function, as we assume the state spaces are of the same
cardinality, and the action spaces are equivalent in this paper. Singh (1992)
and Foster & Dayan (2002) learn multiple tasks by assuming that each goal (or
composite) task is composed of several elemental tasks, and then learning a
set of elemental tasks that can be composed to solve each task of interest.

Solving multiple MDPs has also been approached from the representation
perspective, specifically with the goal of developing a shared representation
that can then be used to solve all tasks. The approach proposed by Asadi &
Huber (2007) focuses on learning a more efficient state-space representation
of the problem that will transfer between multiple tasks, and then learning
options on the new representation. Walsh et al. (2006) use a similar approach,
but rely on the learned state abstraction techniques to transfer between
tasks. Another approach similar to state abstractions is to compare
observations ($\langle s,a,r,s^{\prime}\rangle)$ tuples from previous tasks to
new tasks, then select the best action from the most similar previously
experienced observation. Lazaric (2008) uses this approach in an attempt to
generalize experiences from learned to novel tasks, and then Calandriello et
al. (2014) extend this approach to include sparse representations.

Rusu et al. (2016) introduce progressive neural networks, which is a novel
model architecture which retains a pool of pretrained models throughout
training, and learns lateral connections from these to extract useful features
for new tasks. This architecture leverages transfer learning while avoiding
catastrophic forgetting and allowing for better incorporation of prior
knowledge vs the traditional method of initialization.


# Enhancing Embedding Performance through Large Language Model-based Text
Enrichment and Rewriting

Nicholas Harris  Arizona State University
Tempe, Arizona
<EMAIL_ADDRESS>Anand Butani  MyAutoBio Inc.
Scottsdale, Arizona
<EMAIL_ADDRESS>Syed Hashmy  Arizona State University
Tempe, Arizona
<EMAIL_ADDRESS>
###### Abstract

Embedding models are crucial for various natural language processing tasks but
can be limited by factors such as limited vocabulary, lack of context, and
grammatical errors. This paper proposes a novel approach to improve embedding
performance by leveraging large language models (LLMs) to enrich and rewrite
input text before the embedding process. By utilizing ChatGPT $3.5$ to provide
additional context, correct inaccuracies, and incorporate metadata, the
proposed method aims to enhance the utility and accuracy of embedding models.
The effectiveness of this approach is evaluated on three datasets:
Banking77Classification, TwitterSemEval $2015$, and Amazon Counter-factual
Classification. Results demonstrate significant improvements over the baseline
model on the TwitterSemEval $2015$ dataset, with the best-performing prompt
achieving a score of 85.34 compared to the previous best of $81.52$ on the
Massive Text Embedding Benchmark (MTEB) Leaderboard. However, performance on
the other two datasets was less impressive, highlighting the importance of
considering domain-specific characteristics. The findings suggest that LLM-
based text enrichment has shown promising results to improve embedding
performance, particularly in certain domains. Hence, numerous limitations in
the process of embedding can be avoided.

###### Index Terms:

Large language models, natural language processing, ChatGPT $3.5$

##  I Introduction

Text embeddings are widely adopted in the field of Natural Language Processing
(NLP) that refer to vectorized representation of natural language. An
embedding is a representation of words in a low-dimensional continuous vector
space. They encapsulate the semantic content of the text [1]. These embeddings
find extensive applications across a spectrum of natural language processing
(NLP) endeavors including information retrieval (IR), question answering,
assessing semantic textual similarity, mining bitexts, recommending items,
etc. The researchers are making continuous efforts to improve the accuracy and
reduce the training steps [2].

Furthermore, an efficient technique for creating high-quality text embeddings
using synthetic data and minimal training, avoiding complex pipelines and
extensive labeled datasets, and achieving top results on key benchmarks when
mixed with labeled data [2].

There were early approaches like word2vec [3] and GloVe [4] to more advanced
models such as FastText [5] and BERT [6]. It discusses the strengths and
limitations of each model and their impact on various natural language
processing (NLP) tasks.

Various techniques have been proposed to improve the performance of embedding
models, such as fine-tuning on domain-specific data [7], using ensemble
methods, and incorporating external knowledge sources [8]. Large language
models have been successfully applied to a wide range of NLP tasks, such as
text generation [9], question answering [10], and sentiment analysis [11].
Several studies have explored the use of text enrichment and rewriting
techniques to improve the quality and informativeness of text data. For
example, a method for contextual augmentation of text data using a
bidirectional language model is being proposed [12], while a retrieval-
augmented generation approach for improving the factual accuracy of generated
text was also introduced [13].

Recent research has explored the use of LLMs for text compression to reduce
computational costs in Retrieval-Augmented Generation (RAG) systems and large
LLMs. For instance, RECOMP proposes compressing retrieved documents into
summaries before integrating them with language models, aiming to reduce
computational costs and help LMs identify relevant information more
efficiently [14]. Similarly, TCRA-LLM introduces a token compression scheme
for retrieval-augmented LLMs, employing summarization and semantic compression
techniques to reduce inference costs [15]. Context Tuning for RAG addresses
the limitation of RAG’s tool retrieval step by employing a smart context
retrieval system to fetch relevant information, improving the efficiency and
effectiveness of the generation process [16]. In the domain of prompt
compression, LLMLingua introduces a method for compressing prompts to
accelerate inference in LLMs, achieving up to $20$x compression while
preserving the original prompt’s capabilities [17]. The Natural Language
Prompt Encapsulation (Nano-Capsulator) framework compresses original prompts
into NL formatted Capsule Prompts while maintaining prompt utility and
transferability [18]. Compress-Then-Prompt [18] indicates that the generation
quality in a compressed LLM can be markedly improved for specific queries by
selecting prompts with high efficiency and accuracy trade-offs [19].
LongLLMLingua focuses on improving LLMs’ perception of key information in long
context scenarios through prompt compression, showing that compressed prompts
could derive higher performance with much less cost and reduce the latency of
the end-to-end system. Data Distillation proposes a data distillation
procedure to compress prompts without losing crucial information, addressing
issues related to the efficiency and fidelity of task-agnostic prompt
compression. While these approaches aim to reduce computational costs, the
current study explores the potential of LLMs for text enrichment to enhance
embedding quality.

Embedding models have become an essential component of various natural
language processing (NLP) tasks, such as text classification, clustering, and
retrieval. These models learn dense vector representations of words,
sentences, or documents, capturing semantic and syntactic relationships
between them. The quality of these embeddings directly impacts the performance
of downstream applications.

Despite their widespread use, embedding models face several challenges that
limit their performance. These challenges include limited vocabulary, lack of
context, sensitivity to grammatical errors, data sparsity, and lack of domain-
specific tuning. For example, embedding models may struggle with newer or
domain-specific terms not present in their training data, leading to mis-
classification or poor retrieval performance. Existing approaches to improve
embedding performance often focus on fine-tuning the embedding models on
domain-specific data or using ensemble techniques. However, these methods can
be resource-intensive and may not effectively address the fundamental
limitations of embedding models, such as their inability to capture context or
handle grammatical errors. Large language models (LLMs) have demonstrated
remarkable capabilities in understanding and generating human-like text. By
leveraging the knowledge and contextual understanding of LLMs, it is possible
to enrich and rewrite input text before the embedding process, thereby
addressing the limitations of embedding models and improving their
performance.

##  II Major Contributions

The primary objective of this paper is to propose a novel approach for
enhancing embedding performance by utilizing LLMs for text enrichment and
rewriting. The main contributions of the paper are as follows

  * •

Developing a methodology for leveraging an LLM to enrich and rewrite input
text before embedding

  * •

Identifying and addressing key challenges in embedding models, such as limited
vocabulary, lack of context, and grammatical errors

  * •

Conducting experiments on the TwitterSemEval $2015$ benchmark and others to
demonstrate the effectiveness of the proposed approach

##  III Methodology

The proposed approach involves leveraging the capabilities of ChatGPT $3.5$, a
large language model, to enrich and rewrite input text before the embedding
process. By addressing the limitations of embedding models, such as limited
vocabulary, lack of context, and grammatical errors, the proposed method aims
to improve the performance of embedding models on various NLP tasks. ChatGPT
$3.5$, developed by OpenAI, was chosen as the LLM for this study due to its
strong performance on a wide range of NLP tasks and its ability to generate
human-like text. Its extensive knowledge base and contextual understanding
make it well-suited for text enrichment and rewriting.

The ChatGPT $3.5$ model was used with its default settings and parameters. No
fine-tuning or additional training was performed, ensuring that the
improvements in embedding performance can be attributed solely to the text
enrichment and rewriting process. The text-embedding-3-large model, also
developed by OpenAI, was selected as the embedding model for this study. This
model has demonstrated strong performance on various NLP tasks and serves as a
representative example of state-of-the-art embedding models. The text-
embedding-3-large model was used with its default settings and parameters,
without any fine-tuning or modification. This allows for a fair comparison
between the performance of the embedding model with and without the proposed
text enrichment and rewriting approach. The proposed approach employs several
text enrichment and rewriting techniques to improve the quality and
informativeness of the input text. These techniques include:

###  III-A Context enrichment

ChatGPT 3.5 is used to provide additional context to the input text, making it
more informative and easier for the embedding model to capture the underlying
semantics. This is particularly useful for sparse or list-like entries, where
the LLM can expand the text with relevant descriptions or attributes.

###  III-B Grammatical correction

The LLM identifies and corrects spelling and grammatical errors in the input
text, ensuring that the text conforms to standard language usage. This
improves the quality of the embeddings generated from the text, as the
embedding model can focus on capturing the semantic relationships without
being hindered by grammatical inconsistencies.

###  III-C Terminology normalization

Domain-specific terms, abbreviations, and synonyms are standardized to a
consistent format using the knowledge base of ChatGPT 3.5. This reduces
ambiguity and improves the embedding model’s ability to match related
concepts, even when they are expressed using different terms.

###  III-D Word disambiguation

For polysemous words (words with multiple meanings), the LLM clarifies the
intended meaning based on the surrounding context. This disambiguation helps
the embedding model to capture the correct semantic relationships and improves
the accuracy of downstream tasks.

###  III-E Acronym expansion

ChatGPT 3.5 detects acronyms and abbreviations in the input text and expands
them to their full form. This improves clarity and understanding, enabling the
embedding model to better capture the meaning of the text.

###  III-F Metadata incorporation

Where relevant, the LLM incorporates additional metadata, such as the category
of the text, its intended audience, or domain-specific tags. This contextual
information helps in interpreting the text more accurately and can improve the
performance of the embedding model on domain-specific tasks.

###  III-G Sentence restructuring

The LLM is used to improve the structure of sentences in the input text,
making them clearer, more readable, and coherent. This makes it easier for the
embedding model to process and understand the text, leading to better-quality
embeddings.

###  III-H Inferring missing information

ChatGPT 3.5 uses its contextual understanding to infer missing information
that might be relevant for understanding the text. This can include inferring
the subject of a sentence or the meaning of an unclear reference, thereby
improving the completeness and coherence of the text for the embedding model.

##  IV Prompt Engineering and Optimization

To effectively leverage the capabilities of ChatGPT $3.5$ for text enrichment
and rewriting, a set of prompt design principles were established. These
principles aim to create prompts that clearly communicate the desired tasks
and goals to the LLM, while allowing for flexibility and adaptability to
different types of input text. An iterative prompt refinement process was
employed to identify the most effective prompts for the text enrichment and
rewriting tasks. This process involved creating multiple variations of
prompts, testing their performance on the TwitterSemEval $2015$ dataset, and
analyzing the results to identify areas for improvement. Four main prompt
variations were tested in this study, each focusing on different aspects of
the text enrichment and rewriting process. The prompts ranged from general
instructions for improving text quality to more specific guidance on tasks
such as grammar correction, terminology normalization, and metadata
incorporation.

##  V Numerical Validation

The experimental endeavor was undertaken with the overarching objective of
augmenting the performance of embedding models, particularly in the realms of
classification and clustering tasks, with the aim of securing a prominent
standing on the Massive Text Embedding Benchmark (MTEB) Leaderboard. Central
to this pursuit was the utilization of large language models, notably ChatGPT
3.5, to enhance and refine input text prior to embedding. The proposed
methodology encompasses a multifaceted approach, involving the enrichment of
text with additional contextual information, rectification of grammatical
inaccuracies, standardization of terminology, disambiguation of polysemous
terms, expansion of acronyms, and incorporation of pertinent metadata.
Furthermore, the project endeavors to optimize sentence structures and deduce
missing information, thereby enhancing the overall quality and accuracy of the
resultant embedding. The proposed approach was evaluated on three datasets:
Banking77Classification, TwitterSemEval $2015$, and Amazon Counter Factual
Classification. These datasets cover various domains and have been widely used
as benchmarks for text classification and clustering tasks. The datasets were
preprocessed to remove irrelevant information, such as URLs, hashtags, and
mentions. The text was then tokenized and converted to lowercase to ensure
consistency across the datasets.

The performance of the embedding models was evaluated using the average
precision based on cosine similarity metric in case of TwitterSemEval and
accuracy when evaluated with Banking77Classification data and Amazon Counter
Factual data. This metric assesses the quality of the embeddings by measuring
the similarity between the embedded representations of related texts and
comparing it to the ground truth. The text-embedding-3-large model was used as
a baseline, without any LLM-based text enrichment or rewriting. This allows
for a direct comparison of the performance improvements achieved by the
proposed approach. SFR-Embedding-Mistral model, which was the leading model on
the Massive Text Embedding Benchmark (MTEB) Leaderboard at the time of this
study, was also used as a baseline. This model serves as a representative
example of state-of-the-art embedding models and provides a high-quality
benchmark for comparison. The experimental procedure involved applying the
four prompt variations to the three datasets, using ChatGPT $3.5$ for text
enrichment and rewriting. The enriched and rewritten text was then passed
through the text-embedding-$3$-large model to generate embeddings. The
performance of these embeddings was evaluated using the cosine similarity
metric and accuracy values and then compared to the baseline models.

TABLE I: Performance comparison of the proposed methodology.

Model | B77C | TwitterSemEval | AmazonCF
---|---|---|---
Prompt 1 | 82.24 | 84.84 | 68.9
Prompt 2 | 78.73 | 82.95 | 71.9
Prompt 3 | 75.50 | 83.10 | 76.20
Prompt 4 | 79.71 | 85.34 | 68.00
TE | 85.69 | 77.13 | 78.93
SFR | 88.81 | 81.52 | 77.93
Improvement | -3.45 | 8.21 | -2.73

  * •

Note: TE stands foor text-embedding-3-large (base model) and SFR stands for
SFR-Embedding-Mistral (best performing model on leaderboard). Furthermore,
B77C stands for Banking77Classification, AmazonCF stands for Amazon Counter
Factual data, and improvement is indicated from the baseline. Moreover, the
values corresponding to B77C and AmazonCF are accuracy values whereas for
TwitterSemEval the values indicate the cosine similarities.

The objective was to identify the most effective prompt to achieve the highest
accuracy and average precision based on cosine similarities.

In summary, our MTEB Contextual Rewriting and Optimization project has
delivered significant success, surpassing the performance of the standalone
embedding model and outperforming the current leader in the field. It is worth
noting that due to budgetary constraints, the project was conducted on a
single dataset.

The ChatGPT $3.5$ model was used with its default settings and parameters. No
fine-tuning or additional training was performed, ensuring that the
improvements in embedding performance can be attributed solely to the text
enrichment and rewriting process.

Here are the details of the prompt: -

  * •

Prompt 1: “You are a text enhancer tasked with pre-processing text for
embedding models. Your goals are to enrich the text without losing the
context, correct grammatical inaccuracies, clarify obscure references,
normalize terminology, disambiguate polysemous words, expand acronyms and
abbreviations, incorporate relevant metadata, improve sentence structure for
clarity, and infer missing information where necessary. Your enhancements
should make the text more informative and easier to understand, thereby
improving the performance of embedding models in processing and analyzing the
text. If a user asks a question, then you should return an improved version of
the question. If the user did not ask a question, then you should return an
improved version of an answer.”

  * •

Prompt 2: “You are a text enhancer tasked with preprocessing text for
embedding models. Your goals are to enrich the text with additional context,
correct grammatical inaccuracies, clarify obscure references, normalize
terminology, disambiguate polysemous words, expand acronyms and abbreviations,
incorporate relevant metadata, improve sentence structure for clarity, and
infer missing information where necessary. Your enhancements should make the
text more informative and easier to understand, thereby improving the
performance of embedding models in processing and analyzing the text.”

  * •

Prompt 3: “You are a text enhancer to make better embeddings, your task is to
optimize text for embedding models by enriching, clarifying, and standardizing
it. This involves improving grammar, resolving ambiguities, and inferring
missing information to enhance model performance.”

  * •

Prompt 4: “You are a text enhancer to make better embeddings, your task is to
optimize text for embedding models by enriching, clarifying, and standardizing
it. This involves improving grammar, resolving ambiguities, and inferring
missing information to enhance model performance.”

The results and analysis of using Prompt-1 as input focuses on general
instructions for improving text quality, achieved varying performance across
the three datasets. It performed best on the TwitterSemEval 2015 dataset with
a cosine similarity score of $84.84$, representing a significant improvement
over the baseline text-embedding-3-large model ($77.13$). However, its
performance on Banking77Classification showing an accuracy of $82.24$ and
Amazon Counter Factual Classification with an accuracy of $68.9$ were lower
than the baseline models.

The results and analysis of using Prompt $2$ as input provides more specific
guidance on tasks such as grammar correction and terminology normalization,
also showed mixed results. It achieved a cosine similarity score of $82.95$ on
TwitterSemEval $2015$, outperforming the baseline model but slightly lower
than Prompt $1$. On Banking77Classification ($78.73$) and Amazon Counter
Factual Classification ($71.9$), Prompt $2$ showed better accuracy than for
Prompt $1$ but still fell short of the baseline models.

The insights into the results and analysis of using Prompt $3$, which focused
on concise instructions for optimizing text for embedding models, demonstrated
the best performance on Amazon Counter Factual Classification with an accuracy
of $76.2$, although it still fell short of the baseline models. Its
performance on TwitterSemEval $2015$ with cosine similarity value of $83.1$
was similar to Prompt $2$, while on Banking77Classification with cosine
similarity of $75.5$, it had the lowest score among the prompt variations.

Prompt $4$ was similar to Prompt $3$ but with slight variations in wording,
achieved the highest cosine similarity score on TwitterSemEval $2015$
($85.34$), outperforming all other prompt variations and baseline models.
However, its accuracies when Banking77Classification was used as evaluation
data ($79.71$) and Amazon Counter Factual Classification ($68$) were lower
than the baseline models and some of the other prompt variations.

Comparison with baseline models shows that there is significant improvement
over text-embedding-3-large alone The prompt variations significantly
outperformed the baseline text-embedding-3-large model on the TwitterSemEval
$2015$ dataset, with the best-performing prompt (Prompt $4$) improving upon
the baseline by cosine similarity score of $8.21$. However, on
Banking77Classification and Amazon Counter Factual Classification, the prompt
variations did not surpass the performance (accuracy) of the baseline model.
The best-performing prompt (Prompt $4$) outperformed the leading model on the
MTEB Leaderboard, SFR-Embedding-Mistral, on the TwitterSemEval $2015$ dataset.
However, SFR-Embedding-Mistral maintained its lead on Banking77Classification
and AmazonCounterfactualClassification.

A qualitative analysis of the enriched and rewritten text generated by ChatGPT
$3.5$ revealed several improvements in text quality and informativeness. The
LLM successfully provided additional context, corrected grammatical errors,
normalized terminology, disambiguated polysemous words, expanded acronyms, and
incorporated relevant metadata. These enhancements made the text more
coherent, informative, and easier for the embedding model to process and
understand.

##  VI Conclusion

This paper introduces a novel approach for enhancing embedding performance by
leveraging the capabilities of large language models, specifically ChatGPT
3.5, for text enrichment and rewriting. While recent research has focused on
using LLMs for text compression to reduce computational costs in RAG systems
and large LLMs, this study demonstrates the potential of LLMs for text
enrichment to improve embedding quality. The proposed approach addresses the
limitations of embedding models, such as limited vocabulary, lack of context,
and grammatical errors, by providing additional context, correcting
inaccuracies, normalizing terminology, disambiguating polysemous words,
expanding acronyms, and incorporating metadata. Experimental results on the
TwitterSemEval $2015$ dataset show that the proposed method outperforms the
leading model on the Massive Text Embedding Benchmark (MTEB) Leaderboard.
Hence, the embedding is improved substantially.

## References

  * [1] N. Pittaras, G. Giannakopoulos, G. Papadakis, and V. Karkaletsis, “Text classification with semantically enriched word embeddings,” _Natural Language Engineering_ , vol. 27, no. 4, pp. 391–425, 2021.
  * [2] L. Wang, N. Yang, X. Huang, L. Yang, R. Majumder, and F. Wei, “Improving text embeddings with large language models,” _arXiv preprint arXiv:2401.00368_ , 2023.
  * [3] T. Mikolov, K. Chen, G. Corrado, and J. Dean, “Efficient estimation of word representations in vector space,” _arXiv preprint arXiv:1301.3781_ , 2013.
  * [4] J. Pennington, R. Socher, and C. D. Manning, “Glove: Global vectors for word representation,” in _Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP)_ , 2014, pp. 1532–1543.
  * [5] P. Bojanowski, E. Grave, A. Joulin, and T. Mikolov, “Enriching word vectors with subword information,” _Transactions of the association for computational linguistics_ , vol. 5, pp. 135–146, 2017.
  * [6] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, “Bert: Pre-training of deep bidirectional transformers for language understanding,” _arXiv preprint arXiv:1810.04805_ , 2018.
  * [7] J. Howard and S. Ruder, “Universal language model fine-tuning for text classification,” _arXiv preprint arXiv:1801.06146_ , 2018.
  * [8] Z. Zhang, X. Han, Z. Liu, X. Jiang, M. Sun, and Q. Liu, “Ernie: Enhanced language representation with informative entities,” _arXiv preprint arXiv:1905.07129_ , 2019.
  * [9] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever _et al._ , “Language models are unsupervised multitask learners,” _OpenAI blog_ , vol. 1, no. 8, p. 9, 2019.
  * [10] C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu, “Exploring the limits of transfer learning with a unified text-to-text transformer,” _Journal of machine learning research_ , vol. 21, no. 140, pp. 1–67, 2020.
  * [11] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell _et al._ , “Language models are few-shot learners,” _Advances in neural information processing systems_ , vol. 33, pp. 1877–1901, 2020.
  * [12] S. Kobayashi, “Contextual augmentation: Data augmentation by words with paradigmatic relations,” _arXiv preprint arXiv:1805.06201_ , 2018.
  * [13] K. Guu, K. Lee, Z. Tung, P. Pasupat, and M. Chang, “Retrieval augmented language model pre-training,” in _International conference on machine learning_. PMLR, 2020, pp. 3929–3938.
  * [14] F. Xu, W. Shi, and E. Choi, “Recomp: Improving retrieval-augmented lms with compression and selective augmentation,” _arXiv preprint arXiv:2310.04408_ , 2023.
  * [15] J. Liu, L. Li, T. Xiang, B. Wang, and Y. Qian, “Tcra-llm: token compression retrieval augmented large language model for inference cost reduction,” _arXiv preprint arXiv:2310.15556_ , 2023.
  * [16] R. Anantha, T. Bethi, D. Vodianik, and S. Chappidi, “Context tuning for retrieval augmented generation,” _arXiv preprint arXiv:2312.05708_ , 2023.
  * [17] N. D. Kulkarni and S. Bansal, “Application of generative ai for business analyst role,” _Journal of Artificial Intelligence & Cloud Computing. SRC/JAICC-201. DOI: doi. org/10.47363/JAICC/2023 (2)_, vol. 187, pp. 2–5, 2023.
  * [18] Y.-N. Chuang, T. Xing, C.-Y. Chang, Z. Liu, X. Chen, and X. Hu, “Learning to compress prompt in natural language formats,” _arXiv preprint arXiv:2402.18700_ , 2024.
  * [19] Z. Xu, Z. Liu, B. Chen, Y. Tang, J. Wang, K. Zhou, X. Hu, and A. Shrivastava, “Compress, then prompt: Improving accuracy-efficiency trade-off of llm inference with transferable prompt,” _arXiv preprint arXiv:2305.11186_ , 2023.