Bitune: Leveraging Bidirectional Attention to Improve Decoder-Only LLMs

Authors: Dawid Jan Kopiczko, Tijmen Blankevoort, Yuki M Asano
Primary area: foundation or frontier models, including LLMs
Keywords: Instruction tuning, Parameter-efficient fine-tuning, Transformer, PEFT, LLM

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.

Decision: reject
Reviewer scores: 6, 10, 5, 8

Reviews and discussion:
<thread>
creation date:
1729525767253

thread_type:
official_review

<thread_object>
creator: reviewer_3yh8

<content>
summary:
This paper presents “Bitune”, a method that enables pre-trained causal LMs (unidirectional attention), to use bidirectional attention. The proposed Bitune method makes use of sparse LoRA-based adaptors to create a bidirectional version of the model. Each prompt is passed through both the causal and bidirectional version of the model, and then both sets of features are combined into a single (key, value) cache. This cache is used to generate the rest of the sequence.

strengths:
1. The paper does demonstrate a performance gain from Bitune compared to the baseline pre-trained model and LoRA fine-tuned model on several language modeling benchmarks (table 1,2).
2. Paper does a good job of doing a thorough ablation study to verify the necessity of different components of their proposed method in section 3.3.
3. Paper does discuss the computational cost of adapting the base model to be bi-directional and provides a reasonable solution to mitigate the computational overhead: using parameter efficient fine-tuning. They also highlight that Bitune is not tied to the specific method for updating weights, and instead can be adapted to use different types of weight updating methods.

weaknesses:
Each round of Bitune inference effectively requires two inference passes which can be (time and memory) expensive (especially for really long sequence generation tasks). The authors do highlight this limitation in the discussion section and Appendix 6.4. According to their results, it seems that the most impacted part of the training/inference pipeline is the additional training time cost. Inference time costs did increase but to a much lesser extent. This timing increase could be reasonable if an end user of Bitune feels that the performance gains of the Bitune model warrants the extra computational time. I note that the time analysis is done for a fixed generation length of 200 tokens; I am curious to see how the time scales as generation length increases (e.g., 2000 tokens generated).

questions:
n/a
</content>
</thread_object>

</thread>
<thread>
creation date:
1730588128987

thread_type:
official_review

<thread_object>
creator: reviewer_b9de

<content>
summary:
The paper aims to improve decoder-only LLMs' effectiveness in downstream tasks by enhancing their representation of the input prompts. To achieve that, they introduce bi-directional attention into prompt processing. An LLM learns to represent a given prompt with both causal and bi-directional attention mechanisms and then fuse two features with the learnable combination. Experiments demonstrate the effectiveness of the proposed method in instruction-tuning and question-answering settings. Further ablation study verifies the performance gain comes from the bi-directional attention instead of more parameters or the fusion of two features.

strengths:
1. The proposed idea addresses the innate limitation of decoder-only LLMs while preserving their generation efficiency.
2. The experiments are comprehensive, covering various LLMs and downstream tasks.
3. The ablation study helps identify the contribution of different design choices.

weaknesses:
1. As indicated by the Related Work section, there are several prior works investigating the idea of bi-directional attention in decoder-only LMs. The novelty of this work is thus somewhat limited. It might be helpful if the authors could further discuss how their method stands out.
2. It is unclear why decoder-only LLMs would still face the limitation of representing language using causal attention given that they have undergone extensive pre-training and fine-tuning based on causal attention. Perhaps the authors can provide more evidence besides empirical results such as analysis of the hidden states in cases where the previous context indeed matters for the following one. 
3. The proposed method only applies to the given prompt but not the generated response, which could also benefit from bi-directional attention. Also, it is not convincing that this method would be widely used in practice.

questions:
How could this method be generalized to represent the generated tokens?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730695466966

thread_type:
official_review

<thread_object>
creator: reviewer_e3cz

<content>
summary:
This paper looks at how to add bidriectional attention to pre-trained decoder-only LLMs. This is done by learning new Key and Value matrices used during attention without the causal masking. Instead of just replacing the original matrices, they run the "encoder" part of the network (the parts with bidirectional attention) twice, once with the original weights and causal masking and once with the new weights. The resulting KV values are then mixed with a learnable weight. These mixed values are then attended to during autoregressive generation.

They find this produces strong results on many datasets with many different LLMs. They also include a lot of ablation experiments.

strengths:
The paper includes a lot of ablations that are very convincing of the efficacy of their approach. As I was reading, I keep thinking, I wonder if the see effect is because of X or Y instead of the bidirectional attention, and each time they had an ablation covering that possibility!

weaknesses:
In Figure 2 they show that the initial value of the mixing ratio can cause large differences in the final average ratio once convergence is reached. Figure 3 shows some histograms of mixing ratios at different layers, and Table 6 has experiments that explore how the initial value affects performance. However, the scale of these experiments is small, and the results are inconclusive. Given the initial value causes such a large difference in the ratio, it would have been nice to see a deeper dive into this.

