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

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

We also notice that additional pretraining improves balanced models a little on average but substantially benefits some tasks, like GSM8K and DROP. Besides, 9B-2B performance increases consistently during the adaptation, quickly surpassing Gemma 2 2B and moving towards Gemma 2 9B. This demonstrates the feasibility of encoder-decoder adaptation from varying-sized decoder-only LLMs, as well as its ability to utilize the knowledge from pretrained models.

#### Pretraining objective matters: UL2 and PrefixLM show different characteristics.

Previous study reported the superiority of UL2 over PrefixLM ( Tay et al., 2022 ) , but PrefixLM in our study is enhanced with knowledge distillation, which often improves small models significantly. We compare these two objectives for the adaptation in Table 2 .

We find that PrefixLM and UL2 have their own strengths. Specifically, UL2 delivers stronger contextual representations, outweighing PrefixLM on SuperGLUE across most model scales, resonating with previous findings ( Tay et al., 2022 ) . In contrast, PrefixLM produces more powerful generative LLMs thanks to its generation nature and the knowledge distillation. It surpasses UL2 on PT and IT benchmarks in most cases. Particularly, it outperforms UL2 at 9B-2B on both PT and IT by up to 3.6, a significant margin. Since generative LLMs have become the mainstream, we base our following analysis on PrefixLM. We discuss our attempts to combine PrefixLM and UL2 in the next section.

#### Encoder-decoder LLMs outperform decoder-only LLMs especially after instruction tuning.

Table 2 also shows that the adapted encoder-decoder LLMs achieve comparable or slightly better pretraining performance than their decoder-only counterpart but with substantially improved instruction-tuning performance, echoing with the findings of Wang et al. (2022) . For example, the 9B-9B encoder-decoder LLM surpasses Gemma 2 9B by 1.4 and 4.9 on PT and IT, respectively. The performance gap further increases to 1.8 and 7.1 at 2B-2B scale. We notice that the adaption performs slightly worse at scales below 2B on PT, but the improvements on IT are still promising, e.g. 7.2 at XL-XL.

Regardless of PT or IT models, pretraining objectives, and model scales, encoder-decoder LLMs perform consistently better than decoder-only LLMs on SuperGLUE. This suggests that the contextual representation from encoder-decoder LLMs is often of higher quality, likely due to bidirectional self-attention.

We need to highlight that the above analysis is based on the overall performance, which may not apply when it comes to a specific downstream task. As shown in Table 3 , there are some tasks favoring encoder-decoder models while others favoring decoder-only models especially for PT models. For example, after pretraining, Gemma 2 9B surpasses 9B-9B by 4.1 on ARC-C but underperforms it by 4.4 on Winogrande; while encoder-decoder LLM shows more consistent advantage after instruction tuning, 9B-9B still lags behind Gemma 2 9B by 0.9 on WMT23. This illustrates the complexity when evaluating LLM capability as well as the risk of reaching misleading conclusions when adopting biased evaluation tasks. We reduce such risk by selecting as diverse and broad tasks as possible for evaluation.

#### Encoder-decoder LLMs balance quality and inference efficiency more effectively.

We next analyze different models from the perspective of inference efficiency which becomes increasingly crucial for model deployment. Figure 3 shows that balanced encoder-decoder LLMs have similar inference flops to their decoder-only counterparts, e.g. 2B-2B vs. Gemma 2 2B. As such, encoder-decoder models often dominate the quality-inference efficiency frontier across PT, IT, and SuperGLUE benchmarks.

We acknowledge that inference flops may not correlate well with actual running speed due to factors like inter-device communication, key-value caching, and autoregressive bottleneck. We then provide the latency results measured on GSM8K for 2B and 9B models in Figure 4 , which further verified the above analysis. 9B-9B and 2B-2B show similar latency to Gemma 2 9B and 2B, respectively, but clearly better performance. In particular, 9B-2B, the one pairing large encoder and small decoder, shows similar latency to Gemma 2 2B but significantly better performance than 2B-2B.

Together, these confirm that encoder-decoder adaptation indeed provides a more flexible way for balancing between quality and inference speed.

## 6 Discussion

#### Is the improvement after the adaptation simply due to the extra pretraining compute?

Not really. We also tried to apply more pretraining compute to Gemma 2 2B by going through another 6 trillion tokens, which leads to a PT score of 48.57, still significantly below the encoder-decoder adaptation, 49.7. This indicates that the additional pretraining compute can’t fully explain the improvements from the adaptation and we argue that the inductive bias of encoder-decoder modeling plays a crucial role.

#### Does cross-attention warmup matter for unbalanced encoder-decoder?

Yes. Our preliminary experiments with 9B-2B and UL2 on 800B tokens show that the pretraining performance over Boolq and GSM8K reduces from 62.5 to 61.8 without the warmup. Besides, increasing warmup steps from 1K to 5K further reduces performance to 60.2. An adequate amount of warmup optimization is required to reach the optimal performance.

#### Can we switch from grouped-query attention to multi-head self attention for the encoder?

