#  An experimental sorting method
for improving metagenomic data encoding

Diogo Pratas1,2,3 and Armando J. Pinho1,2
1IEETA/LASI - Institute of Electronics and Informatics Engineering of Aveiro
2DETI - Department of Electronics, Telecommunications and Informatics
University of Aveiro, 3810-193 Aveiro, Portugal<EMAIL_ADDRESS>3DoV -
Department of Virology, University of Helsinki, 00014 Helsinki, Finland

###### Abstract

Minimizing data storage poses a significant challenge in large-scale
metagenomic projects. In this paper, we present a new method for improving the
encoding of FASTQ files generated by metagenomic sequencing. This method
incorporates metagenomic classification followed by a recursive filter for
clustering reads by DNA sequence similarity to improve the overall reference-
free compression. In the results, we show an overall improvement in the
compression of several datasets. As hypothesized, we show a progressive
compression gain for higher coverage depth and number of identified species.
Additionally, we provide an implementation that is freely available at
https://github.com/cobilab/mizar and can be customized to work with other
FASTQ compression tools.

##  1 Introduction

Metagenomics is the study of genetic material collected from environmental
samples, providing insights into the diversity, composition, and functional
potential of microbial communities in various ecosystems. This includes
permeates extreme ecosystems, such as uranium mines [1], intricate landscapes
of both soft and hard tissues [2, 3], vital food crops [4], expansive marine
habitats [5], or in enigmatic depths of subterranean ecosystems [6].

Metagenomics data is the digital material recovered from these samples by a
method called sequencing, specifically using a format called FASTQ. The FASTQ
format contains multiple reads, usually millions or billions. FASTQ reads are
constituted by three channels of information: headers, DNA sequences, and
quality-scores. In each read, DNA sequences and quality-scores must have the
same number of characters, usually 150 or 250 characters. More recently, the
technology of long reads is generating reads with several thousand or tens of
thousand symbols, but at the expense of slightly less accuracy. The DNA
sequence expresses the interpretation of the sequencing machine of a piece of
the genetic material, while the quality-scores represent the quality in which
the sequence has been sequenced (analogous to a confidence value). The headers
(sometimes called identifiers) are specially used to know if the reads are
paired (a correlation between pairs of reads that can extend in some way the
size of the DNA sequences) or single reads and contain an identification to
preserve the uniqueness of the reads.

To speedup the sequencing process, the DNA is sequenced in parallel (cloning
and reading) which generates a FASTQ file with reads in somewhat random order.
These small DNA sequences can overlap given the random sequencing process and
may have non sequenced parts. Fortunately, to reduce these holes and to
increase the certainty in the sequencing process, the same regions are usually
sequenced several times, but not necessarily starting at the same position.
This property creates reads with the same DNA sub-sequences, where the average
total number of a given DNA symbol that has been sequenced at a certain
position of the genome depends on the coverage depth. Usually, the average
coverage depth varies between 5 and 300, where a coverage depth of 30 (for
human sequencing) is considered of high-quality. For smaller genome species,
higher depth coverages are often, specially if enrichment processes have been
applied in the sequencing process.

For some applications using paired-end reads, such as genome reconstruction,
the information of the link between each forward and reverse read is required,
while the order in which each pair appears is not mandatory. However, since
the header contains the same string, with the exception at the end of a digit
to describe if the read is in forward or reverse, grouping these pairs is
trivial and does not require much information than a simple prefix-match
program. Therefore, for most applications, the order of the reads can be
neglected without compromising its purpose.

The compression of these data has widely been addressed using general-purpose
data compressors. These compressors are prepared to deal with all kinds of
finite data and some of them can be very fast and already be included, by
default, in the operating systems. However, for compressing specific data,
they can be much less compression efficient, especially in the case of FASTQ
data. On the other hand, specialized FASTQ data compressors can take into
account the different information channels, as well as to use methods more
appropriate for each of the information channels. Additionally, they can model
specific characteristics of DNA sequences, such as inverted repeats [7].

There are many lossless FASTQ data compressors, for example, SCALCE [8], quip
[9], Fqzcomp [10], DSRC and DSRC2 [11], ORCOM [12], LEON [13], LW-FQZip2 [14],
HARC [15], FaStore [16], SPRING [17], FastqCLS [18], among others. From these,
only a few purely explore clustering or sorting the order of the reads for
improving the compression [19], for example, HARC [15], FaStore [16], and
FastqCLS [18]. However, as far as we know, none of the methods incorporate
metagenomic classification and target metagenomic samples.

In this article, we provide a compression-based method that explores
metagenomic classification and recursive filtering by similarity for sorting
the order of the reads and increase the overall compression. Our main
objective is to decrease the size of FASTQ files from metagenomics, without
loss of information, while neglecting only the information associated to the
order of which the reads appear.

Our main hypothesis stands for the successful ordering of FASTQ reads, through
the approximation of similar classified DNA sequences using compression-based
methods, in order to increase the overall FASTQ compression. Notice that the
FASTQ headers and quality-scores can negatively affect the compression
capacity when changing the order of the reads and, therefore, the overall
compression gain must be superior to this eventual loss. In our second
hypothesis, we argue that the compression gain using the proposed methodology
is superior even if we store the information to transform the data to its
original order (completely lossless). Our third hypothesis stands for the
progressive increase in the compression gain for higher coverage depth and
number of identified species, when our method is used.

In the next section, we present the method along with its implementation.
Then, we present the results for validating our hypothesis and, finally, draw
some conclusions.

##  2 Method

We propose a method that uses metagenomic classification for finding the
sequences that are more probable to be contained in the sample, followed by
recursive read filtering for clustering similar reads according to the select
references, and compression using the sorted file, for improving the
compression of FASTQ metagenomic data. The proposed method is designed for
one-time compression followed by multiple end-users’ decompression. It only
necessitates complete computation of the sorting process during compression,
enabling swift and resource-efficient decompression for all end-users.

Figure 1: Architecture of the proposed methodology.

Figure 1 depicts the architecture of the proposed method, where each component
is explained in the following subsections, namely the process of metagenomic
classification, reads filtering, and reads compression.

### Metagenomic classification

In metagenomics, the DNA from several organisms is sequenced and presented
conjointly. Knowing which organisms constitute a sample requires a database
and a classification task.

The creation of a database requires to download reference genome sequences,
for example, from public repositories, with the disadvantage of requiring
substantial disk space for storing the database, especially if metagenomic
sequencing contains unknown species, where its composition can only be
marginally guessed given frequency patterns from many classification of sets
of FASTQ reads.

The classification task, which requires the previous creation of the database,
is a key process in the proposed methodology. The current reference-free FASTQ
data compressors do not take classification into account, employing only
models that try to cope and adapt to the nature of unknown sequences in the
metagenomic sample.

For increasing the capability of bringing closer similar DNA sequences of
FASTQ reads, that in principle share higher similarity according to higher
coverage depth, we employ compression-based schemes [20, 21].

Consider a set of reference sequences $X$ constituted by
$x_{1},x_{2},...,x_{n}$ and a set of FASTQ read sequences $Y$ constituted by
$y_{1},y_{2},...,y_{m}$. For convenience, we fix the alphabet $\Sigma$ of both
$X$ and $Y$ to $\Sigma=\\{A,C,G,T\\}$. Therefore, symbols that fall outside
this alphabet are considered a random symbol from $\Sigma$ (only in the
analysis phase).

The proposed method uses at its core data compression to estimate the
similarity of the reads according to an existing reference DNA sequence,
$x_{i}$, where reference sequences that have higher similarity are more
probable of being present on the sequences sample. Specifically, the method
loads into several context and substitution tolerant context models of several
depths the DNA sequences present on FASTQ reads of a given file. Then, freezes
these models and, for each DNA reference sequence, estimates the number of
bits required to compress it, $C(x_{i}||Y)$. Notice that $C(x_{i}||Y)$ stands
for the relative compression of $x_{i}$ given exclusively $y$ [22]. Then, the
percentage of relative similarity, $S(x_{i}$), is calculated according to

$S(x_{i})=\left(1-\frac{C(x_{i}||Y)}{|x_{i}|\log_{2}|\Sigma|}\right)\times
100,$ (1)

where $x_{i}$ is the reference sequence to be compressed, $Y$ the sequence
with all the reads concatenated, $|x_{i}|$ the number of symbols in $x_{i}$,
and $|\Sigma|$ the size of the alphabet of $x_{i}$.

Finally, the relative similarities, along with each associated reference
sequence information, are sorted by ascending order. From this classification
information, we take the most similar sequences above a certain threshold
$\mathcal{T}_{1}$, namely $S(x_{i})>\mathcal{T}_{1}$, and provide this
information to the next phase, the reads filtering.

### Reads filtering

In this phase, the objective is to filter input FASTQ reads, $Y$, that have
similarity higher than $\mathcal{T}_{1}$ for each reference sequence, $x_{i}$,
followed by relative compression ($R$) lower than a given threshold,
$\mathcal{T}_{2}$. For filtering these reads, analogous to the metagenomic
classification phase, we apply data compression. Specifically, the proposed
method loads into several context and substitution tolerant context models the
highest similar reference sequence from $X$. Then, for each read, $y_{j}$ of
the FASTQ metagenomic file, the number of bits needed to compress using
exclusively the selected reference $x_{i}$ is calculated according to