The experiments in the paper focus on datasets that have clear single-turn input/output pairs. Given that many LLMs these days are used for chat applications, it would have been nice to see some experiments on how bitune training effects these chat applications. As the text gets longs (more turns into the chat) does the effect of the initial bidirectional attention decrease? Does it help the model follow the initial prompt better? Can their approach be applied to mutli-turn dialogue (all of some of the user text is bidirectional which model generated text is causal?)

questions:
Have you tried training the bitune attention matrices on a prefix language modeling task instead input/output formatted tasks? For example, randomly cutting a sequence of text in two pieces, applying bidirectional attention to the first part and then training on the second? This would also help answer the question of how much data is needed to learn the bitune matrices, which is mentioned but not answered in the paper.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730826252169

thread_type:
official_review

<thread_object>
creator: reviewer_xdhd

<content>
summary:
This paper aims to incorporate decoder-only language models with bidirectional attention. The authors introduce a second set of attention matrices for bidirectional understanding and integrate it with the original causal attention. Experimental results confirm the effectiveness of the proposed method, and comprehensive ablation studies offer insights into the design choices.

strengths:
- The proposed method is simple but effective.
- Comprehensive ablation results with reasonable baselines.

weaknesses:
Some existing works focus on prefix language models (non-causal decoder-only model) that utilize bidirectional attention for processing prefixes (i.e., inputs and instructions). Representative ones include U-PaLM [1] and UniLM [2]. How does the proposed method differ from these? While these models incorporate bidirectional attention during pre-training, the proposed method applies bidirectional attention specifically during instruction tuning. What would happen if we applied simple instruction tuning to these prefix language models (non-causal decoder-only model)? What are the advantages of bidirectional instruction tuning after causal decoder-only pre-training like the proposed method? I would like to see more discussions and even experiments on the comparison.

Another comment would be to visualize bidirectional attention on some examples. This could help highlight cases where bidirectional attention behaves differently from causal attention and therefore bring benefits.

Reference
1. Transcending Scaling Laws with 0.1% Extra Compute, 2022
2. Unified Language Model Pre-training for Natural Language Understanding and Generation, 2019

questions:
Would the proposed bidirectional instruction tuning affect the quality for pure generation tasks?
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces a simple strategy to increase the representation capacity of transformer-based sequence models in situations where generation is to be made conditionally on some input context. Authors argue that in such a situation, similar to the now legacy encoder-decoder architecture, input contexts can be attended to bi-directionally, likely improving the representations via giving attention layers access to additional context.

As pointed out by some of the reviewers, this idea is not novel and models such as Prefix-LM operate similarly. I would also add that recent work has moved back to encoding context with bi-directional models (e.g., CEPE, XC-Cache, You Only Cache Once), though in those cases focus lied mostly on compressing the KV-cache rather than improving the representations. Authors replied arguing that their method has extra improvements over the basic approach in Prefix-LM, mostly related to the fact that they mix representations obtained bi-directionally and with a causal mask. In more detail, the proposed approach fine-tunes models with separate sets of parameters (adapters are used for parameter efficiency) for either causal or bi-directional attention over prefixes. During training, both representations are computed for the prefix part of the input and a convex combination is computed, with learned mixing parameters for each layer. Only the causal layers are used for the remainder of the input.

The idea is simple and results reported in the paper suggest it could be useful. The paper is clearly written and easy to follow. The paper received high scores, although acceptance wasn't unanimous among reviewers. I went through the paper and the discussion, and I have important concerns that prevent me from recommending acceptance. Please see below for a detailed discussion of these concerns.

- Lack of insight and justification of the design: the proposed approach isn't really well justified in the paper. I would expect some justification in the form of some of the following suggestions: theory showing the gap in representation capacity due to causal masks (i.e., moving to bi-directional strictly increases capacity), empirical evidence of the gap in representation capacity due to causal masks (i.e., some kind of simple experiment showing the effect of moving from one setting to the other), or references to work showing one of the previous results.

- Limited evaluation: the evaluation focuses on a single setting, instruction tuning, where instructions are treaded as prefixes. Testing is then carried out in different multiple-choice question answering tasks. There are plenty of settings that were typically tackled in the encoder-decoder era that could have been evaluated here. To name a few examples, it's unclear to what extent the proposal shows benefits in settings such as summarization and translation.

- Inconclusive evidence (main concern): results are reported in a rather inconclusive manner since rankings of models are performed based simply on average accuracy across a number of tasks. Differences in performances are small enough across models that concerns around significance of those differences must be raised. In table 2 for instance, I cannot conclude which method is the best performing one. Same goes for ablation results. 1-2% differences in average accuracy do not really suggest one approach is consistently better than another. Some form of variance would need to be accounted for if all the evidence presented is empirical and especially since authors claim significant improvements.

- The overhead the proposal incurs is not accounted for appropriately: the proposal effectively incurs a 2x overhead in compute and peak memory footprint for prefix encoding. While prefix encoding can be done in parallel to counter overhead in time to first token, that would still require at least twice as much compute. Reviewers brought that up in their reviews, and authors replied with results I believe downplay this limitation. Authors show TTFT results for a setting with a *very short* prefix, but give no insight in terms of how this scales. Memory overhead is not discussed up to a brief note in the limitations section.