Efficient Streaming Language Models with Attention Sinks
Extracted source text · arXiv 2309.17453
\documentclass{article} %
\usepackage{iclr2024_conference,times}
\input{packages}
\input{math_commands}
\input{macro}
\title{Efficient Streaming Language Models \\ with Attention Sinks}
\author{
Guangxuan Xiao$^{1}$\thanks{Part of the work done during an internship at Meta AI.} \quad Yuandong Tian$^2$ \quad Beidi Chen$^3$ \quad Song Han$^{1,4}$ \quad Mike Lewis$^2$ \\
\\ $^1$ Massachusetts Institute of Technology\quad $^2$ Meta AI
\\ $^3$ Carnegie Mellon University \quad $^4$ NVIDIA
\\ \texttt{\url{https://github.com/mit-han-lab/streaming-llm}}
}
\newcommand{\fix}{\marginpar{FIX}}
\newcommand{\new}{\marginpar{NEW}}
\iclrfinalcopy %
\begin{document}
\maketitle
\input{text/0_abstract}
\input{text/1_introduction}
\input{text/2_related}
\input{text/3_method}
\input{text/4_experiment}
\input{text/5_conclusion}
\section*{Reproducibility Statement}
All findings presented in this paper can be reproduced. We have made our code and datasets available in this \href{https://github.com/mit-han-lab/streaming-llm}{github repo}. The models used in this paper are all openly available, and we provide references to access them. Details regarding our experiments, including hyperparameters, training protocols, and evaluation methods, can be found in the Experiments section (Section~\ref{sec:exps}). We are confident that with the provided resources, readers can reproduce the entirety of our presented results.
\section*{Impact Statement}
\method has been widely adopted by various LLM serving solutions including \href{https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/llama#run-llama-with-streamingllm}{NVIDIA TensorRT-LLM}, \href{https://github.com/intel/intel-extension-for-transformers}{Intel Extension for Transformers}, \href{https://huggingface.co/docs/transformers/v4.39.3/en/internal/generation_utils#transformers.SinkCache}{HuggingFace Transformers}, \href{https://github.com/mlc-ai/mlc-llm/pull/1459}{MLC LLM}, etc.
\section*{Acknowledgements}
This work is supported by
MIT-IBM Watson AI Lab,
Amazon and MIT Science Hub,
National Science Foundation.
We thank Angela Li for writing suggestions and demo making, Jingwei Zuo for proofreading, and Xiuyu Li for the suggestion on notations.
\bibliography{reference}
\bibliographystyle{iclr2024_conference}
\clearpage
\input{text/6_appendix}
\end{document}
\usepackage{hyperref}
\newcommand{\theHalgorithm}{\arabic{algorithm}}
\usepackage{times}
\usepackage{latexsym}
\usepackage{amsmath}
\usepackage[T1]{fontenc}
\usepackage[utf8]{inputenc}
\usepackage{microtype}
\usepackage{bbding}
\usepackage{color,xcolor}
\usepackage{epsfig}
\usepackage{graphicx}
\usepackage{adjustbox}
\usepackage{array}
\usepackage{booktabs}
\usepackage{colortbl}
\usepackage{wrapfig}
\usepackage{hhline}
\usepackage{multirow}
\usepackage{subcaption} %
\usepackage{amsmath,amsfonts,amssymb}
\usepackage{bm}
\usepackage{nicefrac}
\usepackage{microtype}
\usepackage{mathtools}
\usepackage{changepage}
\usepackage{extramarks}
\usepackage{fancyhdr}
\usepackage{lastpage}
\usepackage{setspace}
\usepackage{soul}
\usepackage{xspace}
\usepackage{url}
\usepackage{enumerate}
\usepackage{enumitem} %
\usepackage{titlesec}
\usepackage{makecell}
\usepackage{pifont} %
\usepackage{amsthm}
\urlstyle{same}
\usepackage{float}
\usepackage{amsmath,amsfonts,bm}
\newcommand{\figleft}{{\em (Left)}}
\newcommand{\figcenter}{{\em (Center)}}
\newcommand{\figright}{{\em (Right)}}
\newcommand{\figtop}{{\em (Top)}}
\newcommand{\figbottom}{{\em (Bottom)}}
\newcommand{\captiona}{{\em (a)}}
\newcommand{\captionb}{{\em (b)}}
\newcommand{\captionc}{{\em (c)}}
\newcommand{\captiond}{{\em (d)}}
\newcommand{\newterm}[1]{{\bf #1}}
\def\figref#1{figure~\ref{#1}}
\def\Figref#1{Figure~\ref{#1}}
\def\twofigref#1#2{figures \ref{#1} and \ref{#2}}
\def\quadfigref#1#2#3#4{figures \ref{#1}, \ref{#2}, \ref{#3} and \ref{#4}}
\def\secref#1{section~\ref{#1}}
\def\Secref#1{Section~\ref{#1}}
\def\twosecrefs#1#2{sections \ref{#1} and \ref{#2}}
\def\secrefs#1#2#3{sections \ref{#1}, \ref{#2} and \ref{#3}}
\def\eqref#1{equation~\ref{#1}}
\def\Eqref#1{Equation~\ref{#1}}
\def\plaineqref#1{\ref{#1}}
\def\chapref#1{chapter~\ref{#1}}
\def\Chapref#1{Chapter~\ref{#1}}
\def\rangechapref#1#2{chapters\ref{#1}--\ref{#2}}
\def\algref#1{algorithm~\ref{#1}}
\def\Algref#1{Algorithm~\ref{#1}}
\def\twoalgref#1#2{algorithms \ref{#1} and \ref{#2}}
\def\Twoalgref#1#2{Algorithms \ref{#1} and \ref{#2}}
\def\partref#1{part~\ref{#1}}
\def\Partref#1{Part~\ref{#1}}
\def\twopartref#1#2{parts \ref{#1} and \ref{#2}}
\def\ceil#1{\lceil #1 \rceil}
\def\floor#1{\lfloor #1 \rfloor}
\def\1{\bm{1}}
\newcommand{\train}{\mathcal{D}}
\newcommand{\valid}{\mathcal{D_{\mathrm{valid}}}}
\newcommand{\test}{\mathcal{D_{\mathrm{test}}}}
\def\eps{{\epsilon}}
\def\reta{{\textnormal{$\eta$}}}
\def\ra{{\textnormal{a}}}
\def\rb{{\textnormal{b}}}
\def\rc{{\textnormal{c}}}
\def\rd{{\textnormal{d}}}
\def\re{{\textnormal{e}}}
\def\rf{{\textnormal{f}}}
\def\rg{{\textnormal{g}}}
\def\rh{{\textnormal{h}}}
\def\ri{{\textnormal{i}}}
\def\rj{{\textnormal{j}}}
\def\rk{{\textnormal{k}}}
\def\rl{{\textnormal{l}}}
\def\rn{{\textnormal{n}}}
\def\ro{{\textnormal{o}}}
\def\rp{{\textnormal{p}}}
\def\rq{{\textnormal{q}}}
\def\rr{{\textnormal{r}}}
\def\rs{{\textnormal{s}}}
\def\rt{{\textnormal{t}}}
\def\ru{{\textnormal{u}}}
\def\rv{{\textnormal{v}}}
\def\rw{{\textnormal{w}}}
\def\rx{{\textnormal{x}}}
\def\ry{{\textnormal{y}}}
\def\rz{{\textnormal{z}}}
\def\rvepsilon{{\mathbf{\epsilon}}}
\def\rvtheta{{\mathbf{\theta}}}
\def\rva{{\mathbf{a}}}
\def\rvb{{\mathbf{b}}}
\def\rvc{{\mathbf{c}}}
\def\rvd{{\mathbf{d}}}
\def\rve{{\mathbf{e}}}
\def\rvf{{\mathbf{f}}}
\def\rvg{{\mathbf{g}}}
\def\rvh{{\mathbf{h}}}
\def\rvu{{\mathbf{i}}}
\def\rvj{{\mathbf{j}}}
\def\rvk{{\mathbf{k}}}
\def\rvl{{\mathbf{l}}}
\def\rvm{{\mathbf{m}}}
\def\rvn{{\mathbf{n}}}
\def\rvo{{\mathbf{o}}}
\def\rvp{{\mathbf{p}}}
\def\rvq{{\mathbf{q}}}
\def\rvr{{\mathbf{r}}}
\def\rvs{{\mathbf{s}}}
\def\rvt{{\mathbf{t}}}
\def\rvu{{\mathbf{u}}}
\def\rvv{{\mathbf{v}}}
\def\rvw{{\mathbf{w}}}
\def\rvx{{\mathbf{x}}}
\def\rvy{{\mathbf{y}}}
\def\rvz{{\mathbf{z}}}
\def\erva{{\textnormal{a}}}
\def\ervb{{\textnormal{b}}}
\def\ervc{{\textnormal{c}}}
\def\ervd{{\textnormal{d}}}
\def\erve{{\textnormal{e}}}
\def\ervf{{\textnormal{f}}}
\def\ervg{{\textnormal{g}}}
\def\ervh{{\textnormal{h}}}
\def\ervi{{\textnormal{i}}}
\def\ervj{{\textnormal{j}}}
\def\ervk{{\textnormal{k}}}
\def\ervl{{\textnormal{l}}}
\def\ervm{{\textnormal{m}}}
\def\ervn{{\textnormal{n}}}
\def\ervo{{\textnormal{o}}}
\def\ervp{{\textnormal{p}}}
\def\ervq{{\textnormal{q}}}
\def\ervr{{\textnormal{r}}}
\def\ervs{{\textnormal{s}}}
\def\ervt{{\textnormal{t}}}
\def\ervu{{\textnormal{u}}}
\def\ervv{{\textnormal{v}}}
\def\ervw{{\textnormal{w}}}
\def\ervx{{\textnormal{x}}}
\def\ervy{{\textnormal{y}}}
\def\ervz{{\textnormal{z}}}
\def\rmA{{\mathbf{A}}}
\def\rmB{{\mathbf{B}}}
\def\rmC{{\mathbf{C}}}
\def\rmD{{\mathbf{D}}}
\def\rmE{{\mathbf{E}}}
\def\rmF{{\mathbf{F}}}
\def\rmG{{\mathbf{G}}}
\def\rmH{{\mathbf{H}}}
\def\rmI{{\mathbf{I}}}
\def\rmJ{{\mathbf{J}}}
\def\rmK{{\mathbf{K}}}
\def\rmL{{\mathbf{L}}}
\def\rmM{{\mathbf{M}}}
\def\rmN{{\mathbf{N}}}
\def\rmO{{\mathbf{O}}}
\def\rmP{{\mathbf{P}}}
\def\rmQ{{\mathbf{Q}}}
\def\rmR{{\mathbf{R}}}
\def\rmS{{\mathbf{S}}}
\def\rmT{{\mathbf{T}}}
\def\rmU{{\mathbf{U}}}
\def\rmV{{\mathbf{V}}}
\def\rmW{{\mathbf{W}}}
\def\rmX{{\mathbf{X}}}
\def\rmY{{\mathbf{Y}}}
\def\rmZ{{\mathbf{Z}}}
\def\ermA{{\textnormal{A}}}
\def\ermB{{\textnormal{B}}}
\def\ermC{{\textnormal{C}}}
\def\ermD{{\textnormal{D}}}
\def\ermE{{\textnormal{E}}}
\def\ermF{{\textnormal{F}}}
\def\ermG{{\textnormal{G}}}
\def\ermH{{\textnormal{H}}}
\def\ermI{{\textnormal{I}}}
\def\ermJ{{\textnormal{J}}}
\def\ermK{{\textnormal{K}}}
\def\ermL{{\textnormal{L}}}
\def\ermM{{\textnormal{M}}}
\def\ermN{{\textnormal{N}}}
\def\ermO{{\textnormal{O}}}
\def\ermP{{\textnormal{P}}}
\def\ermQ{{\textnormal{Q}}}
\def\ermR{{\textnormal{R}}}
\def\ermS{{\textnormal{S}}}
\def\ermT{{\textnormal{T}}}
\def\ermU{{\textnormal{U}}}
\def\ermV{{\textnormal{V}}}
\def\ermW{{\textnormal{W}}}
\def\ermX{{\textnormal{X}}}
\def\ermY{{\textnormal{Y}}}
\def\ermZ{{\textnormal{Z}}}
\def\vzero{{\bm{0}}}
\def\vone{{\bm{1}}}
\def\vmu{{\bm{\mu}}}
\def\vtheta{{\bm{\theta}}}
\def\va{{\bm{a}}}
\def\vb{{\bm{b}}}
\def\vc{{\bm{c}}}
\def\vd{{\bm{d}}}
\def\ve{{\bm{e}}}
\def\vf{{\bm{f}}}
\def\vg{{\bm{g}}}
\def\vh{{\bm{h}}}
\def\vi{{\bm{i}}}
\def\vj{{\bm{j}}}
\def\vk{{\bm{k}}}
\def\vl{{\bm{l}}}
\def\vm{{\bm{m}}}
\def\vn{{\bm{n}}}
\def\vo{{\bm{o}}}
\def\vp{{\bm{p}}}
\def\vq{{\bm{q}}}
\def\vr{{\bm{r}}}
\def\vs{{\bm{s}}}
\def\vt{{\bm{t}}}
\def\vu{{\bm{u}}}
\def\vv{{\bm{v}}}
\def\vw{{\bm{w}}}
\def\vx{{\bm{x}}}
\def\vy{{\bm{y}}}
\def\vz{{\bm{z}}}
\def\evalpha{{\alpha}}
\def\evbeta{{\beta}}
\def\evepsilon{{\epsilon}}
\def\evlambda{{\lambda}}
\def\evomega{{\omega}}
\def\evmu{{\mu}}
\def\evpsi{{\psi}}
\def\evsigma{{\sigma}}
\def\evtheta{{\theta}}
\def\eva{{a}}
\def\evb{{b}}
\def\evc{{c}}
\def\evd{{d}}
\def\eve{{e}}
\def\evf{{f}}
\def\evg{{g}}
\def\evh{{h}}
\def\evi{{i}}
\def\evj{{j}}
\def\evk{{k}}
\def\evl{{l}}
\def\evm{{m}}
\def\evn{{n}}
\def\evo{{o}}
\def\evp{{p}}
\def\evq{{q}}
\def\evr{{r}}
\def\evs{{s}}
\def\evt{{t}}
\def\evu{{u}}
\def\evv{{v}}
\def\evw{{w}}
\def\evx{{x}}
\def\evy{{y}}
\def\evz{{z}}
\def\mA{{\bm{A}}}
\def\mB{{\bm{B}}}
\def\mC{{\bm{C}}}
\def\mD{{\bm{D}}}
\def\mE{{\bm{E}}}
\def\mF{{\bm{F}}}
\def\mG{{\bm{G}}}
\def\mH{{\bm{H}}}
\def\mI{{\bm{I}}}
\def\mJ{{\bm{J}}}
\def\mK{{\bm{K}}}
\def\mL{{\bm{L}}}
\def\mM{{\bm{M}}}
\def\mN{{\bm{N}}}
\def\mO{{\bm{O}}}
\def\mP{{\bm{P}}}
\def\mQ{{\bm{Q}}}
\def\mR{{\bm{R}}}
\def\mS{{\bm{S}}}
\def\mT{{\bm{T}}}
\def\mU{{\bm{U}}}
\def\mV{{\bm{V}}}
\def\mW{{\bm{W}}}
\def\mX{{\bm{X}}}
\def\mY{{\bm{Y}}}
\def\mZ{{\bm{Z}}}
\def\mBeta{{\bm{\beta}}}
\def\mPhi{{\bm{\Phi}}}
\def\mLambda{{\bm{\Lambda}}}
\def\mSigma{{\bm{\Sigma}}}
\DeclareMathAlphabet{\mathsfit}{\encodingdefault}{\sfdefault}{m}{sl}
\SetMathAlphabet{\mathsfit}{bold}{\encodingdefault}{\sfdefault}{bx}{n}
\newcommand{\tens}[1]{\bm{\mathsfit{#1}}}
\def\tA{{\tens{A}}}
\def\tB{{\tens{B}}}
\def\tC{{\tens{C}}}
\def\tD{{\tens{D}}}
\def\tE{{\tens{E}}}
\def\tF{{\tens{F}}}
\def\tG{{\tens{G}}}
\def\tH{{\tens{H}}}
\def\tI{{\tens{I}}}
\def\tJ{{\tens{J}}}
\def\tK{{\tens{K}}}
\def\tL{{\tens{L}}}
\def\tM{{\tens{M}}}
\def\tN{{\tens{N}}}
\def\tO{{\tens{O}}}
\def\tP{{\tens{P}}}
\def\tQ{{\tens{Q}}}
\def\tR{{\tens{R}}}
\def\tS{{\tens{S}}}
\def\tT{{\tens{T}}}
\def\tU{{\tens{U}}}
\def\tV{{\tens{V}}}
\def\tW{{\tens{W}}}
\def\tX{{\tens{X}}}
\def\tY{{\tens{Y}}}
\def\tZ{{\tens{Z}}}
\def\gA{{\mathcal{A}}}
\def\gB{{\mathcal{B}}}
\def\gC{{\mathcal{C}}}
\def\gD{{\mathcal{D}}}
\def\gE{{\mathcal{E}}}
\def\gF{{\mathcal{F}}}
\def\gG{{\mathcal{G}}}
\def\gH{{\mathcal{H}}}
\def\gI{{\mathcal{I}}}
\def\gJ{{\mathcal{J}}}
\def\gK{{\mathcal{K}}}
\def\gL{{\mathcal{L}}}
\def\gM{{\mathcal{M}}}
\def\gN{{\mathcal{N}}}
\def\gO{{\mathcal{O}}}
\def\gP{{\mathcal{P}}}
\def\gQ{{\mathcal{Q}}}
\def\gR{{\mathcal{R}}}
\def\gS{{\mathcal{S}}}
\def\gT{{\mathcal{T}}}
\def\gU{{\mathcal{U}}}
\def\gV{{\mathcal{V}}}
\def\gW{{\mathcal{W}}}
\def\gX{{\mathcal{X}}}
\def\gY{{\mathcal{Y}}}
\def\gZ{{\mathcal{Z}}}
\def\sA{{\mathbb{A}}}
\def\sB{{\mathbb{B}}}
\def\sC{{\mathbb{C}}}
\def\sD{{\mathbb{D}}}
\def\sF{{\mathbb{F}}}
\def\sG{{\mathbb{G}}}
\def\sH{{\mathbb{H}}}
\def\sI{{\mathbb{I}}}
\def\sJ{{\mathbb{J}}}
\def\sK{{\mathbb{K}}}
\def\sL{{\mathbb{L}}}
\def\sM{{\mathbb{M}}}
\def\sN{{\mathbb{N}}}
\def\sO{{\mathbb{O}}}
\def\sP{{\mathbb{P}}}
\def\sQ{{\mathbb{Q}}}
\def\sR{{\mathbb{R}}}
\def\sS{{\mathbb{S}}}
\def\sT{{\mathbb{T}}}
\def\sU{{\mathbb{U}}}
\def\sV{{\mathbb{V}}}
\def\sW{{\mathbb{W}}}
\def\sX{{\mathbb{X}}}
\def\sY{{\mathbb{Y}}}
\def\sZ{{\mathbb{Z}}}
\def\emLambda{{\Lambda}}
\def\emA{{A}}
\def\emB{{B}}
\def\emC{{C}}
\def\emD{{D}}
\def\emE{{E}}
\def\emF{{F}}
\def\emG{{G}}
\def\emH{{H}}
\def\emI{{I}}
\def\emJ{{J}}
\def\emK{{K}}
\def\emL{{L}}
\def\emM{{M}}
\def\emN{{N}}
\def\emO{{O}}
\def\emP{{P}}
\def\emQ{{Q}}
\def\emR{{R}}
\def\emS{{S}}
\def\emT{{T}}
\def\emU{{U}}
\def\emV{{V}}
\def\emW{{W}}
\def\emX{{X}}
\def\emY{{Y}}
\def\emZ{{Z}}
\def\emSigma{{\Sigma}}
\newcommand{\etens}[1]{\mathsfit{#1}}
\def\etLambda{{\etens{\Lambda}}}
\def\etA{{\etens{A}}}
\def\etB{{\etens{B}}}
\def\etC{{\etens{C}}}
\def\etD{{\etens{D}}}
\def\etE{{\etens{E}}}
\def\etF{{\etens{F}}}
\def\etG{{\etens{G}}}
\def\etH{{\etens{H}}}
\def\etI{{\etens{I}}}
\def\etJ{{\etens{J}}}
\def\etK{{\etens{K}}}
\def\etL{{\etens{L}}}
\def\etM{{\etens{M}}}
\def\etN{{\etens{N}}}
\def\etO{{\etens{O}}}
\def\etP{{\etens{P}}}
\def\etQ{{\etens{Q}}}
\def\etR{{\etens{R}}}
\def\etS{{\etens{S}}}
\def\etT{{\etens{T}}}
\def\etU{{\etens{U}}}
\def\etV{{\etens{V}}}
\def\etW{{\etens{W}}}
\def\etX{{\etens{X}}}
\def\etY{{\etens{Y}}}
\def\etZ{{\etens{Z}}}
\newcommand{\pdata}{p_{\rm{data}}}
\newcommand{\ptrain}{\hat{p}_{\rm{data}}}
\newcommand{\Ptrain}{\hat{P}_{\rm{data}}}
\newcommand{\pmodel}{p_{\rm{model}}}
\newcommand{\Pmodel}{P_{\rm{model}}}
\newcommand{\ptildemodel}{\tilde{p}_{\rm{model}}}
\newcommand{\pencode}{p_{\rm{encoder}}}
\newcommand{\pdecode}{p_{\rm{decoder}}}
\newcommand{\precons}{p_{\rm{reconstruct}}}
\newcommand{\laplace}{\mathrm{Laplace}} %
\newcommand{\E}{\mathbb{E}}
\newcommand{\Ls}{\mathcal{L}}
\newcommand{\R}{\mathbb{R}}
\newcommand{\emp}{\tilde{p}}
\newcommand{\lr}{\alpha}
\newcommand{\reg}{\lambda}
\newcommand{\rect}{\mathrm{rectifier}}
\newcommand{\softmax}{\mathrm{softmax}}
\newcommand{\sigmoid}{\sigma}
\newcommand{\softplus}{\zeta}
\newcommand{\KL}{D_{\mathrm{KL}}}
\newcommand{\Var}{\mathrm{Var}}
\newcommand{\standarderror}{\mathrm{SE}}
\newcommand{\Cov}{\mathrm{Cov}}
\newcommand{\normlzero}{L^0}
\newcommand{\normlone}{L^1}
\newcommand{\normltwo}{L^2}
\newcommand{\normlp}{L^p}
\newcommand{\normmax}{L^\infty}
\newcommand{\parents}{Pa} %
\DeclareMathOperator*{\argmax}{arg\,max}
\DeclareMathOperator*{\argmin}{arg\,min}
\DeclareMathOperator{\sign}{sign}
\DeclareMathOperator{\Tr}{Tr}
\let\ab\allowbreak
\newcommand{\sect}[1]{Sec.~\ref{#1}}
\newcommand{\sectapp}[1]{Appendix~\ref{#1}}
\newcommand{\sectapps}[1]{Appendices~\ref{#1}}
\newcommand{\ssect}[1]{\S~\ref{#1}}
\newcommand{\eqn}[1]{Equation~\ref{#1}}
\newcommand{\fig}[1]{Figure~\ref{#1}}
\newcommand{\figs}[1]{Figs.~\ref{#1}}
\newcommand{\tbl}[1]{Table~\ref{#1}}
\newcommand{\tbls}[1]{Tables~\ref{#1}}
\newcommand{\degree}{\ensuremath{^\circ}\xspace}
\newcommand{\ignore}[1]{}
\newcommand{\norm}[1]{\lVert#1\rVert}
\newcommand{\fcseven}{$\mbox{fc}_7$}
\renewcommand*{\thefootnote}{\fnsymbol{footnote}}
\def\naive{na\"{\i}ve\xspace}
\def\Naive{Na\"{\i}ve\xspace}
\DeclareMathAlphabet{\mathbfit}{OML}{cmm}{b}{it}
\makeatletter
\DeclareRobustCommand\onedot{\futurelet\@let@token\@onedot}
\def\@onedot{\ifx\@let@token.\else.\null\fi\xspace}
\def\eg{e.g\onedot} \def\Eg{E.g\onedot}
\def\ie{i.e\onedot} \def\Ie{I.e\onedot}
\def\cf{\emph{c.f}\onedot} \def\Cf{\emph{C.f}\onedot}
\def\etc{etc\onedot}
\def\vs{\emph{vs}\onedot}
\def\wrt{w.r.t\onedot}
\def\dof{d.o.f\onedot}
\def\etal{et al\onedot}
\def\aka{a.k.a\onedot}
\newcommand{\xgx}[1]{{\color{blue}{[}GX: #1{]}}}
\newcommand{\SH}[1]{{\color{green}{[}SH: #1{]}}}
\newcommand{\bc}[1]{{\color{magenta}{[}BD: #1{]}}}
\newcommand{\yt}[1]{{\color{purple}{[}YT: #1{]}}}
\newcommand{\ml}[1]{{\color{purple}{[}ML: #1{]}}}
\newcommand{\mike}[1]{{\color{purple}{[}ML: #1{]}}}
\newcommand{\xhdr}[1]{{\noindent\bfseries #1}.}
\newcommand{\todo}[1]{{\textcolor{red}{TODO: #1}}}
\newcommand{\red}[1]{{\textcolor{red}{#1}}}
\newcommand{\maroon}[1]{{\textcolor{Maroon}{#1}}}
\newcommand{\CITE}{{\textcolor{red}{[CITE]}}\xspace}
\def\method{StreamingLLM\xspace}
\newcommand{\myparagraph}[1]{\vspace{0pt}\paragraph{#1}}
\renewcommand\theadalign{bc}
\renewcommand\theadfont{\bfseries}
\renewcommand\theadgape{\Gape[4pt]}
\renewcommand\cellgape{\Gape[4pt]}
\begin{abstract}
Deploying Large Language Models (LLMs) in streaming applications such as multi-round dialogue, where long interactions are expected, is urgently needed but poses two major challenges.
Firstly, during the decoding stage, caching previous tokens' Key and Value states (KV) consumes extensive memory.
Secondly, popular LLMs cannot generalize to longer texts than the training sequence length.
Window attention, where only the most recent KVs are cached, is a natural approach --- but we show that it fails when the text length surpasses the cache size.
We observe an interesting phenomenon, namely \emph{attention sink}, that keeping the KV of initial tokens will largely recover the performance of window attention. In this paper, we first demonstrate that the emergence of \emph{attention sink} is due to the strong attention scores towards initial tokens as a ``sink'' even if they are not semantically important.
Based on the above analysis, we introduce \method, an efficient framework that enables LLMs trained with a \textit{finite length} attention window to generalize to \textit{infinite sequence length} without any fine-tuning.
We show that \method can enable Llama-2, MPT, Falcon, and Pythia to perform stable and efficient language modeling with up to 4 million tokens and more.
In addition, we discover that adding a placeholder token as a dedicated attention sink during pre-training can further improve streaming deployment. In streaming settings, \method outperforms the sliding window recomputation baseline by up to 22.2$\times$ speedup.
Code and datasets are provided in the \href{https://github.com/mit-han-lab/streaming-llm}{link}.
\end{abstract}
\section{Introduction}
Large Language Models (LLMs)~\citep{radford2018improving,brown2020language,zhang2022opt,openai2023gpt4,touvron2023llama,touvron2023llama2} are becoming ubiquitous, powering many natural language processing applications such as dialog systems~\citep{schulman2022chatgpt,alpaca,vicuna2023}, document summarization~\citep{goyal-durrett-2020-evaluating,zhang2023benchmarking}, code completion~\citep{chen2021evaluating,rozière2023code} and question answering~\citep{kamalloo2023evaluating}. To unleash the full potential of pretrained LLMs, they should be able to efficiently and accurately perform long sequence generation. For example, an ideal ChatBot assistant can stably work over the content of recent day-long conversations. However, it is very challenging for LLM to generalize to longer sequence lengths than they have been pretrained on, e.g., 4K for Llama-2~\cite{touvron2023llama2}.
The reason is that LLMs are constrained by the attention window during pre-training. Despite substantial efforts to expand this window size~\citep{chen2023extending,kaiokendev,peng2023yarn} and improve training~\citep{dao2022flashattention,dao2023flashattention2} and inference~\citep{pope2022efficiently,xiao2023smoothquant,anagnostidis2023dynamic,wang2021spatten,zhang2023h2o} efficiency for lengthy inputs, the acceptable sequence length remains intrinsically \emph{finite}, which doesn't allow persistent deployments.
In this paper, we first introduce the concept of LLM streaming applications and ask the question:
\begin{center}
\emph{Can we deploy an LLM for infinite-length inputs without sacrificing efficiency and performance?}
\end{center}
When applying LLMs for infinite input streams, two primary challenges arise:
\vspace{-0.5em}
\begin{enumerate}
\item During the decoding stage, Transformer-based LLMs cache the Key and Value states (KV) of all previous tokens, as illustrated in Figure~\ref{fig:scheme} (a), which can lead to excessive memory usage and increasing decoding latency~\citep{pope2022efficiently}.
\item Existing models have limited length extrapolation abilities, i.e., their performance degrades~\citep{alibi,chen2023extending} when the sequence length goes beyond the attention window size set during pre-training.
\end{enumerate}
\vspace{-0.5em}
An intuitive approach, known as window attention~\citep{beltagy2020longformer} (Figure~\ref{fig:scheme} b), maintains only a fixed-size sliding window on the KV states of most recent tokens. Although it ensures constant memory usage and decoding speed after the cache is initially filled, the model collapses once the sequence length exceeds the cache size, \ie, \emph{even just evicting the KV of the first token}, as illustrated in Figure~\ref{fig:nll_curve}.
Another strategy is the sliding window with re-computation (shown in Figure~\ref{fig:scheme} c), which rebuilds the KV states of recent tokens for each generated token. While it offers strong performance, this approach is significantly slower due to the computation of quadratic attention within its window, making this method impractical for real-world streaming applications.
\input{figtab/scheme}
To understand the failure of window attention, we find an interesting phenomenon of autoregressive LLMs: a surprisingly large amount of attention score is allocated to the initial tokens, irrespective of their relevance to the language modeling task, as visualized in Figure~\ref{fig:attention_visualization}. We term these tokens ``\textbf{attention sinks}".
Despite their lack of semantic significance, they collect significant attention scores. We attribute the reason to the Softmax operation, which requires attention scores to sum up to one for all contextual tokens. Thus, even when the current query does not have a strong match in many previous tokens, the model still needs to allocate these unneeded attention values somewhere so it sums up to one. The reason behind \textit{initial} tokens as sink tokens is intuitive: initial tokens are visible to almost all subsequent tokens because of the autoregressive language modeling nature, making them more readily trained to serve as attention sinks.
Based on the above insights, we propose \method, a simple and efficient framework that enables LLMs trained with a finite attention window to work on text of infinite length without fine-tuning. \method exploits the fact that attention sinks have high attention values, and preserving them can maintain the attention score distribution close to normal. Therefore, \method simply keeps the attention sink tokens' KV (with just 4 initial tokens sufficing) together with the sliding window's KV to anchor the attention computation and stabilize the model's performance. With \method, models including Llama-2-[7, 13, 70]B, MPT-[7, 30]B, Falcon-[7, 40]B, and Pythia-[2.9,6.9,12]B can reliably model 4 million tokens, and potentially even more. Compared with the only viable baseline, sliding window with recomputation, \method achieves up to 22.2$\times$ speedup, realizing the streaming use of LLMs.
Furthermore, we confirm our attention sink hypothesis and demonstrate that language models can be pre-trained to require only a single attention sink token for streaming deployment.
Specifically, we suggest that an extra learnable token at the beginning of all training samples can serve as a designated attention sink.
By pre-training 160-million parameter language models from scratch, we demonstrate that adding this single sink token preserves the model's performance in streaming cases. This stands in contrast to vanilla models, which necessitate the reintroduction of multiple initial tokens as attention sinks to achieve the same performance level.
\input{figtab/attention_visualization}
Finally, we emphasize that \method efficiently generates coherent text from tokens within the KV cache without extending the LLMs' context length. It suits continuous operation needs with minimal memory use and past data reliance. Additionally, \method can complement context extension methods to increase the attendable recent context.
\begin{figure}[t]
\centering
\includegraphics[width=\textwidth]{figures/scheme.pdf}
\vspace{-1.5em}
\caption{\small
\textbf{Illustration of \method \vs existing methods.} The language model, pre-trained on texts of length $L$, predicts the $T$th token ($T\gg L$).
(a) Dense Attention has \(O(T^2)\) time complexity and an increasing cache size. Its performance decreases when the text length exceeds the pre-training text length.
(b) Window Attention caches the most recent \(L\) tokens' KV. While efficient in inference, performance declines sharply once the starting tokens' keys and values are evicted. (c) Sliding Window with Re-computation rebuilds the KV states from the \(L\) recent tokens for each new token. While it performs well on long texts, its \(O(TL^2)\) complexity, stemming from quadratic attention in context re-computation, makes it considerably slow. (d) \method keeps the \emph{attention sink} (several initial tokens) for stable attention computation, combined with the recent tokens. It's efficient and offers stable performance on extended texts. Perplexities are measured using the Llama-2-13B model on the first book (65K tokens) in the PG-19 test set.}\label{fig:scheme}
\vspace{-1.5em}
\end{figure}
\begin{figure}[t]
\centering
\includegraphics[width=0.95\textwidth]{figures/attention_weights.pdf}
\vspace{-1em}
\caption{\small Visualization of the \emph{average} attention logits in Llama-2-7B over 256 sentences, each with a length of 16. Observations include: (1) The attention maps in the first two layers (layers 0 and 1) exhibit the "local" pattern, with recent tokens receiving more attention. (2) Beyond the bottom two layers, the model heavily attends to the initial token across all layers and heads.
}
\label{fig:attention_visualization}
\vspace{-2em}
\end{figure}
\vspace{-0.5em}
\section{Related Work}
\vspace{-0.5em}
Extensive research has been done on applying LLMs to lengthy texts, with three main areas of focus: \textbf{Length Extrapolation}, \textbf{Context Window Extension}, and \textbf{Improving LLMs' Utilization of Long Text}. While seemingly related, it's worth noting that progress in one direction doesn't necessarily lead to progress in the other. For example, extending the context size of LLMs doesn't improve the model's performance beyond the context size, and neither approach ensures effective use of the long context. Our \method framework primarily lies in the first category, where LLMs are applied to text significantly exceeding the pre-training window size, potentially even of infinite length. We do not expand the attention window size of LLMs or enhance the model's memory and usage on long texts. The last two categories are orthogonal to our focus and could be integrated with our techniques.
\textbf{Length extrapolation} aims to enable language models trained on shorter texts to handle longer ones during testing. A predominant avenue of research targets the development of relative position encoding methods for Transformer models, enabling them to function beyond their training window. One such initiative is Rotary Position Embeddings (RoPE)~\citep{su2021roformer}, which transforms the queries and keys in every attention layer for relative position integration. Despite its promise, subsequent research~\citep{alibi, chen2023extending} indicated its underperformance on text that exceeds the training window. Another approach, ALiBi~\citep{alibi}, biases the query-key attention scores based on their distance, thereby introducing relative positional information. While this exhibited improved extrapolation, our tests on MPT models highlighted a breakdown when the text length was vastly greater than the training length. Current methodologies, however, have yet to achieve infinite length extrapolation, causing no existing LLMs to fit for streaming applications.
\textbf{Context Window Extension} centers on expanding the LLMs' context window, enabling the processing of more tokens in one forward pass. A primary line of work addresses the training efficiency problem. Given the attention to computation's quadratic complexity during training, developing a long-context LLM is both a computational and memory challenge. Solutions have ranged from system-focused optimizations like FlashAttention~\citep{dao2022flashattention, dao2023flashattention2}, which accelerates attention computation and reduces memory footprint, to approximate attention methods~\citep{zaheer2020bigbird,beltagy2020longformer,wang2020linformer,kitaev2020reformer} that trade model quality for efficiency. Recently, there has been a surge of work on extending pre-trained LLMs with RoPE~\citep{chen2023extending,kaiokendev,blocntkaware,peng2023yarn}, involving position interpolation and fine-tuning. However, all the aforementioned techniques only extend LLMs' context window to a limited extent, which falls short of our paper's primary concern of handling limitless inputs.
\textbf{Improving LLMs' Utilization of Long Text} optimizes LLMs to better capture and employ the content within the context rather than merely taking them as inputs. As highlighted by \citeauthor{liu2023lost} and \citeauthor{longchat2023}, success in the previously mentioned two directions does not necessarily translate to competent utilization of lengthy contexts. Addressing this effective usage of prolonged contexts within LLMs is still a challenge. Our work concentrates on stably harnessing the most recent tokens, enabling the seamless streaming application of LLMs.
\vspace{-0.5em}
\section{\method}
\input{figtab/nll_curve}
\vspace{-0.5em}
\subsection{The Failure of Window Attention and Attention Sinks}
While the window attention technique offers efficiency during inference, it results in an exceedingly high language modeling perplexity. Consequently, the model's performance is unsuitable for deployment in streaming applications. In this section, we use the concept of \emph{attention sink} to explain the failure of window attention, serving as the inspiration behind \method.
\vspace{-1em}
\myparagraph{Identifying the Point of Perplexity Surge.}
Figure~\ref{fig:nll_curve} shows the perplexity of language modeling on a 20K token text. It is evident that perplexity spikes when the text length surpasses the cache size, led by the exclusion of initial tokens. This suggests that the initial tokens, regardless of their distance from the predicted tokens, are crucial for maintaining the stability of LLMs.
\vspace{-0.5em}
\myparagraph{Why do LLMs break when removing \emph{initial} tokens' KV?}
We visualize attention maps from all layers and heads of the Llama-2-7B and models in Figure~\ref{fig:attention_visualization}. We find that, beyond the bottom two layers, the model consistently focuses on the initial tokens across all layers and heads. The implication is clear: removing these initial tokens' KV will remove a considerable portion of the denominator in the SoftMax function (Equation~\ref{equ:softmax}) in attention computation. This alteration leads to a significant shift in the distribution of attention scores away from what would be expected in normal inference settings.
\vspace{-1em}
\begin{equation}
\text{SoftMax}(x)_i = \frac{e^{x_i}}{e^{x_1} + \sum_{j=2}^{N} e^{x_j}}, \quad x_1 \gg x_j, j \in 2, \dots, N
\label{equ:softmax}
\end{equation}
\vspace{-1em}
There are two possible explanations for the importance of the initial tokens in language modeling: (1) Either their semantics are crucial, or
(2) the model learns a bias towards their absolute position. To distinguish between these possibilities, we conduct experiments (Table~\ref{tab:linebreak}), wherein the first four tokens are substituted with the linebreak token ``\textbackslash n". The observations indicate that the model still significantly emphasizes these initial linebreak tokens. Furthermore, reintroducing them restores the language modeling perplexity to levels comparable to having the original initial tokens. This suggests that the absolute position of the starting tokens, rather than their semantic value, holds greater significance.
\vspace{-0.5em}
\myparagraph{LLMs attend to Initial Tokens as Attention Sinks.}
\label{sec:attention_sink}
To explain why the model disproportionately focuses on initial tokens—regardless of their semantic relevance to language modeling, we introduce the concept of ``\emph{attention sink}". The nature of the SoftMax function (Equation~\ref{equ:softmax}) prevents all attended tokens from having zero values. This requires aggregating some information from other tokens across all heads in all layers, even if the current embedding has sufficient self-contained information for its prediction. Consequently, the model tends to dump unnecessary attention values to specific tokens. A similar observation has been made in the realm of quantization outliers~\citep{xiao2023smoothquant,bondarenko2023quantizable}, leading to the proposal of SoftMax-Off-by-One~\citep{softmax1} as a potential remedy.
\input{figtab/start_size_and_linebreak}
Why do various autoregressive LLMs, such as Llama-2, MPT, Falcon, and Pythia, consistently focus on \textit{initial tokens} as their attention sinks, rather than other tokens? Our explanation is straightforward: Due to the sequential nature of autoregressive language modeling, initial tokens are visible to all subsequent tokens, while later tokens are only visible to a limited set of subsequent tokens. As a result, initial tokens are more easily trained to serve as attention sinks, capturing unnecessary attention.
We've noted that LLMs are typically trained to utilize multiple initial tokens as attention sinks rather than just one. As illustrated in Figure~\ref{tab:start_size}, the introduction of four initial tokens, as attention sinks, suffices to restore the LLM's performance. In contrast, adding just one or two doesn't achieve full recovery. We believe this pattern emerges because these models didn't include a consistent starting token across all input samples during pre-training. Although Llama-2 does prefix each paragraph with a ``<s>" token, it's applied before text chunking, resulting in a mostly random token occupying the zeroth position. This lack of a uniform starting token leads the model to use several initial tokens as attention sinks. We hypothesize that by incorporating a stable learnable token at the start of all training samples, it could singularly act as a committed attention sink, eliminating the need for multiple initial tokens to ensure consistent streaming. We will validate this hypothesis in Section~\ref{sec:pretrain-fix}.
\vspace{-1em}
\subsection{Rolling KV Cache with Attention Sinks}
\input{figtab/rolling_kv}
To enable LLM streaming in already trained LLMs, we propose a straightforward method that can recover window attention's perplexity without any model finetuning. Alongside the current sliding window tokens, we reintroduce a few starting tokens' KV in the attention computation. The KV cache in \method can be conceptually divided into two parts, as illustrated in Figure~\ref{fig:rolling_kv}: (1) Attention sinks (four initial tokens) stabilize the attention computation; 2) Rolling KV Cache retains the most recent tokens, crucial for language modeling. \method' design is versatile and can be seamlessly incorporated into any autoregressive language model that employs relative positional encoding, such as RoPE~\citep{su2021roformer} and ALiBi~\citep{alibi}.
When determining the relative distance and adding positional information to tokens, \method focuses on positions \emph{within the cache} rather than those \emph{in the original text}. This distinction is crucial for \method's performance. For instance, if the current cache (Figure~\ref{fig:rolling_kv}) has tokens [0, 1, 2, 3, 6, 7, 8] and is in the process of decoding the 9th token, the positions assigned are [0, 1, 2, 3, 4, 5, 6, 7], rather than the positions in the original text, which would be [0, 1, 2, 3, 6, 7, 8, 9].
For encoding like RoPE, we cache the Keys of tokens \emph{prior to} introducing the rotary transformation. Then, we apply position transformation to the keys in the rolling cache at each decoding phase. On the other hand, integrating with ALiBi is more direct. Here, the contiguous linear bias is applied instead of a 'jumping' bias to the attention scores. This method of assigning positional embedding within the cache is crucial to \method's functionality, ensuring that the model operates efficiently even beyond its pre-training attention window size.
\subsection{Pre-Training LLMs with Attention Sinks}
\label{sec:pretrain-fix}
\input{figtab/pretrain_fix}
As elaborated in Section~\ref{sec:attention_sink}, a significant reason for the model's excessive attention to multiple initial tokens is the absence of a designated sink token to offload excessive attention scores. Due to this, the model inadvertently uses globally visible tokens, primarily the initial ones, as attention sinks. A potential remedy can be the intentional inclusion of a global trainable attention sink token, denoted as a ``Sink Token'', which would serve as a repository for unnecessary attention scores. Alternatively, replacing the conventional SoftMax function with a variant like SoftMax-off-by-One~\citep{softmax1},
\begin{equation}
\small
\text{SoftMax}_1(x)_i = \frac{e^{x_i}}{1+\sum_{j=1}^{N} e^{x_j}},
\label{equ:softmax1}
\end{equation}
which does not require the attention scores on all contextual tokens to sum up to one, may also be effective. Note that $\text{SoftMax}_1$ is equivalent to prepending a token with an all-zero Key and Value features in the attention computation. We denote this method as ``Zero Sink'' to fit our framework.
For validation, we pre-train three language models with 160 million parameters from scratch under identical settings. The first model utilizes the standard SoftMax attention (Vanilla), the second replaced the regular attention mechanism with $\text{SoftMax}_1$ (Zero Sink), and one prepending a learnable placeholder token (Sink Token) in all training samples. As shown in Table~\ref{tab:pretrain_fix}, while the zero sink alleviates the attention sink problem to some extent, the model still relies on other initial tokens as attention sinks. Introducing a sink token is highly effective in stabilizing the attention mechanism. Simply pairing this sink token with recent tokens sufficiently anchors the model's performance, and the resulting evaluation perplexity is even marginally improved. Given these findings, we recommend training future LLMs with a sink token in all samples to optimize streaming deployment.
\begin{figure}
\centering
\includegraphics[width=\textwidth]{figures/nll_curve.pdf}
\vspace{-1.5em}
\caption{\small
Language modeling perplexity on texts with 20K tokens across various LLM. Observations reveal consistent trends:
(1) Dense attention fails once the input length surpasses the pre-training attention window size.
(2) Window attention collapses once the input length exceeds the cache size, i.e., the initial tokens are evicted.
(3) \method demonstrates stable performance, with its perplexity nearly matching that of the sliding window with re-computation baseline.
}
\label{fig:nll_curve}
\vspace{-1.5em}
\end{figure}
\vspace{1em}
\begin{minipage}[t]{\textwidth}
\begin{minipage}[t]{0.41\textwidth}\vspace{0pt}
\centering
\setlength{\tabcolsep}{2pt}
\small
\captionof{table}{\small Window attention has poor performance on long text. The perplexity is restored when we reintroduce the initial four tokens alongside the recent 1020 tokens (4+1020). Substituting the original four initial tokens with linebreak tokens ``\textbackslash n" (4"\textbackslash n"+1020) achieves comparable perplexity restoration.
Cache config x+y denotes adding x initial tokens with y recent tokens. Perplexities are measured on the first book (65K tokens) in the PG19 test set.}
\vspace{-1em}
\label{tab:linebreak}
\begin{tabular}{lc}
\toprule
Llama-2-13B & PPL ($\downarrow$) \\
\midrule
0 + 1024 (Window) & 5158.07 \\
4 + 1020 & 5.40 \\
4"\textbackslash n"+1020 & 5.60 \\
\bottomrule
\end{tabular}
\end{minipage}
\hfill
\begin{minipage}[t]{0.57\textwidth}\vspace{0pt}
\centering
\setlength{\tabcolsep}{2pt}
\small
\captionof{table}{\small Effects of reintroduced initial token numbers on \method. (1) Window attention (0+y) has a drastic increase in perplexity. (2) Introducing one or two initial tokens doesn't fully restore model perplexity, showing that the model doesn't solely use the first token as the attention sink.
(3) Introducing four initial tokens generally suffices; further additions have diminishing returns. Cache config x+y denotes adding x initial tokens to y recent tokens. Perplexities are evaluated on 400K tokens in the concatenated PG19 test set.
}
\vspace{-1em}
\label{tab:start_size}
\begin{tabular}{lccccc}
\toprule
Cache Config & 0+2048& 1+2047 & 2+2046 & 4+2044 & 8+2040 \\
\midrule
Falcon-7B & 17.90 & 12.12 & 12.12 & 12.12 & 12.12 \\
MPT-7B & 460.29 & 14.99 & 15.00 & 14.99 & 14.98 \\
Pythia-12B & 21.62 & 11.95 & 12.09 & 12.09 & 12.02 \\
\midrule
\midrule
Cache Config & 0+4096 & 1+4095 & 2+4094 & 4+4092 & 8+4088 \\
\midrule
Llama-2-7B & 3359.95 & 11.88 & 10.51 & 9.59 & 9.54 \\
\bottomrule
\end{tabular}
\end{minipage}
\end{minipage}
\begin{wrapfigure}{r}{0.4\textwidth}
\centering
\vspace{-4em}
\includegraphics[width=\linewidth]{figures/rolling_kv.pdf}
\vspace{-1.5em}
\caption{\small The KV cache of \method.}
\vspace{-1em}
\label{fig:rolling_kv}
\end{wrapfigure}
\begin{wraptable}{r}{0.5\textwidth}
\centering
\setlength{\tabcolsep}{2pt}
\small
\vspace{-2em}
\caption{\small Comparison of vanilla attention with prepending a zero token and a learnable sink token during pre-training. To ensure stable streaming perplexity, the vanilla model requires several initial tokens. While Zero Sink shows a slight improvement, it still needs other initial tokens. Conversely, the model trained with a learnable Sink Token shows stable streaming perplexity with only the sink token added. Cache config $x$+$y$ denotes adding $x$ initial tokens with $y$ recent tokens. Perplexity is evaluated on the first sample in the PG19 test set.}
\label{tab:pretrain_fix}
\begin{tabular}{lcccc}
\toprule
Cache Config & 0+1024 & 1+1023 & 2+1022 & 4+1020 \\
\midrule
Vanilla & 27.87 & 18.49 & 18.05 & 18.05 \\
Zero Sink & 29214 & 19.90 & 18.27 & 18.01 \\
Learnable Sink & 1235 & \textbf{18.01} & 18.01 & 18.02 \\
\bottomrule
\end{tabular}
\end{wraptable}
\section{Experiments}
\label{sec:exps}
We evaluate \method using four prominent recent model families: Llama-2~\citep{touvron2023llama2}, MPT~\citep{MosaicML2023Introducing}, PyThia~\citep{biderman2023pythia}, and Falcon~\citep{falcon40b}. Notably, Llama-2, Falcon, and Pythia incorporate RoPE~\citep{su2021roformer}, whereas MPT employs ALiBi~\citep{alibi} — two of the most influential position encoding techniques in recent research. Our diverse model selection ensures the validity and robustness of our findings. We benchmark \method against established baselines such as dense attention, window attention, and the sliding window approach with re-computation. In all subsequent experiments with \method, we default to using four initial tokens as attention sinks unless stated otherwise.
\subsection{Language Modeling on Long Texts Across LLM Families and Scales}
We firstly evaluate \method's language modeling perplexity using the concatenated PG19~\citep{Rae2020Compressive} test set, which contains 100 long books. For Llama-2 models, the cache size is set at 2048, while for Falcon, Pythia, and MPT models, it's set at 1024. This is half the pre-training window size chosen to enhance visualization clarity.
Figure~\ref{fig:nll_curve} illustrates that \method can match the oracle baseline (sliding window with re-computation) in terms of perplexity on texts spanning 20K tokens. Meanwhile, the dense attention technique fails when the input length exceeds its pre-training window, and the window attention technique struggles when the input length surpasses the cache size, leading to the eviction of the initial tokens.
In Figure~\ref{fig:long_lm}, we further substantiate that \method can reliably handle exceptionally extended texts, encompassing more than 4 million tokens, across a spectrum of model families and scales. This includes Llama-2-[7,13,70]B, Falcon-[7,40]B, Pythia-[2.8,6.9,12]B, and MPT-[7,30]B.
\input{figtab/long_lm}
\subsection{Results of Pre-Training with a Sink Token}
To validate our suggestion that introducing a sink token to all pre-training samples improves streaming LLMs, we trained two language models, each with 160 million parameters, under identical conditions. While one model adhered to the original training settings, the other incorporated a sink token at the start of every training sample. Our experiments employed the Pythia-160M~\citep{biderman2023pythia} codebase and followed its training recipe. We train the models on an 8xA6000 NVIDIA GPU server using the deduplicated Pile~\citep{pile} dataset. Apart from reducing the training batch size to 256, we retained all Pythia training configurations, including learning rate schedules, model initialization, and dataset permutations. Both models were trained for 143,000 steps.
\input{figtab/pretrain_results}
\myparagraph{Convergence and Normal Model Performance.}
Including a sink token during pre-training has no negative impact on model convergence and subsequent performance on a range of NLP benchmarks.
As depicted in Figure~\ref{fig:curve}, models trained with a sink token exhibit similar convergence dynamics compared to their vanilla counterparts. We evaluate the two models on seven diverse NLP benchmarks, including ARC-[Challenge, Easy]~\citep{arc}, HellaSwag~\citep{hellaswag}, LAMBADA~\citep{paperno-etal-2016-lambada}, OpenbookQA~\citep{OpenBookQA2018}, PIQA~\citep{piqa}, and Winogrande~\citep{sakaguchi2019winogrande}. As shown in Table~\ref{tab:pretrain_eval}, the model pre-trained with a sink token performs similarly to that trained using the vanilla approach.
\myparagraph{Streaming Performance.} As illustrated in Table~\ref{tab:pretrain_fix}, the streaming perplexities differ between models trained using traditional methods and those augmented with a sink token. Remarkably, the vanilla model requires the addition of multiple tokens as attention sinks to maintain stable streaming perplexity. In contrast, the model trained with a sink token achieves satisfactory streaming performance using just the sink token.
\input{figtab/attn_vis_pretrain}
\input{figtab/streameval}
\myparagraph{Attention Visualization.}
Figure~\ref{fig:attn_vis_pretrain} contrasts attention maps for models pre-trained with and without a sink token. The model without the sink token, similar to Llama-2-7B (Figure~\ref{fig:attention_visualization}), shows early-layer local attention and deeper-layer focus on initial tokens. In contrast, models trained with a sink token consistently concentrate on the sink across layers and heads, indicating an effective attention offloading mechanism. This strong focus on the sink, with reduced attention to other initial tokens, explains the sink token's efficacy in enhancing model's streaming performance.
\input{figtab/instruct_qa}
\subsection{Results on Streaming Question Answering with Instruction-tuned Models}
To show \method's real-world applicability, we emulate multi-round question-answering using instruction-tuned LLMs, commonly used in real-world scenarios.
We first concatenate all question-answer pairs from the ARC-[Challenge, Easy] datasets, feed the continuous stream to Llama-2-[7,13,70]B-Chat models, and assess model completions at each answer position using an exact match criterion. As table~\ref{tab:instruct_qa} indicates, dense attention results in Out-of-Memory (OOM) errors, showing it unsuitable for this setting. While the window attention method works efficiently, it exhibits low accuracy due to random outputs when the input length exceeds the cache size. Conversely, \method excels by efficiently handling the streaming format, aligning with the one-shot, sample-by-sample baseline accuracy.
Highlighting a more fitting scenario for \method, we introduce a dataset, StreamEval, inspired by the LongEval~\citep{longchat2023} benchmark. As depicted in Figure~\ref{fig:streameval}, diverging from LongEval's single query over a long-span setup, we query the model every 10 lines of new information. Each query's answer is consistently 20 lines prior, reflecting real-world instances where questions typically pertain to recent information. As illustrated in Figure~\ref{fig:kv_bench}, LLMs employing \method maintain reasonable accuracy even as input lengths approach 120K tokens. In contrast, both dense and window attention fail at the pre-training text length and the KV cache size, respectively. Additionally, we utilize two context-extended models, LongChat-7b-v1.5-32k~\citep{longchat2023} and Llama-2-7B-32K-Instruct~\citep{togetherlong}, to show that \method can complement context extension techniques. Within \method, context extension means broadening the maximum cache size of streaming LLMs, enabling the capture of broader local information.
\input{figtab/kv_bench}
\subsection{Ablation Studies}
\myparagraph{Numbers of Initial Tokens.} In Table~\ref{tab:start_size}, we ablate the effect of adding varying numbers of initial tokens with recent tokens on the streaming perplexity. The results show the insufficiency of introducing merely one or two initial tokens, whereas a threshold of four initial tokens appears enough, with subsequent additions contributing marginal effects. This result justifies our choice of introducing 4 initial tokens as attention sinks in \method.
\myparagraph{Cache Sizes.} In Table \ref{tab:cache_size}, we evaluate cache size's impact on \method's perplexity. Contrary to intuition, increasing the cache size doesn't consistently lower the language modeling perplexity. This inconsistency shows a potential limitation where these models might not maximize the utility of the entire context they receive. Future research efforts should target enhancing these models' capabilities to utilize extensive contexts better.
\input{figtab/cache_size}
\subsection{Efficency Results}
\input{figtab/efficiency}
We benchmark \method's decoding latency and memory usage against the sliding window with re-computation, which is the only baseline with acceptable quality. Both methods are implemented using the Huggingface Transformers library~\citep{wolf2020huggingfaces} and tested on a single NVIDIA A6000 GPU using the Llama-2-7B and Llama-2-13B models.
As shown in Figure~\ref{fig:efficiency}, as the cache size increases, \method's decoding speed has a linear growth. The sliding window with re-computation baseline has a quadratic rise in decoding latency. Thus, \method achieves an impressive speedup, reaching up to 22.2$\times$ per token. Despite its reduced latency, \method sustains a memory footprint consistent with the re-computation baseline.
\begin{figure}
\centering
\includegraphics[width=\textwidth]{figures/long_lm.pdf}
\vspace{-1.5em}
\caption{\small Language modeling perplexity of \method on super long texts with 4 million tokens across various LLM families and scales. The perplexity remains stable throughout. We use the concatenated test set of PG19 (100 books) to perform language modeling, with perplexity fluctuations due to book transitions.}
\label{fig:long_lm}
\vspace{-1em}
\end{figure}
\begin{minipage}[t]{\textwidth}
\begin{minipage}[t]{0.3\textwidth}\vspace{0pt}
\centering
\includegraphics[width=\linewidth]{figures/learnable_sink_curve.png}
\vspace{-2em}
\captionof{figure}{\small Pre-training loss curves of models w/ and w/o sink tokens. Two models have a similar convergence trend.}
\label{fig:curve}
\end{minipage}
\hfill
\begin{minipage}[t]{0.66\textwidth}\vspace{0pt}
\centering
\small
\setlength{\tabcolsep}{3.0pt}
\captionof{table}{\small Zero-shot accuracy (in \%) across 7 NLP benchmarks, including ARC-[Challenge, Easy], HellaSwag, LAMBADA, OpenbookQA, PIQA, and Winogrande. The inclusion of a sink token during pre-training doesn't harm the model performance.}
\label{tab:pretrain_eval}
\begin{tabular}{lccccccc}
\toprule
Methods & {ARC-c} & {ARC-e} & {HS} & {LBD} & {OBQA} & {PIQA} & {WG} \\
\midrule
Vanilla & 18.6 & 45.2 & 29.4 & 39.6 & 16.0 & 62.2 & 50.1 \\
+Sink Token & \textbf{19.6} & \textbf{45.6} & \textbf{29.8} & \textbf{39.9} & \textbf{16.6} & \textbf{62.6} & \textbf{50.8} \\
\bottomrule
\end{tabular}
\end{minipage}
\vspace{-1em}
\end{minipage}
\begin{figure}[t]
\centering
\includegraphics[width=\textwidth]{figures/attention_vis_pretrain.pdf}
\vspace{-1em}
\caption{\small Visualization of average attention logits over 256 sentences, each 16 tokens long, comparing models pre-trained without (left) and with (right) a sink token. Both maps show the same layers and heads. Key observations: (1) Without a sink token, models show local attention in lower layers and increased attention to initial tokens in deeper layers. (2) With a sink token, there is clear attention directed at it across all layers, effectively collecting redundant attention. (3) With the presence of the sink token, less attention is given to other initial tokens, supporting the benefit of designating the sink token to enhance the streaming performance.}
\label{fig:attn_vis_pretrain}
\vspace{-1em}
\end{figure}
\begin{wrapfigure}{r}{0.38\textwidth}
\centering
\vspace{-0.5em}
\includegraphics[width=\linewidth]{figures/streameval_example.pdf}
\vspace{-2em}
\caption{\small The first sample in StreamEval.}
\vspace{-2em}
\label{fig:streameval}
\end{wrapfigure}
\begin{table}
\centering
\setlength{\tabcolsep}{2pt}
\small
\caption{\small Accuracy (in \%) on the ARC-[Easy, Challenge] datasets. Questions were concatenated and answered in a streaming manner to mimic a real-world chat setting. The dense baseline fails due to Out-of-Memory (OOM) errors. Window attention has poor accuracy. \method has comparable results with the one-shot sample-by-sample baseline. Window attention and \method use cache sizes of 1024.}
\label{tab:instruct_qa}
\vspace{-1em}
\begin{tabular}{lcccccc}
\toprule
Model & \multicolumn{2}{c}{Llama-2-7B-Chat} &\multicolumn{2}{c}{Llama-2-13B-Chat} & \multicolumn{2}{c}{Llama-2-70B-Chat} \\
\cmidrule(lr){2-3}\cmidrule(lr){4-5}\cmidrule(lr){6-7}
Dataset & Arc-E & Arc-C & Arc-E & Arc-C & Arc-E & Arc-C \\
\midrule
One-shot & 71.25 & 53.16 & 78.16 & 63.31 & 91.29 & 78.50 \\
\midrule
Dense & \multicolumn{6}{c}{OOM} \\
Window & 3.58 & 1.39 & 0.25 & 0.34 & 0.12 & 0.32 \\
\method & 71.34 & 55.03 & 80.89 & 65.61 & 91.37 & 80.20 \\
\bottomrule
\end{tabular}
\vspace{-1.5em}
\end{table}
\begin{figure}[t]
\centering
\includegraphics[width=\textwidth]{figures/kv_bench.pdf}
\vspace{-1em}
\caption{\small Performance on the StreamEval benchmark. Accuracies are averaged over 100 samples.}
\vspace{-1em}
\label{fig:kv_bench}
\end{figure}
\begin{wraptable}{r}{0.45\textwidth}
\centering
\setlength{\tabcolsep}{2pt}
\small
\caption{\small Effects of cache size on \method's performance.
Increasing the cache size in \method doesn't consistently yield a decrease in perplexity, showing these models may not fully utilize the provided context.
Cache config $x$+$y$ denotes adding $x$ initial tokens with $y$ recent tokens.
Perplexity is evaluated on 400K tokens in the concatenated PG19 test set.
}
\vspace{-0.5em}
\label{tab:cache_size}
\begin{tabular}{lcccc}
\toprule
Cache & 4+252 & 4+508 & 4+1020 & 4+2044 \\
\midrule
Falcon-7B & 13.61 & 12.84 & \textbf{12.34} & 12.84 \\
MPT-7B & \textbf{14.12} & 14.25 & 14.33 & 14.99 \\
Pythia-12B &13.17 & 12.52 & \textbf{12.08} & 12.09 \\
\midrule
\midrule
Cache & 4+508 & 4+1020 & 4+2044 & 4+4092 \\
\midrule
Llama-2-7B & 9.73 & 9.32 & \textbf{9.08} & 9.59 \\
\bottomrule
\end{tabular}
\end{wraptable}
\begin{figure}[t]
\centering
\includegraphics[width=\textwidth]{figures/efficiency.pdf}
\caption{\small Comparison of per-token decoding latency and memory usage between the sliding window approach with re-computation baseline and \method, plotted against the cache size (attention window size) on the X-axis. \method delivers a remarkable speedup of up to 22.2$\times$ per token and retains a memory footprint similar to the re-computation baseline.}
\label{fig:efficiency}
\vspace{-1.5em}
\end{figure}
\section{Conclusion}
Deploying LLMs in streaming applications is urgently needed but comes with challenges due to efficiency limitations and reduced performance with longer texts. Window attention provides a partial solution, but its performance plummets when initial tokens are excluded. Recognizing the role of these tokens as ``attention sinks", we introduced \method—a simple and efficient framework that enables LLMs to handle unlimited texts without fine-tuning. By adding attention sinks with recent tokens, \method can efficiently model texts of up to 4 million tokens. We further show that pre-training models with a dedicated sink token can improve the streaming performance. \method firstly decouples the LLM's pre-training window size and its actual text generation length, paving the way for the streaming deployment of LLMs.
\appendix
\section{Discussions}
\textbf{Applications.}
\method is particularly suited for streaming applications, such as multi-round dialogues, where continuous operation without heavy reliance on extensive memory or historical data is crucial. For instance, in a daily assistant application based on LLMs, StreamingLLM enables the model to function seamlessly over extended periods. It bases its responses on recent interactions, thus avoiding the need for frequent cache refreshes. Traditional methods might require resetting the cache when the conversation length surpasses the training length, leading to a loss of recent context, or they might need to recompute key-value (KV) states from recent text history, which can be inefficient.
\textbf{Limitations.}
While \method improves the efficiency of LLMs in streaming contexts, it does not extend the models' context window or enhance their long-term memory capabilities. As detailed in Section~\ref{sec:kv_lines}, the model is limited to operating within the confines of its current cache. Consequently, \method is not suitable for tasks that demand long-term memory and extensive data dependency, such as long document question-answering (QA) and summarization. However, it excels in scenarios only requiring short-term memory, like daily conversations and short document QA, where its strength lies in generating coherent text from recent context without the need for cache refreshment.
\textbf{Broader Societal Impacts.}
\method significantly enhances the efficiency and accessibility of LLMs, democratizing their use across various sectors. By enabling nonstop and rapid interactions in applications like conversational agents, \method improves user experiences, especially in scenarios requiring fixed-length models. This advancement allows for more seamless and contextually aware dialogues, potentially benefiting sectors like education, healthcare, and customer service.
Additionally, \method's efficiency in processing reduces the computational load, aligning with the need for environmentally sustainable AI technologies. This aspect is crucial in making advanced AI tools more accessible in regions with limited technological resources.
However, the potential negative impacts of \method mirror those associated with general language models, such as misinformation and biased content generation risks. It's essential to address these risks with robust ethical guidelines and safeguards.
In summary, while \method shares some risks common to language models, its positive contributions towards enhancing user experience, democratizing AI access, and promoting sustainability are noteworthy. These benefits underscore the importance of responsible deployment and ethical use of this technology.
\section{Additional Related Works}
\textbf{Sparse Transformers.} The literature on efficient Transformer models primarily focuses on reducing the computational and memory complexity of the self-attention mechanism. A relevant line of work involves sparsifying the attention matrix by restricting the field of view to fixed, predefined patterns, such as local windows or block patterns with fixed strides~\citep{tay2022efficient}.
Sparse Transformer~\citep{child2019generating} introduces sparse factorizations of the attention matrix, reducing the computational complexity of attention to $O(n\sqrt{n})$. LongFormer~\citep{beltagy2020longformer} combines dilated local windowed attention with task-motivated global attention. Extended Transformer Construction (ETC)~\cite{ainslie2020etc} presents a novel global-local attention mechanism, incorporating four types of attention patterns: global-to-global, local-to-local, local-to-global, and global-to-local. Building on ETC, BigBird~\citep{zaheer2020big} proposes another linear complexity attention alternative, utilizing global tokens, local sliding window attentions, and random attention.
However, these methods have several limitations. First, Sparse Transformer and ETC require custom GPU kernels for a specific block-sparse variant of matrix-matrix multiplication. Second, LongFormer, ETC, and BigBird all rely on a global attention pattern, which is unsuitable for autoregressive language models. Third, these methods are incompatible with pre-trained models, necessitating retraining from scratch.
In contrast, our method offers ease of implementation using standard GPU kernels and is compatible with pre-trained autoregressive language models using dense attention, which are prevalent in the NLP community. This compatibility provides a significant advantage, allowing for the leveraging of existing pre-trained models without any fine-tuning.
\textbf{Concurrent Works.} Our research coincides with the work of \citeauthor{han2023lminfinite}, who conducted a theoretical study on the length generalization failure of language models, identifying three out-of-distribution factors. Their approach, inspired by this analysis, involves employing a ``$\Lambda$"-shaped attention pattern and reconfiguring position encoding distances to enhance length generalization in LLMs. This approach bears a resemblance to our methodology. However, our work uncovers the ``attention sink" phenomenon, wherein Transformer models tend to assign high attention scores to initial tokens with small semantics. This phenomenon extends beyond the scope of length generalization failure, indicating a more pervasive issue in Transformer models. We observe this ``attention sink" behavior not only in auto-regressive language models but also in encoder Transformers such as BERT (see Section~\ref{sec:bert_attn}), and Vision Transformers (ViTs)~\citep{Darcet2023Registers}, suggesting its broader prevalence in Transformer architectures. To mitigate the ``attention sink" phenomenon, we propose the introduction of a learnable sink token during pre-training, and we support our findings with extensive ablation studies.
In parallel, \citeauthor{Darcet2023Registers} observed similar attention concentration on random background patch tokens in Vision Transformers, termed as "registers." These registers act as repositories for global image information. Their solution, adding dedicated "register" tokens, aims to balance attention distribution. Our finding of "attention sinks" parallels this concept. In our paper, the ``attention sinks" are initial tokens that disproportionately attract attention from subsequent tokens. Introducing a dedicated sink token during pre-training prevents the model from inappropriately using content tokens as attention sinks, leading to more effective attention distribution. However, a key difference exists: "registers" in Vision Transformers function as global information holders within intermediate layers, whereas our "attention sinks" are positioned as initial tokens in autoregressive models. This positional variance suggests that the softmax function in attention computation might play a more fundamental role in the emergence of attention sinks.
\section{Accuracy on StreamEval with Increasing Query-Answer Line Distance}
\label{sec:kv_lines}
\input{figtab/kv_lines.tex}
To assess \method's handling of extended inputs, we evaluated the Llama-2-7B-32K-Instruct model on StreamEval, focusing on different query-answer line distances under various cache configurations. In StreamEval, each line consists of 23 tokens, making the line distances equivalent to token distances of $23 \times \text{line distances}$. Accuracy was calculated by averaging results over 100 samples, with each sample comprising 100 queries. Table~\ref{tab:kv_lines} illustrates that \method retains accuracy when the token distance between the query and answer is within the cache size. However, accuracy diminishes as this distance increases and eventually drops to zero when it surpasses the cache capacity.
These results demonstrate that while \method is effective in generating coherent text based on recent context, it cannot extend the context length of language models. These results also emphasize a broader challenge in current language models: their inability to fully utilize context information within the cache, a finding that aligns with the observations made by \citeauthor{liu2023lost}.
\section{Long-Range Benchmark Evaluation}
\input{figtab/longbench.tex}
We evaluated \method using the Llama-2-7B-chat model (max context length 4k) on LongBench~\citep{bai2023longbench}, which encompasses three key NLP tasks: single-document QA (NarrativeQA~\citep{narrativeqa} and Qasper~\citep{dasigi2021qasper}), multi-document QA (HotpotQA~\citep{yang2018hotpotqa} and 2WikiMQA~\cite{ho-etal-2020-constructing}), and summarization (GovReport~\citep{huang2021govreport}, MultiNews~\citep{fabbri2019multinews}). LongBench sets a default max sequence length of 3,500 tokens for the Llama-2-7B-chat model, truncating from the middle to preserve beginning and end information (1,750 tokens each). Table~\ref{tab:longbench} shows that \method with a 4+3496 cache configuration underperforms compared to the truncation baseline, likely due to the loss of crucial initial input prompt information. However, aligning the attention sink number to 1750 restores performance to the level of the text truncation baseline. These results corroborate the findings in Section~\ref{sec:kv_lines}, demonstrating that \method's effectiveness is contingent on the information within its cache, with in-cache performance comparable to the text truncation baseline.
\section{Llama-2-7B Attention Visualization on Longer Sequences}
\input{figtab/llama-2-7b-128-attn.tex}
Figure~\ref{fig:attention_visualization} visualizes the attention map of Llama-2-7B using short sequences (length of 16) for clarity. We further visualize the attention of Llama-2-7B on longer sequences (length of 128) in Figure~\ref{fig:llama-2-7b-128-attn}. We find the observations on short sequences also hold on longer sequences, where the attention scores of the initial tokens are much higher than the rest of the tokens in most layers, regardless of the distance between the initial tokens and the tokens in the rest of the sequence. Because the longer the sequence, the thinner the attention sinks' scores are visualized on the heatmap. We further analyze the attention distribution on longer sequences (length of 4096) using a different method in Section~\ref{sec:quantitative_analysis}.
\section{Quatitative Analysis of Attention Sinks in Long Inputs}
\label{sec:quantitative_analysis}
\input{figtab/first_token_attn.tex}
Figures~\ref{fig:attention_visualization} and \ref{fig:llama-2-70b-attn} illustrate the attention sink phenomenon using short sequences for clarity. Extending this analysis, Figure~\ref{fig:first_token_attn} demonstrates the distribution of attention scores (after SoftMax) towards the first token in lengthy inputs (sequence length of 4096). We average attention scores across 256 sequences, with each sequence comprising 4096 tokens. The plotted data represent the attention allocated by the 4096th token to the initial token in every layer. Notably, the attention scores for the first token are significantly high, often exceeding half of the total attention, except for the two bottom layers. This observation empirically substantiates the preferential focus on the first token by the majority of layers and heads, irrespective of other tokens' distances within the sequence. Such a trend underscores the critical role of the initial tokens in a sequence, as their removal has a huge impact on language model performance due to a large portion of the denominator in the SoftMax function being removed.
\section{Llama-2-70B Attention Visualization}
\input{figtab/llama-2-70b-avg-attn.tex}
Figure~\ref{fig:attention_visualization} shows the attention visualization of Llama-2-7B, we further visualize the attention of Llama-2-70B in Figure~\ref{fig:llama-2-70b-attn}. We find the observation on Llama-2-7B also holds on Llama-2-70B, where the attention scores of the initial tokens are much higher than the rest of the tokens in most layers.
\section{Attention Sinks in Encoder Transformers}
\label{sec:bert_attn}
\input{figtab/bert_attn.tex}
In this paper, we mainly explore the attention sink phenomenon observed in autoregressive, decoder-only language models like GPT and Llama. Building upon the insights from Section~\ref{sec:attention_sink}, we propose that this phenomenon likely extends to other Transformer architectures, including encoder models such as BERT~\citep{Devlin2019BERTPO} and ViT~\citep{dosovitskiy2021image}. This assumption stems from the fact that these models share a similar Transformer structure and utilize SoftMax attention mechanisms. To substantiate our hypothesis, we analyze the attention patterns of BERT-base-uncased, as depicted in Figure~\ref{fig:bert-attn}. Our findings reveal that BERT-base-uncased exhibits the attention sink phenomenon, characterized by disproportionately high attention scores assigned to the \texttt{[SEP]} token in most layers. This indicates that the model consistently relies on the omnipresent \texttt{[SEP]} token as a focal point for attention. Furthermore, concurrent research by~\citeauthor{Darcet2023Registers} identifies similar attention spikes in Vision Transformers, attributed to random background patch tokens acting as "registers" for global image information. We contend that these "registers" are analogous to the attention sink phenomenon we observed, suggesting that this is a universal characteristic across all Transformer models.
\section{Using More Sink Tokens in the Pre-Training Stage}
Section~\ref{sec:pretrain-fix} illustrated that incorporating a single dedicated sink token in the pre-training stage doesn't affect model performance but enhances streaming performance by centralizing attention sinks to one token. This section delves into whether adding additional sink tokens during pre-training could further optimize the performance of pre-trained language models.
As depicted in Figure~\ref{fig:curve2}, our experiments show that incorporating either one or two sink tokens during pre-training results in pre-training loss curves that closely resemble those of the baseline (vanilla) model. However, as detailed in Table~\ref{tab:pretrain_eval2}, the introduction of a second sink token does not yield substantial improvements in performance across most benchmark tasks.
Further analysis, as shown in Table~\ref{tab:pretrain_fix2}, reveals that the inclusion of additional sink tokens does not enhance streaming performance. Interestingly, the model appears to rely on both sink tokens to maintain stable streaming performance. These findings suggest that while a single sink token is adequate for improving streaming performance, adding more sink tokens does not lead to further enhancements in overall language model performance. This contrasts with findings in Vision Transformers (ViT)~\citep{Darcet2023Registers}, where multiple "registers" have been found to be beneficial.
\begin{minipage}[t]{\textwidth}
\begin{minipage}[t]{0.3\textwidth}\vspace{0pt}
\centering
\includegraphics[width=\linewidth]{figures/two_learnable_sink_curve.png}
\captionof{figure}{\small Pre-training loss curves of models with 0, 1, and 2 sink tokens.}
\label{fig:curve2}
\end{minipage}
\hfill
\begin{minipage}[t]{0.66\textwidth}\vspace{0pt}
\centering
\small
\setlength{\tabcolsep}{3.0pt}
\captionof{table}{Zero-shot accuracy (in \%) across 7 NLP benchmarks, including ARC-[Challenge, Easy], HellaSwag, LAMBADA, OpenbookQA, PIQA, and Winogrande.}
\label{tab:pretrain_eval2}
\begin{tabular}{lccccccc}
\toprule
Methods & {ARC-c} & {ARC-e} & {HS} & {LBD} & {OBQA} & {PIQA} & {WG} \\
\midrule
Vanilla & 18.6 & 45.2 & 29.4 & 39.6 & 16.0 & 62.2 & 50.1 \\
+ 1 Sink Token & \textbf{19.6} & \textbf{45.6} & \textbf{29.8} & \textbf{39.9} & \textbf{16.6} & 62.6 & \textbf{50.8} \\
+ 2 Sink Tokens & 18.7 & 45.6 & 29.6 & 37.5 & 15.8 & \textbf{64.3} & 50.4 \\
\bottomrule
\end{tabular}
\end{minipage}
\end{minipage}
\begin{table}
\centering
\setlength{\tabcolsep}{3pt}
\small
\caption{Comparison of vanilla attention with prepending a zero token and a learnable sink token during pre-training. Cache config $x$+$y$ denotes adding $x$ initial tokens with $y$ recent tokens. Perplexity is evaluated on the first sample in the PG19 test set.}
\label{tab:pretrain_fix2}
\begin{tabular}{lcccc}
\toprule
Cache Config & 0+1024 & 1+1023 & 2+1022 & 4+1020 \\
\midrule
Vanilla & 27.87 & 18.49 & 18.05 & 18.05 \\
+ 1 Sink Token & 1235 & \textbf{18.01} & 18.01 & 18.02 \\
+ 2 Sink Tokens & 1262 & 25.73 & 18.05 & 18.05 \\
\bottomrule
\end{tabular}
\end{table}
\begin{table}
\centering
\caption{\small Accuracy (in \%) on StreamEval with increasing query-answer distance. Each line in StreamEval contains 23 tokens. Accuracies are averaged over 100 samples, and each sample contains 100 queries.}
\label{tab:kv_lines}
\begin{tabular}{llcccc}
\toprule
\multicolumn{2}{c}{Llama-2-7B-32K-Instruct} & \multicolumn{4}{c}{Cache Config} \\
\cmidrule(lr){1-2} \cmidrule(lr){3-6}
Line Distances & Token Distances & 4+2044 & 4+4092 & 4+8188 & 4+16380 \\
\midrule
20 & 460 & 85.80 & 84.60 & 81.15 & 77.65 \\
40 & 920 & 80.35 & 83.80 & 81.25 & 77.50 \\
60 & 1380 & 79.15 & 82.80 & 81.50 & 78.50 \\
80 & 1840 & 75.30 & 77.15 & 76.40 & 73.80 \\
100 & 2300 & 0.00 & 61.60 & 50.10 & 40.50 \\
150 & 3450 & 0.00 & 68.20 & 58.30 & 38.45 \\
200 & 4600 & 0.00 & 0.00 & 62.75 & 46.90 \\
400 & 9200 & 0.00 & 0.00 & 0.00 & 45.70 \\
600 & 13800 & 0.00 & 0.00 & 0.00 & 28.50 \\
800 & 18400 & 0.00 & 0.00 & 0.00 & 0.00 \\
1000 & 23000 & 0.00 & 0.00 & 0.00 & 0.00 \\
\bottomrule
\end{tabular}
\end{table}
\begin{table}
\centering
\small
\setlength{\tabcolsep}{3pt}
\caption{\small Performance comparison of \method against the default truncation baseline in LongBench~\citep{bai2023longbench}. The baseline truncates inputs to 1750 initial and 1750 final tokens. \method 4+3496 uses 4 attention sink tokens and 3496 recent tokens, while \method 1750+1750 uses 1750 tokens for both initial and recent segments.}
\label{tab:longbench}
\begin{tabular}{lcccccc}
\toprule
\multirow{2}{*}{Llama2-7B-chat} & \multicolumn{2}{c}{Single-Document QA} & \multicolumn{2}{c}{Multi-Document QA} & \multicolumn{2}{c}{Summarization} \\
& NarrativeQA & Qasper & HotpotQA & 2WikiMQA & GovReport & MultiNews \\
\midrule
Truncation 1750+1750 & 18.7 & 19.2 & 25.4 & 32.8 & 27.3 & 25.8 \\
\midrule
\method 4+3496 & 11.6 & 16.9 & 21.6 & 28.2 & 23.9 & 25.5 \\
\method 1750+1750 & 18.2 & 19.7 & 24.9 & 32.0 & 26.3 & 25.9 \\
\bottomrule
\end{tabular}
\end{table}
\begin{figure}[htbp]
\centering
\small
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-7b-128-attn/0.png}
\caption{\small Layer 0 Head 0}
\end{subfigure}
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-7b-128-attn/6.png}
\caption{\small Layer 6 Head 0}
\end{subfigure}
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-7b-128-attn/12.png}
\caption{\small Layer 12 Head 0}
\end{subfigure}
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-7b-128-attn/18.png}
\caption{\small Layer 18 Head 0}
\end{subfigure}
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-7b-128-attn/24.png}
\caption{\small Layer 24 Head 0}
\end{subfigure}
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-7b-128-attn/30.png}
\caption{\small Layer 30 Head 0}
\end{subfigure}
\caption{\small Visualization of the \emph{average} attention logits in Llama-2-7B over 256 sentences, each with a length of 128.}
\label{fig:llama-2-7b-128-attn}
\end{figure}
\begin{figure}
\centering
\includegraphics[width=\textwidth]{figures/Llama-2-7B-first-token-attn.png}
\caption{Visualization of attention scores (after SoftMax) on the first token across layers in Llama-2-7B. Attention Scores are the 4096th token's attention towards the first token in each layer. The error bars are the standard deviation of the first token's attention scores across different heads in one layer. Results are averaged over 256 sentences, each having a length of 4096 tokens.}
\label{fig:first_token_attn}
\end{figure}
\begin{figure}[htbp]
\centering
\small
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/0_0.png}
\caption{\small Layer 0 Head 0}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/0_1.png}
\caption{\small Layer 0 Head 1}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/8_0.png}
\caption{\small Layer 8 Head 0}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/8_1.png}
\caption{\small Layer 8 Head 1}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/16_0.png}
\caption{\small Layer 16 Head 0}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/16_1.png}
\caption{\small Layer 16 Head 1}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/24_0.png}
\caption{\small Layer 24 Head 0}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/24_1.png}
\caption{\small Layer 24 Head 1}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/32_0.png}
\caption{\small Layer 32 Head 0}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/32_1.png}
\caption{\small Layer 32 Head 1}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/40_0.png}
\caption{\small Layer 40 Head 0}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/40_1.png}
\caption{\small Layer 40 Head 1}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/48_0.png}
\caption{\small Layer 48 Head 0}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/48_1.png}
\caption{\small Layer 48 Head 1}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/56_0.png}
\caption{\small Layer 56 Head 0}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/56_1.png}
\caption{\small Layer 56 Head 1}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/64_0.png}
\caption{\small Layer 64 Head 0}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/64_1.png}
\caption{\small Layer 64 Head 1}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/72_0.png}
\caption{\small Layer 72 Head 0}
\end{subfigure}
\begin{subfigure}{0.24\textwidth}
\includegraphics[width=\linewidth]{figures/llama-2-70b-avg-attn/72_1.png}
\caption{\small Layer 72 Head 1}
\end{subfigure}
\caption{\small Visualization of the \emph{average} attention logits in Llama-2-70B over 256 sentences, each with a length of 16.}
\label{fig:llama-2-70b-attn}
\end{figure}
\begin{figure}[htbp]
\centering
\small
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/bert-attn/0.png}
\end{subfigure}
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/bert-attn/1.png}
\end{subfigure}
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/bert-attn/2.png}
\end{subfigure}
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/bert-attn/3.png}
\end{subfigure}
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/bert-attn/4.png}
\end{subfigure}
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/bert-attn/5.png}
\end{subfigure}
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/bert-attn/6.png}
\end{subfigure}
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/bert-attn/7.png}
\end{subfigure}
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/bert-attn/8.png}
\end{subfigure}
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/bert-attn/9.png}
\end{subfigure}
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/bert-attn/10.png}
\end{subfigure}
\begin{subfigure}{0.32\textwidth}
\includegraphics[width=\linewidth]{figures/bert-attn/11.png}
\end{subfigure}
\caption{\small Visualization of attention maps for sentence \emph{``StreamingLLM can work on infinite-length texts without compromising efficiency and performance.''} in BERT-base-uncased.}
\label{fig:bert-attn}
\end{figure}
\begin{thebibliography}{63}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
\providecommand{\doi}[1]{doi: #1}\else
\providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi
\bibitem[Ainslie et~al.(2020)Ainslie, Ontanon, Alberti, Cvicek, Fisher, Pham,
Ravula, Sanghai, Wang, and Yang]{ainslie2020etc}
Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher,
Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li~Yang.
\newblock Etc: Encoding long and structured inputs in transformers, 2020.
\bibitem[Almazrouei et~al.(2023)Almazrouei, Alobeidli, Alshamsi, Cappelli,
Cojocaru, Debbah, Goffinet, Heslow, Launay, Malartic, Noune, Pannier, and
Penedo]{falcon40b}
Ebtesam Almazrouei, Hamza Alobeidli, Abdulaziz Alshamsi, Alessandro Cappelli,
Ruxandra Cojocaru, Merouane Debbah, Etienne Goffinet, Daniel Heslow, Julien
Launay, Quentin Malartic, Badreddine Noune, Baptiste Pannier, and Guilherme
Penedo.
\newblock {Falcon-40B}: an open large language model with state-of-the-art
performance.
\newblock 2023.
\bibitem[Anagnostidis et~al.(2023)Anagnostidis, Pavllo, Biggio, Noci, Lucchi,
and Hofmann]{anagnostidis2023dynamic}
Sotiris Anagnostidis, Dario Pavllo, Luca Biggio, Lorenzo Noci, Aurelien Lucchi,
and Thomas Hofmann.
\newblock Dynamic context pruning for efficient and interpretable
autoregressive transformers, 2023.
\bibitem[Bai et~al.(2023)Bai, Lv, Zhang, Lyu, Tang, Huang, Du, Liu, Zeng, Hou,
Dong, Tang, and Li]{bai2023longbench}
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang,
Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and
Juanzi Li.
\newblock Longbench: A bilingual, multitask benchmark for long context
understanding.
\newblock \emph{arXiv preprint arXiv:2308.14508}, 2023.
\bibitem[Beltagy et~al.(2020)Beltagy, Peters, and Cohan]{beltagy2020longformer}
Iz~Beltagy, Matthew~E. Peters, and Arman Cohan.
\newblock Longformer: The long-document transformer, 2020.
\newblock arXiv:2004.05150.
\bibitem[Biderman et~al.(2023)Biderman, Schoelkopf, Anthony, Bradley, O'Brien,
Hallahan, Khan, Purohit, Prashanth, Raff, Skowron, Sutawika, and van~der
Wal]{biderman2023pythia}
Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle
O'Brien, Eric Hallahan, Mohammad~Aflah Khan, Shivanshu Purohit, USVSN~Sai
Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van~der
Wal.
\newblock Pythia: A suite for analyzing large language models across training
and scaling, 2023.
\bibitem[Bisk et~al.(2020)Bisk, Zellers, Bras, Gao, and Choi]{piqa}
Yonatan Bisk, Rowan Zellers, Ronan~Le Bras, Jianfeng Gao, and Yejin Choi.
\newblock Piqa: Reasoning about physical commonsense in natural language.
\newblock In \emph{Thirty-Fourth AAAI Conference on Artificial Intelligence},
2020.
\bibitem[bloc97(2023)]{blocntkaware}
bloc97.
\newblock {NTK-Aware Scaled RoPE allows LLaMA models to have extended (8k+)
context size without any fine-tuning and minimal perplexity degradation.},
2023.
\newblock URL
\url{https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/}.
\bibitem[Bondarenko et~al.(2023)Bondarenko, Nagel, and
Blankevoort]{bondarenko2023quantizable}
Yelysei Bondarenko, Markus Nagel, and Tijmen Blankevoort.
\newblock Quantizable transformers: Removing outliers by helping attention
heads do nothing, 2023.
\bibitem[Brown et~al.(2020)Brown, Mann, Ryder, Subbiah, Kaplan, Dhariwal,
Neelakantan, Shyam, Sastry, Askell, et~al.]{brown2020language}
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared~D Kaplan, Prafulla
Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell,
et~al.
\newblock Language models are few-shot learners.
\newblock \emph{Advances in neural information processing systems},
33:\penalty0 1877--1901, 2020.
\bibitem[Chen et~al.(2021)Chen, Tworek, Jun, Yuan, de~Oliveira~Pinto, Kaplan,
Edwards, Burda, Joseph, Brockman, Ray, Puri, Krueger, Petrov, Khlaaf, Sastry,
Mishkin, Chan, Gray, Ryder, Pavlov, Power, Kaiser, Bavarian, Winter, Tillet,
Such, Cummings, Plappert, Chantzis, Barnes, Herbert-Voss, Guss, Nichol,
Paino, Tezak, Tang, Babuschkin, Balaji, Jain, Saunders, Hesse, Carr, Leike,
Achiam, Misra, Morikawa, Radford, Knight, Brundage, Murati, Mayer, Welinder,
McGrew, Amodei, McCandlish, Sutskever, and Zaremba]{chen2021evaluating}
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique~Ponde
de~Oliveira~Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph,
Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy
Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder,
Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens
Winter, Philippe Tillet, Felipe~Petroski Such, Dave Cummings, Matthias
Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert-Voss,
William~Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor
Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher
Hesse, Andrew~N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa,
Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter
Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and
Wojciech Zaremba.
\newblock Evaluating large language models trained on code, 2021.
\bibitem[Chen et~al.(2023)Chen, Wong, Chen, and Tian]{chen2023extending}
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian.
\newblock Extending context window of large language models via positional
interpolation, 2023.
\newblock arXiv: 2306.15595.
\bibitem[Chiang et~al.(2023)Chiang, Li, Lin, Sheng, Wu, Zhang, Zheng, Zhuang,
Zhuang, Gonzalez, Stoica, and Xing]{vicuna2023}
Wei-Lin Chiang, Zhuohan Li, Zi~Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin
Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph~E. Gonzalez, Ion Stoica, and
Eric~P. Xing.
\newblock Vicuna: An open-source chatbot impressing gpt-4 with 90\%* chatgpt
quality, March 2023.
\newblock URL \url{https://lmsys.org/blog/2023-03-30-vicuna/}.
\bibitem[Child et~al.(2019)Child, Gray, Radford, and
Sutskever]{child2019generating}
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever.
\newblock Generating long sequences with sparse transformers.
\newblock 2019.
\bibitem[Clark et~al.(2018)Clark, Cowhey, Etzioni, Khot, Sabharwal, Schoenick,
and Tafjord]{arc}
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa
Schoenick, and Oyvind Tafjord.
\newblock Think you have solved question answering? try arc, the ai2 reasoning
challenge.
\newblock \emph{arXiv:1803.05457v1}, 2018.
\bibitem[Dao(2023)]{dao2023flashattention2}
Tri Dao.
\newblock Flash{A}ttention-2: Faster attention with better parallelism and work
partitioning.
\newblock 2023.
\bibitem[Dao et~al.(2022)Dao, Fu, Ermon, Rudra, and Ré]{dao2022flashattention}
Tri Dao, Daniel~Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré.
\newblock {FlashAttention}: Fast and memory-efficient exact attention with
{IO}-awareness, 2022.
\newblock arXiv:2205.14135.
\bibitem[Darcet et~al.(2023)Darcet, Oquab, Mairal, and
Bojanowski]{Darcet2023Registers}
Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski.
\newblock Vision transformers need registers, 2023.
\bibitem[Dasigi et~al.(2021)Dasigi, Lo, Beltagy, Cohan, Smith, and
Gardner]{dasigi2021qasper}
Pradeep Dasigi, Kyle Lo, Iz~Beltagy, Arman Cohan, Noah~A. Smith, and Matt
Gardner.
\newblock A dataset of information-seeking questions and answers anchored in
research papers, 2021.
\bibitem[Devlin et~al.(2019)Devlin, Chang, Lee, and
Toutanova]{Devlin2019BERTPO}
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.
\newblock Bert: Pre-training of deep bidirectional transformers for language
understanding.
\newblock In \emph{North American Chapter of the Association for Computational
Linguistics}, 2019.
\newblock URL \url{https://api.semanticscholar.org/CorpusID:52967399}.
\bibitem[Dosovitskiy et~al.(2021)Dosovitskiy, Beyer, Kolesnikov, Weissenborn,
Zhai, Unterthiner, Dehghani, Minderer, Heigold, Gelly, Uszkoreit, and
Houlsby]{dosovitskiy2021image}
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn,
Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg
Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.
\newblock An image is worth 16x16 words: Transformers for image recognition at
scale, 2021.
\bibitem[Fabbri et~al.(2019)Fabbri, Li, She, Li, and
Radev]{fabbri2019multinews}
Alexander Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir Radev.
\newblock Multi-news: A large-scale multi-document summarization dataset and
abstractive hierarchical model.
\newblock In \emph{Proceedings of the 57th Annual Meeting of the Association
for Computational Linguistics}, 2019.
\bibitem[Gao et~al.(2020)Gao, Biderman, Black, Golding, Hoppe, Foster, Phang,
He, Thite, Nabeshima, Presser, and Leahy]{pile}
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles
Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser,
and Connor Leahy.
\newblock The {P}ile: An 800gb dataset of diverse text for language modeling.
\newblock \emph{arXiv preprint arXiv:2101.00027}, 2020.
\bibitem[Goyal \& Durrett(2020)Goyal and
Durrett]{goyal-durrett-2020-evaluating}
Tanya Goyal and Greg Durrett.
\newblock Evaluating factuality in generation with dependency-level entailment.
\newblock In \emph{Findings of the Association for Computational Linguistics:
EMNLP 2020}, Online, 2020. Association for Computational Linguistics.
\bibitem[Han et~al.(2023)Han, Wang, Xiong, Chen, Ji, and
Wang]{han2023lminfinite}
Chi Han, Qifan Wang, Wenhan Xiong, Yu~Chen, Heng Ji, and Sinong Wang.
\newblock {LM}-{I}nfinite: Simple on-the-fly length generalization for large
language models, 2023.
\bibitem[Ho et~al.(2020)Ho, Duong~Nguyen, Sugawara, and
Aizawa]{ho-etal-2020-constructing}
Xanh Ho, Anh-Khoa Duong~Nguyen, Saku Sugawara, and Akiko Aizawa.
\newblock Constructing a multi-hop {QA} dataset for comprehensive evaluation of
reasoning steps.
\newblock In \emph{Proceedings of the 28th International Conference on
Computational Linguistics}, December 2020.
\bibitem[Huang et~al.(2021)Huang, Cao, Parulian, Ji, and
Wang]{huang2021govreport}
Luyang Huang, Shuyang Cao, Nikolaus Parulian, Heng Ji, and Lu~Wang.
\newblock Efficient attentions for long document summarization, 2021.
\bibitem[kaiokendev(2023)]{kaiokendev}
kaiokendev.
\newblock {Things I'm learning while training superhot.}, 2023.
\newblock URL \url{https://kaiokendev.github.io/til#extending-context-to-8k}.
\bibitem[Kamalloo et~al.(2023)Kamalloo, Dziri, Clarke, and
Rafiei]{kamalloo2023evaluating}
Ehsan Kamalloo, Nouha Dziri, Charles L.~A. Clarke, and Davood Rafiei.
\newblock Evaluating open-domain question answering in the era of large
language models, 2023.
\bibitem[Kitaev et~al.(2020)Kitaev, Kaiser, and Levskaya]{kitaev2020reformer}
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya.
\newblock Reformer: The efficient transformer.
\newblock In \emph{8th International Conference on Learning Representations,
ICLR 2020}. OpenReview.net, April 2020.
\bibitem[Kočiský et~al.(2017)Kočiský, Schwarz, Blunsom, Dyer, Hermann,
Melis, and Grefenstette]{narrativeqa}
Tomáš Kočiský, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl~Moritz
Hermann, Gábor Melis, and Edward Grefenstette.
\newblock The narrativeqa reading comprehension challenge, 2017.
\bibitem[Li et~al.(2023)Li, Shao, Xie, Sheng, Zheng, Gonzalez, Stoica, Ma, ,
and Zhang]{longchat2023}
Dacheng Li, Rulin Shao, Anze Xie, Ying Sheng, Lianmin Zheng, Joseph~E.
Gonzalez, Ion Stoica, Xuezhe Ma, , and Hao Zhang.
\newblock How long can open-source llms truly promise on context length?, June
2023.
\newblock URL \url{https://lmsys.org/blog/2023-06-29-longchat}.
\bibitem[Liu et~al.(2023)Liu, Lin, Hewitt, Paranjape, Bevilacqua, Petroni, and
Liang]{liu2023lost}
Nelson~F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua,
Fabio Petroni, and Percy Liang.
\newblock Lost in the middle: How language models use long contexts, 2023.
\bibitem[Mihaylov et~al.(2018)Mihaylov, Clark, Khot, and
Sabharwal]{OpenBookQA2018}
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal.
\newblock Can a suit of armor conduct electricity? a new dataset for open book
question answering.
\newblock In \emph{EMNLP}, 2018.
\bibitem[Miller(2023)]{softmax1}
Evan Miller.
\newblock Attention is off by one, 2023.
\newblock URL \url{https://www.evanmiller.org/attention-is-off-by-one.html}.
\bibitem[OpenAI(2023)]{openai2023gpt4}
OpenAI.
\newblock Gpt-4 technical report, 2023.
\bibitem[Paperno et~al.(2016)Paperno, Kruszewski, Lazaridou, Pham, Bernardi,
Pezzelle, Baroni, Boleda, and Fern{\'a}ndez]{paperno-etal-2016-lambada}
Denis Paperno, Germ{\'a}n Kruszewski, Angeliki Lazaridou, Ngoc~Quan Pham,
Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel
Fern{\'a}ndez.
\newblock The {LAMBADA} dataset: Word prediction requiring a broad discourse
context.
\newblock In \emph{Proceedings of the 54th Annual Meeting of the Association
for Computational Linguistics (Volume 1: Long Papers)}, pp.\ 1525--1534,
Berlin, Germany, August 2016. Association for Computational Linguistics.
\newblock \doi{10.18653/v1/P16-1144}.
\newblock URL \url{https://aclanthology.org/P16-1144}.
\bibitem[Peng et~al.(2023)Peng, Quesnelle, Fan, and Shippole]{peng2023yarn}
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole.
\newblock Yarn: Efficient context window extension of large language models,
2023.
\bibitem[Pope et~al.(2022)Pope, Douglas, Chowdhery, Devlin, Bradbury, Levskaya,
Heek, Xiao, Agrawal, and Dean]{pope2022efficiently}
Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury,
Anselm Levskaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean.
\newblock Efficiently scaling transformer inference.
\newblock \emph{arXiv preprint arXiv:2211.05102}, 2022.
\bibitem[Press et~al.(2022)Press, Smith, and Lewis]{alibi}
Ofir Press, Noah Smith, and Mike Lewis.
\newblock Train short, test long: Attention with linear biases enables input
length extrapolation.
\newblock In \emph{International Conference on Learning Representations}, 2022.
\newblock URL \url{https://openreview.net/forum?id=R8sQPpGCv0}.
\bibitem[Radford et~al.(2018)Radford, Narasimhan, Salimans, Sutskever,
et~al.]{radford2018improving}
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et~al.
\newblock Improving language understanding by generative pre-training.
\newblock 2018.
\bibitem[Rae et~al.(2020)Rae, Potapenko, Jayakumar, Hillier, and
Lillicrap]{Rae2020Compressive}
Jack~W. Rae, Anna Potapenko, Siddhant~M. Jayakumar, Chloe Hillier, and
Timothy~P. Lillicrap.
\newblock Compressive transformers for long-range sequence modelling.
\newblock In \emph{International Conference on Learning Representations}, 2020.
\bibitem[Rozière et~al.(2023)Rozière, Gehring, Gloeckle, Sootla, Gat, Tan,
Adi, Liu, Remez, Rapin, Kozhevnikov, Evtimov, Bitton, Bhatt, Ferrer,
Grattafiori, Xiong, Défossez, Copet, Azhar, Touvron, Martin, Usunier,
Scialom, and Synnaeve]{rozière2023code}
Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat,
Xiaoqing~Ellen Tan, Yossi Adi, Jingyu Liu, Tal Remez, Jérémy Rapin, Artyom
Kozhevnikov, Ivan Evtimov, Joanna Bitton, Manish Bhatt, Cristian~Canton
Ferrer, Aaron Grattafiori, Wenhan Xiong, Alexandre Défossez, Jade Copet,
Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usunier, Thomas Scialom,
and Gabriel Synnaeve.
\newblock Code {L}lama: Open foundation models for code, 2023.
\bibitem[Sakaguchi et~al.(2019)Sakaguchi, Bras, Bhagavatula, and
Choi]{sakaguchi2019winogrande}
Keisuke Sakaguchi, Ronan~Le Bras, Chandra Bhagavatula, and Yejin Choi.
\newblock Winogrande: An adversarial winograd schema challenge at scale.
\newblock \emph{arXiv preprint arXiv:1907.10641}, 2019.
\bibitem[Schulman et~al.(2022)Schulman, Zoph, Kim, Hilton, Menick, Weng, Uribe,
Fedus, Metz, Pokorny, et~al.]{schulman2022chatgpt}
John Schulman, Barret Zoph, Christina Kim, Jacob Hilton, Jacob Menick, Jiayi
Weng, Juan Felipe~Ceron Uribe, Liam Fedus, Luke Metz, Michael Pokorny, et~al.
\newblock Chatgpt: Optimizing language models for dialogue.
\newblock \emph{OpenAI blog}, 2022.
\bibitem[Su et~al.(2021)Su, Lu, Pan, Murtadha, Wen, and Liu]{su2021roformer}
Jianlin Su, Yu~Lu, Shengfeng Pan, Ahmed Murtadha, Bo~Wen, and Yunfeng Liu.
\newblock Roformer: Enhanced transformer with rotary position embedding.
\newblock \emph{arXiv preprint arXiv:2104.09864}, 2021.
\bibitem[Taori et~al.(2023)Taori, Gulrajani, Zhang, Dubois, Li, Guestrin,
Liang, and Hashimoto]{alpaca}
Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos
Guestrin, Percy Liang, and Tatsunori~B. Hashimoto.
\newblock Stanford alpaca: An instruction-following llama model.
\newblock \url{https://github.com/tatsu-lab/stanford_alpaca}, 2023.
\bibitem[Tay et~al.(2022)Tay, Dehghani, Bahri, and Metzler]{tay2022efficient}
Yi~Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler.
\newblock Efficient transformers: A survey.
\newblock \emph{ACM Computing Surveys}, 55\penalty0 (6), dec 2022.
\newblock ISSN 0360-0300.
\bibitem[Team(2023)]{MosaicML2023Introducing}
MosaicML~NLP Team.
\newblock Introducing mpt-7b: A new standard for open-source, commercially
usable llms, 2023.
\newblock URL \url{www.mosaicml.com/blog/mpt-7b}.
\newblock Accessed: 2023-05-05.
\bibitem[Together(2023)]{togetherlong}
Together.
\newblock Llama-2-7b-32k-instruct — and fine-tuning for llama-2 models with
together api, June 2023.
\newblock URL \url{https://together.ai/blog/llama-2-7b-32k-instruct}.
\bibitem[Touvron et~al.(2023{\natexlab{a}})Touvron, Lavril, Izacard, Martinet,
Lachaux, Lacroix, Rozi{\`e}re, Goyal, Hambro, Azhar,
et~al.]{touvron2023llama}
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne
Lachaux, Timoth{\'e}e Lacroix, Baptiste Rozi{\`e}re, Naman Goyal, Eric
Hambro, Faisal Azhar, et~al.
\newblock Llama: Open and efficient foundation language models.
\newblock \emph{arXiv preprint arXiv:2302.13971}, 2023{\natexlab{a}}.
\bibitem[Touvron et~al.(2023{\natexlab{b}})Touvron, Martin, Stone, Albert,
Almahairi, Babaei, Bashlykov, Batra, Bhargava, Bhosale,
et~al.]{touvron2023llama2}
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine
Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale,
et~al.
\newblock Llama 2: Open foundation and fine-tuned chat models.
\newblock \emph{arXiv preprint arXiv:2307.09288}, 2023{\natexlab{b}}.
\bibitem[Wang et~al.(2021)Wang, Zhang, and Han]{wang2021spatten}
Hanrui Wang, Zhekai Zhang, and Song Han.
\newblock Spatten: Efficient sparse attention architecture with cascade token
and head pruning.
\newblock \emph{HPCA}, 2021.
\bibitem[Wang et~al.(2020)Wang, Li, Khabsa, Fang, and Ma]{wang2020linformer}
Sinong Wang, Belinda~Z Li, Madian Khabsa, Han Fang, and Hao Ma.
\newblock Linformer: Self-attention with linear complexity.
\newblock 2020.
\bibitem[Wolf et~al.(2020)Wolf, Debut, Sanh, Chaumond, Delangue, Moi, Cistac,
Rault, Louf, Funtowicz, Davison, Shleifer, von Platen, Ma, Jernite, Plu, Xu,
Scao, Gugger, Drame, Lhoest, and Rush]{wolf2020huggingfaces}
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue,
Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe
Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien
Plu, Canwen Xu, Teven~Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest,
and Alexander~M. Rush.
\newblock Huggingface's transformers: State-of-the-art natural language
processing, 2020.
\bibitem[Xiao et~al.(2023)Xiao, Lin, Seznec, Wu, Demouth, and
Han]{xiao2023smoothquant}
Guangxuan Xiao, Ji~Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han.
\newblock {S}mooth{Q}uant: Accurate and efficient post-training quantization
for large language models.
\newblock In \emph{Proceedings of the 40th International Conference on Machine
Learning}, 2023.
\bibitem[Yang et~al.(2018)Yang, Qi, Zhang, Bengio, Cohen, Salakhutdinov, and
Manning]{yang2018hotpotqa}
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William~W. Cohen, Ruslan
Salakhutdinov, and Christopher~D. Manning.
\newblock {HotpotQA}: A dataset for diverse, explainable multi-hop question
answering.
\newblock In \emph{Conference on Empirical Methods in Natural Language
Processing ({EMNLP})}, 2018.
\bibitem[Zaheer et~al.(2020{\natexlab{a}})Zaheer, Guruganesh, Dubey, Ainslie,
Alberti, Ontanon, Pham, Ravula, Wang, Yang, and Ahmed]{zaheer2020big}
Manzil Zaheer, Guru Guruganesh, Kumar~Avinava Dubey, Joshua Ainslie, Chris
Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li~Yang,
and Amr Ahmed.
\newblock {Big Bird}: {T}ransformers for longer sequences.
\newblock In \emph{Proc. of NeurIPS}, volume~33, 2020{\natexlab{a}}.
\bibitem[Zaheer et~al.(2020{\natexlab{b}})Zaheer, Guruganesh, Dubey, Ainslie,
Alberti, Ontañón, Pham, Ravula, Wang, Yang, and Ahmed]{zaheer2020bigbird}
Manzil Zaheer, Guru Guruganesh, Kumar~Avinava Dubey, Joshua Ainslie, Chris
Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang,
Li~Yang, and Amr Ahmed.
\newblock Big bird: Transformers for longer sequences.
\newblock In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina
Balcan, and Hsuan-Tien Lin (eds.), \emph{Advances in Neural Information
Processing Systems 33: Annual Conference on Neural Information Processing
Systems 2020, NeurIPS 2020}. Curran Associates, Inc., 2020{\natexlab{b}}.
\bibitem[Zellers et~al.(2019)Zellers, Holtzman, Bisk, Farhadi, and
Choi]{hellaswag}
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi.
\newblock Hellaswag: Can a machine really finish your sentence?
\newblock \emph{CoRR}, abs/1905.07830, 2019.
\newblock URL \url{http://arxiv.org/abs/1905.07830}.
\bibitem[Zhang et~al.(2022)Zhang, Roller, Goyal, Artetxe, Chen, Chen, Dewan,
Diab, Li, Lin, Mihaylov, Ott, Shleifer, Shuster, Simig, Koura, Sridhar, Wang,
and Zettlemoyer]{zhang2022opt}
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui
Chen, Christopher Dewan, Mona Diab, Xian Li, Xi~Victoria Lin, Todor Mihaylov,
Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit~Singh Koura, Anjali
Sridhar, Tianlu Wang, and Luke Zettlemoyer.
\newblock Opt: Open pre-trained transformer language models, 2022.
\bibitem[Zhang et~al.(2023{\natexlab{a}})Zhang, Ladhak, Durmus, Liang, McKeown,
and Hashimoto]{zhang2023benchmarking}
Tianyi Zhang, Faisal Ladhak, Esin Durmus, Percy Liang, Kathleen McKeown, and
Tatsunori~B. Hashimoto.
\newblock Benchmarking large language models for news summarization,
2023{\natexlab{a}}.
\bibitem[Zhang et~al.(2023{\natexlab{b}})Zhang, Sheng, Zhou, Chen, Zheng, Cai,
Song, Tian, Ré, Barrett, Wang, and Chen]{zhang2023h2o}
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai,
Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, and
Beidi Chen.
\newblock H$_2$o: Heavy-hitter oracle for efficient generative inference of
large language models, 2023{\natexlab{b}}.
\end{thebibliography}