$R(y_{j})=\frac{C(y_{j}||x_{i})}{|y_{j}|\log_{2}|\Sigma|}.$ (2)

According to Figure 1, the reads where $R(y_{j})\leq\mathcal{T}_{2}$ are
redirected to a FASTQ file called filtered reads, otherwise to a FASTQ file
called unfiltered reads.

Then, the unfiltered reads are used as input to repeat the described process
for the second highest similar reference sequence determined in the
metagenomic classification phase. The process is again repeated until the
number of reference sequences above $\mathcal{T}_{1}$ is reached or the number
of unfiltered reads is zero.

The last sub-phase is constituted by the concatenation of all the filtered
reads and the last set of unfiltered reads (if exists) is sent to a single
file. In theory, reads with higher similarity should be more close, according
to the identified reference sequence in the metagenomic classification phase.
Notice that this final file contains the same number of characters than the
original file.

### Reads compression

In principle, data compressors with larger memory model capacity will require
larger FASTQ files for noticing compression gains. Additionally, this type of
compressors require more computational resources, including time, which for a
scenario of one compression on a server side and multiple downloads in client
sides significantly reduces its usage. In a FASTQ file sorted by similarity,
these large memory data compressors show less effectiveness because when
similar blocks are closer the need for a larger memory model is reduced.

For this purpose, we choose three data compressors to incorporate the proposed
method that are available by parameter selection, namely LZMA, Fqzcomp,
JARVIS. LZMA 111Available for free download at https://tukaani.org/lzma is a
general purpose data compressor. Fqzcomp [10] is a specific FASTQ data
compressor that uses low computational resources and good compression
capabilities. JARVIS (Version 3) 222Development version is available for free
download at https://github.com/cobilab/jarvis3 is an experimental genomic
compressor based on JARVIS2 [23]. Both Fqzcomp and JARVIS use different models
for compressing headers, DNA sequences, and quality-scores.

### Implementation

The implementation of the proposed method (named MizaR) is done using the
C/C++ language and Bash. This implementation is available for free download
(GPv3 license) at https://github.com/cobilab/MizaR along with the scripts to
generate all the results presented in this article.

##  3 Results

In this section, we evaluate the performance of the proposed method using
several datasets. These datasets are composed by real viral sequences in which
the process of sequencing is simulated using different conditions. The reason
to use sequencing simulation is to understand the dynamics that affect
different sequencing coverage depth and number of input sequences.

For simulating the process of sequencing, we used the ART simulator [24]. The
ART simulator generates reads using the information of genome sequences that
are feed using a multi-FASTA format. Notice that this simulator includes
alterations to better approximate the sequencing process. The size of the
reads was set to 150 base pairs, the mean size (-m) to 200 and the standard
deviation (-s) to 10 while set for paired-end reads similar to a Illumina
HiSeq 2500. The genome sequences were picked pseudo-randomly for each run from
the viral database in use (also available at the repository). For each ART
simulation, the depth sequencing coverage was set to a specific value and
MizaR run with the default threshold parameters ($\mathcal{T}_{1}=50$ and
$\mathcal{T}_{2}=0.5$).

All the computations have been performed using a laptop computer running Linux
Ubuntu 20.04 with Intel® Core™ i7-8650U CPU @ 1.90GHz × 8, 15.5 GiB RAM, and a
disk of approximately 512 GB.

Figure 2 depicts the compression gain for different compressors (Fqzcomp,
JARVIS, LZMA) with and without preserving the order of the reads using the
simulation of sequencing from real viral reference sequences while using a
coverage depth of 50 and varying the number of reference sequences from 100 to
700.

Figure 2: Compression gain in Megabytes (original compressed number of bytes
minus the sorted compressed number of bytes) for the Fqzcomp, JARVIS, and LZMA
compressors while varying the number of reference sequences used in the
sequencing simulation.

As it can be seen, there is a progressive increase in the compression gain as
the number of reference sequences increases. By compression gain, we mean the
Megabytes of the compressed original file minus the compressed sorted file
using the same data compressor. There is one point below the average increase
(for 600 reference sequences) because the pseudo-random selection was
performed from nearly 3k sequences and the higher values have higher
probability of contain similar reference sequences which harm the compression
gain.

Figure 2 also depicts that the general-purpose compressor LZMA obtains higher
compression gains compared to specialized FASTQ compressors (Fqzcomp and
JARVIS). Moreover, LZMA compression gain was similar to the specialized
compressor Fqzcomp if we would need to account for the information required to
represent the order of the reads and using the worst model, approximated by
the Stirling’s formula, $y^{n}\log_{2}(y^{n})-y^{n}\log_{2}e$, where $y^{n}$
stands for the number of reads of $y$.

Globally, these gains represent, for example, in 700 reference sequences,
approximately 3.7% of improvement using the sorting methodology (computational
time of sorting with 8 threads was 93 minutes) compared to not using the
sorting methodology; assuming the percentage gain as $(1-C_{s}/C_{u})\times
100$, where $C_{s}$ and $C_{u}$ stand for the compression of the sorted and
unsorted files, respectively. For the same example, considering the worst case
for preserving the order, the compression percentage gain is approximately
2.2%.

The JARVIS compressor seems to achieve a lower gain than Fqzcomp and LZMA.
This is justified by the memory model that is larger than in the other
compressors. For starting to notice higher gains in JARVIS, the size of the
file must be larger. Usually, larger files are the most frequent case, where
files have one order of magnitude large than those used here. However, because
there were many computations to perform, we used smaller examples.

Figure 3 depicts the gains for each of the information channels (headers, DNA
sequences, and quality-scores) while using Fqzcomp. We can notice a higher
gain in the DNA sequence channel (as expected) and a small negative gain in
the headers, while the quality-scores remain approximately the same.

Figure 3: Compression gain in Megabytes (original compressed number of bytes
minus the sorted compressed number of bytes) for the Fqzcomp compressor while
varying the number of reference sequences used in the sequencing simulation.

Accordingly, these values correspond to an average coverage depth of 50, which
for enriched viral metagenomic samples is very low. For studying the impact of
the coverage depth, we applied the same experience but now fixing the number
of reference sequences to 500 and varying the coverage depth from 1 to 91.
Figure 4 depicts this analysis, showing a gain increase for higher coverage
depth.

Figure 4: Compression gain in Megabytes (original compressed number of bytes
minus the sorted compressed number of bytes) for the Fqzcomp compressor while
varying the sequencing coverage depth used in the sequencing simulation.

Notice that MizaR depends on the diversity and quality of the database for
metagenomic classification. The recursive filter will only group reads that
share similarity above the defined threshold with the reference sequence in
the database. Therefore, building a good reference sequence database is of
extreme importance for increasing the compression gain capability of MizaR.

Additionally, MizaR is very flexible because it can be applied in other
contexts, such as human genome sequencing. For this case, besides the usual
bacterial and viral contamination in the database, the human genome can also
be included in the reference sequences database. However, for improved
compression, instead of having a reference for each chromosome, these
chromosomal references can be separated into junks of, say, 5 MB. MizaR will
assume different references sequences and will provide better resolution for
such a large genome.

The final considerations are for the parameters, namely the $\mathcal{T}_{1}$
and $\mathcal{T}_{2}$ thresholds. Although MizaR already provides default
values, $\mathcal{T}_{1}=50$ and $\mathcal{T}_{2}=0.5$, these may require
adaptation according to the application. This flexibility can be an advantage,
but it can also be a disadvantage because without expertise in the area it is
difficult to know which parameters better work for each case. Nevertheless, an
intuition can be given. For example, for $\mathcal{T}_{1}<50$ and a high
number of similar references in the database, MizaR will filter some of these
sequences and collect only a few reads, increasing the computational time
without increasing much the compressing gain. For different values of
$\mathcal{T}_{2}$, the reads, which share similarity with other reads
corresponding to other genome species, will be grouped or missed.

The choice of these parameters calls for a study which target different case
scenarios. Moreover, the tests to compute Figure 2 used at most a file size
less than 2 GB, which are only indicative of the true potential of MizaR,
because higher gains are associated with larger files (for example $>20$ GB).

##  4 Conclusions

In this paper, we presented MizaR, a compression-based method that uses
metagenomic classification and recursive filtering for improving the
compression of FASTQ metagenomic files, by changing the order of the reads
through DNA sequence similarity. We have shown that our methodology is able to
provide additional gains as the sequencing coverage depth increases and the
number of reference sequences. We have shown that general-purpose compression
methods, such as LZMA, are able to benefit more with this methodology.

## Acknowledgements

D.P. is funded by national funds through FCT – Fundação para a Ciência e a
Tecnologia, I.P., under the Scientific Employment Stimulus - Institutional
Call - reference CEECINST/00026/2018. We also acknowledge national funds
through the FCT in the context of the project UIDB/00127/2020.

##  5 References

