Encoder-Decoder Gemma: Improving the Quality-Efficiency Trade-Off via Adaptation — part 1/2

Cleaned-text code points [0, 15743). Bibliography omitted here; source.md includes it.

# Encoder-Decoder Gemma: Improving the Quality-Efficiency Trade-Off via Adaptation

###### Abstract

While decoder-only large language models (LLMs) have shown impressive results, encoder-decoder models are still widely adopted in real-world applications for their inference efficiency and richer encoder representation. In this paper, we study a novel problem: adapting pretrained decoder-only LLMs to encoder-decoder, with the goal of leveraging the strengths of both approaches to achieve a more favorable quality-efficiency trade-off. We argue that adaptation not only enables inheriting the capability of decoder-only LLMs but also reduces the demand for computation compared to pretraining from scratch. We rigorously explore different pretraining objectives and parameter initialization/optimization techniques. Through extensive experiments based on Gemma 2 (2B and 9B) and a suite of newly pretrained mT5-sized models (up to 1.6B), we demonstrate the effectiveness of adaptation and the advantage of encoder-decoder LLMs. Under similar inference budget, encoder-decoder LLMs achieve comparable (often better) pretraining performance but substantially better finetuning performance than their decoder-only counterpart. For example, Gemma 2B-2B outperforms Gemma 2B by ∼ \sim 7% after instruction tuning. Encoder-decoder adaptation also allows for flexible combination of different-sized models, where Gemma 9B-2B significantly surpasses Gemma 2B-2B by > > 3%. The adapted encoder representation also yields better results on SuperGLUE. We will release our checkpoints to facilitate future research.

## 1 Introduction

Neural network architectures are often designed to incorporate certain assumptions or inductive biases regarding the input data, leading to either improved model performance or better computational efficiency, if not both. Unlike the popular decoder-only architecture used for large language model (LLM) ( Brown et al., 2020 ) , the encoder-decoder architecture adopts separate modeling modules – an encoder for input understanding and a decoder for output generation ( Vaswani et al., 2017 ) . This separation decouples parameters for different functionalities and thus enjoys higher freedom in handling contextual representation and challenging tasks ( Tay et al., 2022 ; Wang et al., 2022 ) . It also offers high flexibility in changing the encoder and decoder size (e.g., a large encoder paired with a small decoder) to control the quality-efficiency trade-off ( Kasai et al., 2020 ; Zhang et al., 2022 ) , an increasingly important aspect for LLM deployment ( Gemini et al., 2024 ) . Despite these benefits, however, the study on encoder-decoder LLMs receive little to no attention nowadays.

In this paper, we revisit this classical architecture by exploring the following question: can we get strong(er) encoder-decoder LLMs by adapting from existing pretrained decoder-only LLMs? We consider the adaptation more significantly than pretraining new models from scratch since pretraining is resource-intensive and powerful decoder-only models at different sizes are already widely available ( Dubey et al., 2024 ; Team et al., 2024 ; Liu et al., 2024a ; Yang et al., 2024 ; Jiang et al., 2024 ) . Our hypothesis is that, by reusing parameters from decoder-only models, we can accelerate training and effectively transfer their internal knowledge to encoder-decoder, preserving (even enhancing) their capabilities. Note adaptation also allows for pairing varying-sized decoder-only models to achieve specific quality-efficiency considerations. Yet, the optimal method for such adaptation and the extent to which performance can be improved remain open questions, which we aim to address rigorously.

We employ Gemma 2 ( Team et al., 2024 ) as the testbed. As shown in Figure 1 , the encoder-decoder architecture follows the original Transformer ( Vaswani et al., 2017 ) but equipped with Gemma 2 modifications. The key idea behind the adaptation is to initialize the parameters of the encoder-decoder model from pretrained decoder-only model(s) as a warmup and then pretrain or adapt all parameters with self-supervised learning. Depending on whether the encoder and the decoder share the same configuration, we propose different initialization and optimization strategies for the cross-attention layer. We also compare different pretraining objectives, including prefix language modeling with knowledge distillation ( Hinton et al., 2015 ) and UL2 ( Tay et al., 2022 ) . Apart from Gemma 2 2B and 9B, we pretrain a series of small models to better understand the adaptation at different scales.

To thoroughly evaluate model performance, we adopt different benchmarks for pretrained and instruction-tuned models respectively, each covering a range of established academic evaluations. In addition, we use SuperGLUE ( Wang et al., 2019a ) to measure the quality of the learned contextual representations. Our main findings are below: • Leveraging pretrained decoder-only LLMs is an effective way to build powerful encoder-decoder LLMs, which yields substantially improved downstream performance particularly after instruction tuning under similar inference flops.

• Our adaptation method is highly flexible, allowing for pairing large encoder with small decoder, such as 9B-2B, with significant quality gains over Gemma 2 2B but similar generation latency.

• Adaptation is not only more compute efficient but also more effective than pretraining from scratch.

