Dodo: Dynamic Contextual Compression for Decoder-only LM s (2024)

Abstract: Transformer-based language models (LMs) are inefficient in long contexts. We propose Dodo, a solution for context compression. Instead of one vector per token in a standard transformer model, Dodo represents text with a dynamic number of hidden states at each layer, reducing the cost of self-attention to a fraction of typical time and space. Moreover, off-the-shelf models such as LLaMA can be adapted to Dodo by efficient parameter tuning methods such as LoRA. In use, Dodo can act as either an autoregressive LM or a context compressor for downstream tasks. We demonstrate through experiments in language modeling, question answering, and summarization that Dodo retains capabilities in these tasks, while drastically reducing the overhead during decoding. For example, in the autoencoding task, Dodo shrinks context at a 20x compression ratio with a BLEU score of 98% for reconstruction, achieving nearly lossless encoding.

## introduction transformer-based lms @xcite suffer from quadratic computational complexity w.r.t. sequence length, making it challenging to scale to long sequences. proposed solutions @xcite include sparsifying attention patterns @xcite or approximating the attention computation with kernel methods @xcite . however, not all these approaches are proven effective for nlp tasks @xcite , and very few of them are applied to large language models (llms), such as llama @xcite . we propose dodo , a solution for dynamic contextual compression for decoder-only lms. while a standard transformer represents a text with vector sequences of the same length as tokens, nlp is concerned with human language, … q: what is nlp? a: it is … compression w/ figure 1 : dodo efficiently maps long inputs into a compressed set of vectors named nuggets , which can then be attended to when processing a query. the intuition of dodo is to use a smaller, variable number of vectors as a contextual representation. past research indicates that a subset of token embeddings, named nuggets , in an encoder with global attention may carry enough information to reconstruct surrounding context @xcite , and upon inspection those authors observed these nuggets tended to account for preceding text. this suggests a decoder-only model might be dynamically capable of deriving such a representation online (fig. 1 ). to enable dodo requires addressing a selection process that is not differentiable: we adopt the straight-through estimator @xcite to make the model end-to-end trainable. past work on context compression, such as @xcite and @xcite , appends fixed additional tokens. dodo grows the representation with sequence length and re-uses existing token embeddings. moreover, unlike pattern-based methods that evenly chunk the text @xcite , experiments show that dodo spontaneously learns to use textual delimiters as nuggets , naturally splitting the text into subsentential units (section 4.3). dodo supports causal masking and can be naturally used as an autoregressive lm. we experimentally demonstrate that dodo can achieve a perplexity score lower than the original lm with restricted memory, outperforming the baseline model of @xcite . for tasks with a fixed context, e.g. long-form qa, dodo works as a context compressor: it encodes a token sequence into a shorter vector sequence, achieving a configurable compression ratio. in experiments on autoencoding, we demonstrate that dodo can achieve near lossless encoding with a compression ratio as high as 20x, a marked improvement over icae @xcite . after fine-tuning, dodo is effective in downstream nlp tasks such as question answering (qa) and summarization, where it performs on par with or even better than the original lms while achieving a compression ratio as high as 10x. in summary, we propose dodo for contextual compression for decoder-only transformers. it learns to subselect a fractional number of tokens as context representation. a straight-through estimator ensures that dodo is differentiable and can be trained with the next-token prediction objective. dodo achieves a remarkable compression ratio of up to 20x and is shown to be effective in tasks such as autoencoding, language modeling, and applications including qa and summarization. ## approach in this paper, we study the language modeling problem p(w t | @xmath0, where w @xmath1. the common transformer @xcite approach encodes a token sequence w 1:n into a sequence of vectors and then predicts the next token: ## overall experiment setup we adopt the decoder-only transformer architecture of llama @xcite as our base model. for the autoencoding experiment, we use the checkpoint of llama-7b following the baseline model icae @xcite . we use the checkpoint of llama-2-7b for the autoregressive language modeling experiments (section 5) and llama-2-7b-chat (section 6) for the downstream nlp tasks. we adopt lora @xcite with a rank of 32 to fine-tune the parameters of the lm, namely θ and ϕ. we adopt the implementation of huggingface/peft packakge @xcite . more specifically, we fix the original parameters of llamaand add two lora adapters for θ and ϕ respectively. different adapters are activated for the computation of compressing and decoding of dodo . we disable the adapters to produce the features to scorer. we employ mixed precision to save gpu memory. the training is scaled up to 16 nvidia v100 cards with deepspeed @xcite . see appendix b for further training details, including hyperparameters, and parameter counts. in this section, we use dodo as a context compressor (section 2.3) and apply it to the autoencoding task. as a comparison, we use in-context autoencoder @xcite as a baseline model. in this task, a model is asked to reconstruct the input text from a compressed representation. following icae, we fine-tune the llama-7b model on the pile @xcite dataset. we manually split the corpus into train, dev, and test splits, and train the model until convergence. as stated in section 2.3, we use dodo to compress the input text into fewer hidden states z, and then use the lm to decode the input sequence. the size of hidden states z, i.e. k, is set to be proportional to the length of the input sequence: @xmath2, and we set @xmath3. we prepend a trainable soft token to the decoding sequence to signal the model to reconstruct inputs @xcite . the key idea of icae is to append 128 tokens to the input sequence as "memory slots," and train the decoder to reconstruct the input from the memories: we measure using bleu @xcite score on pairs of input and decoded texts. foot_3 ## downstream task experiments we pick downstream tasks where a document as context is followed by a query. the model is asked to encode the document and decode the answer conditioned on the document encoding and question. in these tasks, we use dodo as a context compressor (section 2.3), and we set the compression @xmath4. to train dodo to perform these tasks, we consider 2 scenarios. a) fine-tuning: dodo is trained on the training set of the downstream tasks. b) zero-shot: dodo is trained on normal texts randomly sampled from the pile and directly tested on the downstream task. in this case, each text is split into 2 parts, containing up to 512 and 128 tokens, and the model is asked to decode the second part conditioned on the encoding of the first part. we consider the tasks of question answering and summarization. datasets used in this section are squad @xcite and cnn/dailymail v3.0.0 @xcite for summarization. their statistics are listed in we use the following baseline methods: • full : results of the original lm. compressed tokens: the number of compressed tokens that precede context tokens: the uncompressed context immediately before the token to be predicted. this adds up to total state, which is directly comparable between systems, using three settings (256, 128, and 64). dodo trades off explicit context for larger history, with better perplexity results. ## conclusion in this work, we propose dodo , a method for contextual compression for decoder-only transformers. in language modeling (section 5) and summarization (section 6.2), dodo is shown to generate a highly condensed representation of the context, while the results in autoencoding (section 4) and question answering (section 6.1) reflect that the details of the contexts can be recovered from nuggets . moreover, in section 6.1 we show that dodo trained with text continuation preserves the capability of instruction following. this demonstrates llms can encapsulate more of their input into fewer hidden states than previously realized, suggesting a new direction for efficient foundation models. future work will explore more specialized versions of this proposal for optimizing results on individual applications, such as in dialog, supervised fine-tuning, reinforcement learning with human feedback, and in-context learning.