## References

  * [1] Cláudia Covas, Tânia Caetano, Andreia Cruz, Tiago Santos, Liliana Dias, Guenter Klein, Amir Abdulmawjood, Luis M Rodríguez-Alcalá, Lígia L Pimentel, Ana Gomes, et al.,  “Pedobacter lusitanus sp. nov., isolated from sludge of a deactivated uranium mine,”  International journal of systematic and evolutionary microbiology, vol. 67, no. 5, pp. 1339–1348, 2017.
  * [2] Lari Pyöriä, Diogo Pratas, Mari Toppinen, Klaus Hedman, Antti Sajantila, and Maria F Perdomo,  “Unmasking the tissue-resident eukaryotic DNA virome in humans,”  Nucleic Acids Research, vol. 51, no. 7, pp. 3223–3239, 2023.
  * [3] Mari Toppinen, Diogo Pratas, Elina Väisänen, Maria Söderlund-Venermo, Klaus Hedman, Maria F Perdomo, and Antti Sajantila,  “The landscape of persistent human DNA viruses in femoral bone,”  Forensic Science International: Genetics, vol. 48, pp. 102353, 2020\.
  * [4] Weihong Qi, Yi-Wen Lim, Andrea Patrignani, Pascal Schläpfer, Anna Bratus-Neuenschwander, Simon Grüter, Christelle Chanez, Nathalie Rodde, Elisa Prat, Sonia Vautrin, et al.,  “The haplotype-resolved chromosome pairs of a heterozygous diploid African cassava cultivar reveal novel pan-genome and allele-specific transcriptome features,”  GigaScience, vol. 11, pp. giac028, 2022.
  * [5] Stelios Katsanevakis, Fernando Tempera, and Heliana Teixeira,  “Mapping the impact of alien species on marine ecosystems: the Mediterranean Sea case study,”  Diversity and Distributions, vol. 22, no. 6, pp. 694–707, 2016\.
  * [6] Don A Cowan, Jean-Baptiste Ramond, Thulani P Makhalanyane, and Pieter De Maayer,  “Metagenomics of extreme environments,”  Current opinion in microbiology, vol. 25, pp. 97–102, 2015.
  * [7] Morteza Hosseini, Diogo Pratas, and Armando J Pinho,  “A survey on data compression methods for biological sequences,”  Information, vol. 7, no. 4, pp. 56, 2016.
  * [8] Faraz Hach, Ibrahim Numanagić, Can Alkan, and S Cenk Sahinalp,  “SCALCE: boosting sequence compression algorithms using locally consistent encoding,”  Bioinformatics, vol. 28, no. 23, pp. 3051–3057, 2012.
  * [9] Daniel C Jones, Walter L Ruzzo, Xinxia Peng, and Michael G Katze,  “Compression of next-generation sequencing reads aided by highly efficient de novo assembly,”  Nucleic acids research, vol. 40, no. 22, pp. e171–e171, 2012.
  * [10] James K Bonfield and Matthew V Mahoney,  “Compression of FASTQ and SAM format sequencing data,”  PloS one, vol. 8, no. 3, pp. e59190, 2013.
  * [11] Łukasz Roguski and Sebastian Deorowicz,  “DSRC 2—Industry-oriented compression of FASTQ files,”  Bioinformatics, vol. 30, no. 15, pp. 2213–2215, 2014.
  * [12] Szymon Grabowski, Sebastian Deorowicz, and Łukasz Roguski,  “Disk-based compression of data from genome sequencing,”  Bioinformatics, vol. 31, no. 9, pp. 1389–1395, 2015.
  * [13] Gaëtan Benoit, Claire Lemaitre, Dominique Lavenier, Erwan Drezen, Thibault Dayris, Raluca Uricaru, and Guillaume Rizk,  “Reference-free compression of high throughput sequencing data with a probabilistic de Bruijn graph,”  BMC bioinformatics, vol. 16, no. 1, pp. 1–14, 2015.
  * [14] Zhi-An Huang, Zhenkun Wen, Qingjin Deng, Ying Chu, Yiwen Sun, and Zexuan Zhu,  “LW-FQZip 2: a parallelized reference-based compression of FASTQ files,”  BMC bioinformatics, vol. 18, pp. 1–8, 2017.
  * [15] Shubham Chandak, Kedar Tatwawadi, and Tsachy Weissman,  “Compression of genomic sequencing reads via hash-based reordering: algorithm and analysis,”  Bioinformatics, vol. 34, no. 4, pp. 558–567, 2018.
  * [16] Łukasz Roguski, Idoia Ochoa, Mikel Hernaez, and Sebastian Deorowicz,  “FaStore: a space-saving solution for raw sequencing data,”  Bioinformatics, vol. 34, no. 16, pp. 2748–2756, 2018.
  * [17] Shubham Chandak, Kedar Tatwawadi, Idoia Ochoa, Mikel Hernaez, and Tsachy Weissman,  “SPRING: a next-generation compressor for FASTQ data,”  Bioinformatics, vol. 35, no. 15, pp. 2674–2676, 2019.
  * [18] Dohyeon Lee and Giltae Song,  “FastqCLS: a FASTQ compressor for long-read sequencing via read reordering using a novel scoring model,”  Bioinformatics, vol. 38, no. 2, pp. 351–356, 2022.
  * [19] Raymond Wan and Kiyoshi Asai,  “Sorting next generation sequencing data improves compression effectiveness,”  in 2010 IEEE International Conference on Bioinformatics and Biomedicine Workshops (BIBMW). IEEE, 2010, pp. 567–572.
  * [20] Diogo Pratas and Armando J Pinho,  “Metagenomic composition analysis of sedimentary ancient DNA from the Isle of Wight,”  in 2018 26th european signal processing conference (EUSIPCO). IEEE, 2018, pp. 1177–1181.
  * [21] Diogo Pratas, Morteza Hosseini, Gonçalo Grilo, Armando J Pinho, Raquel M Silva, Tânia Caetano, João Carneiro, and Filipe Pereira,  “Metagenomic composition analysis of an ancient sequenced polar bear jawbone from Svalbard,”  Genes, vol. 9, no. 9, pp. 445, 2018.
  * [22] Armando J Pinho, Diogo Pratas, and Paulo JSG Ferreira,  “Authorship attribution using relative compression,”  in 2016 Data Compression Conference (DCC). IEEE, 2016, pp. 329–338.
  * [23] Diogo Pratas and Armando J Pinho,  “JARVIS2: a data compressor for large genome sequences,”  in 2023 Data Compression Conference (DCC). IEEE, 2023, pp. 288–297.
  * [24] Weichun Huang, Leping Li, Jason R Myers, and Gabor T Marth,  “ART: a next-generation sequencing read simulator,”  Bioinformatics, vol. 28, no. 4, pp. 593–594, 2012.


# Layer-Condensed KV Cache for Efficient Inference of
Large Language Models

Haoyi Wu  Kewei Tu
School of Information Science and Technology, ShanghaiTech University
Shanghai Engineering Research Center of Intelligent Vision and Imaging
{wuhy1<EMAIL_ADDRESS>Corresponding author.

###### Abstract

Huge memory consumption has been a major bottleneck for deploying high-
throughput large language models in real-world applications. In addition to
the large number of parameters, the key-value (KV) cache for the attention
mechanism in the transformer architecture consumes a significant amount of
memory, especially when the number of layers is large for deep language
models. In this paper, we propose a novel method that only computes and caches
the KVs of a small number of layers, thus significantly saving memory
consumption and improving inference throughput. Our experiments on large
language models show that our method achieves up to 26$\times$ higher
throughput than standard transformers and competitive performance in language
modeling and downstream tasks. In addition, our method is orthogonal to
existing transformer memory-saving techniques, so it is straightforward to
integrate them with our model, achieving further improvement in inference
efficiency. Our code is available at https://github.com/whyNLP/LCKV.

Layer-Condensed KV Cache for Efficient Inference of
Large Language Models

Haoyi Wu and Kewei Tu††thanks:  Corresponding author. School of Information
Science and Technology, ShanghaiTech University Shanghai Engineering Research
Center of Intelligent Vision and Imaging {wuhy1<EMAIL_ADDRESS>

##  1 Introduction

High throughput and low latency are essential for deploying large language
models (LLMs) in real-world applications Tillet et al. (2019); Kwon et al.
(2023). However, the huge memory consumption of LLMs has been a major
bottleneck, preventing a large batch size and high throughput generation.
Among the memory-consuming components, the key-value (KV) cache is one of the
most significant parts Pope et al. (2023); Zhang et al. (2023) that takes over
30% of the GPU memory during deployment Kwon et al. (2023). The KV cache is
used to store the keys and values in each transformer layer during generation
to avoid re-computation. Its memory consumption is proportional to both the
sequence length and the number of layers.

There have been substantial works on reducing the memory consumption of the KV
cache in LLMs. Most of them focus on compressing the KV cache by reducing the
length of the cached KV sequence. For example, Jiang et al. (2023a); Li et al.
(2023); Mu et al. (2023) compress the prompts to save the memory consumption.
Ren et al. (2023) incrementally compress a specified span of tokens into
compact ones to reduce the KV cache length. Xiao et al. (2024); Han et al.
(2023) propose to store only the KVs of initial and recent tokens in the KV
cache. Zhang et al. (2023) propose a dynamic KV cache eviction policy to only
keep a small portion of the KV cache in memory.

(a) Standard transformer

(b) Our model

Figure 1: Illustration of a standard transformer decoder and our model. Each
node represents one layer of transformer computation of one token. Each
horizontal edge $a\rightarrow b$ denotes that the queries at $b$ are paired
with the KVs at $a$.

In this paper, we propose to reduce the memory consumption of the KV cache
from a novel perspective that is orthogonal to previous efforts: reducing the
number of layers. Specifically, we present a new variant of transformer
decoders in which queries of all layers are paired with keys and values of
just the top layer, as illustrated in Figure 1. In this way, we only need to
cache the keys and values of one layer (vs. tens of layers in a typical LLM),
significantly saving memory consumption while introducing no computation
overheads during inference. In fact, since we only need the keys and values of
the top layer, we can omit the key-value computation and discard the key-value
parameters for all the other layers, further improving the throughput and
reducing the memory consumption as well as the mode size.

We draw our inspiration from the interpretation of the stacking layer
structure of a transformer as an iterative process of improving token
representation Wu and Tu (2023). In this interpretation, the representation at
the top layer is the most informative, so it makes sense to attend only to the
top layer. We also note the similarity of our idea to the cross-attention
mechanism in a standard transformer encoder-decoder, in which all the decoder
layers attend to the top encoder layer. However, applying this idea to a
decoder has never been attempted before as far as we know.

Although our model presented above achieves very efficient inference, its
performance in language modeling and downstream tasks degrades in comparison
with standard transformers. Therefore, we further propose to retain standard
attention for a small number of layers in our model, which slightly diminishes
our saving of the KV cache memory consumption but leads to almost no
performance degradation.

Another challenge faced by our model is training. When training a standard
transformer decoder, the computation of all the tokens can be fully
parallelized. In our model, however, the computation at each token depends on
the top layer of the previous token, creating sequential dependencies that
spoil parallel training. We address the challenge by deriving an approximate
training method that supports parallel training.

Our experiments on Llama Touvron et al. (2023) show that our model achieves up
to 32$\times$ larger batch sizes and up to 26$\times$ higher throughput than
standard transformers for LLMs of 1B–30B parameters; at the same time, our
model has competitive performance to standard transformers in language
modeling and downstream tasks. We further empirically demonstrate that it is
straightforward to integrate our model with other memory-saving techniques
like StreamingLLM Xiao et al. (2024), achieving further improvements in
inference efficiency.

We summarize our contributions as follows: 1) we propose a new variant of
transformer decoders that reduces the memory consumption of the KV cache by
dramatically reducing the number of cached layers; 2) we make parallel
training of our model feasible by designing a novel approximate training
method; 3) we conduct extensive experiments to verify and analyze the
effectiveness and efficiency of our method.