Yes but with mixed results. Gemma 2 adopts grouped-query attention (GQA) to improve its decoding efficiency. However, unlike the decoder, the encoder can be fully parallelized during inference, making the use of multi-head attention (MHA) reasonable. We tried to expand GQA in Gemma 2 2B to MHA by replicating head parameters for the encoder self-attention. Under PrefixLM, this improves PT performance to 50.2 by 0.5 at 2B-2B but reduces IT performance to 43.5 by 2.9. We thus still stick to GQA when adapting Gemma 2 2B and 9B for the encoder.

#### Does bidirectional self-attention matter for the encoder?

Yes. A crucial difference between encoder-decoder and decoder-only LLMs is the use of bidirectional self-attention. We also tested keeping the encoder self-attention causal at 2B-2B, which achieves a PT and IT score of 45.6 and 41.7, lagging behind its bidirectional counterpart substantially by 4.1 and 4.7, respectively. Note, the causal 2B-2B model surpasses Gemma 2 2B on IT by 2.7, although it performs worse on PT. This suggests that bidirectional self-attention contributes greatly to the success of our adaptation, but is not the only factor.

#### Would pretraining encoder-decoder LLMs from scratch yield better performance?

Not really. Pretraining from scratch is a common method for developing new LLMs. We also pretrained encoder-decoder LLMs from scratch on 8 trillion tokens with PrefixLM. Table 4 summarizes the results. Despite using more pretraining tokens, encoder-decoder LLMs pretrained from scratch only perform better at small scales, such as S-S and B-B, beyond which adaptation shows clear superiority. As such, adaptation is a more computationally efficient way of developing powerful encoder-decoder LLMs.

#### Is IT/SuperGLUE score predicable from PT score?

Mixed. A general assumption in LLM development is that PT performance can be used as an indicator for downstream applications. We summarize all our ablations and put them in Figure 6 . Over all data points and across all model sizes, the correlation is pretty strong: a Spearman’s ρ \rho of 0.97 and 0.89 for IT vs. PT and SuperGLUE vs. PT, respectively. However, when considering data points within each model size separately, the averaged Spearman’s ρ \rho reduces to 0.42 and 0.05, respectively and is not significant anymore.

In practice, we also noticed that PT checkpoints with weaker performance sometimes yield significantly better IT or SuperGLUE performance. When selecting PT checkpoints for a specific model size, it’s better to also examine their IT performance apart from PT results to avoid some biases or overfitting.

#### Can we get the best of both worlds from PrefixLM and UL2?

This is non-trivial. Our first attempt is to merge checkpoints trained from PrefixLM and UL2 with uniform weighting. Unfortunately, the merged model results in either similar or much worse performance. We argue that PrefixLM and UL2 lead to different training dynamics and converge to very different local minima. Directly merging their weights doesn’t work right out of the box.

We next explore a two-stage optimization, where we first adapt with PrefixLM and then shift to UL2 for the last 10% of training, and vice versa. Figure 5 shows very mixed results. Switching from PrefixLM to UL2 generally hurts performance. In contrast, switching from UL2 to PrefixLM improves IT performance, but suffers from reduction in PT and SuperGLUE performance.

Another direction is to jointly optimize the model on PrefixLM and UL2, which we leave for future work.

## 7 Conclusion and Future Work

In this paper, we presented methods for building powerful, general purpose encoder-decoder LLMs by adapting from pretrained decoder-only LLMs. Such adaptation offers high flexibility in leveraging different types/families of pretrained decoder-only models as well as combining different-sized models. Through extensive experiments based on Gemma 2, we demonstrated the feasibility and effectiveness of the adaptation: the adapted encoder-decoder LLMs outperform their decoder-only counterparts substantially after instruction tuning, dominating the quality-inference efficiency frontier. Besides, encoder-decoder LLMs also provide better contextual representations as evaluated on SuperGLUE.

We hope our findings inspire more researchers from academia and industry to revisit the encoder-decoder paradigm for LLM development. To facilitate the research, we will release the code and checkpoints at XXX (coming soon).

Our work still suffers from several limitations. Particularly, we only experimented with Gemma 2 models up to 9B, although the proposed approach could apply to other LLM families. In the future, we are interested in scaling the model size (e.g, to 27B), exploring other LLMs (such as LLaMA), examining more unbalanced setups, and testing the combination of dense and MoE LLMs. As mentioned above, we will also investigate better ways to leverage PrefixLM, knowledge distillation, and UL2. Extending our adapted encoder-decoder LLM to cross/multi-modality modeling (e.g., vision-language and speech-language) would be another intriguing direction.

## Acknowledgements

We’d like to thank Enrique Alfonseca, Tris Warkentin, Xiaodan Song, Sugato Basu, Inderjit Dhillon, Alexander Grushetsky, Pandu Nayak, Ramakrishnan Srikant, and Slav Petrov for their constructive feedback on the manuscript. We are grateful to Srinivasan Venkatachary for supporting this project.