Table of Contents
Fetching ...

Scalable Multi-Stage Influence Function for Large Language Models via Eigenvalue-Corrected Kronecker-Factored Parameterization

Yuntai Bao, Xuhong Zhang, Tianyu Du, Xinkui Zhao, Jiang Zong, Hao Peng, Jianwei Yin

TL;DR

This paper tackles training data attribution for large language models under full-parameter fine-tuning by introducing a scalable multi-stage influence function that traces downstream behavior to pre-training data. It achieves scalability through Eigenvalue-corrected Kronecker-Factored Approximation Curvature (EK-FAC) and a semantic-similarity–based candidate-filtering strategy, enabling efficient iHVP computations. The authors show that MLP parameters contribute more to influence estimates than MHA in large models and validate the approach on a fact-tracing benchmark and a real-world model (dolly-v2-3b), with publicly available code. The work demonstrates practical pathways to calibrate trust in generative AI by grounding model outputs in pre-training data, while outlining limitations and directions for broader applicability to other architectures and components.

Abstract

Pre-trained large language models (LLMs) are commonly fine-tuned to adapt to downstream tasks. Since the majority of knowledge is acquired during pre-training, attributing the predictions of fine-tuned LLMs to their pre-training data may provide valuable insights. Influence functions have been proposed as a means to explain model predictions based on training data. However, existing approaches fail to compute ``multi-stage'' influence and lack scalability to billion-scale LLMs. In this paper, we propose the multi-stage influence function to attribute the downstream predictions of fine-tuned LLMs to pre-training data under the full-parameter fine-tuning paradigm. To enhance the efficiency and practicality of our multi-stage influence function, we leverage Eigenvalue-corrected Kronecker-Factored (EK-FAC) parameterization for efficient approximation. Empirical results validate the superior scalability of EK-FAC approximation and the effectiveness of our multi-stage influence function. Additionally, case studies on a real-world LLM, dolly-v2-3b, demonstrate its interpretive power, with exemplars illustrating insights provided by multi-stage influence estimates. Our code is public at https://github.com/colored-dye/multi_stage_influence_function.

Scalable Multi-Stage Influence Function for Large Language Models via Eigenvalue-Corrected Kronecker-Factored Parameterization

TL;DR

This paper tackles training data attribution for large language models under full-parameter fine-tuning by introducing a scalable multi-stage influence function that traces downstream behavior to pre-training data. It achieves scalability through Eigenvalue-corrected Kronecker-Factored Approximation Curvature (EK-FAC) and a semantic-similarity–based candidate-filtering strategy, enabling efficient iHVP computations. The authors show that MLP parameters contribute more to influence estimates than MHA in large models and validate the approach on a fact-tracing benchmark and a real-world model (dolly-v2-3b), with publicly available code. The work demonstrates practical pathways to calibrate trust in generative AI by grounding model outputs in pre-training data, while outlining limitations and directions for broader applicability to other architectures and components.

Abstract

Pre-trained large language models (LLMs) are commonly fine-tuned to adapt to downstream tasks. Since the majority of knowledge is acquired during pre-training, attributing the predictions of fine-tuned LLMs to their pre-training data may provide valuable insights. Influence functions have been proposed as a means to explain model predictions based on training data. However, existing approaches fail to compute ``multi-stage'' influence and lack scalability to billion-scale LLMs. In this paper, we propose the multi-stage influence function to attribute the downstream predictions of fine-tuned LLMs to pre-training data under the full-parameter fine-tuning paradigm. To enhance the efficiency and practicality of our multi-stage influence function, we leverage Eigenvalue-corrected Kronecker-Factored (EK-FAC) parameterization for efficient approximation. Empirical results validate the superior scalability of EK-FAC approximation and the effectiveness of our multi-stage influence function. Additionally, case studies on a real-world LLM, dolly-v2-3b, demonstrate its interpretive power, with exemplars illustrating insights provided by multi-stage influence estimates. Our code is public at https://github.com/colored-dye/multi_stage_influence_function.
Paper Structure (55 sections, 50 equations, 4 figures, 6 tables)

This paper contains 55 sections, 50 equations, 4 figures, 6 tables.

Figures (4)

  • Figure 1: Spearman correlation coefficient of influence scores versus wall-clock time. Hollow markers and solid markers share the same correlation values. Hollow markers only accounts for the time on obtaining pair-wise influence estimates, while solid ones additionally account for the overhead.
  • Figure 2: Distribution of $||\boldsymbol{\theta}^{\text{ft}}-\boldsymbol{\theta}^{\text{pt}}||_2$ and $||\boldsymbol{\theta}^{\text{ft}}-\boldsymbol{\theta}^{\text{pt}}||_2 / ||\boldsymbol{\theta}^{\text{pt}}||_2.$
  • Figure 3: Fact-tracing results (in percentage). Round markers and solid lines denote MRR results, while triangular markers and dotted lines denote Recall@10 results.
  • Figure 4: Pearson correlation of influence scores versus wall-clock time.

Theorems & Definitions (2)

  • proof
  • proof