##  2 Layer-Condensed KV Cache

###  2.1 Model

As shown in Figure 1(b), we pair the queries of all layers with KVs of only
the top layer, so that we do not have to cache or even compute KVs for layers
other than the top layer, saving both memory consumption and computation.
Furthermore, since we no longer need to compute KVs for these layers, nor do
we need to keep the weights $W_{K},W_{V}$ that map hidden representations to
KVs for these layers, thus also saving model parameters.

One problem with this method is that, since each token also attends to itself,
we need its top-layer KVs for its attention computation at lower layers, but
the top-layer cannot be computed until we finish the computation of all the
lower layers. A straightforward solution to this cyclic dependency problem is
to drop the attention of each token to itself, which is equivalent to masking
the diagonal of the attention matrix. Now the first token of the sequence has
nothing to attend to, so we just use zero vectors as dummy KVs in its
attention computation. Note that even without self-attention of each token,
its information can still be incorporated in its bottom-up computation thanks
to residual connections. Empirically, we find that the diagonal mask of the
attention matrix does not affect the performance of the model.

It has been previously observed that transformers tend to attend to syntactic
information in lower layers and semantic information in higher layers Clark et
al. (2019b). Intuitively, applying KVs of the same layer to queries of all
layers might break this pattern. Empirically, we do find that our method as
described above underperforms standard transformers in language modeling and
downstream tasks. A simple yet effective solution to this problem is to retain
standard attention for a small number of layers, which we call _warmup layers_
, and only apply our method to the rest of the layers. Inspired by the
sandwich module of Reid et al. (2021), we propose to keep the top $w/2$ layers
and the bottom $w/2$ layers as warmup layers. We empirically find that such a
sandwich configuration outperforms alternative configurations and leads to
almost no performance degradation compared with standard transformers.

###  2.2 Training

Though the inference process of our model is straightforward and almost
identical to that of a standard transformer, i.e., decoding one token at a
time from left to right, the training process of our model is more
complicated. Since each token relies on the KVs of the top layer of previous
tokens, it is impossible to train on a sequence of tokens in parallel as in a
standard transformer. Below we derive a parallel training process of our model
in three steps. For simplicity, we assume no warmup layers (i.e., $w=0$). It
is straightforward to add warmup layers in the derived training process.

####  2.2.1 From Sequential to Parallel Training

During training, we minimize the cross entropy loss for each token. In our
model, the computation of each token depends on the top-layer KVs of its
previous tokens. Therefore, for a token sequence of length $n$, training is
done sequentially on a computation graph consisting of $n$ passes of bottom-up
transformer computation, as shown in Figure 2(a).

In the following, we present a different training computation graph and show
its equivalence to the original computation graph. Specifically, we perform
$n$ iterations of bottom-up transformer computation on all tokens in parallel,
and in each iteration, we pair the queries of all layers with KVs of the top
layer from the previous iteration, thus breaking the sequential dependencies
within the iteration. We compute the cross entropy loss only after the last
iteration. Note that the first iteration has no previous iteration, and we
pair its queries with dummy KVs which are zero vectors. Figure 2(b) shows the
new computation graph.

######  Theorem 1.

The two computation graphs are equivalent in terms of model training.

(a)

(b)

Figure 2: Two equivalent computation graphs for training our model with two
layers. (a) Sequential over $n$ tokens. (b) Parallel over $n$ tokens with $n$
iterations. Sub-graphs with the same color (except grey) represent identical
data flow. Sub-graphs in grey are unused in loss computation.

We leave the proof of the theorem to Appendix A. Here, we provide an intuitive
explanation. We say the computation of a token (w.r.t. its hidden
representations and top-layer KVs) is correct in an iteration of the second
computation graph if and only if it is identical to that in the first
computation graph. Since we have masked the diagonal of the attention matrix,
the computation of the first token does not rely on any key or value (except
for dummy zero vectors) and thus is correct in every iteration. For the second
token, its computation relies on the KVs of the first token and thus is
correct starting from the second iteration. In general, the computation of the
$i$-th token relies on the KVs of the first $i-1$ tokens and by induction, it
is correct starting from the $i$-th iteration. Therefore, after $n$
iterations, all the tokens are correctly computed. As a result, the
computation sub-graphs of the cross-entropy losses of all the tokens are
identical in the two graphs, and hence training processes following the two
graphs are equivalent.

Essentially, the second computation graph replaces horizontal dependencies in
the first graph with vertical dependencies, and it does not change the length
of the longest dependency chain. Consequently, although the second computation
graph supports parallel training over all the tokens, it still requires the
same $n$ iterations as the first graph. Next, we will trim the iterations
first in terms of backpropagation and then in terms of forward propagation.

####  2.2.2 Backpropagation: Gradient Stopping

We compute the cross entropy loss after the last iteration, which
backpropagates through $n$ iterations, resulting in a large computation graph
that is impossible to fit into GPU memory for large $n$. To solve the issue,
we follow the practice of gradient stopping in Transformer-XL Dai et al.
(2019) and backpropagate the loss only through the last $b$ iterations ($b\ll
n$).

Notice that the KVs used in the last iteration come from the second-to-last
iteration, so if we set $b=1$, then backpropagation would not reach the KVs
and hence the model parameters used to calculate the KVs are not trained at
all, which would result in a large performance degradation. We empirically
find that with $b\geq 2$, the performance of our model is comparable with that
of a standard transformer. To reduce memory consumption, we set $b=2$ by
default, which means we only backpropagate through two iterations.

####  2.2.3 Forward Propagation: Fast Convergence of KV

With gradient stopping, the first $n-b$ iterations are solely used in forward
propagation to compute KVs that are fed into the last $b$ iterations. When $n$
is large, it is still too costly to run $n-b$ iterations of forward
propagation. Fortunately, we observe that the values of KVs converge very fast
over iterations and hence we do not have to run $n-b$ iterations to obtain the
final KVs. Figure 3 shows the convergence of KVs of a randomly initialized
model with the same configuration as TinyLlama Zhang et al. (2024). The input
is a randomly picked segment of text from the MiniPile Kaddour (2023) dataset
with 2048 tokens (i.e., $n=2048$). We measure the change of KVs over
consecutive iterations using the mean squared error. As can be seen, KVs
converge after only a few tens of iterations, with more warmup layers leading
to faster convergence. Therefore, we use $m$ iterations ($m\ll n$) to
approximate the KVs of $n-b$ iterations. We empirically find that $m=7$ is
sufficient for model training and larger values of $m$ do not further improve
the performance.

Figure 3: MSE of the KV before and after the $i$th iteration. The model is
randomly initialized and tested with 2048 tokens.

###  2.3 Inference with Prompts

