Bitune: Leveraging Bidirectional Attention to Improve Decoder-Only LLMs — part 1/3
Cleaned-text code points [0, 15979). Bibliography omitted here; source.md includes it.
# Bitune: Leveraging Bidirectional Attention to Improve Decoder-Only LLMs
###### Abstract
Decoder-only large language models typically rely solely on masked causal attention, which limits their expressiveness by restricting information flow to one direction. We propose Bitune, a method that enhances pretrained decoder-only LLMs by incorporating bidirectional attention into prompt processing. We evaluate Bitune in instruction-tuning and question-answering settings, showing significant improvements in performance on commonsense reasoning, arithmetic, and language understanding tasks. Furthermore, extensive ablation studies validate the role of each component of the method, and demonstrate that Bitune is compatible with various parameter-efficient finetuning techniques and full model finetuning.
## 1 Introduction
Large Language Models (LLMs) are being deployed in numerous practical applications where humans engage with them through various forms of natural language interaction. In use cases such as general purpose assistants ( OpenAI, 2024 ) , medical diagnosticians ( Thirunavukarasu et al., 2023 ) , game-conversation generation ( Cox and Ooi, 2023 ) or coding-assistants ( Roziere et al., 2023 ) , the ability for an LLM to precisely interpret and respond to user inputs is of primary concern.
Correspondingly, Instruction-Tuning (IT) ( Chung et al., 2024 ; Ouyang et al., 2022a ) is the prevailing paradigm for finetuning LLMs after their self-supervised pretraining phase to improve them for such tasks. Here, the model is trained on a dataset comprised of pairs of instructions and corresponding responses. Given the instruction-with-response structure of IT data, the generation of an LLM response can be divided into two phases: first, converting the instruction into key and value embeddings, which we refer to as instruction features; second, using these features to autoregressively generate an answer. Due to this task’s inherently conditional nature, the instruction features’ effectiveness is crucial for obtaining high-quality model outputs.
In the past, bidirectional attention ( Schuster and Paliwal, 1997 ) has been a key technique for obtaining stronger features for words or tokens. This is because the meaning of a word depends greatly on its context. In particular, for some words in a sentence, the information that comes later might be far more informative for generating a meaningful representation and resolving ambiguities. With only uni-directional causal attention, where the representation of each word is restricted to depend solely on the words that came before, this cannot be achieved. This is the reason why many previous transformers such as encoder-only BERT ( Devlin et al., 2019 ) and encoder-decoder T5 ( Raffel et al., 2020 ) employed bidirectional attention to improve the encoding of the input and why tasks like text retrieval ( Lewis et al., 2020 ; Li and Li, 2023 ) still rely on this.
However, in the context of LLMs, architectures utilizing bidirectional attention have fallen out of favor, as decoder-only models such as GPT ( OpenAI, 2024 ) and Llama ( AI@Meta, 2024 ) have focused on and vastly improved the generative performance of language models. These architectures are trained by large volumes of data with next-token prediction, eschewing any look-ahead mechanism for the sake of better autoregressive modeling. As there is simply more unlabeled data available for pretraining, training a decoder-only architecture on unlabeled data, and then finetuning it for tasks with instruction-tuning, is the best modus operandi of today ( Wang et al., 2022 ) . However, with this switch to decoder-only architectures, we lost bidirectional attention in the process. As we know this can improve feature representations for instructions, we set out to re-introduce bidirectional attention, such that it can be integrated into pretrained decoder-only LLMs.
Our new method Bitune adds bidirectional attention to decoder-only architectures and combines it with causal attention to generate two sets of instruction features, using two different sets of weights. These features are then integrated, utilizing learnable mixing coefficients, and later used as the KV-cache for response generation. Notably, the autoregressive response generation process remains unaffected by the bidirectional attention and continues to be causal. By realizing these adaptations with parameter-efficient finetuning methods, we introduce only a minimal set of new parameters.
Overall, our contributions are as follows: • We propose a novel method, Bitune, that improves the performance of pretrained decoder-only LLMs in instruction-following and question-answering settings.
• We evaluate the method on multiple downstream tasks, showing consistent improvements over the baselines.
• We conduct an extensive ablation study investigating the necessity of each component of the method, and showing the method’s PEFT-agnosticism, as well as its effectiveness in full finetuning scenarios.
## 2 Bidirectional Instruction-Tuning
In the instruction-tuning setting ( Ouyang et al., 2022b ; Zhang et al., 2024 ) , a dataset 𝒟 \mathcal{D} consists of instruction-answer pairs that are used to adapt the model in a supervised fashion. Formally, a dataset of size N N can be described as 𝒟 = { q , a } i = 1 N \mathcal{D}=\{q,a\}^{N}_{i=1} , where q q and a a are instructions and answers. The training objective is to model p ( a | q ) p(a|q) in an autoregressive manner: This means the answer is generated one token at a time, such that token a i a_{i} at position i i has access to all earlier tokens: p ( a | q ) = Π i = 1 | a | p ( a i | a 1 , … a i − 1 , q ) , p(a|q)=\Pi^{|a|}_{i=1}p(a_{i}|a_{1},\dots a_{i-1},q), where | a | |a| denotes the length of the answer. Note how compared to the regular language modeling objective, the response is already conditional (on the instruction q q ) even for the first generated token.
This naturally leads response-generation to be divided into two phases: prefilling and decoding. During the prefilling phase, the entire instruction – also often called a prompt – is processed concurrently to generate a series of features to be stored. For a Transformer architecture ( Vaswani et al., 2017 ) , these features are those of the key and value vectors, which can be stored in a KV-cache to avoid costly recomputations. During the subsequent decoding phase, the model generates output tokens sequentially, one token at a time, based on the KV-cache of the instruction and the already generated tokens.
In this work, we introduce Bitune, a method to leverage this two-phase process to improve instruction-tuning of language models. In our approach, the model processes the instruction with both causal and bidirectional attention using separate sets of parameters, leading to an enhanced KV-cache that is then used to condition the answer. Figure 1 provides an overview of the method, while Algorithm 1 presents pseudocode for the inference process.
#### Two Sets of Features.
In Bitune, the model performs two passes on the instruction to obtain two kinds of features for every transformer block. Namely, a set of c ausal features that the model was originally trained to process and utilize, K c = X c W k c , V c = X c W v c , K_{\color[rgb]{1,0,0}c}=X_{\color[rgb]{1,0,0}c}W_{k{\color[rgb]{1,0,0}c}},\quad V_{\color[rgb]{1,0,0}c}=X_{\color[rgb]{1,0,0}c}W_{v{\color[rgb]{1,0,0}c}}, (1) and a set of b idirectional features encoding the instruction without the constraints of causal masking, K b = X b W k b , V b = X b W v b . K_{\color[rgb]{0,0,1}b}=X_{\color[rgb]{0,0,1}b}W_{k{\color[rgb]{0,0,1}b}},\quad V_{\color[rgb]{0,0,1}b}=X_{\color[rgb]{0,0,1}b}W_{v{\color[rgb]{0,0,1}b}}. (2) To allow the model to learn how to process the causal and bidirectional features differently, we introduce two sets of weights: one for the bidirectional pass on the instruction ( W k b W_{k{\color[rgb]{0,0,1}b}} , W v b W_{v{\color[rgb]{0,0,1}b}} ) and another for the causal pass on the instruction, which is also used for the causal generation of answer tokens ( W k c W_{k{\color[rgb]{1,0,0}c}} , W v c W_{v{\color[rgb]{1,0,0}c}} ).
In the case of the first block of the model, representations X c X_{\color[rgb]{1,0,0}c} , X b X_{\color[rgb]{0,0,1}b} are the initial token embeddings. In other cases, they are the output of the preceding block and were processed by different components including the self-attention mechanism, which can be defined as: Attn ( Q , K , V , M ) = σ ( Q K T / d k + M ) V , \text{Attn}(Q,K,V,M)=\sigma({QK^{T}}/{\sqrt{d_{k}}}+M)V, (3) where σ \sigma is the softmax, Q Q are the queries, M M is the attention mask, and d k d_{k} is the dimension of keys and queries. For the causal pass, the mask M c M_{\color[rgb]{1,0,0}c} enforces causality by masking future tokens, such that tokens j j can only attend to earlier tokens i ≤ j i\leq j , while for the bidirectional pass, no masking is applied: M c ( i , j ) \displaystyle M_{{{\color[rgb]{1,0,0}c}}}(i,j) = { 0 if i ≤ j − ∞ if i > j \displaystyle=\begin{cases}0&\text{if }i\leq j\\ -\infty&\text{if }i>j\end{cases} (4) M b ( i , j ) \displaystyle M_{\color[rgb]{0,0,1}b}(i,j) = 0 \displaystyle=0 (5) The final KV-cache is obtained by a learnable convex combination of causal and bidirectional features, K Bitune \displaystyle K_{\texttt{Bitune}} = K c ⋅ ( 1 − α k ) + K b ⋅ α k \displaystyle=K_{\color[rgb]{1,0,0}c}\cdot(1-\alpha_{k})+K_{\color[rgb]{0,0,1}b}\cdot\alpha_{k} (6) V Bitune \displaystyle V_{\texttt{Bitune}} = V c ⋅ ( 1 − α v ) + V b ⋅ α v , \displaystyle=V_{\color[rgb]{1,0,0}c}\cdot(1-\alpha_{v})+V_{\color[rgb]{0,0,1}b}\cdot\alpha_{v}, (7) where α \alpha represents the bidirectional-to-causal ratio of features. This ratio is parameterised as α j = | θ j | / ( θ init + | θ j | ) , j ∈ { k , v } \alpha_{j}={|\theta_{j}|}/({\theta_{\text{init}}+|\theta_{j}|}),\quad j\in\{k,v\} (8) where θ j \theta_{j} is a learnable mixing coefficient per transformer block, and θ init \theta_{\text{init}} is a hyperparameter defining the initial value of θ j \theta_{j} . The mixing coefficients are learnable to allow each block to independently adjust the balance between bidirectional and causal features throughout the training.
#### Parameter Efficient Fine-tuning
Note that the components of the model, other than the key and value projections, can have their own separate sets of weights as well. In the case of full finetuning, this approach would require an additional set of full weights, which is impractical for large models.
Instead, we adapt our model using parameter-efficient finetuning methods. These introduce only a fraction of trainable parameters, making it viable to have two modified variants of the model within a single forward pass. In the default configuration of our method, we utilize the Low-Rank Adaptation (LoRA) of Hu et al. (2022) to adapt the model. However, Bitune can utilize different methods for updating the weights, including full model finetuning and other parameter-efficient techniques, as demonstrated in our ablations section.
## 3 Experiments
### 3.1 Instruction-Tuning
Our core experiments involve training pretrained language models on an instruction-tuning dataset and zero-shot evaluating them on downstream tasks. We evaluate Bitune on multiple models, comparing results to standard finetuning with LoRA, and zero-shot results of pretrained models without finetuning.
Specifically, we use a subset of the cleaned UltraFeedback ( Cui et al., 2023 ) dataset, which contains instructions and corresponding answers generated by various LLMs. From this dataset, we select completions generated by GPT-4 ( OpenAI, 2024 ) , ensuring high-quality responses for training. To fit every model on a single GPU, we filter out samples longer than 512 tokens, which leaves us with roughly 10,000 samples for training. For results on another instruction-tuning dataset, please see the Appendix A.6 .
We test the method on pretrained decoder-only language models of two different scales of approximately 2 billion and 7 billion parameters. The specific models used in our experiments are: Gemma 2B and 7B ( Gemma Team et al., 2024 ) , Llama2 7B ( Touvron et al., 2023 ) and Llama3 8B ( AI@Meta, 2024 ) , and Phi-2 ( Li et al., 2023 ) , which has 2.7 billion parameters. We use HuggingFace Transformers ( Wolf et al., 2020 ) implementation of these models.
For updating the weights we use the HuggingFace PEFT ( Mangrulkar et al., 2022 ) implementation of LoRA, with the default rank of 8, and apply it to all linear layers of MLP and self-attention components of the model. We compare Bitune with the following three baselines: Pretrained - initial model without any finetuning; LoRA - model finetuned with LoRA without Bitune-specific modifications, using rank of 8 as used in our method; and LoRA 16 {}_{\text{16}} - model finetuned with LoRA, using a rank of 16 to provide a fair comparison in terms of the number of parameters, as our method introduces two sets of weights.
For each model, we tune the learning rate on the LoRA baseline using steps on the approximate logarithmic scale ( 1 e − 4 1\mathrm{e}{-4} , 3 e − 4 3\mathrm{e}{-4} , 1 e − 3 1\mathrm{e}{-3} , 3 e − 3 3\mathrm{e}{-3} ), and then apply the same rate to the other approaches. Note that this potentially puts our method at a disadvantage compared to the LoRA baseline. All hyperparameters are reported in the Appendix A.2 .
Models are evaluated zero-shot on multiple-choice tasks to assess their performance. For commonsense reasoning, we use the PIQA ( Bisk et al., 2020 ) , CommonsenseQA ( Talmor et al., 2019 ) , ARC-Challenge ( Clark et al., 2018 ) , and SIQA ( Sap et al., 2019 ) datasets, while for language understanding, we use the MMLU ( Hendrycks et al., 2021 ) benchmark. Each task consists of a series of questions, each with multiple choices, where only one answer is correct. As the tasks follow the question-answer pattern, they are compatible with the instruction-tuning setting.
For evaluation, we use the Language Model Evaluation Harness framework ( Gao et al., 2023 ) . This framework formats each question using a predefined template, tokenizes the question-choice pairs, runs them through the model, and compares the log-likelihoods of the choices to determine the selected answer. For each model and approach configuration, we conduct experiments using three different random seeds, and average the results.
Models are loaded and trained using bfloat16 precision, except for the mixing part, which operates in the full 32-bit floating-point format. This high level of precision for the mixing of features is important, as minor numerical inaccuracies in the learnable coefficients and intermediate results of the mixing operation may lead to significant deviations in the model’s behavior.
In the decoding phase of the inference with Bitune, to initiate generation, the model requires at least a single token to obtain a set of attention queries , in addition to the keys and values extracted from the instruction. To facilitate this, one can introduce a new learnable <sep> token that would be placed at the beginning of modeled answer, or utilize an existing token. For our experiments, we opted to move the last token of the instruction template to the beginning of the modeled answer. For details on the instruction template used, please refer to the section A.10 of the Appendix.
#### Results.
Table 1 shows consistent and significant gains after instruction-tuning with Bitune, with the highest gains seen on the Gemma-2B model, showing a 4 percentage point (pp) improvement over the baseline LoRA and a 9.3 pp improvement over the pretrained model. For the other models, the average gains over baseline finetuning are equal to 1.8, 1.4, and 0.9 pp, for Llama3-8B, Llama2-7B, and Phi-2 respectively.