Bitune: Leveraging Bidirectional Attention to Improve Decoder-Only LLMs — part 2/3

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

It is worth noting that the Gemma-7B model shows the lowest average improvement across all tasks, with merely 0.1 pp gain over the baseline finetuning. It is also a single case where the baseline pretrained model achieved the highest score on a task, MMLU, with degraded performance in all fine-tuning approaches. However, this is not an issue with the model’s scale, as significant gains are observed with the Llama2-7B and Llama3-8B models.

### 3.2 Downstream Task Training

This complementary experiment verifies whether Bitune increases the capacity of the model within the narrow scope of a single task. It follows the setup from the instruction-tuning experiments with a few changes. Namely, models are not instruction-tuned but trained separately for each evaluation task using the corresponding training set. We use PIQA, ARC, CSQA, and SIQA introduced earlier, and an additional arithmetic task, GSM8K ( Cobbe et al., 2021 ) .

GSM8K differs from the other tasks, where we compare log-likelihoods of predefined answers, as it requires the model to generate a full answer token-by-token, including the intermediate step-by-step reasoning. The final answer follows a specific pattern, making it feasible to extract the answer using methods such as regular expressions as the model learns to adhere to this pattern during training.

#### Results.

Table 2 presents the results, demonstrating improvements when finetuning on the downstream tasks with Bitune, similar to those seen with instruction-tuning. While there are a few cases where the baseline finetuning achieves better results on specific tasks, when considering the average gains, applying our method is beneficial across all models. Most importantly, on the GSM8K dataset, we see consistent high gains, suggesting that our method improves the model’s reasoning ability in generative tasks. We present additional results on GSM8K with a 22B parameter model in the Appendix A.7 .

Similar to the instruction-tuning results, the highest gains are observed on the Gemma-2B model, while the lowest on the Gemma-7B. This indicates that the effectiveness of our method depends on the specific model used.

### 3.3 Chain-of-Thought Reasoning

To further test the effectiveness of our method on generative tasks that require explicit reasoning, we created four additional training sets for PIQA, ARC, CSQA, and GSM8K. Each set was distilled from the GPT-4.1 model ( OpenAI, 2025 ) by prompting it to solve every training-split question step-by-step; the resulting reasoning traces were then used to finetune Gemma-2B and Llama-3-8B on each benchmark separately. At evaluation time the models generated a full answer, including intermediate reasoning, and scores were averaged over three random seeds.

The original versions of these benchmarks supply only the final label, forcing a model trained on them to perform all reasoning in latent space and yielding short generations that do not reflect real-world usage. Our distilled chain-of-thought data instead mirrors the way large language models are typically invoked, producing longer, multi-step explanations before the answer, thereby providing a more faithful test bed. Moreover, because the public test sets may already appear in pre-training corpora, training the models to answer the same questions in a custom, step-by-step format reduces the risk that apparent gains stem from data contamination. The prompts used to distill the datasets, along with example completions, are available in Appendix A.9 .

#### Results.

Table 3 shows that Bitune improves performance on chain-of-thought reasoning tasks, raising average accuracy by 2.4 pp for Gemma-2B and 1.3 pp for Llama-3-8B.

### 3.4 Inference Speed Comparison

During inference, Bitune performs two forward passes over the instruction sequence and merges the resulting features, which affects prefilling runtime. We compare inference time for Gemma-2B on a long-context instruction of 2000 tokens and report the time required to prefill this instruction and to generate 2000 subsequent tokens on a single A100 GPU.

Table 4 shows that the extra latency for processing the instruction is negligible because most of the computation time is spent on autoregressive answer generation.

### 3.5 Ablations

We conduct an ablation study on Bitune using the same experimental setup as in the instruction-tuning experiment. For this purpose, two models are used: Gemma-2B and Llama3-8B, representing different size scales and model families.

#### Component Removal

To verify the necessity of each component of the method, we remove selected parts to answer the following questions:

• Can we simply modify the attention mask to apply bidirectional attention on the prompt, without using separate weights and mixing? - We test this simplest variant, which we refer to as Naive Bidir.

• Do we need two sets of features? Is it sufficient to obtain bidirectional features from the prompt using different weights than those used for causal answer generation? - We remove the part responsible for generating the set of causal features, and therefore also the mixing component; we refer to this as No Mixing .