It is straightforward to employ our model for autoregressive generation.
However, our model cannot do parallel encoding of prompts like standard
transformers for the same reason as it cannot do parallel training.
Fortunately, since iterative computation of KVs is fast to converge, we can
just perform iterative computation for the prompts for $m+b$ iterations.
Typically, $m+b$ is far less than the number of tokens to generate, and thus
the extra time spent in encoding the prompts is negligible.

##  3 Experiments

We empirically verify the effectiveness of our method on the Llama model
Touvron et al. (2023). We show that our method achieves significant memory
reduction and throughput improvement as well as competitive performance in
language modeling and downstream tasks compared with standard transformers. We
also show that our method could effectively integrate with other memory-saving
techniques.

###  3.1 Generation Throughput

GPU | Model Size | Seq. Length | Batch Size | Throughput (tokens/s)
---|---|---|---|---
Llama |  | Ours
---
$w=2$
| Ours
---
$w=10$
Llama |  | Ours
---
$w=2$
| Ours
---
$w=10$
RTX 3090 | 1.1B | 5+8187 | 48 | 384 (8$\times$) | 119 (2.5$\times$) | 1424.96 | 4113.37 (2.9$\times$) | 2374.05 (1.7$\times$)
5+2043 | 239 | 1150 (4.8$\times$) | 289 (1.2$\times$) | 5142.86 | 10033.40 (2.0$\times$) | 7239.92 (1.4$\times$)
7B | 5+8187 | 1 | 12 (12$\times$) | 4 (4$\times$) | 32.02 | 151.91 (4.7$\times$) | 83.80 (2.6$\times$)
2048+2048 | 2 | 23 (11.5$\times$) | 8 (4$\times$) | 56.98 | 171.65 (3.0$\times$) | 119.68 (2.1$\times$)
5+2043 | 5 | 64 (12.8$\times$) | 16 (3.2$\times$) | 140.88 | 534.02 (3.8$\times$) | 315.38 (2.2$\times$)
512+512 | 9 | 95 (10.6$\times$) | 32 (3.6$\times$) | 225.31 | 378.89 (1.7$\times$) | 380.60 (1.7$\times$)
512+1024 | 7 | 72 (10.3$\times$) | 16 (2.3$\times$) | 174.11 | 401.92 (2.3$\times$) | 310.05 (1.8$\times$)
| 30B
---
(CPU-offload)
512+1024 | 4 | 83 (20.8$\times$) | 23 (5.8$\times$) | 0.23 | 5.99 (26.0$\times$) | 1.63 (7.1$\times$)
A100 | 7B | 2048+2048 | 15 | 128 (8.5$\times$) | 42 (2.8$\times$) | 141.10 | 421.02 (3.0$\times$) | 315.09 (2.2$\times$)
30B | 2048+2048 | 1 | 32 (32$\times$) | 8 (8$\times$) | 14.10 | 108.29 (7.7$\times$) | 77.65 (5.5$\times$)
Table 1: Maximum generation batch size and throughput on an RTX 3090 (24GB)
and an A100 (80GB) GPU respectively with different sequence lengths. Following
Zhang et al. (2023), we use “$x+y$” to denote a prompt length of $x$ and a
generation length of $y$.

We test our method with 1.1B, 7B, and 30B parameters on an NVIDIA GeForce RTX
3090 (24GB) GPU and an NVIDIA A100 (80GB) GPU respectively. The 1.1B model
configuration follows that of TinyLlama Zhang et al. (2024) and the 7B and 30B
model configuration follows that of the original Llama Touvron et al. (2023).
We set $m=7,b=2$ and $w=\\{2,10\\}$. Our implementation is based on
HuggingFace Transformers Wolf et al. (2020) with kernel replacement with
FlashAttention 2 Dao (2023), fused RMS norm, fused cross-entropy and fused
SwiGLU.

Following FlexGen Sheng et al. (2023), the evaluation is conducted in an end-
to-end fashion. For a prompt of length $s$, we let the model generate output
sequences of length $n$ with batch size $b$. The latency $t$ is defined as the
total time in seconds spent in processing the prompts and generating all the
$bn$ tokens. The generation throughput is defined as $bn/t$ tokens per second.

Table 1 compares the maximum batch sizes and throughput of standard Llama
models and our models on the two types of GPUs. Note that some of the sequence
lengths exceed the maximum input lengths that the models are trained on, but
that does not affect batch size and throughput measurement. We still benchmark
the batch sizes and throughput to show the potential of our method on other
models allowing larger sequence lengths. It can be seen that our method
achieves significantly larger batch sizes and higher throughput on all of the
settings.

Figure 4: Throughput of 7B Llama and our model w.r.t. the batch size.

Notice that the maximum throughput is not necessarily achieved using the
maximum batch size. Figure 4 shows the throughput of 7B Llama and our model on
an A100 GPU w.r.t. the batch size. The prompt length and generation length are
both 2048. We find that when the batch size grows larger than 32, the
throughput no longer increases and even decreases with a batch size of 128.
This may indicate that the model operations are turning from memory-bound into
compute-bound Dao et al. (2022) and the throughput is limited by the
computation power. Further improvement in throughput may require more
efficient computation kernels instead of larger batch sizes.

We also would like to mention that the increased throughput is not necessarily
due to the increased batch size. As shown in Figure 4, our model has much
higher throughput than Llama even with the same batch sizes. We leave the
detailed analysis to Appendix C.6.

###  3.2 Model Performance

Model | HellaSwag | Obqa | WinoGrande | ARC-c | ARC-e | BoolQ | PIQA | Avg
---|---|---|---|---|---|---|---|---
TinyLlama | 44.58 | 30.2 | 50.99 | 25.00 | 46.38 | 60.46 | 68.93 | 46.65
Ours ($w=2$) | 42.22 | 30.6 | 49.64 | 24.74 | 43.10 | 61.38 | 66.49 | 45.45
Ours ($w=10$) | 44.74 | 31.0 | 51.70 | 24.83 | 46.38 | 61.38 | 67.90 | 46.84
Table 2: Zero-shot accuracy on commonsense reasoning tasks. Model | Dev ppl.
---|---
TinyLlama | 9.219
Ours ($w=2$) | 9.746
Ours ($w=10$) | 9.265
Table 3: Perplexity on a 10M subset of the development set of SlimPajama.

To evaluate the performance of our model in language modeling and downstream
tasks, we pre-train from scratch two 1.1B models with $m=7$, $b=2$ and
$w=\\{2,10\\}$. We use TinyLlama as our baseline, whose size is also 1.1B. We
pre-train the models on a 100B subset of the SlimPajama dataset Soboleva et
al. (2023). The training details are consistent with those of TinyLlama Zhang
et al. (2024). All models are trained with AdamW Loshchilov and Hutter (2019)
with $\beta_{1}=0.9$ and $\beta_{2}=0.95$. The batch size is 2M tokens. We use
a cosine learning rate schedule with a maximum learning rate of $4.0\times
10^{-4}$ and a warmup of 200 steps. The final learning rate is $4.0\times
10^{-5}$. We use a weight decay of 0.1 and gradient clipping of 1.0. The
models are trained on 128 NVIDIA A800 (80GB) GPUs.

During evaluation, we perform inference in the standard left-to-right fashion
instead of using the method of Section 2.3. We report the perplexity on a 10M
subset of the development set of SlimPajama. We also test the zero-shot
performance on commonsense reasoning tasks following Zhang et al. (2024),
including Hellaswag Zellers et al. (2019), OpenBookQA Mihaylov et al. (2018),
WinoGrande Sakaguchi et al. (2021), ARC-Easy and ARC-Challenge Clark et al.
(2018), BoolQ Clark et al. (2019a), and PIQA Bisk et al. (2020). The tests are
conducted using the lm-eval-harness framework Gao et al. (2023). For these
tasks, we encode the prompts with the same number of iterations ($m+b=9$) as
in training.

Table 2 and 3 show the results. The performance of our models is comparable to
that of TinyLlama. In particular, our model with $w=10$ has almost no
performance degradation, while achieving significantly higher generation
throughput as evaluated in Section 3.1. Our model with $w=2$ has a small but
noticeable decrease in performance for most of the tasks, but it achieves even
higher increase in throughput.

Despite the competitive performance and higher inference efficiency of our
models, we note that pre-training our model costs about 3 times the time of
pre-training TinyLlama with the same amount of data due to the iterative
training process. Nevertheless, we believe that in most scenarios, a speedup
in inference is worth a slowdown in training which is a one-time process.

###  3.3 Integration with StreamingLLM

Figure 5: Comparison of latency per token and memory consumption of
StreamingLLM and our model ($w=10$) integrated with StreamingLLM w.r.t.
different cache sizes.

We have previously mentioned that our method is orthogonal to other memory-
saving techniques and can be easily integrated with them. Here we integrate
our method with StreamingLLM Xiao et al. (2024). StreamingLLM employs an
attention sink that only preserves the KV cache of the first few tokens (four
by default) and recent tokens, which empowers LLMs to process infinite-length
inputs.

As shown in Figure 5, the integration of StreamingLLM and our model ($w=10$)
achieves lower latency and memory consumption compared to the original
StreamingLLM on different cache sizes (numbers of cached tokens).