• Pretraining objective matters. Models trained with prefix language modeling and knowledge distillation are generally better at generative tasks, while UL2 models have better encoder representations.

## 2 Related Work

While the decoder-only architecture has become the de facto standard for LLMs, the debate between encoder-decoder and decoder-only modeling is still not conclusive. Many prior studies proposed different approaches to pretrain strong encoder-decoder models, e.g., MASS ( Song et al., 2019 ) , T5 ( Raffel et al., 2020 ) , mT5 ( Xue et al., 2021 ) , byT5 ( Xue et al., 2022 ) , BART ( Lewis et al., 2020 ) , and OpenBA ( Li et al., 2023 ) . Tay et al. (2022) compared different pretraining objectives, highlighting the superiority of UL2 and encoder-decoder modeling. Zhang et al. (2022) systematically examined the scaling behavior of both architectures on machine translation, showing their similarity when adequate objectives are applied. Wang et al. (2022) thoroughly explored different modeling choices and training objectives with a focus on LLM zero-shot generalization. They discovered that encoder-decoder LLMs after instruction tuning achieve the best performance, echoing with our experiments. They also studied adaptation, but it is between different pretraining objectives rather than from decoder-only LLMs to encoder-decoder LLMs.

Leveraging pretrained models for encoder-decoder modeling has been extensively explored. In the BERT era ( Devlin et al., 2019 ) , researchers developed different ways of utilizing it to enhance encoder-decoder performance on downstream tasks, such as machine translation ( Zhu et al., 2020 ; Clinchant et al., 2019 ; Yang et al., 2020 ) , grammatical error correction ( Kaneko et al., 2020 ) , summarization ( Liu & Lapata, 2019 ) , and text generation ( Chen et al., 2019 ) . Our work follows a similar spirit but is based on pretrained decoder-only LLMs and focuses on developing general-purpose encoder-decoder LLMs.

Another related direction is the development of inference friendly LLMs. Techniques for improving inference efficiency are many, ranging from quantization ( Dettmers & Zettlemoyer, 2023 ) , key-value cache optimization ( Corallo & Papotti, 2024 ) , recurrent modeling ( Gu & Dao, 2023 ; Botev et al., 2024 ) , to strong small LLMs with improved pretraining ( Abdin et al., 2024 ; Liu et al., 2024b ) , to name a few. While these techniques offer significant efficiency gains, their focus is fundamentally distinct and complementary to our proposed encoder-decoder adaptation, i.e., both approaches can be used in conjunction to realize greater overall efficiency.

## 3 Approach: Encoder-Decoder Adaptation

### 3.1 Architecture

Pretraining LLMs is both compute and time intensive. To reduce the amount of training required, we propose to adapt existing decoder-only LLMs to encoder-decoder and leverage pretrained decoder-only checkpoints for initialization, as shown in Figure 1 . Due to this, we keep the encoder-decoder architecture as similar as possible to original decoder-only model, only introducing changes when necessary. This results in the following architecture: 1. Encoder has exactly the same architecture as the decoder-only model, but self-attention is switched from causal to bidirectional. We provide ablations in Section 6 that illustrate the critical effect of bidirectional attention on downstream performance.

2. In each Decoder block, FFN and self-attention parts are identical to the corresponding parts in decoder-only models, and cross-attention has the same number of heads and head dimension as self-attention, but attends to the whole output of the encoder.

We base our study on Gemma 2 ( Team et al., 2024 ) . But note our approach is highly flexible and isn’t restricted to specific decoder-only architectures. We can easily apply our method to other model families, such as LLaMA ( Dubey et al., 2024 ) , QWen ( Yang et al., 2024 ) , and DeepSeek ( Liu et al., 2024a ) . In theory, we can also adapt decoder-only models from different families, such as pairing LLaMA models with QWen models.

In addition, our approach allows for unbalanced encoder-decoder models, where the decoder is significantly smaller than the encoder. This provides better support for applications where input processing capabilities are more important than generative capacity. For example, for summarization, deep understanding of the input text is often more important than the generation part, as it doesn’t need to generate any new information. As a result, generation time is significantly reduced, while providing competitive quality.

### 3.2 Initialization

When initializing an encoder-decoder model from a decoder-only checkpoint, we try to map every layer to the most similar weight in the decoder-only checkpoint. In particular, the encoder is fully initialized from the decoder-only checkpoint, as it doesn’t introduce any new weights. In the decoder, FFN and self-attention subblocks are initialized from the FFN and self-attention weights from the corresponding layers in the decoder-only checkpoint.

Cross-attention is initialized from self-attention weights in the balanced setup where encoder and decoder have the same configuration. Otherwise, we first initialize cross-attention from scratch and then finetune it for the first K K steps as a warmup while freezing other model parameters. After K K steps, all model parameters are tuned.

### 3.3 Pretraining Objective

