기술 요약: LoRIF – 확장 가능한 훈련 데이터 귀속을 위한 저랭크 영향 함수

문제 제기

훈련 데이터 귀속 (Training Data Attribution, TDA) 은 주어진 쿼리에 대한 모델의 예측에 가장 큰 영향을 미친 특정 훈련 예시들을 식별하는 것을 목표로 합니다. 영향 함수 (Influence Functions, IF) 는 그라디언트와 역 헤시안 행렬을 통해 leave-one-out 재훈련을 근사함으로써 TDA 에 이론적으로 근거한 접근 방식을 제공하지만, 수백만 개의 예시로 훈련된 수십억 개의 파라미터를 가진 현대 규모의 모델에서는 사실상 접근이 불가능한 상태였습니다.

기존의 확장 가능한 변형들인 TRAK, LoGRA, TrackStar 등은 투영된 예시별 그라디언트를 저장하고 감쇠된 가우스-뉴턴 헤시안 근사를 사용하여 계산 병목 현상을 해결합니다. 그러나 이러한 방법들은 귀속 품질과 계산 실현 가능성 사이의 트레이드오프를 초래하는 두 가지 중요한 확장성 병목 현상에 직면해 있습니다:

1. 그라디언트 저장 및 I/O: $N$개의 훈련 예시에 대한 투영된 예시별 그라디언트를 저장하는 데는 $O(ND)$의 공간이 필요합니다. 여기서$D$는 유효 투영 차원입니다. 쿼리 시점에 이러한 그라디언트를 로드하는 작업이 지연 시간의 대부분 (쿼리 시간의 약 96%) 을 차지합니다.
2. 역 헤시안 확장: $D \times D$ 역 헤시안 근사를 구성하고 저장하는 데는 $O(D^2)$의 메모리가 필요합니다.

귀속 품질을 향상시키기 위해서는 투영 차원 $D$를 증가시켜야 하지만, 그렇게 하면 저장/I/O 및 메모리 병목 현상이 모두 악화되어 최첨단 규모의 모델에 대한 고품질 귀속을 비현실적으로 만듭니다.

방법론: LoRIF

LoRIF (Low-Rank Influence Functions) 는 신경망 그라디언트 내에 내재된 저랭크 구조를 활용하여 이러한 병목 현상을 해결합니다. 이 방법은 다음 세 단계의 파이프라인을 따릅니다:

1. 저랭크 그라디언트 저장 (I/O 병목 현상 해결)

LoRIF 는 전체 투영된 예시별 그라디언트 행렬 $\tilde{G}_i \in \mathbb{R}^{d_1 \times d_2}$ (여기서 $D = d_1 d_2$) 을 저장하는 대신, 랭크-$c$ 분해법을 저장합니다:
$$\tilde{G}_i \approx u_i v_i^\top$$
여기서 $u_i \in \mathbb{R}^{d_1 \times c}$이고 $v_i \in \mathbb{R}^{d_2 \times c}$입니다.

• 메커니즘: 이 분해는 블록 파워 반복 (block power iterations) 을 통해 계산됩니다.
• 영향: 샘플당 계층당 저장 및 쿼리 시점 I/O 가 $O(D)$에서 $O(c(\sqrt{D}))$로 감소합니다. 실험 결과에 따르면 $c=1$일지라도 근사가 의미 있는 귀속 품질을 유지하며, 높은 $c$보다 더 큰 $D$를 우선시하는 것이 비용 효율성이 더 뛰어납니다.

2. 잘린 SVD 곡률 근사 (메모리 병목 현상 해결)

조밀한 $D \times D$ 역 헤시안 $(G^\top G + \lambda I)^{-1}$을 구성하는 것을 피하기 위해, LoRIF 는 집계 그라디언트 행렬 $G$의 랭크-$r$ 잘린 특이값 분해 (SVD) 를 계산합니다:
$$G \approx U_r \Sigma_r V_r^\top$$
**우드버리 행렬 항등식 (Woodbury matrix identity)**을 사용하여 역 헤시안 항을 $r$차원 부분 공간에서 근사합니다:
$$(V_r \Sigma_r^2 V_r^\top + \lambda I)^{-1} = \frac{1}{\lambda}I - \frac{1}{\lambda^2} V_r \left( \Sigma_r^{-2} + \frac{1}{\lambda}I_r \right)^{-1} V_r^\top$$

• 메커니즘: 상위-$r$개의 특이 벡터 ($V_r$) 와 특이값 ($\Sigma_r$) 만 저장됩니다.
• 영향: 곡률 항에 대한 메모리 복잡도가 $O(D^2)$에서 $O(Dr)$로 감소합니다. 이를 통해 고정된 메모리 예산 내에서 훨씬 더 큰 유효 투영 차원$D$를 사용할 수 있게 됩니다.

3. 영향 점수 계산