Figure 6: Comparison of StreamingLLM and our model integrated with
StreamingLLM w.r.t. the cache size. We use 4 initial tokens for all settings.
The results are collected on the first text sample of PG19 Rae et al. (2020).
Figure 7: Language modeling perplexity of our model integrated with
StreamingLLM on texts with 4M tokens. Following Xiao et al. (2024), we use the
concatenated test set of the PG19 dataset Rae et al. (2020) as the input.

We further showcase that integration with our method does not hinder the
ability of StreamingLLM to process infinite-length tokens. Specifically, we
integrate StreamingLLM into our model ($w=10$) trained in Section 3.2 with
different cache sizes and find that the integrated model achieves even lower
perplexities than StreamingLLM as shown in Figure 6. We also let the model
handle inputs with a sequence length of four million tokens. As shown in
Figure 7, the integrated model can effectively process the input with the
perplexity remaining stable.

##  4 Analyses

In this section, we empirically analyze design choices in our method. For
experiment details, please refer to Appendix B.

###  4.1 The Sandwich Configuration

In Section 2.1, we propose to add some warmup layers to improve model
performance. Note that inference efficiency is determined by the number of
warmup layers $w$ and not by their placement. Here we ask the following
question: with the number of warmup layers fixed, where should we place the
warmup layers to achieve the best performance?

Model Size | Dev ppl.
---|---
all-bottom | all-top | sandwich
50M | 14.556 | 221.850 | 14.069
1.1B | 7.668 | 9.098 | 7.381
Table 4: Model performance with different warmup layer placements ($w=2$).

Table 4 shows the model performance in language modeling with two warmup
layers (i.e., $w=2$) that are placed at the bottom, at the top, and with the
default sandwich configuration. It can be seen that the sandwich style
performs the best. A possible explanation is that top and bottom layers serve
different functionalities (e.g., semantic vs. syntactic) and it makes more
sense to warm up both than just one of them.

###  4.2 Number of Warmup Layers

Warmup layers serve as a bridge between the standard transformer and our
model. The more warmup layers we keep, the more similar it is to the standard
transformer, and the less memory we save. In this section, we ask the
following question: how does the number of warmup layers $w$ affect the model
performance and throughput?

Figure 8: Effect of the number of warmup layers on model performance and
throughput. The right-most point ($w=22$) denotes that all the layers are
warmup layers and hence our model becomes exactly the standard transformer
(the baseline). The throughput is tested on an RTX 3090 GPU with prompt length
5 and generation length 2043.

We test the model with 1.1B (22 layers) parameters and different numbers of
warmup layers (Figure 8). Surprisingly, we find that the log dev perplexity
does not monotonically decrease with the number of warmup layers. Without
warmup layers (the red region), the model performance is significantly worse
than that of the standard transformer. With only a few warmup layers (the
yellow region), the model performance is greatly improved and close to that of
the standard transformer, but its throughput is decreased at the same time.
With more warmup layers (the green region), the model even outperforms the
standard transformer, while suffering from further decrease in throughput.

The results point to a trade-off between model performance and throughput that
is controllable by the number of warmup layers. If pursuing a high throughput,
one could keep only a few warmup layers (the yellow region). If good
performance is crucial, one could keep more warmup layers (the left part of
the green region).

###  4.3 Convergence of KV

In Section 2.2.3, we have shown that KVs converge very fast for a random model
and hence we use $m\ll n$ iterations to compute KVs. Here, we ask the
following questions: how fast do KVs converge for a trained model and what
value shall we pick for $m$?

Figure 9: (Top) Experiment on a Llama (50M) with different values of $m$
during training. (Bottom) MSE of the KV before and after the $i$th iteration.
The models are tested with 1024 tokens.

We measure the convergence of KVs of a 50M Llama (Figure 9, bottom) and the
1.1B model ($w=2$) pre-trained in Section 3.2 (Figure 10). It can be seen that
while a random model requires 15–20 iterations to converge, a trained model
requires far fewer iterations. This hints at a small value of $m$ especially
during late stages of training. We then evaluate model performance when
trained with different values of $m$ (Figure 9, top). It can be seen that the
model performance converges with $m\geq 6$, with more warmup layers leading to
faster convergence. This justifies our default choice of $m=7$.

Figure 10: MSE of the KV before and after the $i$th iteration. The models are
tested with 2048 tokens.

##  5 Related Work

Extensive research has been done on reducing KV cache for efficient inference
of LLMs. With the exception of vLLM Kwon et al. (2023), which proposes paged
attention to reduce memory fragmentation of the KV cache from a system
perspective, most recent works focus on compressing the KV cache by reducing
the length of the cached KV sequence. Jiang et al. (2023a, b) accelerate model
inference by compressing document-level prompts into short prompts. Li et al.
(2023) remove the redundancy in the input context. Mu et al. (2023) train gist
tokens to replace the reusable system prompts. Ren et al. (2023) incrementally
compress a specified span of tokens into compact ones to reduce the KV cache
length. Liu et al. (2023) find that only pivotal tokens have a significant
influence at a future step, so pruning unimportant tokens does not affect the
performance. Ge et al. (2023) argue that the attention heads could be
classified into different types and propose to apply different KV pruning
strategies to different types of attention heads. Xiao et al. (2024); Han et
al. (2023) find that only initial tokens and recent tokens are crucial and
propose to store only the KVs of these tokens to enable infinite-length
context for LLMs. Zhang et al. (2023) propose a KV cache eviction policy based
on the summation of attention scores to only keep a small portion of the KV
cache in memory. Unlike these previous methods, our method reduces the memory
consumption of the KV cache by reducing the number of layers, which is
orthogonal to these methods and can potentially be combined with these methods
to further reduce the KV cache and improve inference efficiency.

Feedback Transformers Fan et al. (2020) aggregate hidden representations from
all layers and use them as token memory. This is followed by KV projections to
obtain the key-value pairs for all layers. Their experimental results show
improved performance with this strategy, even when using only the hidden
representation from the top layer. However, their sequential training process
is time-costly and not practical for large models. Our method supports
parallel training, which is more efficient and scalable.

##  6 Conclusion

In this paper, we propose a novel method to reduce the memory consumption and
improve the throughput of LLMs by reducing the number of layers whose keys and
values need to be computed and cached. We empirically show that our method
achieves significant memory reduction and throughput improvement with
negligible performance degradation. We also show that our method could
effectively integrate with other memory-saving techniques like StreamingLLM.
Future work includes designing more efficient training approaches, developing
large-batch-friendly kernels, and verifying our method on larger and more
complex LLMs. We hope that our work could provide a new perspective for
improving inference efficiency of LLMs and inspire more research in this
direction.

## Limitations

Though our method achieves impressive memory reduction and throughput
improvement, we would like to point out its limitations from the following
aspects:

  * •

Due to the iterative training, our model requires about $3\times$ the time to
pre-train a model with the same amount of data. In other words, our method
improves the inference efficiency at the cost of the training efficiency. A
potential remedy is that if one has a pre-trained model, one could use it to
initialize our model, which is empirically found to speed up the process of
training.

  * •

Since our method requires iteratively processing the prompts, the throughput
degrades when the prompts are much longer than the generation length, e.g., in
document summarization. Generally, our method is more suitable for tasks with
a large generation length, such as translation, dialogue, question answering,
CoT problem solving, etc.