Decoder-only pretraining often adopts causal language modeling on a single sequence. In contrast, encoder-decoder adaptation requires separate input and target sequences to be fed to the encoder and decoder separately. We explore two classical pretraining objectives for encoder-decoder modeling: prefix language modeling (PrefixLM) and UL2 ( Tay et al., 2022 ; Wang et al., 2022 ) .

PrefixLM behaves similar to causal language modeling except for its prefix condition. To simplify the preprocessing, we split a sequence equally into two halves, the first half used as input and the second one as target. This also eases the adoption of knowledge distillation from decoder-only models. UL2 is more complicated. It is composed of several denoising tasks at different levels of complexity. We prepare UL2 data following Tay et al. (2022) . We compare their performance in experiments.

## 4 Setup

#### Data Setting

Our data for pretraining and instruction tuning – including supervised finetuning (SFT) and reinforcement learning from human feedback (RLHF) – follow Gemma 2 ( Team et al., 2024 ) . For the adaptation, we preprocess the Gemma 2 pretraining data (8 trillion tokens) with PrefixLM and UL2. Note Gemma 2 pretraining data comes with knowledge distillation. We preserve this information for PrefixLM while adopting ground-truth targets for UL2 as mapping the teacher logits to UL2 is non-trivial. The preprocessed data has an input-output sequence length of 4096-4096 and 8192-8192 for PrefixLM and UL2, respectively. We adapt our models on up to 2 trillion tokens.

#### Model Setting

We use Gemma 2 (2B and 9B) as the base decoder-only LLM. We also pretrain several smaller models (Small, Base, Large, and XL) following mT5 configurations ( Xue et al., 2021 ) under the Gemma 2 framework, and then adapt them to encoder-decoder LLMs. Detailed model configurations are given in Table 1 .

#### Evaluation

We employ diverse academic evaluation datasets to evaluate different capabilities of LLMs. Concretely, we use the following benchmarks: • Pretraining (PT) benchmark: Boolq ( Clark et al., 2019 ) , SIQA ( Sap et al., 2019 ) , PIQA ( Bisk et al., 2020 ) , ARC-c&ARC-e ( Clark et al., 2018 ) , MMLU ( Hendrycks et al., 2021 ) , MMLU Pro ( Wang et al., 2024 ) , HellaSwag ( Zellers et al., 2019 ) , Winogrande ( Sakaguchi et al., 2021 ) , TruthfulQA ( Lin et al., 2021 ) , AGIEval ( Zhong et al., 2023 ) , BBH ( Suzgun et al., 2022 ) , DROP ( Dua et al., 2019 ) , GPQA ( Rein et al., 2023 ) , GSM8K ( Cobbe et al., 2021 ) , HumanEval ( Chen et al., 2021 ) , Lambada ( Paperno et al., 2016 ) , MATH-500 ( Hendrycks et al., 2021 ) , MBPP ( Austin et al., 2021 ) , NQ ( Kwiatkowski et al., 2019 ) , TriviaQA ( Joshi et al., 2017 ) , and WMT23 ( Kocmi et al., 2023 ) . We perform zero/few-shot prompting for pretrained LLMs, and report the averaged result as PT score .

• Instruction-tuning (IT) benchmark: GSM8K, MMLU, MMLU Pro, MBPP, HumanEval, MATH-500, BBH, GPQA (Diamond), WMT23, and MGSM ( Shi et al., 2022 ) . We perform zero/few-shot prompting with task-specific instruction for instruction-tuned models, and report the averaged result as IT score .

• SuperGLUE ( Wang et al., 2019b ) : we use this benchmark to examine the learned contextual representation. We stack a task-specific head on the representation of the last token in the encoder (decoder) of the encoder-decoder (decoder-only) LLM, and finetune all parameters on the training set. Learning rate, batch size, and dropout are grid-searched for each task. We reformulate all tasks as classification tasks and report averaged dev-set accuracy over COPA, WIC, WSC, RTE, MultiRC, CB, and Boolq.

For generative tasks, we always apply greedy sampling. We perform pretraining, SFT, and RLHF based on the Gemma 2 recipe except for the learning rate which we tune empirically for encoder-decoder LLMs. In unbalanced encoder-decoder adaptation, e.g. 9B-2B, we set the cross-attention warmup step K K to 1000.

## 5 Results

#### The encoder-decoder adaptation converges rapidly, particularly for balanced architectures.

While adaptation leverages pretrained parameters for initialization, whether and how this benefits model convergence is still questionable. Figure 2 shows the change of PT performance with respect to the amount of pretrained tokens. Obviously, adaptation is very computationally efficient, converging quickly and achieving similar performance to its decoder-only counterpart after only tens of billions of tokens. Balanced architectures (2B-2B and 9B-9B) converge much faster than the unbalanced ones (9B-2B) since all parameters in the former are initialized from pretrained decoder-only models while the cross-attention in the latter is randomly initialized.