• Are the gains solely from mixing two sets of features generated with different weights, or is bidirectional attention necessary? - Here we keep the attention mask causal to generate both sets of features, which we refer to as Only Causal .

• Do we need separate weights, or can the same weights be used to generate both causal and bidirectional features? - To answer this question, we do not introduce the second set of weights and use the same LoRA for both passes on the prompt, calling it Shared Weights .

The results, averaged over three seeds and presented in Table 5 , indicate that all variants of Bitune lead to gains over the baseline LoRA finetuning. However, the highest gains are observed in the full variant of Bitune, demonstrating that each component contributes to the method’s effectiveness.

#### Different PEFT Methods

To verify the impact of different PEFT methods on the performance of our method, we compare Bitune in combination with the following techniques: LoRA ( Hu et al., 2022 ) , that reparametrizes weight updates as a multiplication of two low-rank matrices; DoRA ( Liu et al., 2024 ) , which decomposes these weight updates into direction and magnitude; and IA3 ( Liu et al., 2022 ) , that instead rescales activations with learnable vectors.

The results are shown in Table 6 . We find consistent gains across all three PEFT methods we analyze, with gains ranging from + 1.6 % +1.6\% to + 4.0 % +4.0\% for averaged accuracy. This demonstrates that Bitune is PEFT-agnostic and can be combined with existing and future innovations in PEFT methods.

#### Full Finetuning

Additionally, we test whether Bitune leads to gains with full finetuning ( Full-FT ), by optimizing two sets of full model’s parameters. We conduct experiments on Gemma-2B model, and compare results with standard Full-FT baseline. The results in Table 7 demonstrate that Bitune improves the model’s performance even in full finetuning scenarios.

#### Attention Mask of Second Pass

We test another option for the attention mask of the second pass on the instruction. We transpose the causal attention mask, blocking information flow from the past tokens, and allowing from the future tokens - we call it anti-causal attention mask.

Results shown in Table 8 indicate that the instruction has to be processed with full bidirectional attention to achieve the highest gains. Combining causal and anti-causal features independently does not lead to the same high performance.

## 4 Related Work

Our approach shares similarities with the concept of "prefix language modeling", which enables a decoder-only model to handle bidirectional context within a prefix (instruction) while maintaining causal generation for the output sequence. The prefix-LM architecture was introduced by Liu et al. (2018) and further explored and popularized by Raffel et al. (2020) . In their work on T5, Raffel et al. (2020) pretrained the prefix-LM architecture alongside other architectures, such as encoder-decoder and decoder-only models, demonstrating that prefix-LM outperforms decoder-only models on both training objectives: denoising and language modeling.

The prefix-LM approach has been used in UniLM ( Dong et al., 2019 ) , which trains a single transformer on three types of language modeling tasks: unidirectional, bidirectional, and sequence-to-sequence prediction. UniLM employs a shared Transformer network and utilizes specific self-attention masks to control the context that predictions are conditioned on, where the sequence-to-sequence task is equivalent to the prefix-LM approach.

Additionally, UL2 ( Tay et al., 2023 ) introduces a pretraining objective called "Mixture of Denoisers", which combines various denoising strategies, including the prefix-LM approach. Lastly, XLNet ( Yang et al., 2019 ) also allows for non-causal word ordering by allowing random permutations to be used with a next-token prediction objective.

All these works focused on the model pretraining . As for the utilization of pretrained causal language models, Springer et al. (2024) show in their work that simply repeating the input to these models improves the quality of token embeddings for text-retrieval. This work addresses the limitation that token embeddings in autoregressive models cannot contain information from tokens appearing later in the input. By repeating the input twice, the early tokens are allowed to encode information about later tokens, thereby improving the quality of the embeddings. Another approach, LLM2Vec ( BehnamGhader et al., 2024 ) , demonstrates that pretrained causal LLMs can be effectively converted to BERT-like encoders. It can be done by enabling bidirectional attention, training the model on the objective of masked token prediction, and applying unsupervised contrastive learning.

#### Bitune vs. Prefix-LM.