## References

  * Adnan et al. (2024) Muhammad Adnan, Akhil Arunkumar, Gaurav Jain, Prashant J Nair, Ilya Soloveychik, and Purushotham Kamath. 2024.  Keyformer: Kv cache reduction through key tokens selection for efficient generative inference.  _arXiv preprint arXiv:2403.09054_.
  * Bisk et al. (2020) Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. 2020.  Piqa: Reasoning about physical commonsense in natural language.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 34, pages 7432–7439.
  * Clark et al. (2019a) Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. 2019a.  BoolQ: Exploring the surprising difficulty of natural yes/no questions.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 2924–2936, Minneapolis, Minnesota. Association for Computational Linguistics.
  * Clark et al. (2019b) Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. 2019b.  What does BERT look at? an analysis of BERT’s attention.  In _Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_ , pages 276–286, Florence, Italy. Association for Computational Linguistics.
  * Clark et al. (2018) Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. 2018.  Think you have solved question answering? try arc, the ai2 reasoning challenge.  _arXiv preprint arXiv:1803.05457_.
  * Dai et al. (2019) Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. 2019.  Transformer-XL: Attentive language models beyond a fixed-length context.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pages 2978–2988, Florence, Italy. Association for Computational Linguistics.
  * Dao (2023) Tri Dao. 2023.  Flashattention-2: Faster attention with better parallelism and work partitioning.  _arXiv preprint arXiv:2307.08691_.
  * Dao et al. (2022) Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022.  FlashAttention: Fast and memory-efficient exact attention with IO-awareness.  In _Advances in Neural Information Processing Systems_.
  * Fan et al. (2020) Angela Fan, Thibaut Lavril, Edouard Grave, Armand Joulin, and Sainbayar Sukhbaatar. 2020.  Addressing some limitations of transformers with feedback memory.  _arXiv preprint arXiv:2002.09402_.
  * Gao et al. (2023) Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. 2023.  A framework for few-shot language model evaluation.
  * Ge et al. (2023) Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. 2023.  Model tells you what to discard: Adaptive KV cache compression for LLMs.  In _Workshop on Advancing Neural Network Training: Computational Efficiency, Scalability, and Resource Optimization (WANT@NeurIPS 2023)_.
  * Han et al. (2023) Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. 2023.  Lm-infinite: Simple on-the-fly length generalization for large language models.  _arXiv preprint arXiv:2308.16137_.
  * Jiang et al. (2023a) Huiqiang Jiang, Qianhui Wu, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. 2023a.  LLMLingua: Compressing prompts for accelerated inference of large language models.  In _Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing_ , pages 13358–13376, Singapore. Association for Computational Linguistics.
  * Jiang et al. (2023b) Huiqiang Jiang, Qianhui Wu, Xufang Luo, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. 2023b.  Longllmlingua: Accelerating and enhancing llms in long context scenarios via prompt compression.  _arXiv preprint arXiv:2310.06839_.
  * Kaddour (2023) Jean Kaddour. 2023.  The minipile challenge for data-efficient language models.  _arXiv preprint arXiv:2304.08442_.
  * Kwon et al. (2023) Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. 2023.  Efficient memory management for large language model serving with pagedattention.  In _Proceedings of the 29th Symposium on Operating Systems Principles_ , pages 611–626.
  * Li et al. (2023) Yucheng Li, Bo Dong, Frank Guerin, and Chenghua Lin. 2023.  Compressing context to enhance inference efficiency of large language models.  In _Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing_ , pages 6342–6353, Singapore. Association for Computational Linguistics.
  * Liu et al. (2023) Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. 2023.  Scissorhands: Exploiting the persistence of importance hypothesis for LLM KV cache compression at test time.  In _Thirty-seventh Conference on Neural Information Processing Systems_.
  * Loshchilov and Hutter (2019) Ilya Loshchilov and Frank Hutter. 2019.  Decoupled weight decay regularization.  In _International Conference on Learning Representations_.
  * Merity et al. (2017) Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2017.  Pointer sentinel mixture models.  In _International Conference on Learning Representations_.
  * Mihaylov et al. (2018) Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. 2018.  Can a suit of armor conduct electricity? a new dataset for open book question answering.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing_ , pages 2381–2391, Brussels, Belgium. Association for Computational Linguistics.
  * Mu et al. (2023) Jesse Mu, Xiang Lisa Li, and Noah Goodman. 2023.  Learning to compress prompts with gist tokens.  In _Thirty-seventh Conference on Neural Information Processing Systems_.
  * Pope et al. (2023) Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. 2023.  Efficiently scaling transformer inference.  _Proceedings of Machine Learning and Systems_ , 5.
  * Rae et al. (2020) Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. 2020.  Compressive transformers for long-range sequence modelling.  In _International Conference on Learning Representations_.
  * Reid et al. (2021) Machel Reid, Edison Marrese-Taylor, and Yutaka Matsuo. 2021.  Subformer: Exploring weight sharing for parameter efficiency in generative transformers.  In _Findings of the Association for Computational Linguistics: EMNLP 2021_ , pages 4081–4090, Punta Cana, Dominican Republic. Association for Computational Linguistics.
  * Ren et al. (2023) Siyu Ren, Qi Jia, and Kenny Zhu. 2023.  Context compression for auto-regressive transformers with sentinel tokens.  In _Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing_ , pages 12860–12867, Singapore. Association for Computational Linguistics.
  * Sakaguchi et al. (2021) Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. 2021.  Winogrande: An adversarial winograd schema challenge at scale.  _Communications of the ACM_ , 64(9):99–106.
  * Sheng et al. (2023) Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Re, Ion Stoica, and Ce Zhang. 2023.  FlexGen: High-throughput generative inference of large language models with a single GPU.  In _Proceedings of the 40th International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pages 31094–31116. PMLR.
  * Soboleva et al. (2023) Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. 2023.  SlimPajama: A 627B token cleaned and deduplicated version of RedPajama.
  * Tillet et al. (2019) Philippe Tillet, Hsiang-Tsung Kung, and David Cox. 2019.  Triton: an intermediate language and compiler for tiled neural network computations.  In _Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages_ , pages 10–19.
  * Touvron et al. (2023) Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. 2023.  Llama: Open and efficient foundation language models.  _arXiv preprint arXiv:2302.13971_.
  * Wolf et al. (2020) Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi 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 Rush. 2020.  Transformers: State-of-the-art natural language processing.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations_ , pages 38–45, Online. Association for Computational Linguistics.
  * Wu and Tu (2023) Haoyi Wu and Kewei Tu. 2023.  Probabilistic transformer: A probabilistic dependency model for contextual word representation.  In _Findings of the Association for Computational Linguistics: ACL 2023_ , pages 7613–7636, Toronto, Canada. Association for Computational Linguistics.
  * Xiao et al. (2024) Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. 2024.  Efficient streaming language models with attention sinks.  In _The Twelfth International Conference on Learning Representations_.
  * Zellers et al. (2019) Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. 2019.  HellaSwag: Can a machine really finish your sentence? In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pages 4791–4800, Florence, Italy. Association for Computational Linguistics.
  * Zhang et al. (2024) Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. 2024.  Tinyllama: An open-source small language model.  _arXiv preprint arXiv:2401.02385_.
  * Zhang et al. (2023) Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Re, Clark Barrett, Zhangyang Wang, and Beidi Chen. 2023.  H2o: Heavy-hitter oracle for efficient generative inference of large language models.  In _Thirty-seventh Conference on Neural Information Processing Systems_.

##  Appendix A Proof of The Training Theorem

Here we formally prove Theorem 1 in Section 2.2. Remember that we need to
compute the loss for each token sequentially, but we propose to train all the
tokens in parallel with $n$ iterations.

See 1

To prove the theorem, we first introduce the following lemma.

######  Lemma 1.

In the first computation graph, denote the final hidden representation of the
$i$-th token as $h_{i}=f_{1,i}(x_{1},x_{2},\dots,x_{i})$. In the second
computation graph, denote the final hidden representation of the $i$-th token
in iteration $t$ as $h^{(t)}_{i}=f^{(t)}_{2,i}(x_{1},x_{2},\dots,x_{i})$. We
have $f_{1,i}=f^{(t)}_{2,i},\forall t\geq i$, no matter what the initial KVs
are.

###### Proof.

In the first computation graph, denote the KVs of the $i$-th token as
${KV}_{i}$ and that under parallel training with iteration $t$ as
${KV}^{(t)}_{i}$.

Since the basic networks are the same for the two computation graphs, we
further denote $g_{i}$ as
$h_{i}=f_{1,i}(x_{1},x_{2},\dots,x_{i})=g_{i}(x_{1},x_{2},\dots,x_{i},{KV}_{1},{KV}_{2},\dots,{KV}_{i-1})$,
as well as
$h^{(t)}_{i}=f^{(t)}_{2,i}(x_{1},x_{2},\dots,x_{i})=g_{i}(x_{1},x_{2},\dots,x_{i},{KV}^{(t-1)}_{1},\dots,{KV}^{(t-1)}_{i-1})$.
The only difference in the computation of the final hidden representations
lies in the KVs.

We prove the lemma by induction. For the base case, since we have removed the
diagonal of the attention matrix, $h_{1}$ and $h^{(t)}_{1}$ does not rely on
any key or value. So we have
$h_{1}=f_{1,1}(x_{1})=g_{i}(x_{1})=f^{(t)}_{2,1}(x_{1})=h^{(t)}_{1},\forall
t$. That is $f_{1,1}=f^{(t)}_{2,1},\forall t$. Since the computation graphs of
the first token are the same, ${KV}_{1}$ and ${KV}^{(t)}_{1}$ are the same for
all $t$.

For the inductive step, we assume that $\forall i\leq
T,f_{1,i}=f^{(T)}_{2,i}=g_{i},{KV}_{i}={KV}^{(T)}_{i}$. Then for iteration
$T+1$:

$\forall i\leq T+1$, the computation of the $i$-th token relies on the KVs of
the first $i-1$ tokens. That is,
$h^{(T+1)}_{i}=f^{(T+1)}_{2,i}(x_{1},x_{2},\dots,x_{i})=g_{i}(x_{1},x_{2},\dots,x_{i},{KV}^{(T)}_{1},\dots,{KV}^{(T)}_{i-1})$.

Since $i\leq T+1$, we have $i-1\leq T$. By induction we have
${KV}^{(T)}_{i-1}={KV}_{i-1}$. Thus, we have
$g_{i}(x_{1},x_{2},\dots,x_{i},{KV}^{(T)}_{1},\dots,{KV}^{(T)}_{i-1})=g_{i}(x_{1},x_{2},\dots,x_{i},{KV}_{1},\dots,{KV}_{i-1})=f_{1,i}(x_{1},x_{2},\dots,x_{i})=h_{i}$.
The computation graphs are the same for the first $T+1$ tokens, then we have
${KV}_{i}={KV}^{(T+1)}_{i}$.

Thus, we have $f_{1,i}=f^{(t)}_{2,i},\forall i\leq t$. The proof does not rely
on initial KVs, so the conclusion holds for any initial KVs. ∎

Let $t=n$, we have the entire computation graphs are equivalent. Therefore, we
have proved Theorem 1.

