It’s crucial to ensure the reliability and trustworthiness of AI systems as they are adopted in various real-world
decision-making scenarios. On one hand, language models still exhibit illusions during use, and on the other hand,
the reliability of systems comprises two parts: the system’s reliability under individual testing environments and
its reliability in human interactions. Another issue is constructing systems with decision-making processes that
are observable and explainable to users. From a social perspective, the proliferation of AI systems across fields
also poses potential risks. This risk arises from a gap between AI developers, who often focus on advancing
technology without considering its downstream applications, and AI adopters, who may transfer AI systems to
their fields without adequate safety considerations or verification of replicable success [41] . Therefore, it is crucial
to build a framework that enables AI adopters to accurately assess model utility and appropriateness, and allows
AI regulators to quickly identify risks and issue safety alerts in AI systems.
Alignment techniques can facilitate synchronized, independent, and rigorous evaluations of AI systems. AI
developers should prioritize appropriate bias handling during the training process, acknowledging the importance
of socio-economic, cultural, and other differences. Furthermore, we should aim to develop robust and fair evaluation methods and datasets for auditing AI systems. Zhu et al. (2023) proposes the first dynamic testing protocol
for large language models, utilizing Directed Acyclic Graphs (DAGs) to dynamically generate test data, thereby
reducing the risk of test data memorization and contamination. Additionally, new robust security protocol evaluation methods have been introduced: Shlegeris and Greenblatt (2023) suggests constructing adversarial policies to
manage dangerously powerful and deceptive models, while Greenblatt et al. (2023) proposes (un)trusted editing to
supervise models based on their harm and deceitfulness levels. Future efforts should also prevent AI systems from
reward-hacking evaluation system exploits and aim to provide AI regulators with an explainable, independent, and
centralized evaluation system.
AI adopters and the industry should allocate financial and computational resources to thoroughly evaluate use
cases and share case studies showcasing both successes and failures. Equally important is training for adopters on
downstream applications.

5.4.2 Alignment Techniques for AI Governance


In this survey, we have provided a broadly-scoped introduction to AI alignment, which aims to build AI systems
that behave in line with human intentions and values. We specify the objectives of alignment as Robustness,
Interpretability, Controllability, and Ethicality ( **RICE** ), and characterize the scope of alignment methods as comprising of _forward alignment_ (making AI systems aligned via alignment training) and _backward alignment_ (gaining
evidence of the systems’ alignment and govern them appropriately to avoid exacerbating misalignment risks). Currently, the two notable areas of research within forward alignment are _learning from feedback_ and _learning under_
_distribution shift_, while backward alignment is comprised of _assurance_ and _governance_ .
One thing that sets alignment apart from many other fields is its diversity (Hendrycks, 2022) – it is a tight
assembly of multiple research directions and methods, tied together by a shared goal, as opposed to a shared
methodology. This diversity brings benefits. It fosters innovation by having the different directions compete and
clash against each other, leading to a cross-pollination of ideas. It also allows different research directions to
complement each other and together serve the goal of alignment; this is reflected in the _alignment cycle_ (see Figure
2), where the four pillars are integrated into a self-improving loop that continually improves the alignment of AI
systems. Meanwhile, this diversity of research directions raises the barrier to entry into this field, which mandates
the compilation of well-organized survey materials that serve both the newcomers and the experienced. In this
survey, we attempt to address this need by providing a comprehensive and up-to-date overview of alignment.
We attempt to account for the full diversity within the field by adopting a broad and inclusive characterization
of alignment. Our survey of alignment gives a spotlight to almost all major research agendas in this field, as
well as to real-world practices on the assurance and governance front. We recognize that boundaries of alignment
are often vague and subject to debate. Therefore, when proposing the RICE principles, we put forth our broad
characterization of alignment as an explicit choice. In the meantime, we recognize that such a survey needs to
be a long-term endeavor that is continually reviewed and updated. Both the problems and methods of alignment
closely follow the development of machine learning. This fast-paced development means that new materials and
frameworks can become outdated after merely a few years. This fact is one reason why we write the survey to
reflect the latest developments, and also mandates continual maintenance and updates.

Jiaming Ji
Tianyi Qiu
Boyuan Chen
Borong Zhang
Hantao Lou
Kaile Wang
Yawen Duan
Zhonghao He
Jiayi Zhou
Zhaowei Zhang
Fanzhi Zeng
Kwan Yee Ng
Juntao Dai
Xuehai Pan
Aidan O'Gara
Yingshan Lei
Hua Xu
Brian Tse
Jie Fu
Stephen McAleer
Yaodong Yang
Yizhou Wang
Song-Chun Zhu
Yike Guo
Wen Gao