Both Bitune and the classic prefix-LM masking scheme grant full bidirectional access to the instruction prefix. The key difference is that Bitune is a hybrid approach designed for the optimal adaptation of decoder-only models. Existing large language models are almost always pre-trained as pure decoder-only networks ; simply replacing the causal mask with the prefix-LM masking scheme before finetuning therefore discards the causal features encoded during pretraining and forces the model to relearn the mechanics of bidirectional attention over the prefix tokens. Bitune keeps the original causal stream intact and adds a second bidirectional pass, injecting fresh capacity into the model. Finetuning then proceeds in a single architecture that unifies the two attention regimes. Our ablation study (Section 3.5 ) confirms the payoff: Bitune yields larger gains than straightforward “prefix-LM conversion” (ablation Naive Bidir ). Whether the same advantage holds when both schemes are trained from scratch remains an open question.

## 5 Conclusion

This work proposes a method that exploits the inherent instruction–answer structure of IT datasets to incorporate bidirectional attention into pretrained decoder-only models. Bitune demonstrates general applicability across different models and scales, and it consistently improves performance across commonsense reasoning, arithmetic, and language understanding tasks. We further demonstrate that the method is compatible with different existing PEFT methods and will likely benefit from further developments in this area.

## Limitations

During standard instruction-tuning training, the instruction and the answer are processed in a single forward pass. In our method, this processing is explicitly split into phases - extracting instruction features with causal attention, bidirectional attention, and answer modeling, which results in increased both training time and memory usage.

However, this is a minor limitation in the context of instruction-tuning, since typically smaller datasets are used compared to pretraining, leading to relatively short training times. Furthermore, as shown in Section 3.4 , the additional latency during prefilling at inference time is negligible, since most time is spent on sequential token generation.

## Appendix A Appendix

### A.1 Pseudocode for Bitune Training Step

### A.2 Hyperparameters

### A.3 Datasets

For all experiments we used HuggingFace Datasets ( Lhoest et al., 2021 ) library to obtain necessary datasets.

### A.4 Training Speed & Memory Usage

As the method introduces two additional forward passes during training, both the training speed and the memory usage are impacted. Here we present average training times and GPU memory usage on the instruction-tuning setup with 3000 update steps (30000 actual steps, due to gradient accumulation), on a single A100 GPU, for models of two different scales - Gemma-2B & Llama3-8B. Our implementation has not been optimized, which means that e.g. training times could be improved via parallelization of two passes on the prompt. Table shows average training time, peak GPU memory usage during training, and average accuracy on downstream tasks.

Using these values, one can approximate required compute to reproduce results on a given tasks, as all experiments shared the same batch size and many other hyperparameters.

### A.5 Initialization of Mixing Coefficient

The initial value of the mixing coefficient θ \theta is a hyperparameter in our method. To evaluate its impact on the performance and the training dynamics of the bidirectional-to-causal ratio of features, we conduct experiments on the instruction-tuning setup with the following values: 0.1 0.1 , 0.01 0.01 , and 0.001 0.001 .

Table 15 demonstrates that the initial value of the mixing coefficient impacts the performance, with 0.01 0.01 being the most optimal value for both models, regardless of their scale. Figure 2 shows that the initial value substantially affects the rate of change of the mixing ratio, with the higher value leading to nearly no change in the ratio, while the lower value results in sharp changes at the very beginning of the training. In Figure 3 , we also observe that after training, all layers utilize the bidirectional attention.

### A.6 Instruction-tuning on Alpaca Dataset

We tested Bitune on another, larger instruction-tuning dataset - cleaned Alpaca dataset 1 1 1 https://huggingface.co/datasets/yahma/alpaca-cleaned ( Taori et al., 2023 ) with 50,000 samples. Similarly to previous experimental settings, first we tested different learning rates for the baseline LoRA finetuning, picked the best one, and then used the same learning rate for other approaches used in the experiment - Bitune, and Naive Bidir. (introduced in the ablation study section 3.5 ). The results demonstrate that our method’s benefits extend to larger datasets as well.

### A.7 Finetuning Larger Model on GSM8K

In order to verify whether improvements hold for larger, already highly capable models, we finetune Codestral ( Mistral AI, 2024 ) with 22B parameters on the GSM8K dataset. The results, averaged over 3 random seeds, show a substantial 4.3 percentage point improvement over the baseline LoRA finetuning, indicating that Bitune can provide considerable gains even for larger models, which already have strong performance on a given task.

### A.8 Multi-Turn Setting

In our experiments, we focus on single-turn QA settings, but the approach can be extended to multi-turn chat applications in the following two ways:

(A) “Recompute” The most straightforward way is to treat the entire chat history as a prefix or query and recompute the KV cache at every new round. This allows bidirectional attention over all tokens, including previously generated answers.