From the lemma, we could learn one more thing: The parallel training has the
same computation graph as the sequential training when the iteration $t\geq
n$. This indicates that the parallel training is theoretically guaranteed to
converge to the same solution as the sequential training. Once it is
converged, it will not diverge. Our work finds that the KVs converge much
faster than the theoretical $n$ iterations, which significantly reduces the
training time.

##  Appendix B Model and Training Details

Model Size | 50M | 1.1B | 7B | 30B
---|---|---|---|---
Hidden Size | 512 | 2048 | 4096 | 6656
Intermediate Size | 1024 | 5632 | 11008 | 17920
Max Trained Length | 1024 | 2048 | – | –
# Layers | 8 | 22 | 32 | 60
# Attention Heads | 8 | 32 | 32 | 52
# KV Heads | 4 | 4 | 32 | 52
RMS Norm eps | 1e-5 | 1e-5 | 1e-6 | 1e-6
Vocab Size | 32000
Table 5: Model configurations. Section | 4.1 | 4.2 | 4.3
---|---|---|---
Model Size | 50M | 1.1B | 1.1B | 50M
$m$ | 7 | 7 | 7 | –
$b$ | 2 | 2 | 2 | 2
$w$ | 2 | 2 | – | 2
lr scheduler | cosine
max. lr | 3e-4 | 3e-4 | 4e-4 | 3e-4
min. lr | 0 | 0 | 4e-5 | 0
optimizer | AdamW
$\beta_{1}$ | 0.9 | 0.9 | 0.9 | 0.9
$\beta_{2}$ | 0.999 | 0.999 | 0.95 | 0.999
batch size (tokens) | 16K | 256K | 256K | 16K
warmup ratio | 0.015 | 0.015 | 0.024 | 0.015
weight decay | 6.6e-6 | 6.6e-6 | 1e-1 | 6.6e-6
gradient clipping | 1.0 | 1.0 | 1.0 | 1.0
initialize from pre-trained | yes | yes | no | no
epochs | 3 | 1 | 2B tokens | 3
Data | WikiText-103 | MiniPile | SlimPajama | WikiText-103
GPU | RTX 3090x1 | A100x8 | A800x16 | RTX 3090x1
Table 6: Training details for Section 4.

We provide the model configurations and training details in Table 5 and 6. The
7B and 30B model configurations are consistent with those of the original
Llama Touvron et al. (2023). The 1.1B model configuration follows that of
TinyLlama Zhang et al. (2024). We use the WikiText-103 Merity et al. (2017)
(licensed under CC-BY-SA 3.0), MiniPile Kaddour (2023) (licensed under MIT)
and SlimPajama Soboleva et al. (2023) (various licenses depending on the data
source) as our datasets. Our use of the datasets is consistent with their
intended use.

The training of TinyLlama in Section 3 takes 14:42:59, while our models take 1
day, 16:44:16 (2.77x, $w=2$) and 1 day, 15:52:38 (2.71x, $w=10$),
respectively. The training took place before we optimize our codes, so the
training time could be further reduced, especially for the model with $w=10$.

##  Appendix C More Analyses

In this section, we provide more analyses beyond those in Section 4.

###  C.1 Initialize with Pre-trained Models

Model | Dev ppl. | Avg.
---|---|---
TinyLlama | 9.219 | 46.65
Ours ($w=2$) | 9.746 | 45.45
TinyLlama-2.5T | 7.822 | 53.97
TinyLlama-500B | 9.046 | 48.28
Ours ($w=2$, init w/ 2.5T) | 8.514 | 49.55
Table 7: Dev ppl: Perplexity on a 10M subset of the validation set of
SlimPajama. Avg: Average zero-shot accuracy on commonsense reasoning tasks.
Models are trained on a 100B subset of SlimPajama. The models with italic
fonts are from TinyLlama checkpoints.

Since our model structure resembles that of the standard transformer, we could
initialize our model with pre-trained models. For $W_{K},W_{V}$ of the middle
layers, we just ignore the parameters. Experiments show that initialization
with pre-trained models could effectively speed up the training process (Table
7). Though our model has a different computation graph from the standard
transformer, it could still benefit from the pre-trained models. Our model
initialized with a TinyLlama checkpoint trained on 2.5T tokens achieves a
perplexity of 8.514, which is much better than the randomly initialized model.
It is even better than the TinyLlama checkpoint trained on 500B tokens. Thus,
if the pre-trained models are available, initializing our model with them
could save a lot of training time.

###  C.2 Iterations with Gradients

In Section 2.2.2 we set $b=2$ by default. What if we use larger $b$? To make
sure that the inference process is equivalent (specifically, encoding the
prompts), we fix $m+b=9$ and experiment with different values of $b$.

Figure 11 shows the perplexity of a 50M Llama with different values of $b$.
Though from $b=2$ to $b=3$ the perplexity decreases, for $b=4$ the perplexity
increases. We further confirm that for $b=4$ the KVs do not converge as fast
as for $b=2$.

Experiments on a 1.1B model (Table 8) show that the perplexity increases with
$b$ and the training time also increases. From the training curve, we find
that larger $b$ leads to more unstable training, thus the model is harder to
converge. Therefore, we set $b=2$ by default.

Figure 11: Effect of $b$ on a 50M model with different number of warmup layers. Model | Dev ppl. | Train Time
---|---|---
Ours ($b=2$) | 10.390 | 8h
Ours ($b=3$) | 10.476 | 10h
Ours ($b=4$) | 10.885 | 13h
Table 8: Effect of $b$ on a 1.1B model.

###  C.3 KV Loss for Less Iterations

In Section 4.3 we have shown that the KVs converge very fast for a trained
model, yet we still want to make it converge even faster, saving both training
and inference time costs. An intuitive idea is to add an MSE loss to the KVs
before and after the last iteration to force the KVs to converge. We call this
term the “KV Loss”. Our experiments (Table 9) show that for small data with
small $w$, the KV loss could lead to better performance. While for large data
or large $w$, the KV loss hurts the performance. This is probably because the
KVs are not converged at the beginning of training and the KV loss helps the
KVs converge. However, when the KVs are already converged, the KV loss could
slow down the training process. In our method, we do not use the KV loss.

Model Size | $w$ | KV Loss | Dev ppl.
---|---|---|---
50M | 0 | no | 15.965
0 | yes | 15.610
2 | no | 15.004
2 | yes | 15.065
1.1B | 2 | no | 9.746
2 | yes | 10.073
Table 9: Effect of the KV loss on a 50M and a 1.1B model. The 50M model is
trained on WikiText-103 with 3 epochs and the 1.1B model is trained on a 100B
subset of SlimPajama.

###  C.4 Encode Prompts with Different Number of Iterations

Though we set $m=7,b=2$ during training, it does not necessarily mean that we
have to encode the prompts with 9 iterations. Is it possible to encode the
prompts with less iterations? What if we encode the prompts with more
iterations?

We treat the token segments as prompts and test the model trained in Section
3.2 with different numbers of iterations during encoding. As shown in Figure
12, the perplexity increases when the number of iterations is reduced, but
still in a reasonable scale if we only reduce one or two iterations. The more
warmup layers there are, the more stable the performance is. Increasing the
number of iterations does not noticeably affect the performance. Thus, one
could make a trade-off to set the proper number of iterations during inference
to balance the time encoding prompts and the quality of generation texts.

Figure 12: Performance of the models with different numbers of iterations for prompt encoding. Model Size | Model | Batch Size | Latency (s) | Throughput (token/s)
---|---|---|---|---
7B | Llama | 15 | 217.71 | 141.10
Ours ($w=2$) | 16 | 89.15 | 367.56
Ours ($w=10$) | 16 | 131.43 | 249.32
30B | Llama | 1 | 145.23 | 14.10
Ours ($w=2$) | 1 | 101.12 | 20.25
Ours ($w=10$) | 1 | 106.99 | 19.14
Table 10: Inference latency and throughput of the models with different model
sizes. The models are tested on an A100 (80GB) GPU with prompt length 2048 and
generation length 2048. The batch size is (approximately) the largest batch
size for standard Llama model.

###  C.5 Model Performance With Respect to Token Position

The long-context performance of LLMs highly relies on KV cache Xiao et al.
(2024); Han et al. (2023); Adnan et al. (2024). To verify that the performance
of our model does not degrade under long context, we test the perplexity of
our model with different token positions on PG19 Rae et al. (2020). Despite
the fact that we only compute the KVs of a few layers, the performance of our
model does not degrade with the token position (Figure 13) and is comparable
to that of TinyLlama. Due to the limitation of computational resources, we
only train models with context length 2048.

Figure 13: Perplexity of the models with different token positions on PG19. We
use the concatenated test set of the PG19 dataset as the input.

###  C.6 Improvement of Throughput Not Necessarily Due to Larger Batch Size

In Section 3.1, we show that our model achieves higher generation throughput
than standard transformers. While one might assume that this improvement is
solely due to a larger batch size, Figure 4 demonstrates that our model
outperforms the standard transformer even with the same batch size.
Additionally, Table 10 reports reduced latency for our model when compared
with the standard transformer with the same batch size. This suggests that our
method can also benefit scenarios that require fast response. The exact reason
for this phenomenon is not yet clear and we speculate that it could be
attributed to factors such as the reduced calculation of KVs, the decreased
memory consumption enabling faster memory transfer and access, and various
implementation details.