쿼리 시점에 영향 점수는 쿼리 그라디언트를 $r$차원 부분 공간으로 투영하고 훈련 그라디언트에 대해 저장된 저랭크 인자를 활용하여 계산됩니다. 영향 점수 $I(x_{tr}, x_{te})$는 다음과 같이 근사됩니다:
$$I(x_{tr}, x_{te}) \approx \frac{1}{\lambda} g_{te}^\top g_{tr} - \frac{1}{\lambda^2} (g'_{te})^\top \left( \Sigma_r^{-2} + \frac{1}{\lambda}I_r \right)^{-1} g'_{tr}$$
여기서 $g'$는 특이 벡터 $V_r$로 투영된 그라디언트를 나타냅니다. 첫 번째 항은 랭크-$c$ 인자를 사용하여 효율적으로 계산되며, 두 번째 항은 중간 행렬의 대각선 특성으로 인해 $O(r)$ 연산만 포함합니다.

주요 기여

1. 병목 현상 식별: 이 논문은 투영 차원 $D$가 확장 가능한 IF 스타일 TDA 의 주요 제약 사항이며, 정확한 귀속 (큰 $D$) 이 현재 비현실적인 품질 - 확장성 트레이드오프를 생성한다고 규명했습니다.
2. 랭크-$c$ 분해: 투영된 그라디언트를 랭크-$c$ 인자로 저장하여 예시당 계층당 저장을 $O(D)$에서 $O(c\sqrt{D})$로 줄이는 방안을 제안했습니다.
3. 랭크-$r$ 곡률 근사: 잘린 SVD 와 우드버리 항등식을 사용하여 곡률 항을 근사하고, 계층당 메모리를 $O(D^2)$에서 $O(Dr)$로 줄이는 방안을 제안했습니다.
4. 확장성 입증: 수백만 개의 훈련 예시를 가진 최대 700 억 파라미터 (Apertus-70B) 모델에 대한 효과적인 데이터 귀속을 최초로 입증하여 실용적인 비용과 지연 시간을 유지했습니다.

실험 결과

LoRIF 는 124M 에서 70B 파라미터에 이르는 모델 (GPT2-small, OLMo-3-7B, Apertus-70B) 과 최대 380 만 개의 예시를 포함하는 데이터셋에서 평가되었습니다.

• 저장 및 지연 시간: LoGRA 와 비교하여 LoRIF 는 귀속 품질을 유지하거나 초과하면서 그라디언트 저장을 2.3–20 배 감소시키고 쿼리 지연 시간을 1.3–20 배 감소시켰습니다.
  • GPT2-small 에서 LoRIF 는 중간 저장 regimes 에서 저장을 최대 9 배, 지연 시간을 5.6 배 줄였습니다.
  • OLMo-3-7B 에서 LoRIF 는 저장을 20.3 배, 지연 시간을 22.0 배 줄이면서 LoGRA 와 거의 동일한 tail-patch 점수를 달성했습니다.
  • Apertus-70B 에서 LoRIF 는 더 낮은 저장 및 지연 시간으로 LoGRA 대비 tail-patch 점수를 1.5–2.2 배 향상시켰습니다.
• 귀속 품질: Linear Datamodeling Score (LDS) 와 Tail-patch 점수를 통해 측정된 결과, LoRIF 는 일관되게 파레토 프론티어를 개선했습니다. 메모리 제약으로 인해 이전에는 비현실적이었던 더 큰 유효 투영 차원 $D$를 가능하게 함으로써, 더 작은 $D$로 제한된 베이스라인보다 더 높은 품질을 달성했습니다.
• 정성적 평가: 상위 1 개 검색 예시에 대한 LLM-as-a-judge 평가 결과, LoRIF 는 LoGRA 에 비해 더 높은 의미적 관련성 (OLMo-3-7B 에서 평균 점수 3.04 대 2.49) 을 가진 예시를 검색하고 완전히 관련 없는 검색이 더 적었습니다.

중요성과 주장

이 논문은 LoRIF 가 최첨단 규모에서 그라디언트 기반 훈련 데이터 귀속을 실용적으로 만든다고 주장합니다. 그라디언트의 저랭크 구조를 활용함으로써 LoRIF 는 대규모 언어 모델에 영향 함수를 적용하는 것을 방해해 온 품질 - 확장성 트레이드오프를 해결합니다.

저자들은 LoRIF 가 새로운 저랭크 원시 연산에 의존하는 것이 아니라, 예시별 및 집계 그라디언트에서 모두 관찰되는 기존 저랭크 구조가 투영된 그라디언트 영향 방법에서 체계적으로 어떻게 활용될 수 있는지를 보여준다고 강조합니다. 이러한 결과는 귀속 인덱스가 많은 쿼리에 걸쳐 재사용될 때, 이전 방법들의 prohibitive 한 저장 또는 I/O 비용 없이 대규모 LLM 감사, 디버깅 및 안전 점검에 그라디언트 기반 TDA 를 사용할 수 있음을 시사합니다.

이 논문은 한계를 인정하며, 귀속 인덱스를 구성하는 것에는 여전히 코퍼스 전체에 대한 그라디언트 패스 (전처리 비용) 가 필요하며, 현재 대규모 실험은 전체 사전 훈련 코퍼스보다는 지도형 미세 조정 데이터에 초점을 맞추고 있음을 지적합니다. 그러나 이 연구는 확장 가능한 귀속을 위한 새로운 기준을 설정하여, 이전에는 이러한 방법의 적용이 불가능하다고 간주되었던 모델들의 분석을 가능하게 합니다.