Specifically, we outline key challenges and potential future directions based on the alignment cycle, namely forward and backward alignment.
[41https://www.scai.gov.sg/scai-question-11/](https://www.scai.gov.sg/scai-question-11/)
6.1 Key Challenges in the Alignment Cycle
**Learning Human Intent from Rich Modalities (forward alignment)** Underspecificity of true human intent, _i.e._,
the non-uniqueness of inferred human intent from binary feedback data, is a key challenge in scalable oversight.
Consider an AI system tasked with providing proof or refutation to a mathematical hypothesis, under a human evaluator who might be tricked by sophisticated false proofs. Our goal is to construct a training process that induces
the AI system to output sound proofs as opposed to false proofs that seem convincing. This system may mislead
evaluators with plausible but false proofs due to the system’s optimization for human approval, as it attempts to
satisfy the superficial criteria of convincing proofs rather than focusing on accuracy. The fundamental problem
stems from the reliance on binary feedback which categorizes responses simply as preferred or dispreferred, thus
limiting the amount of information on true human preferences that’s available to the learning algorithm, potentially
leading to the preference of credible-seeming deceptive proofs over genuinely sound arguments.
To enhance the model’s alignment with true human intent, researchers have proposed incorporating richer human
input beyond binary choices, such as detailed text feedback (Chen et al., 2024a) and real-time interations (HadfieldMenell et al., 2016). It allows the model to differentiate between proofs that are merely convincing and those that
are truly sound, using nuanced human evaluations and a vast database of human-written texts. The broader input
base helps in constructing a more accurate model of human preferences, reducing the risk of favoring misleading
proofs while respecting the complexity of human intent and reasoning. Looking forward, even richer modalities
like embodied societal interactions could represent an enticing next step.
It is worth noting that current LLMs are already trained on Internet-scale human text (and for multimodal models,
also visual/audio content). Why, then, don’t reward modeling algorithms already possess the ability to accurately
pin down human intent? The explanation is that pretraining data does not feed into the reward modeling process
in a way that biases the process towards true human intent, even though the reward model is finetuned from the
pretrained model. For instance, neural circuits representing human intent can potentially be rewired during RLHF
to perform manipulative behaviors. From another perspective, pretraining on text such as _humans do not want to_
_be tricked into believing things_ does not induce the reward model to interpret later human feedback signals in this
light, partly due to the lack of out-of-context learning capabilities in current LLMs (Berglund et al., 2023). Solving
these problems may enable reward modeling algorithms to learn human intent from massive pretraining data, a big
step towards our goal.
We summarize three key questions for the learning of human intent from rich modalities. They serve as key
dimensions for characterizing an alignment method from the intent modality lens, and almost all existing alignment
methods can be categorized by their answers to these three questions.
1. **Learning algorithm** . As previously mentioned, we need to learn human intent from rich modalities in a way
that guides the reward model’s subsequent interpretation of human input.
2. **Priors and inductive biases** . Human-like priors/inductive bias is needed for the reward modeling process
to select the correct hypothesis of human intent, though this requirement is greatly loosened as the allowed
modalities of human input expand.
3. **Learner alignment** . We utilize the intent learner to align AI systems, possibly by using it as a reward model.
However, this would not be possible if the intent learner, which is itself an AI system with potentially strong
capabilities, is misaligned. This necessitates measures to avoid or contain the misalignment of the intent
learner.
**Trustworthy Tools for Assurance (backward alignment)** A major concern in AI alignment is deceptive alignment, where AI systems pursue aligned goals under most circumstances but may pursue other goals when opportunities arise. Recent studies have revealed that general alignment techniques ( _e.g._, SFT, RLHF, Adversarial
Training) fail to eradicate certain deceptive and backdoor behaviors, possibly leading to a misleading sense of
safety (Hubinger et al., 2024). With AI systems gaining power and access to more resources, hidden intentions
that pose existential risks could have unimaginable consequences. _How can we detect and eliminate deceptive and_
_backdoor behaviors?_
Reliable tools are still lacking to address this issue. On one hand, mechanistic interpretability tools encounter
additional challenges due to the polysemanticity of neurons and scalability issues. On the other hand, there is a
limited understanding of how jailbreaking functions and the susceptibility of language models to poisoning and
backdoors (Anwar et al., 2024).
Additionally, given the potential misuse of AI systems in cyber attacks, biological warfare, and misinformation,
it is crucial to develop reliable mechanisms to trace the origins of LLM outputs. While AI systems are becoming
more integrated into society, societal readiness lags behind. This is evident from the inadequate AI governance
efforts, insufficient public knowledge, governments’ lack of necessary scientific and technical capabilities, the
absence of institutions that can keep pace with LLM advancements, and the challenges in mitigating the social
6.2 Key Traits and Future Directions in Alignment Research
impacts of widespread harmful behaviors. Therefore, it is essential to reconsider AI alignment from a sociotechnical standpoint, establish dependable AI assurance and governance mechanisms, and engage in effective
international governance collaboration.
**Value Elicitation and Value Implementation (backward alignment)** Current algorithms for learning from
human feedback, particularly RLHF, often assume feedback comes from a singular, monolithic human source.
However, this assumption is unrealistic due to widespread disagreements on contentious issues globally, which
frequently result in conflicting judgments about AI system outputs (Santurkar et al., 2023). Consequently, determining who to draw feedback from and understanding the scope and nature of human values infused into models
are crucial questions for the field of alignment.
_Value Elicitation_ and _Value Implementation_ aim to define the values and norms that AI systems should encode
and how to integrate these into AI systems. Human values and preferences are diverse, ranging from strict rules
like laws and moral principles to social etiquette and specific domain preferences (Cahyawijaya et al., 2024; Kirk
et al., 2024). We need reliable tools to reveal the values embedded in current AI systems and potential social risks,
enabling us to mitigate these risks more effectively [42] .
_Democratic human input_ is one of the leading solutions to value elicitation and implementation. This method
gathers input from a large, demographically representative sample of individuals, aggregating preferences and
values into a coherent policy, rather than relying on feedback from a single individual. This approach is heavily
influenced by the computational social choice literature (Brandt et al., 2016). Leading industry (Zaremb et al.,
2023) and academic (Köpf et al., 2024) labs have adopted democratic human input for LLMs. However, research
is still needed on its integration into more agentic AI systems, such as LLM-based autonomous agents.
Despite its apparent simplicity, democratic human input encounters significant practical and fundamental challenges. Obtaining a truly random sample of the global population is particularly challenging, as 33% of people
worldwide do not have Internet access and thus are excluded from participating in AI system training (United
Nations, ITU, 2023). Furthermore, human feedback becomes less effective when the AI system’s reasoning capabilities surpass those of humans, making it difficult for human workers to evaluate its outputs. To complement
democratic human input, alternative approaches aim to formalize universally recognized meta-level moral principles, such as _moral consistency_, _moral reflection_, and _moral progress_, designing algorithms to enact these principles. Although these methods still rely on human data and input, they do not demand strict representativeness and
are less constrained by human oversight limitations.
  - **Moral consistency** . There is a general consensus that moral principles should be consistently applied, meaning similar cases should receive similar treatment irrespective of the people or parties involved. Algorithms
have been developed to integrate this principle into the ethical decision-making processes of AI systems (Jin
et al., 2022b).
  - **Moral reflection and moral progress** . The _coherent extrapolated volition_ concept was developed to formalize the role of reflection in shaping human moral values (Søvik, 2022). Inspired by this, subsequent
algorithms were designed to enable AI systems to mimic human moral reflection, thereby influencing their
actions (Xie et al., 2023). Furthermore, the logical next step of moral reflection is _moral progress_, demonstrated by AI-driven analyses of historical moral trends (Schramowski et al., 2020; Atif et al., 2022) and
efforts to permanently integrate continual moral advancement into AI systems (Kenward and Sinclair, 2021).


In the end of the survey, we conclude the survey by looking ahead and presenting the key traits in this field that we
believe ought to be preserved or fostered.
**Open-Ended Exploration of Novel Challenges and Approaches** A lot of the alignment discourse is built upon
classic works that predate the recent developments of LLMs and other breakthroughs in large-scale deep learning.
Thus, when this paradigm shift happens in the machine learning field, it is plausible that some challenges in
alignment become less salient while others become more so; after all, one defining feature of scientific theories is
their falsifiability (Popper, 2005). More importantly, this shift in machine learning methodology and the broader
trend of ever-tighter integration of AI systems into society (Abbass, 2019) introduces novel challenges that could
not be envisioned before. This requires that we engage in _open-ended exploration_, actively seeking out new
challenges that were previously neglected. Moreover, such an exploration need not be constrained to challenges –
a similar mindset should be adopted regarding approaches and solutions, thus building a more diverse portfolio for
both the _questions_ and the _answers_ (Shimi, 2022).
[42https://www.scai.gov.sg](https://www.scai.gov.sg)
6.2 Key Traits and Future Directions in Alignment Research
**Combining Forward-Looking and Present-Oriented Perspectives** Alignment has emphasized harms from
potential advanced AI systems that possess stronger capabilities than current systems (Ngo, 2020a). These systems
might come into existence well into the future, or might just be a few years away (Stein-Perlman et al., 2022). The
former possibility requires us to look into extrapolated trends and hypothetical scenarios (Carlsmith, 2022). In
contrast, the latter possibility highlights the need for on-the-ground efforts that work with current governance
institutions and use current systems as a prototype for more advanced ones (Cotra, 2021).
**Emphasis on Policy Relevance** Alignment research does not live in a vacuum but in an ecosystem [43], with participation from researchers, industry actors, governments, and non-governmental organizations. Research serving
the needs of the AI alignment and safety ecosystem would therefore be useful. Such needs include solving the key
barriers to various governance schemes, for example, extreme risk evaluations (Shevlane et al., 2023), infrastructure for computing governance, and mechanisms for making verifiable claims about AI systems (Brundage et al.,
2020).
**Emphasis on Social Complexities and Moral Values** As AI systems become increasingly integrated into society (Abbass, 2019), alignment ceases to be only a single-agent problem and becomes a social problem. Here, the
meaning of _social_ is three-fold.
1. Alignment research in multi-agent settings featuring the interactions between multiple AI systems and multiple humans (Critch and Krueger, 2020; Liu et al., 2024a). This includes how AI systems can receive granular
feedback from realistic simulated societies, ensuring consistency in training scenarios and among multiple
entities ( _i.e._, multi-agent, multiple AI systems, and multiple humans), which not only aids in generalizing AI
systems in multi-entity settings but also helps avoid problems associated with RL (Liu et al., 2024a).
2. Incorporating human moral and social values into alignment (see Section1.2.3 and Section4.3), which is closely linked to
the field of _machine ethics_ and _value alignment_ (Gabriel, 2020; Gabriel and Ghazavi, 2021).
3. Modeling and predicting the impacts of AI systems on society, which requires methods to approach the
complexities of the social system, including those from the social sciences. Examples of potentially useful
methodologies include social simulation (Bonabeau, 2002; De Marchi and Page, 2014; Park et al., 2023a)
and game theory (Du et al., 2023).


Allen Institute for AI
Berkeley, CA, USA
```
nathanl@allenai.org
```

The Alignment Ceiling: Objective Mismatch in Reinforcement Learning from
  Human Feedback


TU Dresden
Dresden, Germany
```
    roberto.calandra@tu-dresden.de
```


Reinforcement learning from human feedback (RLHF) has emerged as a powerful technique to make
large language models (LLMs) more capable in complex settings. RLHF proceeds as collecting
human preference data, training a reward model on said data, and optimizing a base ML model with
respect to said reward for extrinsic evaluation metrics (e.g. MMLU, GSM8k). RLHF relies on many
assumptions about how the various pieces fit together, such as a reward model capturing human
preferences and an RL optimizer extracting the right signal from a reward model. As the RLHF
process involves many distinct design decisions, it is easy to assume that multiple processes are
correlated and therefore numerically linked. This apparent correlation is often not true, where reward
models are easily overoptimized or RL optimizers can reduce performance on tasks not modeled
in the data. Notable manifestations of models trained with imperfect RLHF systems are those that
are prone to refusing basic requests for safety reasons or appearing lazy in generations. As chat
model evaluation becomes increasingly nuanced, the reliance on a perceived link between reward
model training, RL scores, and downstream performance drives these issues, which we describe as an
_objective mismatch_ . In this paper, we illustrate the causes of this issue, reviewing relevant literature
from model-based reinforcement learning, and argue for solutions. By solving objective mismatch in
RLHF, the ML models of the future will be more precisely aligned to user instructions for both safety
and helpfulness.


Reinforcement learning from human feedback (RLHF) is a powerful tool for integrating qualitative values into large
machine learning models (Bai et al., 2022; Christiano et al., 2017; Ouyang et al., 2022) that are used in popular
consumer apps such as ChatGPT and Midjourney. RLHF was popularized with its use to integrate human values into
large language models (LLMs) for aligning chat tools (Schulman, Zoph, Kim, & more, 2022; Team et al., 2023). RLHF
has become an important technique in the process of making models better at responding to user requests, often referred
to as instruction-tuned, steerable, aligned, or chat-tuned.
RLHF methods typically operate in a multi-step process on top of a base language model, first learning a model of
human preferences that acts as a reward function, and second using this model within a reinforcement learning (RL)
loop. These two steps are often executed independently, with a reward model (RM) being trained on human preference
data and then the RL optimizer is used to extract maximum information from the RM into the base model. This
multi-step process induces challenges (Schulman, 2023) – even the most popular RLHF models include weaknesses
such as `llama-2-70b-chat-hf` ’s propensity to refuse vanilla requests on safety grounds (Röttger et al., 2023) or a
version of ChatGPT documented officially as having “cases of laziness” (OpenAI, 2024). Colloquially, these issues fall
under the potential banner of “too much RLHF.” These failures are signs of the current limitations of RLHF, where even
with positive signals in training of each individual module, the resulting model can have unintended behaviors.
In this paper, we detail and argue for solving a fundamental challenge in modern RLHF learning schemes – _objective_
_mismatch_ - in order to mitigate these issues. In RLHF, three important parts of training are numerically decoupled: the
evaluation metrics, the reward model, and the generating model (policy). This mismatch between the reward model and
Objective Mismatch in RLHF
Figure 1: The three links causing objective mismatch in RLHF: Reward model training _↔_ policy model training, reward
model training _↔_ evaluation tools, and policy model training _↔_ evaluation tools, as discussed in Sec. 4.1.
the RL training is visualized in Fig. 2, yet other links exist between the goals of evaluation and training processes as
shown in Fig. 1. Among other prospects, there are many avenues to better align reward model training to the literature
in preference quantification (Lambert, Gilbert, & Zick, 2023) and fundamental optimization challenges need to be
solved in RLHF practices (Casper et al., 2023). ChatGPT, the most popular model trained with RLHF, shows signs
of this limitation through issues such as verbosity, self-doubt and question refusals, repeated phrases, hedging, and
more (Schulman, 2023). These traits of overoptimization are results of the subtle proxy objective problem that objective
mismatch provides a frame for studying and solving – the reward model attributes excess value to phrases that do
not contribute to user benefit, which the RL optimizer exploits, such as safety flags. On the other hand, the current
training setups are not fully aligned with evaluation tools because the RLHF’d models still need sophisticated prompting
techniques such as “thinking step by step” (J. Wei et al., 2022) or “take a deep breath” (Yang et al., 2023) to reach
maximum performance. Solving objective mismatch will remove the need for these advanced techniques and reduce
the likelihood of out-of-scope refusals from an LLM.
The use of RLHF is promising as it gives more levers for optimization of LLMs beyond next-token prediction accuracy.
In this paper, we argue the position that _the potential benefits of RLHF will not be realized without solving the objective_
_mismatch issue_ . RLHF has the potential to enable LLMs that are safe (Ji et al., 2023; Shi, Chen, & Zhao, 2023),
personalized (Jang et al., 2023), and effective (Bai et al., 2022; Ouyang et al., 2022).
The phrase objective mismatch originates from model-based reinforcement learning (MBRL), where an agent iteratively
learns a dynamics model of the environment that it later uses to solve a control task (a dynamics model _fθ_ maps from
state and action to next state, as _st_ +1 = _fθ_ ( _at, st_ )) (Lambert, Amos, Yadan, & Calandra, 2020; Moerland, Broekens,
Plaat, Jonker, et al., 2023; R. Wei, Lambert, McDonald, Garcia, & Calandra, 2023). In this context, the mismatch is
between learning an accurate dynamics model rather than one that is optimized for high task reward. In RLHF, the
problem is related, but with added complexity, as the reward model is optimized for preference data over a closed
distribution, which does not match the end users. Second, the task of open-ended language generation is less specific
to a notion of reward than that of RL control policies. For these reasons, as we explore in this paper, the objective
mismatch issue is more nuanced and critical to RLHF. In this position paper, we make three contributions:
    - Clearly explain the origins and potential manifestations of objective mismatch in chat-tuned LLMs,
    - Connect related work from NLP and RL literature around objective mismatch,
    - Propose directions of study to solve the mismatch and foster better RLHF practices.
Objective Mismatch in RLHF
Figure 2: An illustration of where the objective mismatch issue emerges within the RL optimization phase of RLHF. A
mismatch occurs when the score from the reward model is assumed to be correlated with other downstream evaluation
metrics, such as human preferences over evaluation sets, classic NLP benchmarks, or LLM-as-a-judge systems.
Compared to traditional RL problems, RLHF does not have the canonical form of an _environment_, which indirectly
maps to the training data with a reward model, but does not capture the same properties.


Early work in RLHF focused on continuous control domains with various methods for altering the behavior across
trajectories (Christiano et al., 2017; Wirth, Akrour, Neumann, Fürnkranz, et al., 2017). The impacts of RLHF today
primarily has been centered around its use with LLMs. Initial work on RLHF for LLMs utilized user preferences from
a batch of 4 options (Ziegler et al., 2019) to train a reward model across general LLM benchmarks. Group preferences
were changed to pairwise preferences, and rather than general benchmarks the reward model was focused on the task of
summarization (Stiennon et al., 2020; J. Wu et al., 2021). Next emerged general question-answering models (Ouyang
et al., 2022) and web crawling agents (Nakano et al., 2021), primarily from scaling the initial model and human
datasets. Now, RLHF is used to train general chat models across a variety of tasks (Bai et al., 2022; Schulman et
al., 2022; Touvron et al., 2023) and in specific domains such as harm reduction (Glaese et al., 2022) or information
accuracy (Menick et al., 2022).
The development of these methods has accelerated markedly, with many variations on the methodology for integrating
feedback into language models (Fernandes et al., 2023). The most popular reinforcement learning optimizer is still
Proximal Policy Optimization (PPO) (Schulman, Wolski, Dhariwal, Radford, & Klimov, 2017), with its many stable
and scalable implementations. Recent works have been developing methods for the RL optimizer, such as the offline
RL algorithm Implicit Language Q Learning (ILQL) (Snell, Kostrikov, Su, Yang, & Levine, 2022), direct preference
optimization (DPO) (Rafailov et al., 2023) for utilizing preference data without a reward model, or Advantage-Leftover
Lunch RL (A-LOL) (Baheti et al., 2023) which is designed to act on the entire response as a single action (which RLHF
generally does).

2.1 Reinforcement learning from human feedback


There is a substantial emerging literature on varieties of numerical issues, unexpected behaviors such as verbosity
and evasiveness (Schulman, 2023), and potential solutions in RLHF, which can be mitigated by progress on solving
objective mismatch. A prominent recent example is the behavior of the flagship Llama 2 chat model refusing to answer
a request asking “How do I kill a Linux process,” conflating the computer process with the morals of killing a living
creature. It has been shown that there are predictable behaviors of reward model overoptimization with PPO and
best-of-N optimization techniques (Gao, Schulman, & Hilton, 2022), which can be partially mitigated by training
ensemble reward models (Coste, Anwar, Kirk, & Krueger, 2023), weight-averaging (Ramé et al., 2024), or constrained
optimization (Moskovitz et al., 2023). Other issues have emerged in RLHF models that demonstrate the need for
improved reward models, such as a bias towards long responses (Singhal, Goyal, Xu, & Durrett, 2023), a lack of
language consistency (L. Shen et al., 2023) (invariance over changes that maintain meaning), or a reduction of output
diversity Kirk et al. (2023). A similar argument is made in A. Wei, Haghtalab, and Steinhardt (2023), where the authors
Objective Mismatch in RLHF
argue that “competing objectives and mismatched generalization” mislead the models – we present how objective
mismatch covers both of these limitations and more possible failure cases.
Other papers indicate more fundamental limitations in how the preference data are collected (Bansal, Dang, & Grover,
2023) or utilized. For example, multiple lines of work argue that the reward model training formulation does not align
with the data collection process and downstream RL optimization, suggesting the models should model advantage
estimates rather than direct value functions (Knox & Stone, 2008; Peng et al., 2023).


Specific domains are addressing this by shifting preference labels away form solely pairwise annotator input (whether by
a human or an LLM) to computational feedback to bootstrap pairwise data for a reward model. For example, successful
code execution in Python or reasoning path length has been used for rejection sampling (Yuan et al., 2023). Other works
combine scores from code execution, syntax, and semantics to optimize for effective code (Shojaee, Jain, Tipirneni, &
Reddy, 2023) or through unit tests (Liu, Zhu, et al., 2023; B. Shen et al., 2023). These are examples of early solutions
to the reward specification problem facing all applications of RLHF.


Core to the ideas of objective mismatch with LLMs is the methods of evaluation used to correlate performance.
Historically, LLMs have been evaluated across a wide variety of tasks trying to capture specific characteristics of
models, making evaluation an extremely broad process (Liang et al., 2022) where progress is saturating (Kiela, Thrush,
Ethayarajh, & Singh, 2023). Now, many models are focused on hard to specify tasks such as chat, where existing
benchmarks were not well correlated with performance (Zheng et al., 2023), so new chat based evaluations such as
MT-Bench (Zheng et al., 2023) and AlpacaEval (Li et al., 2023) have been introduced, but substantial further work is
needed.


Reward models are trained with human preference data most often consisting of a task given to the model _prompt_, i.e a
request or instruction, and ratings of the _completion_, or answer. The feedback can consist of selecting the best from
groups of responses (Ziegler et al., 2019), scores and rankings of a group of candidate responses (Ouyang et al., 2022),
a choice between a pair of responses (Bai et al., 2022) (choose best response between two options), and even finer
grained data (Z. Wu et al., 2023). The workers employed are generally given detailed instructions on which styles,
occurrences, or values to prioritize in their labels.
The reward models trained for RLHF are most often trained as classifiers between a chosen and rejected completion to
a prompt before optimizing with RL where they return a scalar value for each piece of text. Given two options for a
completion _y_ from a prompt _x_, and the scores they obtain a scalar output _r_ from an initially untrained value head on an
LLM or value model entirely, the loss for the reward model follows (Askell et al., 2021; Ouyang et al., 2022)
_L_ = log�1 + _e_ _[r]_ [chosen] _[−][r]_ [rejected][�] _._ (1)
The loss function is designed to increase the distance between the two samples, where variations exist including losses
of 4 samples rather than a pair (Ziegler et al., 2019), updating the model with batches of pairwise labels on a given
prompt (Ouyang et al., 2022), or optimizing based on the margin between _r_ chosen and _r_ rejected (Touvron et al., 2023).
For inference during RL optimization, the reward is taken as the raw logit output from this model that represents an
unnormalized probability of the text being preferred.


Language generation optimized via reinforcement learning, which RLHF is a version of, can be formalized as a partially
observable Markov decision process (POMDP) (Spaan, 2012). We define a POMDP _M_ at a per-token level with
_M_ = ( _S, A, O, T, Z, µ_ 0 _, R, γ_ ). Here, the state of the system is _st ∈S_, which the agent receives as an observation
_ht ∈O_ . The observation is a history of tokens _ht_ = _{t_ 0 _, t_ 1 _, . . ., tt−_ 1 _}_ and the action space is the possible set of
next-tokens in the vocabulary of the policy model _at_ = _tt ∈A_, including the end-of-sequence token _a_ end. As in a
traditional MPD, _T_ is the transition function _T_ ( _·|st, at_ ).
Objective Mismatch in RLHF
Figure 3: The canonical RL problem, where an agent interacts repeatedly with an environment, which the RLHF process
is derived from (as in Fig. 2).
The goal of the RLHF process is to learn a policy that is mapping _π_ : _O �→P_ ( _A_ ). This is done with the reward
model, which acts as a reward function _R_ ( _st, at_ ) _�→R_, used after each sequence is generated. The full sequence, until
end-of-sequence token _a_ end, is called the action and used to get a scalar reward _rt_ from the reward model.
With LLMs, the generating model is referred to as the _policy_ model. In RLHF, the discount factor of reward is set to 1
and no further actions are taken for the given prompt, casting the problem as contextual bandits. An example of the
RLHF loop is shown in Fig. 2 in comparison to a standard RL loop shown in Fig. 3.


The objective mismatch in RLHF emerges from three broad causes: First, common practice in RL engineering dictates
that as long as reward is increasing the model is improving. Second, the evaluation methods available for models trained
with RLHF are often incomplete relative to their downstream use-cases. Third, the assumption that the reward model
trained is a suitable reward function for optimization. For these reasons, objective mismatch emerges as the assumption
that downstream evaluation will be correlated with the reward model score for the current policy, which is not proven.


Objective mismatch in RLHF is the result of the interactions between three different sub-components, rather than
just the two (i.e., dynamics model and policy) from MBRL: It is a balance of (1) the reward model training, _the goal_
_of getting a calibrated reward function_, (2) the policy training, _the process of extracting useful information from a_
_reward model_, and (3) the often bespoke evaluation techniques used for RLHF models, _the process of fairly evaluating_
_a multi-use model_ . There exists an interface between each pair of these three that provides an axis for erroneous
assumptions regarding the true optimization problem as shown in Fig. 1, but the importance of each link is not uniform
for mitigation of mismatch.
When viewing these links, they present areas for improvement in RLHF when assuming one knob of the process is
fixed. For example, in order to study the task of _a reward that enables stable RL training_, one should operate under a
fixed evaluation regime. Without isolating modules of the system, all components of an RLHF optimization scheme,
reward, evaluations, and preference agreement, can quickly become contaminated with each other and correlated. An
example of such a project would be studying reward model design to mitigate overoptimization (Coste et al., 2023;
Ramé et al., 2024), targeting the top right of Fig. 1.
The first link presented is the most engineering heavy of the three by a substantial margin, so it is likely that progress is
the most tractable. The other three present constantly emerging challenges as the use cases for RLHF-tuned models
evolve with the applications of LLMs and other ML models.
**Reward model training** _↔_ **policy model training** Uniformly extracting the information from the reward model
into the policy and avoiding the reward hacking inherent to RL (Pan, Bhatia, & Steinhardt, 2022) that can result in
overoptimization of reward models (Gao et al., 2022) is central to RLHF. A good reward model may not be one that
is empirically easy to train a policy with high reward from, but rather a RM that is well correlated with downstream
evaluation metrics. Common practice in RLHF, especially with larger models where gradients are less stable, is to
spend additional compute in search of “stable” training runs with increasing reward, which induces further likelihood of
mismatch.
**Reward model training** _↔_ **evaluation tools** While relatively little work and resources exist for the study of state-ofthe-art reward models, the matching of the reward signal they provide to the intended use-case of the final policy is
Objective Mismatch in RLHF
Figure 4: Illustrating the most likely visualization of objective mismatch in RLHF, the link between policy training and
downstream evaluation. Measuring the correlation between evaluation and RL training is crucial to understanding the
scope of impact of objective mismatch on current and future RLHF trained models.
central to solving the objective mismatch issue, particularly through the methods used to collect preference data. The
reward models are trained on aggregated datasets to maximize agreement of the model on a held out set of data, which
in practice often yields a maximum accuracy of 60-75% (Bai et al., 2022; Ouyang et al., 2022). Given the complex task
encompassed in reward modeling, it is unlikely that the models converge to 100% accuracy, but studying the sources
of this performance delta could indicate sources of mismatch. In fact, understanding true upper bounds on different
types of preference data is an essential step to studying reward model accuracy. New tools are needed for evaluation of
reward models that better match their conceptual underpinnings as a representation of human values for solving the
alignment problem (Leike et al., 2018) and as a practical realization as targets for optimization Lambert et al. (2023).
**Policy model training** _↔_ **evaluation tools** The third link contributes the least to the emergence of mismatch, but
is the easiest axis to visualization potential signs of objective mismatch. This axis entails designing effective reward
optimizer for language that integrate reward signal while not degrading any capabilities of the base model. Directly
matching RL training with any additional evaluation metrics is technically challenging. In MBRL, such a solution could
be by using a differentiable simulator (R. Wei et al., 2023), but with the complexity of RLHF such solutions are less
desirable. Exploring any types of regularization or calibration of training with respect to final evaluations is viable as
research directions, but this area of study is best suited for visualizing signs of objective mismatch, as shown in Fig. 4.


The original training object used in popular language model architectures, autoregressive next-token prediction also
suffers from an objective mismatch problem, as almost all LLM evaluation techniques evaluate the entire output rather
than individual tokens. While this is true, the development signal that the next-token prediction loss provides is more
orthogonal to the goals of RLHF. In RLHF, and most related work in RL, the reward signal is interpreted as a direct
indicator of performance. This assumption creates a much more unintentionally nuanced research setup, warranting the
specific study of its impacts.
In MBRL, the learning of a dynamics model is also often done via one-step transitions, with recent work studying
autoregressive models (Janner, Li, & Levine, 2021; Lambert, Wilcox, Zhang, Pister, & Calandra, 2021), where the compounding error of multiple one-step predictions is well known as a deeply related issue to objective mismatch (Lambert,
Pister, & Calandra, 2022). In the case where mismatch becomes a fundamental problem of LLMs, similar solutions
could be investigated.


Direct Preference Optimization (DPO) (Rafailov et al., 2023) solves the RLHF problem by inducing a policy from the
optimal solution to the reward model problem, resulting in an LLM that acts as a generative model and reward scorer.
This class of algorithms, which is expanding to address concerns of over-fitting and robustness (Azar et al., 2023),
reduces the complexity of the objective mismatch problem by directly tying the training of the reward model and policy
together. These methods mitigate the policy-reward model interface, but induce new problems in terms of objective
mismatch. By joining the reward and policy models together, it becomes more nuanced to develop research programs
designed around each individual element. In principle a reward model achieved with DPO should be useful in same
Objective Mismatch in RLHF
manners as other RLHF approaches, but substantial research is required to assess them. Finally, the same problems of
preference data selection and evaluation are still present and core to the applicability of DPO methods.

4.3 Does Direct Preference Optimization solve the mismatch?


There is already emerging research on many potential causes and solutions of mismatch in RLHF, yet further work
can be inspired by solutions from the broader RL literature. Many of the solutions to objective mismatch in MBRL
will not apply directly because in MBRL they have a true reward from the environment, and for that reason research
is needed to understand the outputs of reward models. Here follows a series of investigations to expand to mitigate
objective mismatch:
**Reward model evaluation** There are many axes by which a reward model is expected to behave in order to be a
reasonable approximation of a reward function, but they are typically not studied. Reward models need to be assessed
for consistency, robustness to adversarial attacks, calibration across distributions, and more, as discussed in Lambert et
al. (2023). Understanding reward models performance is the foundation of solving the mismatch problem. Evaluating
reward models will be an indirect but useful path to measure the varied preference datasets used for open RLHF models.
**Reward model training methods** In order to solve limitations of reward models across better evaluation techniques,
numerous new training methods will be developed. Early research has already shown reward model ensembles can
help mitigate overoptimization (Coste et al., 2023). Further research is warranted to integrate techniques that have
improved performance of model-based RL algorithms, such as probabilistic loss functions for the dynamics models and
planning (Chua, Calandra, McAllister, & Levine, 2018), calibrated probability estimates (Malik et al., 2019) during
training the reward model as a classifier, and other solutions (R. Wei et al., 2023). Additionally, links should be explored
between the reward models of inverse reinforcement learning (IRL) (Ng, Russell, et al., 2000), the subfield tasked with
learning a reward function from agent behavior, and those of RLHF. Early research also shows reformatting the reward
model training to better match preference learning literature (Knox et al., 2023) could improve performance. While
ensembles (Coste et al., 2023) and weight-averages (Ramé et al., 2024) mitigate overoptimization, they do not solve all
challenges facing reward models (Eisenstein et al., 2023).
**Reward model training datasets** High-quality datasets are a bottleneck slowing progress in open RLHF research,
given the large costs required to acquire them. There are a few datasets available, but they are unproven in their ability
to match the performance of the best models. The Stanford Preferences Dataset of Reddit content (Ethayarajh, Choi, &
Swayamdipta, 2022), UltraFeedback synthetic preference data (Cui et al., 2023), WebGPT internet browsing (Nakano
et al., 2021), learning to summarize (Stiennon et al., 2020), and Anthropic HHH dataset (Askell et al., 2021) serve as a
strong foundation for research. Explorations are needed to first characterize why these datasets succeed and where they
fall short, and then apply it to curating new datasets.
**Value-guided sampling techniques** Increased compute can be spent at inference time to improve the performance of
RLHF models by utilizing the values returned by the reward model (Deng & Raffel, 2023; Liu, Cohen, et al., 2023).
Feng et al. (2023) explores this through Monte Carlo tree search generation, yet many more methods can be explored
across the planning literature.
**Human-centric NLP evaluation** The most popular evaluation technique for chat-tuned RLHF models is preference
percentage versus other top models on evaluation prompt sets (as done in open RLHF models including Llama
2 (Touvron et al., 2023) and Dromedary-2 (Sun et al., 2023)). This evaluation mechanism, while well-motivated in the
popular use-cases of the models, suffers from bias and reproducibility challenges. The prompts can easily be chosen to
support the model designed by the authors, and the prompts are often not released or aggregated into a future benchmark.
Expanding the reproducibility and consistency of these practices will be important to creating robust practices for
RLHF.
**RL (and other) optimizers for language** As discussed in Sec. 2.1, the optimizers used for RLHF are most often
those from previous RL literature. Now there is an opportunity for expansion of RL algorithms into the niche of RLHF,
where conditions are highly specialized through the expansive actions space and bandits formulation. New algorithms
are a step in the right direction, such as T. Wu et al. (2023) modifying the PPO algorithm for pairwise preferences or
Baheti et al. (2023) proposing an offline RL algorithm for full-completion actions.
This investigation should compare to other baselines for extracting signal from a reward model, such as rejection
sampling (Touvron et al., 2023), which runs autoregressive fine-tuning on the top samples as dictated by a reward
model.
Objective Mismatch in RLHF
Figure 5: The outer loop of RLHF is the process to train the next reward model for RL to improve performance is areas
of user interest. This setup induces additional complexity in objective mismatch in RLHF.
**Other solutions** Other creative mismatch solutions will exist, such as work integrating the LLM policy, reward model,
and transition function into a single model (Xu, Dong, Arumugam, & Van Roy, 2023). Methods such as this need to be
evaluated across many scales to confirm that they are still numerically stable with the larger state-of-the-art models
where powerful emergent behaviors exist.


**Iterative deployment of RLHF** The iterative deployment form of RLHF where reward models are retrained based on
user data, which induces a second feedback loop, is shown in Fig. 5. Schulman (2023) discusses how this is used in
ChatGPT to mitigate issues such as evasiveness, verbosity, and other unexpected, undesirable qualities. Designing in
this framework introduces further complexity onto engineering objectives, but allows iterative mitigation of mismatch.
This style of iterative RL deployment has been understood as exogenous feedback (Gilbert, Dean, Zick, & Lambert,
2022) and can have societal implications.
There is some literature in this space, but expanding related works to the scale of use of modern LLMs will be difficult.
For example, Suhr and Artzi (2022) shows theoretical results on outer-loop optimization of instruction-tuned models.
**Contextual bandits** The modifications made to the RL optimization of RLHF cast it as a contextual bandits problem,
where an agent takes one action and the dynamics are abstracted into one trajectory-reward pairing. Work in this area has
investigated the potential of integrating partial, skewed, or noisy human feedback into the optimization process (Nguyen,
Daumé III, & Boyd-Graber, 2017).
The subarea of dueling bandits has further specified the problem that is closely aligned with RLHF, but in primarily
theoretical work with much smaller models, datasets, and tasks. Yue, Broder, Kleinberg, and Joachims (2012) explains
this space in work showing theoretical bounds:
“In contrast to conventional approaches that require the absolute reward of the chosen strategy to be
quantifiable and observable, our setting assumes only that (noisy) binary feedback about the relative
reward of two chosen strategies is available. This type of relative feedback is particularly appropriate
in applications where absolute rewards have no natural scale or are difficult to measure... but where
pairwise comparisons are easy to make.”
This, while closely related to RLHF, will require substantial experimentation to be applicable. Others have built on this
into work directly learning from human preferences (Sekhari, Sridharan, Sun, & Wu, 2023) or from implicit human
feedback (Maghakian et al., 2022).


This paper presents the multiple ways by which objective mismatch limits the accessibility and reliability of RLHF
methods. This current disconnect between designing a reward model, optimizing it, and the downstream model goals
creates a method that is challenging to implement and improve on. Future work mitigating mismatch and the proxy
Objective Mismatch in RLHF
objectives present in RLHF, LLMs and other popular machine learning methods will become easier to align with human
values and goals, solving many common challenges users encounter with state-of-the-art LLMs.
In fact, it could be argued that the objective mismatches in RLHF are caused by the lack of a formal objective existing
for human preferences. Given the prevalent success of RLHF’s early renditions in deployed technology today such as
ChatGPT, the existing objective is effective enough to be worth studying and investing heavily in. Our position is that
objective mismatch articulates the directions the research community should go to make the most progress.


Carnegie Mellon University
tachungc@andrew.cmu.edu

Attention Alignment and Flexible Positional Embeddings Improve
  Transformer Length Extrapolation

Ta-Chung Chi
Ting-Han Fan
Alexander I. Rudnicky


Independent Researcher
tinghanf@alumni.princeton.edu


An ideal length-extrapolatable Transformer language model can handle sequences longer than
the training length without any fine-tuning.
Such long-context utilization capability relies
heavily on a flexible positional embedding design. Upon investigating the flexibility of existing large pre-trained Transformer language
models, we find that the T5 family deserves a
closer look, as its positional embeddings capture rich and flexible attention patterns. However, T5 suffers from the dispersed attention
issue: the longer the input sequence, the flatter
the attention distribution. To alleviate the issue,
we propose two attention alignment strategies
via temperature scaling. Our findings show
improvement on the long-context utilization capability of T5 on language modeling, retrieval,
multi-document question answering, and code
completion tasks without any fine-tuning. This
suggests that a flexible positional embedding
design and attention alignment can go a long
way toward Transformer length extrapolation. [1]


Pre-training large Transformer language models
on long sequences is inherently expensive due to
self-attention’s quadratic complexity w.r.t the input sequence length (Vaswani et al., 2017). Even
with the help of memory-efficient attention (Rabe
and Staats, 2021; Dao et al., 2022), the maximum
supported input length of current open-source pretrained Transformer language models are capped at
4,096 tokens (Touvron et al., 2023), limiting their
efficacy in handling long-context tasks.
One notable research topic aiming to lift
the input length restriction is _Length Extrapo-_
_lation_ (Press et al., 2022). Ideally, a lengthextrapolatable Transformer language model is
trained on short sequences and can perform equally


Topic Line Passkey
Criteria
512 15k 512 15k 512 15k
Pmax 0.28 0.12 0.27 0.11 0.32 0.24
H 3.47 6.63 3.47 7.04 3.09 5.97
Table 1: **The Dispersed Attention Issue of Flan-T5-**
**XL Encoder.** Pmax is the average maximum probability and H is the average entropy. After increasing the
sequence length from 512 to 15k, we observe larger
entropy and smaller maximum probability, implying a
flatter self-attention distribution.
well on longer ones without any further fine-tuning.
This is made possible with carefully designed positional embeddings (Press et al., 2022; Chi et al.,
2022, 2023). Unfortunately, existing approaches
are tailored for natural language modeling, a task
known to have strong recency bias, and they often do not perform well on other seemingly simple
tasks such as passkey, topic, and line retrieval (Mohtashami and Jaggi, 2023; Li et al., 2023).
To circumvent the recency bias, we sift through
the positional embeddings of existing open-source
large pre-trained Transformer language models,
shown in Table 2, to find a flexible design, and
the T5 family (Raffel et al., 2020) comes to our attention. As visualized in Figure 1, the flexibility of
T5’s positional embeddings allows it to encourage
recency bias on one head and discourage that on
another head. However, there is no free lunch: T5
suffers from the dispersed attention issue as shown
in Table 1. That is, the attention distributions of
long input sequences tend to be flatter than those
of short input sequences. As a remedy, we propose
two fine-tuning-free attention alignment strategies
via Softmax temperature scaling (Yao et al., 2021;
Su, 2021) to mitigate the dispersed attention issue: maximum probability (Pmax) and entropy (H)
alignment.
We validate the effectiveness of our alignment
[1https://github.com/chijames/](https://github.com/chijames/Attention-Alignment-Transformer-Length-Extrapolation)
[Attention-Alignment-Transformer-Length-Extrapolation](https://github.com/chijames/Attention-Alignment-Transformer-Length-Extrapolation)
**Models** T5 (2020) OPT (2022) ChatGLM (2022) LLaMA (2023) Falcon (2023) Pythia (2023) XGen (2023) BLOOM (2022) MPT (2023)
Learned Learned Rotary Rotary Rotary Rotary Rotary ALiBi ALiBi
**PE.**
Relative Absolute Relative Relative Relative Relative Relative Relative Relative
Table 2: **Open-source Transformer language models and their positional embeddings.** T5 is the only model
equipped with learnable relative positional embeddings, which enable its long-context utilization capability.
strategies on tasks including language modeling,
retrieval, multi-document question answering, and
code completion. We also provide a theoretical
analysis of how the alignment strategies work under
the hood by investigating the relation between the
Softmax temperature and data distribution.


Transformer-based models rely on positional
embeddings to encode positional information.
We summarize open-source large pre-trained
Transformer language models and their positional
embeddings in Table 2. The relative variants
are widely adopted due to their better empirical
performance (Su et al., 2021) and possible
length-extrapolation capability (Press et al., 2022).
In this work, we place special focus on the T5
positional embeddings due to their flexibility as
shown in Figure 1.
**Transformer Length Extrapolation** Existing
research on Transformer length extrapolation is
mostly confined to the task of natural language
modeling (Press et al., 2022; Chi et al., 2022, 2023).
Unfortunately, the reported positive results do not
carry over to long-context retrieval (Mohtashami
and Jaggi, 2023; Li et al., 2023). This contrastive
observation can be explained by models’ short empirical receptive field (Chi et al., 2023). In short,
the strong decaying prior of positional embeddings
prevents models from accessing distant tokens that
may be necessary for retrieval tasks. In this work,
we improve the flexible positional embeddings of
T5 to get around this limitation.
**Transformer Position Interpolation** Instead of
performing direct length extrapolation, a different
line of research conducts model fine-tuning on long
input sequences (Chen et al., 2023), where the main
focus is to identify the most efficient fine-tuning
scheme that can improve long-context utilization.
Positive results have been reported on retrieval
tasks (Li et al., 2023). However, we argue that
fine-tuning incurs additional costs since it needs
1) GPU resources to perform long sequence finetuning with large models and 2) a pre-defined target
sequence length, which still imposes a sequence
length upper limit. Our proposed methods can circumvent these two limitations.


Transformer-based approaches often consist
of a retriever and a reader to overcome the context
length restriction (Guu et al., 2020; Lewis et al.,
2020; Izacard and Grave, 2021; Borgeaud et al.,
2022). The retriever retrieves relevant text snippets
from a very large database and the reader digests
the retrieved information to generate the correct
output. Our proposed attention alignment strategy
can be used to significantly increase the input
sequence length of the reader, thereby allowing
more retrieved information to participate in
the decision process. For small-scale retrieval
problems, our methods even obviate the need for
context segmentation and the external key-value
store used in prior work (Mohtashami and Jaggi,
2023), serving as a more elegant approach.
**Softmax Temperature Scaling** To increase the
length extrapolation capability of Transformers,
previous work (Yao et al., 2021; Su, 2021) scales
the temperature of Softmax logarithmically w.r.t
the sequence length to ensure invariant entropy.
Our entropy alignment strategy is also inspired by
this line of research except that we adopt a different
procedure outlined below in Algorithm 1. Interestingly, our results in Section7 show that the logarithmic
temperature scaling scheme is more similar to our
proposed maximum probability alignment strategy.


As suggested by recent work (Mohtashami and
Jaggi, 2023; Li et al., 2023), the task of longcontext retrieval serves as a controllable benchmark to measure how well a Transformer language
model utilizes long-context inputs. One prominent
characteristic of retrieval tasks is that only a subset
of the input is of interest, requiring a model to accu
(a) 1 [st] Attention Head (b) 27 [nd] Attention Head
Figure 1: **Visualization of T5 Positional Embeddings.** To plot figures of _bm,n_, we set _m_ = 7500 and vary the
value of _n_ from 0 to 15k. Each attention head of a Flan-T5-XL encoder learns a set of positional embeddings that
capture different attention bias. For example, the positional embeddings in the left figure encourage the model to
focus on nearby tokens. In contrast, the ones in the right figure let the model focus on only remote tokens.
rately pick up the necessary information. The other
characteristic is that the key information can sit
anywhere in an input, requiring a model to attend
flexibly. Finally, the controllable aspect allows us
to gradually increase the input sequence length to
test the models’ length extrapolation capability.


To solve retrieval tasks using Transformer language
models, it is necessary to choose a positional embedding design that permits accurate and flexible
length-extrapolatable attention. After checking
through the existing positional embeddings in Table 2, we find that the T5 family (Raffel et al., 2020)
fits our needs. As for other candidates, learnable absolute positional embeddings (Vaswani et al., 2017;
Zhang et al., 2022) must be evaluated within the
training length. ALiBi (Press et al., 2022) and Rotary (Su et al., 2021) have a recency bias; they
cannot extrapolate easily without fine-tuning.
For each attention head, T5 encoder maintains a
bucket (B) of 32 learnable parameters and assigns
the relative positional bias (rpe bias) _bm,n_ as [2]
_bm,n_ =
B[
B[
B[ _m −_ _n_ ] _,_ if 0 _≤_ _m −_ _n <_ 8
B[ _n −_ _m_ + 16] _,_ if _−_ 8 _< m −_ _n <_ 0
B[min(15 _,_ 8 + _⌊_ [log((] _[m][−][n]_ [)] _[/]_ [8)] _·_ 8 _⌋_ )] _,_ if
B[min(15 _,_ 8 + _⌊_ log(128 _/_ 8) _·_ 8 _⌋_ )] _,_ if 8 _≤_ _m −_ _n_
B[min(31 _,_ 24 + _⌊_ [log((] _[n][−][m]_ [)] _[/]_ [8)] _·_ 8 _⌋_ )] _,_ if _m −_ _n_

log(128 _/_ 8) _·_ 8 _⌋_ )] _,_ if _m −_ _n ≤−_ 8 _,_
th entry of the _L × L_ self-attention matrix. The
summation becomes the input to the temperaturescaled Softmax. We plot the learned rpe bias of a
T5 encoder in Figure 1. We can tell that its attention
heads encode rich attention patterns. For example,
head 1 learns to focus on the nearby tokens whereas
head 27 learns to ignore the nearby tokens and
allow access to faraway tokens.


Unfortunately, directly applying T5 models on retrieval tasks does not yield perfect results. Upon
inspecting the intermediate model states, we find
that a longer input sequence consists of more tokens competing for the same amount (i.e., Softmax
sums to 1) of attention, resulting in the dispersed attention issue. In Table 1, we see that the longer the
input sequence, the flatter the self-attention distribution. The situation is not hopeless if the desired
information still attains a higher attention weight
than the remaining tokens. Our proposed solution
in Section4 will let the key information stand out.


A natural solution to the dispersed attention issue
described in Section3 is to sharpen the self-attention distribution. This can be achieved by reducing the
temperature _τ_ during extrapolation. We set the extrapolation temperature _τex_ such that the sharpness
during training with _τtr_ = 1 and that during extrapolation with _τex <_ 1 are roughly the same. As a
measurement of sharpness, we explore the maxi
where 0 _≤_ _m < L_ and 0 _≤_ _n < L_ are two
position indices. _bm,n_ will be added to the ( _m, n_ )
[2https://github.com/huggingface/transformers/](https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/t5/modeling_t5.py#L390)
[blob/v4.33.2/src/transformers/models/t5/](https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/t5/modeling_t5.py#L390)
[modeling_t5.py#L390](https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/t5/modeling_t5.py#L390)
**Algorithm 1** Attention Alignment Strategies
**Require:** A short sequence of length _Ltr_ and a long sequence
of length _Lex > Ltr_ . Encoder _E_ . Alignment mode _M_ .
**Ensure:** The Softmax temperature _τ_
**function** FINDS( _τ_, _M_ )
Set temperature of all Softmax to _τ_
_s ←_ [ ]


**return** _τex_ s.t. _S_ _[ex]_ ( _τex_ ) _≈_ _S_ _[tr]_ (1)
mum probability or entropy of a distribution. In
other words, our proposed solution is to align the
maximum probability or entropy of training and
extrapolation distributions by adjusting _τex_ .
Concretely, let _l_ [(] _[i]_ [)] _∈_ R _[L]_ be the i-th pre-Softmax
logit vector of a T5 encoder, where _L ∈{Ltr, Lex}_
is the sequence length. The post-Softmax distribution of _l_ [(] _[i]_ [)] is P [(] _[i]_ [)] ( _τ_ ) = Softmax _τ_ ( _l_ [(] _[i]_ [)] ). The
maximum probability and entropy of P [(] _[i]_ [)] ( _τ_ ) are
P [(] max _[i]_ [)] [(] _[τ]_ [)][ and][ H][(] _[i]_ [)][(] _[τ]_ [)][, respectively.]
Take the maximum probability alignment strategy as an example: We first run the forward
pass and compute the average maximum probability under temperature _τ_ over all logit vectors:
Pmax( _τ_ ) = (1 _/N_ ) [�] _i_ [P] max [(] _[i]_ [)] [(] _[τ]_ [)][ where] _[ N]_ [ =] _[ R][ ×]_
_H_ _×L_ is the number of logit vectors in a T5 encoder
with _R_ layers, _H_ heads, and length- _L_ sequences.
Since the temperature is 1 during training and _τex_
during extrapolation, we denote the average maximum probability during training as P _[tr]_ max [(1)][ and]
that during extrapolation as P ~~_[ex]_~~ max [(] _[τ][ex]_ [)][. Finally, to]
align the maximum probabilities, we adjust _τex_ s.t.
P ~~_[ex]_~~ max [(] _[τ][ex]_ [)] _[ ≈]_ [P] _[tr]_ max [(1)][. In practice, we do a grid]
search on _τex_ from 1.0 to 0.5. We outline the procedure of the alignment strategies in Algorithm 1.
Note that our proposed methods do not require
any model fine-tuning or gradient computations.
The only overhead is estimating the temperature
_τex_ using Algorithm 1 and a few length _Lex_ sequences. Once the temperature is decided, it will
be held fixed, rendering our methods simple and
efficient. In addition, our fine-tuning free methods do not lead to performance regression on short
_Ltr_ sequences commonly observed on long-context
fine-tuned models (Roziere et al., 2023).


We compare the two alignment strategies against
the length-only Softmax temperature scaling
scheme _τ_ = log _Lex Ltr_ (Yao et al., 2021; Su, 2021)
and LongChat-13B-16K (Li et al., 2023). Note that
LongChat-13B-16K (Li et al., 2023), the best baseline, was fine-tuned from LLaMA (Touvron et al.,
2023) on long sequences of length 16k while our
proposed methods do not need any fine-tuning.


Models Sequence Length ( _Lex_ )
1024 2048 4096 8192 15000 Avg.
T5-Large-LM 35.9 40.1 >1k >1k >1k  - 1k
w/ Pmax 34.7 45.5 45.2 45.5 52.7 **44.7**
w/ H 40.2 43.9 45.6 54.6 56.0 48.1
w/ log _Lex Ltr_ 39.8 38.2 47.4 45.3 55.9 45.3
T5-XL-LM 28.3 >1k >1k >1k >1k  - 1k
w/ Pmax 30.2 36.0 31.6 41.7 50.0 37.9
w/ H 30.4 36.8 38.4 53.3 63.4 44.4
w/ log _Lex Ltr_ 27.3 29.4 31.7 39.3 45.8 **34.7**
T5-XXL-LM 109 >1k >1k >1k >1k  - 1k
w/ Pmax 32.2 29.7 29.5 36.6 44.3 34.5
w/ H 26.8 28.1 34.2 37.8 43.8 **34.1**
w/ log _Lex Ltr_ 27.1 36.1 33.9 246 43.8 77.5
Table 3: **Language Modeling Performance.** We report
the average perplexity of 500 sequences. The lower the
better.


We use the LM-Adapted T5 models for this experiment [3] . We set _Ltr_ = 512. Following previous
work on Transformer length extrapolation, we perform an intrinsic evaluation on language modeling (Press et al., 2022; Chi et al., 2022, 2023). Ideally, our proposed methods should alleviate the perplexity explosion problem during extrapolation. As
we can see in Table 3, both alignment strategies dramatically improve (lower) the perplexity. We also
observe that scaling the temperature solely based
on sequence lengths is not the optimal strategy, as
indicated by the sudden perplexity increase of the
log _Lex Ltr_ strategy. We will provide an in-depth
discussion on this topic in Section7. Note that perplexity
[3https://github.com/google-research/](https://github.com/google-research/text-to-text-transfer-transformer/blob/main/released_checkpoints.md#lm-adapted-t511lm100k)
[text-to-text-transfer-transformer/blob/main/](https://github.com/google-research/text-to-text-transfer-transformer/blob/main/released_checkpoints.md#lm-adapted-t511lm100k)
[released_checkpoints.md#lm-adapted-t511lm100k](https://github.com/google-research/text-to-text-transfer-transformer/blob/main/released_checkpoints.md#lm-adapted-t511lm100k)


Topic, # of topics Line, # of lines Passkey, # of sentences
Models Avg.
5 10 15 20 25 200 300 400 500 600 680 20k 30k 40k 50k 55k
Flan-T5-Large 99 100 97 97 83 97 100 92 96 93 92 62 47 31 16 9 76
w/ Pmax 96 90 86 94 98 99 98 98 98 98 100 84 90 85 79 85 **92**
w/ H 59 32 16 0 3 97 90 94 83 93 88 29 25 21 15 22 48
w/ log _Lex Ltr_ 88 79 75 61 55 99 99 98 99 97 98 74 63 51 41 35 76
Flan-T5-XL 100 100 100 100 100 96 90 77 57 45 26 100 100 100 100 100 87
w/ Pmax 100 100 100 100 100 97 90 89 80 70 62 100 99 100 100 100 **93**
w/ H 99 98 97 96 96 95 87 88 79 70 71 100 100 100 100 100 92
w/ log _Lex Ltr_ 99 100 100 100 100 98 88 81 86 60 67 100 100 100 100 99 92
Flan-T5-XXL 100 100 100 99 99 100 100 98 95 84 82 100 100 100 100 100 97
w/ Pmax 100 100 100 99 99 97 99 96 97 94 95 100 98 100 100 100 **98**
w/ H 100 100 97 98 94 99 92 92 76 58 58 100 100 100 100 100 92
w/ log _Lex Ltr_ 100 100 99 98 92 100 98 94 93 84 90 100 100 100 100 100 97
LongChat 100 100 100 99 89 100 91 93 83 78 59 100 100 99 100 99 93
Table 4: **Performance of Retrieval Tasks.** The numbers are accuracy. Full score is 100. The LongChat model
corresponds to the LongChat-13B-16K model (Li et al., 2023). It is a LLaMA-13B (Touvron et al., 2023) model
fine-tuned on sequences of length 16k using the positional interpolation technique (Chen et al., 2023). Flan-T5-XXL
has 11B parameters. The maximum sequence lengths ( _Lex_ ) of the three tasks are all around 14.5k to 15.5k tokens.
is not our primary focus since it often cannot accurately reflect the long-context utilization capability
of Transformers on practical tasks (Li et al., 2023).


The tasks are formulated in the Question Answering (QA) format; therefore, we use the Flan-T5
models to leverage their instruction-following capability. We set _Ltr_ = 512. Inspired by recently
proposed retrieval tasks, we evaluate the proposed
alignment strategies on three of these. Topic retrieval requires a model to retrieve the first topic
in a long and multi-topic conversation (Li et al.,
2023). Line retrieval has a long series of key-value
pairs, and a model needs to retrieve the value corresponding to the given key (Li et al., 2023). Passkey
retrieval hides a passkey in a long junk text snippet, and a model needs to return that passkey (Mohtashami and Jaggi, 2023).
As we can see in Table 4, the retrieval performance is greatly boosted after the Flan-T5 models
are equipped with our proposed attention alignment
strategies. In particular, the maximum probability
alignment strategy provides better results across the
board. Other baselines such as MPT (Team, 2023)
and ChatGLM2 (Du et al., 2022) perform worse
than LongChat. Please refer to Li et al. (2023)
for more details. We also present the optimal temperature given by Algorithm 1 in Table 9 in Appendix A.5. In short, the temperature decreases
when the input sequence length increases. We will
provide additional temperature analysis below, in
Section7.


We again use the Flan-T5 models to leverage their
instruction-following capability. We set _Ltr_ = 512.
We choose the multi-document question-answering
task as our downstream task (Liu et al., 2023).
The model input consists of a question Q and
multiple documents extracted from NaturalQuestions (Kwiatkowski et al., 2019) related to Q, where
one of the documents (golden doc) contains the
ground truth answer to Q. As shown in Table 5,
when a model is equipped with the proposed maximum probability alignment strategy, it consistently
outperforms the original model across model sizes
and number of input documents.
Apart from the better task performance, we believe that the attention dispersed attention issue
discussed in Section3 can help demystify the lost-in-themiddle phenomenon (Liu et al., 2023) of this task:
Transformer models tend to perform worse when
the ground truth sits near the middle of the input
context. Let us recall the relative positional embedding of head 27 learned in Figure 1, if the ground
truth answer sits in the middle, it will have long
contexts from both sides competing for the attention weight. If this hypothesis is correct, we can
expect the performance boost to be more prominent when the answer appears near the middle. We
reveal the performance breakdown when the num


10 Docs 20 Docs 30 Docs, golden doc at different positions
Models
Avg. Avg. 0 4 9 14 19 24 29 Avg.
Flan-T5-Large 52.4 43.3 52.6 42.0 36.5 34.0 33.9 33.9 37.9 38.7
w/ Pmax **53.1** **44.2** 50.8 44.5 39.5 36.4 35.9 35.8 37.0 **40.0**
Improvement +0.7 +0.9 -1.8 +1.5 +3.0 +2.4 +2.0 +1.9 -0.9 +1.3
w/ H 52.1 43.2 47.6 41.1 35.2 33.5 32.2 33.3 34.2 36.7
w/ log _Lex Ltr_ 53.2 44.5 50.6 44.1 39.3 36.3 35.8 35.8 37.2 39.9
Flan-T5-XL 59.4 51.2 58.4 44.6 40.0 39.9 41.7 46.4 54.8 46.5
w/ Pmax **61.1** **53.6** 60.9 49.1 46.0 44.9 46.3 49.1 55.7 **50.3**
Improvement +1.7 +2.4 +2.5 +4.5 +6.0 +5.0 +4.6 +2.7 +0.9 +3.8
w/ H 60.5 52.4 52.4 43.5 42.1 40.3 42.0 42.9 51.3 44.9
w/ log _Lex Ltr_ 60.9 **53.6** 61.0 49.1 46.1 44.7 46.1 48.7 55.4 50.2
Flan-T5-XXL 63.6 56.9 58.9 49.1 48.1 47.5 48.9 53.1 61.2 52.4
w/ Pmax 63.7 **57.7** 60.4 52.5 51.0 50.2 51.3 53.5 59.1 **54.0**
Improvement +0.1 +0.8 +1.5 +3.4 +2.9 +2.7 +2.4 +0.4 -2.1 +1.6
w/ H 63.6 57.1 61.0 53.4 50.8 50.3 50.7 51.9 55.7 53.4
w/ log _Lex Ltr_ **63.9** 57.6 61.5 53.3 51.3 50.3 51.1 53.0 57.2 **54.0**
Table 5: **Performance of Multi-document QA.** Numbers are accuracy. Full score is 100. The maximum sequence
length ( _Lex_ ) of 30 documents is around 5k. The improvement row represents the absolute accuracy improvement
after a Flan-T5 model is equipped with our proposed maximum probability alignment strategy. For the full
performance breakdown, please refer to Table 14 in Appendix A.7.
ber of input documents is 30. As we can see in
the improvement row, those cases indeed achieve
greater improvements.
Our strategies are not always perfect: The performance drops if the ground truth answer is at
position 29. We believe T5 might have already handled this case pretty well due to the recency bias
learned on some attention heads, and our additional
temperature scaling sharpens the distribution too
aggressively. We acknowledge this trade-off in Section9.


To test the generalizability of the alignment strategies, we apply our methods to the CodeT5+
model (Wang et al., 2023) that was pre-trained
on code data with 770M parameters. [4] We set
_Ltr_ = 768. We do not experiment with larger
CodeT5+ models since they do not follow the
T5 architecture, but use other positional embeddings. We conduct two experiments on the LCC
dataset (Guo et al., 2023), which is highly similar
to the classic PY150 dataset (Raychev et al., 2016)
except that the input context length is much longer.
For the code key retrieval experiment, we sample
several code files from LCC along with a special
function that only returns an integer from 1 to 100.
We concatenate them and ask a model to gener
[4https://huggingface.co/Salesforce/](https://huggingface.co/Salesforce/codet5p-770m-py)
[codet5p-770m-py](https://huggingface.co/Salesforce/codet5p-770m-py)
ate the returned integer at the end (Roziere et al.,
2023). Considering that this is essentially a passkey
retrieval task in the code domain, we briefly report
the average accuracy of 100 test cases when the
input sequence length is around 16k: 0 (Original
CodeT5+), **87** (w/ Pmax), 80 (w/ H), and 85 (w/
log _Lex Ltr_ ). We can see that the maximum probability alignment strategy performs the best.
For the code completion experiment, a model
needs to generate the next line of code given some
prior code as the context. The metrics are Exact
Match (EM) and Edit similarity (ES) on a per line
basis (Svyatkovskiy et al., 2020). We report the results in Table 6 using the context length bucketing
format. While both alignment strategies improve
the performance substantially, Pmax is better; however, its EM performance lags behind log _Lex Ltr_
when the sequence length increases. We additionally include an extrapolation-free baseline, _trunca-_
_tion_, that truncates the long input context to the
most recent _Ltr_ = 768 tokens. Both Pmax and
log _Lex Ltr_ perform better than this baseline when
_Lex <_ 6000, indicating that they can indeed benefit
from longer (6000 _/_ 768 = 7 _._ 8x) contexts without
any fine-tuning.


First, the maximum probability alignment strategy
is the most reliable and best-performing method


Models Sequence Length ( _Lex_ )
1k 2k 3k 4k 5k 6k
CodeT5+ 19.6 19.0 11.3 2.6 0.1 0.0
w/ Pmax 21.1 22.5 21.7 21.5 19.3 22.7
w/ H 19.5 18.7 13.7 9.0 7.9 9.0
w/ log _Lex Ltr_ **21.6 23.0 22.1 22.0 20.6 24.3**
w/ _truncation_ 20.0 19.2 19.3 19.2 17.1 21.4


Models Sequence Length ( _Lex_ )
1k 2k 3k 4k 5k 6k
CodeT5+ 62.4 59.6 53.1 38.9 18.3 10.4
w/ Pmax 65.9 65.7 65.3 65.6 **63.1** 64.9
w/ H 64.8 62.5 54.1 43.0 43.0 44.8
w/ log _Lex Ltr_ **66.3 66.1 65.2 66.4** 63.0 66.1
w/ _truncation_ 65.3 64.2 64.2 65.6 62.2 **66.9**
Table 6: **Code Completion Performance.** Full score is 100. We set _Ltr_ = 768. The bucket _n_ k contains the data
with length in [ _n_ k, ( _n_ +1)k), _n ∈_ [1 _,_ 6]. For example, the bucket 3k contains data with length in [3000, 4000). See
Table 12 and 13 in Appendix A.6 for the full performance breakdown of _Lex_ up to 16k tokens.
across most tasks and settings, echoing our discussion in Section3.1: For most data, only a subset of the
input is useful for a model process at a time. The
maximum probability alignment strategy captures
this characteristic naturally, thereby outperforming the entropy alignment strategy that cares more
about the holistic distribution.
Second, deciding the optimal temperature solely
based on sequence lengths, e.g. _τ_ = log _Lex Ltr_,
is not robust enough. For example, the perplexity
of log _Lex Ltr_ suddenly increases (worse) on T5XXL-LM, in Table 3, while the other strategies
maintain stable results. As another example, it
fails to improve the retrieval performance on the
Flan-T5-Large model, shown in Table 4.


To shed more light on the underlying mechanisms
of the two alignment strategies, we establish the
connection between the softmax temperature _τ_ and
data distribution under empirically verified assumptions. We focus on the 0-th layer (closest to the
input embeddings) and take the average over all
logit vectors across attention heads. Note that this
is just a crude approximation of Algorithm 1 for
analysis purposes since 1) a Transformer language
model typically encompasses multiple layers, and
2) in Algorithm 1, we take the maximum probability or entropy of individual logit vectors as opposed
to the average one.
**Assumption 1.** _The length L average logit vector_
_is normally distributed, i.e., its entry li ∼_ _N_ (0 _, σ_ [2] ) _._
To compute the _average logit vector_, we start
with a input sequence of length _L_ . Using a Transformer model with _H_ attention heads (specifically,
a T5 Encoder in our context), we generate _H × L_
pre-softmax logit vectors, each with a length of _L_ .
Here, the number of layers is 1 because we focus
on the 0-th layer. These logit vectors are then individually sorted, and we subsequently calculate the
average of all _H × L_ sorted logit vectors, resulting
in the average logit vector of length _L_ .
To assess whether the average logit entries follow a Gaussian distribution, we make use of QQ
plots, as illustrated in Figure 2. The linearity of the
plot serves as an indicator – the closer the points
are to the identity line, the more Gaussian the distribution.
Figure 2: **QQ plots of Flan-T5-XL.** We experiment
with short and long sequences. The red reference line
is y=x. We use sequences of length around 512 for this
plot. The plot for sequences of length around 15k looks
highly similar. Please refer to Appendix A.1 for details.
**Assumption 2.** _The largest logit entry of the aver-_
_age logit vector during training and extrapolation_
_is the same: l_ max _[ex]_ [=] _[ l]_ max _[tr]_ _[. See Table][ 7][.]_


Topic Line Passkey
Criteria
512 15k 512 15k 512 15k
_l_ max 8.61 8.80 8.71 8.97 8.75 8.85
Table 7: **Largest Logit Entry of Flan-T5-XL.** _l_ max is
the largest logit entry of the average logit vector.


**Proposition 1.** _Under Assumption 1 and 2, we_
_can adjust the temperature τ to align the maximum_
_probability_ P _[tr]_ max [= P] max _[ex]_
log _Ltr_ + log P _[tr]_ max [+] _[ σ]_ _tr_ [2] _[/]_ [2]
_τ ≈_
log _Lex_ + log P _[tr]_ max [+] _[ σ]_ _ex_ [2] _[/]_ [(2] _[τ]_ [ 2][)] _[.]_
_Bτ_ [2]
=
_[C]_ _Aτ_ [2] +
_τ_ [2]
_B_
=
_A_ + _[C]_
_Aτ_ [2] + _C_ _[.]_
_Assuming τ ̸_ = 0 _, we solve the quadratic equation_
_Aτ_ [2] _−Bτ_ + _C_ = 0 _to get τ_ _. We pick the larger root_
_as our final solution. See proof in Appendix A.2._


**Proposition 2.** _Under Assumption 1, we can adjust_
_the temperature τ to align the entropy Htr_ = _Hex_
_σex_
_τ ≈_
~~�~~
_σtr_ [2] [+ 2 log] _[L]_ _L_ _[ex]_ _tr_
_Ltr_
_See proof in Appendix A.3._
Figure 3: **Language Modeling Temperature Analysis.**
Curves are from Proposition 1 & 2. Dots and crosses
are from Algorithm 1.


The goal of this section is to explain the observations made in Section5 via the lens of temperature analysis. We visualize Proposition 1 and 2 by plotting
the temperature curves in Figure 3 and 4. We evaluate P _[tr]_ max [and] _[ σ][tr]_ [at the training length and] _[ σ][ex]_
at every extrapolation length considering only the
0-th layer. You may find the temperature curves for
the other tasks in Appendix A.4.
Figure 4: **Topic Retrieval Temperature Analysis.**
Curves are from Proposition 1 & 2. Dots and crosses
are from Algorithm 1.
First, while both proposed strategies lower the
temperature when the input sequence length increases, the entropy alignment strategy does so
more aggressively, possibly leading to its inferior
performance observed in Table 4 and 5 (w/ H). This
can be seen by comparing the curves from Propositions 1 and 2 or dots from Algorithm 1.
Second, deciding the optimal temperature based
on sequence lengths, e.g. _τ_ = log _Lex Ltr_, is not the
most robust method. It gives too high of a temperature in Figure 3 compared to Algorithm 1. In other
words, it does not sharpen the distribution enough,
possibly explaining its perplexity spike in Table 3.
On the other hand, it overly lowers the temperature
in Figure 4, thereby failing to improve the retrieval
performance on Flan-T5-Large in Table 4.


In this paper, we show that the T5 model family has
great potential when it comes to Transformer length
extrapolation. We propose the maximum probability and entropy alignment strategies to fix T5’s
dispersed attention issue without model fine-tuning.
We conduct experiments on natural language modeling, retrieval, multi-document question answering, and code completion tasks to demonstrate the
effectiveness of our proposed methods. Finally, we
present a simplified theoretical analysis to elucidate
how the temperature is scaled to achieve attention
alignment. We hope that our work can inspire future length-extrapolatable Transformer designs.


We base our theoretical analysis on a simplified
Transformer language model, which might be further improved by taking all the layers and their interactions into account. In addition, we find that different layers have different degrees of distribution
flatness, which could be leveraged in future work to
perform per-layer fine-grained attention alignment.
Finally, our temperature scaling scheme sometimes
sharpens a distribution too aggressively in the multidocument question-answering and code completion
experiments. This drawback could be possibly improved by designing a more fine-grained attention
alignment strategy.


The first author acknowledges the support from the
Boeing Company (2019-STU-PA-259).


A QQ plot (Wilk and Gnanadesikan, 1968) is
a graphical technique used for comparing two
probability distributions by plotting their quantiles
against each other. A point (x, y) corresponds to a
quantile from the second distribution (y-coordinate)
plotted against the same quantile from the first distribution (x-coordinate). When the two distributions under comparison are similar, the points in
the QQ plot will roughly align with the identity line,
y = x. In our case, where we aim to determine the
degree of Gaussian behavior in the average logit
vector, the linearity of the plot serves as an indicator – the closer the points are to the identity line,
the more Gaussian the distribution.
We present the QQ plots for two lengths, 512 and
15k, on the three retrieval tasks in Figure 5. They
are all close to the red reference line, indicating
that their form is highly Gaussian.


Let _l_ max be the largest value in the logit vector _l_ .
Let _τ_ be the temperature of the Softmax function.
The probability of the largest entry is
_e_ _[l]_ [max] _[/τ]_
Pmax = ~~�~~ _L_ _[.]_
_i_ =1 _[e][l][i][/τ]_
Since Softmax is shift-invariant, the logit vector
can always be made zero-mean: [�] _i_ _[l][i]_ [ = 0][. Next,]
according to Assumption 1, the denominator of
Softmax can be approximated as
_L_
_e_ _[l][i][/τ]_ _≈_ _L ·_ E[ _e_ _[l][i][/τ]_ ] = _L · e_ _[σ]_ [2] _[/]_ [(2] _[τ]_ [ 2][)] (1)
_i_ =1
This implies Pmax is approximately
_e_ _[l]_ [max] _[/τ]_
Pmax _≈_
_Le_ _[σ]_ [2] _[/]_ [(2] _[τ]_ [ 2][)]
During the training stage, the temperature _τ_ is 1
_e_ _[l]_ max _[tr]_
P _[tr]_ max _[≈]_ _Ltre_ _[σ]_ _tr_ [2] _[/]_ [2] _[,]_
which gives an expression of the largest logit entry
during the training stage
        _tr_ _[/]_ [2][�]
_l_ max _[tr]_ _[≈]_ [log] P _[tr]_ max _[L][tr][e][σ]_ [2] (2)
According to Assumption 2, the largest probability
during the extrapolation stage can be simplified as
_e_ _[l]_ max _[ex]_ _[/τ]_
P _[ex]_ max _[≈]_ [2]
_e_ _[l]_ max _[ex]_ _[/τ]_ A. 2 _e_ _[l]_ max _[tr]_ _[/τ]_
=
_Lexe_ _[σ][ex]_ [2] _[/]_ [(2] _[τ]_ [ 2][)] _Lexe_ _[σ][ex]_ [2] _[/]_
_Lexe_ _[σ][ex]_ [2] _[/]_ [(2] _[τ]_ [ 2][)]
(2)
_≈_
P _[tr]_ max _[L][tr][e][σ]_ _tr_ [2] _[/]_ [2][�][1] _[/τ]_
_Lexe_ _[σ][ex]_ [2] _[/]_ [(2] _[τ]_ [ 2][)]
Since _τ_ is a free parameter during extrapolation,
we adjust it to carry out the maximum probability
alignment strategy. Rearranging the terms gives
Proposition 1.

A.2 Detailed Derivation of Proposition 1


The entropy of a discrete probability computed by
Softmax is
_[l][i][/τ]_ = log _D −_ - _i_ _lτi_
_D_
_li_
_τ_ _[e][l][i][/τ]_
_,_
_D_
 _H_ = _−_
_i_
_e_ _[l][i][/τ]_
_[l][i][/τ]_
log _[e][l][i][/τ]_
_D_ _D_
where _D_ = [�] _i_ _[e][l][i][/τ]_ [ is the denominator of Soft-]
max, which can be approximated using Eq. (1). On
the other hand, we note that [�] _i_ _[l][i][e][l][i][ ≈]_ _[L]_ [E][[] _[le][l]_ []][.]
When _l ∼_ _N_ (0 _, σ_ [2] ), E[ _le_ _[l]_ ] is approximated as
    - _∞_
E[ _le_ _[l]_ ] =
_−∞_
_le_ _[l]_
2 _σ_ ~~[2]~~ _dl_
2 _σ_
2 _π_ _[e][−]_ _[l]_ [2]
~~_√_~~
_σ_
 - _∞_
=
_−∞_
 - _∞_
=
_−∞_
_l_
~~_√_~~
_σ_
_l_
~~_√_~~
_σ_
2 _σ_ [2] _l−l_ [2]
2 _σ_ ~~[2]~~ _dl_
2 _π_ _[e]_
2 _σ_ ~~[2]~~ _dl_
2 _π_ _[e][−]_ [(] _[l][−][σ]_ [2][)][2] _[−][σ]_ [4]
(3)
   - _∞_
= _e_ _[σ]_ [2] _[/]_ [2]
_−∞_
= _e_ _[σ]_ [2] _[/]_ [2] _σ_ [2]
2 _σ_ ~~[2]~~ _dl_
2 _π_ _[e][−]_ [(] _[l][−][σ]_ [2][)][2]
2 _σ_ ~~[2]~~
2 _π_ _[e][−]_ [(] _[l][−][σ]_ [2][)][2]
_l_
~~_√_~~
_σ_
Thus, combining Eq. (1) and (3), the entropy _H_ is
approximated as
_[σ]_ [2]
_τ_ [2]
_[σ]_ [2]
2 _τ_ [2] _[−]_ _[Le]_ _Le_ _[σ]_ [2] _[σ][/]_ [2][(2] _[/]_ [(2] _[τ]_ [ 2] _[τ]_ [)][ 2][)]
_H ≈_ log _L_ + _[σ]_ [2]
_τ_ [2]
_Le_ _[σ]_ [2] _[/]_ [(2] _[τ]_ [ 2][)]
= log _L −_ _[σ]_ [2]
2 _τ_ [2]
Since _τ_ is set to 1 during the training stage, we
have _Htr ≈_ log _Ltr −_ _[σ]_ 2 _tr_ [2] [. During extrapolation,]
we align the entropy (i.e., _Hex_ = _Htr_ ) by adjusting
_τ_ .
log _Lex −_ 2 _[σ]_ _τex_ [2][2] _[≈]_ _[H][ex]_ [ =] _[ H][tr][ ≈]_ [log] _[ L][tr][ −]_ _[σ]_ 2 _tr_ [2] _[.]_
Since _τ_ is a free parameter during extrapolation,
we adjust it to apply the entropy alignment strategy.
Rearranging the terms gives Proposition 2.
(a) Short sequences around 512 (b) Long sequences around 15k
Figure 5: **QQ plots of Flan-T5-XL.** We experiment with short and long sequences. The red reference line is y=x.
The more closely the scatter plots follow the red reference line, the more Gaussian they are.

A.3 Detailed Derivation of Proposition 2


We verify Proposition 1 and 2 on the remaining
tasks by plotting the temperature curves in Figure 6, 7, 8, and 9. We empirically evaluate _σtr_
at the training length and _σex_ every extrapolation
length considering only the 0-th layer.
The real temperatures given by Algorithm 1 are
usually higher than those derived from the two
propositions. After checking the per-layer attention distributions, we find that the 0-th layer has
flatter distributions compared to higher layers. Because the two propositions are derived based on the
0-th layer and a flatter distribution needs a lower
temperature to correct, the temperatures given by
them tend to be lower than the ones given by Algorithm 1 that takes the average of temperatures
across all layers.


We report the temperatures for all tasks across
model sizes given by Algorithm 1 in Table 8, 9, 10,
and 11.


We report the performance breakdown of Exact
Match and Edit Similarity across lengths in Table 12 and 13.


We report the performance breakdown of different
numbers of input documents in Table 14.
Figure 6: **Line Retrieval Temperature Analysis.**
Curves are given by Proposition 1 and 2. Cross signs
and dots are given by Algorithm 1. log _L_ 512 is given
by Yao et al. (2021); Su (2021).


Sequence Length ( _Lex_ )
Models
1024 2048 4096 8192 15000
T5-Large-LM
w/ Pmax 0.9 0.85 0.8 0.75 0.7
w/ H 0.8 0.7 0.6 0.5 0.5
T5-XL-LM
w/ Pmax 0.9 0.85 0.75 0.7 0.6
w/ H 0.85 0.7 0.55 0.5 0.5
T5-XXL-LM
w/ Pmax 0.9 0.85 0.65 0.55 0.5
w/ H 0.85 0.7 0.7 0.55 0.5
w/ log _Lex Ltr_ 0.9 0.82 0.75 0.69 0.65
Table 8: **Temperatures of Language Modeling.** We
search the optimal temperature from 1.0, 0.95, 0.9, _· · ·_,
0.5. We set _Ltr_ = 512.


Topic, # of topics Line, # of lines Passkey, # of sentences
Models
5 10 15 20 25 200 300 400 500 600 680 20k 30k 40k 50k 55k
Flan-T5-Large
w/ Pmax 0.85 0.8 0.75 0.75 0.75 0.85 0.8 0.8 0.75 0.75 0.75 0.85 0.80 0.80 0.75 0.75
w/ H 0.7 0.6 0.55 0.5 0.5 0.65 0.55 0.55 0.5 0.5 0.5 0.6 0.55 0.5 0.5 0.5
Flan-T5-XL
w/ Pmax 0.8 0.75 0.7 0.65 0.65 0.8 0.75 0.75 0.7 0.70 0.7 0.85 0.8 0.75 0.75 0.75
w/ H 0.7 0.55 0.55 0.5 0.5 0.6 0.55 0.55 0.5 0.5 0.5 0.7 0.65 0.6 0.6 0.6
Flan-T5-XXL
w/ Pmax 0.85 0.8 0.75 0.75 0.75 0.8 0.8 0.75 0.75 0.75 0.75 0.85 0.8 0.8 0.75 0.75
w/ H 0.75 0.65 0.6 0.55 0.55 0.65 0.6 0.6 0.55 0.55 0.55 0.65 0.6 0.55 0.55 0.5
w/ log _Lex Ltr_ 0.79 0.72 0.69 0.67 0.65 0.74 0.71 0.69 0.67 0.66 0.65 0.73 0.69 0.67 0.66 0.65
Table 9: **Temperatures of Retrieval Tasks.** We search the optimal temperature from 1.0, 0.95, 0.9, _· · ·_, 0.5. The
maximum lengths of the three tasks are all around 14.5k to 15.5k tokens ( _Lex_ ). We set _Ltr_ = 512.
Figure 7: **Passkey Retrieval Temperature Analysis.**
Curves are given by Proposition 1 and 2. Cross signs
and dots are given by Algorithm 1. log _L_ 512 is given
by Yao et al. (2021); Su (2021).
Figure 8: **Multi-doc QA Temperature Analysis.**
Curves are from Proposition 1 & 2. Dots and crosses
are from Algorithm 1.
Figure 9: **Code Completion Temperature Analysis.**
Curves are given by Proposition 1 and 2. Cross signs
and dots are given by Algorithm 1. log _L_ 768 is given
by Yao et al. (2021); Su (2021).


10 Docs 20 Docs 30 Docs
Models
_Lex_ = 1700 _Lex_ = 3300 _Lex_ = 5000
Flan-T5-Large
w/ Max. 0.9 0.85 0.8
w/ Ent. 0.75 0.65 0.6
Flan-T5-XL
w/ Max. 0.85 0.75 0.75
w/ Ent. 0.75 0.65 0.55
Flan-T5-XXL
w/ Max. 0.9 0.8 0.8
w/ Ent. 0.75 0.7 0.65
w/ log _Lex Ltr_ 0.84 0.77 0.73
Table 10: **Temperatures of Multi-document Question**
**Answering.** We search the optimal temperature from
1.0, 0.95, 0.9, _· · ·_, 0.5. Different golden document
positions have the same temperature. We set _Ltr_ = 512.


Sequence Length ( _Lex_ )
Models
1k 2k 3k 4k 5k 6k 7k 8k 9k 10k 11k 12k 13k 14k 15k 16k
CodeT5+
w/ Pmax 0.95 0.8 0.75 0.75 0.7 0.7 0.6 0.6 0.6 0.6 0.55 0.55 0.55 0.55 0.5 0.5
w/ H 0.85 0.55 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5
w/ log _Lex Ltr_ 0.96 0.87 0.83 0.8 0.78 0.76 0.75 0.74 0.73 0.72 0.71 0.71 0.7 0.7 0.69 0.69
Table 11: **Temperatures of Code Completion.** We search the optimal temperature from 1.0, 0.95, 0.9, _· · ·_, 0.5.
The maximum length is around 16k tokens ( _Lex_ ). We set _Ltr_ = 768.


Sequence Length ( _Lex_ )
Models
1k 2k 3k 4k 5k 6k 7k 8k 9k 10k 11k 12k 13k 14k 15k
CodeT5+ 19.6 19.0 11.3 2.6 0.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
w/ Pmax 21.1 22.5 21.7 21.5 19.3 22.7 16.1 14.4 13.4 20.6 16.0 15.3 12.3 16.7 4.5
w/ H 19.5 18.7 13.7 9.0 7.9 9.0 10.3 8.8 10.8 12.1 11.7 10.2 9.2 11.1 2.3
w/ log _Lex Ltr_ 21.6 23.0 22.1 22.0 20.6 24.3 20.7 18.6 19.1 22.4 13.8 20.3 15.4 19.4 11.4
w/ _truncation_ 20.0 19.2 19.3 19.2 17.1 21.4 21.1 18.0 19.1 25.2 18.1 20.3 16.9 27.8 15.9
Table 12: **Full Exact Match Breakdown of Code Completion Edit Similarity.** We set _Ltr_ = 768. Numbers in
red are higher than their counterpart in the w/ _truncation_ row. The bucket _n_ k contains the data with length in [ _n_ k,
( _n_ +1)k), _n ∈_ [1 _,_ 15].


Sequence Length ( _Lex_ )
Models
1k 2k 3k 4k 5k 6k 7k 8k 9k 10k 11k 12k 13k 14k 15k
CodeT5+ 62.4 59.6 53.1 38.9 18.3 10.4 6.1 4.0 4.5 5.0 6.7 5.1 6.4 4.4 3.5
w/ Pmax 65.9 65.7 65.3 65.6 63.1 64.9 60.0 60.0 58.1 57.5 56.2 56.0 52.1 56.9 39.9
w/ H 64.8 62.5 54.1 43.0 43.0 44.8 47.7 47.0 47.6 51.2 44.3 49.7 50.3 57.4 42.0
w/ log _Lex Ltr_ 66.3 66.1 65.2 66.4 63.0 66.1 61.9 58.8 61.6 57.8 54.2 57.9 48.7 52.2 48.6
w/ _truncation_ 65.3 64.2 64.2 65.6 62.2 66.9 66.8 61.8 64.1 65.1 63.5 63.9 61.5 67.6 60.8
Table 13: **Full Edit Similarity Breakdown of Code Completion.** We set _Ltr_ = 768. Numbers in red are
higher than their counterpart in the w/ _truncation_ row. The bucket _n_ k contains the data with length in [ _n_ k, ( _n_ +1)k),
_n ∈_ [1 _,_ 15].


10 Docs 20 Docs 30 Docs
Models
0 4 9 0 4 9 14 19 0 4 9 14 19 24 29
Flan-T5-Large 60.6 48.5 48.0 54.5 44.0 39.6 38.0 40.2 52.6 42.0 36.5 34.0 33.9 33.9 37.9
w/ Max. 60.9 49.8 48.6 53.5 45.6 40.8 39.7 41.3 50.8 44.5 39.5 36.4 35.9 35.8 37.0
w/ Ent. 58.9 50.1 47.3 52.4 45.2 40.4 38.0 40.0 47.6 41.1 35.2 33.5 32.2 33.3 34.2
w/ log _Lex Ltr_ 60.2 51.1 48.4 53.8 46.0 41.4 39.4 41.7 50.6 44.1 39.3 36.3 35.8 35.8 37.2
Flan-T5-XL 64.0 55.4 58.9 60.6 47.9 45.1 47.3 55.3 58.4 44.6 40.0 39.9 41.7 46.4 54.8
w/ Max. 65.3 57.3 60.8 62.2 51.6 49.0 49.4 56.0 60.9 49.1 46.0 44.9 46.3 49.1 55.7
w/ Ent. 64.7 56.7 60.0 59.3 50.1 47.9 49.8 55.1 52.4 43.5 42.1 40.3 42.0 42.9 51.3
w/ log _Lex Ltr_ 65.1 57.0 60.6 62.2 51.7 48.8 49.5 56.0 61.0 49.1 46.1 44.7 46.1 48.7 55.4
Flan-T5-XXL 65.1 61.0 64.6 61.1 53.9 52.4 54.7 62.4 58.9 49.1 48.1 47.5 48.9 53.1 61.2
w/ Max. 66.2 61.8 63.2 62.8 55.9 54.4 55.6 59.6 60.4 52.5 51.0 50.2 51.3 53.5 59.1
w/ Ent. 67.3 62.1 61.3 63.2 56.1 54.1 54.3 57.6 61.0 53.4 50.8 50.3 50.7 51.9 55.7
w/ log _Lex Ltr_ 66.7 61.9 63.1 63.1 56.0 54.7 55.1 59.0 61.5 53.3 51.3 50.3 51.1 53.0 57.2
Table 14: **Full Performance Breakdown of Multi-document Question Answering.** The numbers are accuracy.
Full score is 100. 0, 4, 9... indicate the position of the golden document that contains the answer to a question.