Exploiting Contextual Knowledge in LLMs through V-usable Information based Layer Enhancement

TL;DR

The paper tackles context-faithful generation in LLMs by shifting focus from decoding strategies to internal state processing. It introduces CaLE, a layer-aware intervention guided by $I_V(h_l -> Y)$, to identify and enhance context-rich layers via amplification or residual connections, with theoretical support for reducing $H_V(Y|h_f)$. CaLE employs both supervised and unsupervised (KL-based) layer identification, enabling practical layer selection without labeled data. Empirical results on CounterFact, NQ, SQuAD, and StrategyQA across multiple model families and decoding methods demonstrate robust improvements in context-faithful generation, particularly under unknown or conflicting contextual knowledge. The work offers a versatile, architecture-agnostic approach that complements decoding strategies and highlights the importance of internal state dynamics in leveraging contextual information.

Abstract

Large Language Models (LLMs) have demonstrated remarkable capabilities in various tasks, yet they often struggle with context-faithfulness generations that properly reflect contextual knowledge. While existing approaches focus on enhancing the decoding strategies, they ignore the fundamental mechanism of how contextual information is processed within LLMs' internal states. As a result, LLMs remain limited in their ability to fully leverage contextual knowledge. In this paper, we propose Context-aware Layer Enhancement (CaLE), a novel intervention method that enhances the utilization of contextual knowledge within LLMs' internal representations. By employing V-usable information analysis, CaLE strategically amplifies the growth of contextual information at an optimal layer, thereby enriching representations in the final layer. Our experiments demonstrate that CaLE effectively improves context-faithful generation in Question-Answering tasks, particularly in scenarios involving unknown or conflicting contextual knowledge.

Figures (6)

• Figure 1: An illustration of CaLE Method.
• Figure 2: Visualization of Information Flow. The vertical axis represents the variation in $\mathcal{V}$-information, as reflected by the $-H_\mathcal{V}$ metric. The horizontal axis denotes the information content across different layers, while the shaded region indicates one standard deviation from the mean.
• Figure 3: Variation of the KL divergences across layers in different models. The $\mathrm{KL}_q$ quantifies the impact of question conditioning on layer representations by measuring their distributional divergence, while $\mathrm{KL}_c$ captures the incremental influence of context conditioning given the question on these representations. The shaded region represents the confidence interval.
• Figure 4: Validation set size impact on supervised layer selection and comparative layer selection with unsupervised CaLE. The selected layers are detailed in Table \ref{['tab:layer']}.
• Figure 5: Visualization of Analysis on the CounterFact Dataset for Llama models.

Theorems & Definitions (2)

• Proposition 3.1
• proof