F i DO : Fusion-in-Decoder optimized for stronger performance and faster inference (2023)

Abstract: Fusion-in-Decoder (FiD) is a powerful retrieval-augmented language model that sets the state-of-the-art on many knowledge-intensive NLP tasks. However, the architecture used for FiD was chosen by making minimal modifications to a standard T5 model, which our analysis shows to be highly suboptimal for a retrieval-augmented model. In particular, FiD allocates the bulk of FLOPs to the encoder, while the majority of inference time results from memory bandwidth constraints in the decoder. We propose two simple changes to the FiD architecture to alleviate memory bandwidth constraints, and speed up inference by 7x. This allows us to use a much larger decoder at modest cost. We denote FiD with the above modifications as FiDO, and show that it strongly improves performance over existing FiD models for a wide range of inference budgets. For example, FiDO-Large-XXL performs faster inference than FiD-Base and achieves better performance than FiD-Large.

## introduction a large body of work has demonstrated that language model performance on downstream tasks can be improved by augmenting the model with relevant retrieved text @xcite @xcite . in particular, the fusion-in-decoder (fid) architecture @xcite stands out for strong performance, even outperforming much larger models on many knowledge-intensive tasks @xcite . however, fid uses a standard t5 encoder-decoder architecture @xcite which was not designed for use as a retrieval-augmented model. in this work we propose fido, a modified fid architecture optimized for the retrieval-augmented setting. the fid decoder is responsible for a difficult task, assimilating information from many passages and reasoning over the information to generate an output. however, because the encoder and decoder are similar size and the encoder is applied to a large number of retrieved passages, fid devotes an order of magnitude more floating point operations (flops) to the encoder than the decoder. in spite of this, the majority of inference time is actually spent in the decoder, as has been observed in prior work @xcite . this surprising result is shown in figure 1 . our analysis finds that for typical inference settings the fid decoder is memory-bandwidth bound @xcite due to using multi-head cross-attention @xcite over a large input sequence. based on this analysis, we propose two sets of architectural changes. we first propose to reduce the cost of cross-attention over retrieved passages by removing most cross-attention layers from the decoder. this reduces cost and yields much smaller losses in performance than fid-light (hofstätter et al., 2022), the best previously-proposed approach for optimizing fid. we also replace multi-head attention with multi-query attention @xcite . with these modifications the memory-bandwidth bottleneck is eliminated: decoder inference is now orders of magnitude faster and most inference time is spent in the encoder, consistent with the balance of flops between components. finally, we propose to partially rebalance compute towards the decoder by massively scaling decoder size, using a smaller encoder to extract information from retrieved passages and a larger decoder to assimilate the information and reason about the desired output. we refer to the resulting series of models as fido (fusion in decoder optimized) and show that fido strongly outperforms standard fid models on the question-answering datasets natural questions @xcite , triviaqa @xcite and webquestions @xcite for a wide range of inference budgets and settings. figure 2 summarizes some of these results. ## analysis retrieval-augmented models generally read many context tokens relative to the number of question or answer tokens, such that processing retrieved text consumes the bulk of flops. however, past work has shown that most inference time for fusion-in-decoder (fid) is spent in the decoder @xcite . our own experiments support this (figure 1 ). this section investigates fid's computational structure and decoder inference speed, and finds the slower decoder speed to be the result of memory bandwidth constraints, exacerbated by attention over retrieved documents. ## method we have shown that the encoder accounts for the bulk of fid flops and training cost, while fid spends the majority of inference time in the decoder due to low operational intensity of cross-attention layers. next we propose several ways to alleviate the decoder bottleneck. this allows us to efficiently allocate more compute to the decoder by scaling decoder size without significantly increasing the inference speed. we denote fusion-in-decoder with the proposed optimizations as fido (fusionin-decoder optimized). ## related work retrieval-augmented models there exists a large body of retrieval-augmented approaches. some particularly well known models are realm @xcite , rag (lewis et al., 2020), retro @xcite and fusion-in-decoder @xcite . fid in particular has achieved state-of-the-art performance on a wide variety of tasks @xcite @xcite and in this work we focus on improving the performanceefficiency trade-offs for fid. retro is another closely related retrieval-augmented model, as it uses a small encoder for retrieved context and a larger primary decoder like fido does. unlike retro, fido's efficiency improvements allow it to tractably attend to many retrieved passages with a much larger decoder. ing factor for computations on modern devices @xcite @xcite . @xcite shows that autoregressive transformers are particularly bandwidth bound during inference, and proposes multi-query attention as a partial solution. we find that this is exacerbated by the fid setting, and adopt multi-query attention for fido to ameliorate the problem. pope et al. (2022) also investigates multi-query attention, primarily in the context of efficient inference and parallelization for very large language models, whereas we focus on performance/cost trade-offs for the retrieval-augmented setting. ## conclusion we perform analysis of the performance-inference speed tradeoff for fid, showing that the encoder uses more flops but most time is spent in the decoder due to memory bandwidth constraints. we propose fido, an extension of fid which removes most cross-attention layers and employs multiquery attention to vastly reduce the cost of the decoder. the resulting model spends most time in the encoder, consistent with compute analysis, which fido takes advantage of by strongly increasing the size of the decoder. we show that fido achieves much stronger performance for the same inference budget relative to existing fid models. the second difference to the training recipe arises from the observation that fido xl-xxl is unstable for the standard training regimen. we solve the instability by restarting from a recent healthy checkpoint with a 10x decreased learning rate, which happened once. during fine-tuning, we load not only model weights but also second moment estimates, which we find leads to better fine-tuning in general and particularly for asymmetric models. we finetune with learning rate 0.001 and batch size 64 for all datasets. for evaluation on test sets we select the checkpoint with the best validation performance.