MuDAF: Long-Context Multi-Document Attention Focusing through Contrastive Learning on Attention Heads

TL;DR

MuDAF enhances long-context question answering in LLMs by explicitly optimizing attention heads with contrastive learning, reducing distractions and improving focus on relevant information.

Contribution

The paper introduces MuDAF, a novel contrastive learning method that explicitly optimizes attention heads to improve focus and performance in long-context multi-document tasks.

Findings

Significant improvement in long-context question answering performance.

Enhanced attention focus on relevant information demonstrated by visualization.

MuDAF reduces attention distractions effectively.

Abstract

Large Language Models (LLMs) frequently show distracted attention due to irrelevant information in the input, which severely impairs their long-context capabilities. Inspired by recent studies on the effectiveness of retrieval heads in long-context factutality, we aim at addressing this distraction issue through improving such retrieval heads directly. We propose Multi-Document Attention Focusing (MuDAF), a novel method that explicitly optimizes the attention distribution at the head level through contrastive learning. According to the experimental results, MuDAF can significantly improve the long-context question answering performance of LLMs, especially in multi-document question answering. Extensive evaluations on retrieval scores and attention visualizations show that MuDAF possesses great potential in making attention heads more focused on relevant information and reducing attention distractions.