그라디언트 원자: 훈련 그래디언트의 희소 분해로 모델 행동을 자동 발견·귀속·조정

본 논문은 대형 언어 모델(LLM)의 파인튜닝 과정에서 학습된 행동을 이해하고 제어하기 위한 새로운 프레임워크인 “Gradient Atoms”를 제안한다. 기존 훈련 데이터 귀속(TDA) 연구는 특정 행동에 대해 사전 정의된 쿼리를 기반으로 각 훈련 문서가 해당 행동에 기여하는 정도를 점수화한다. 이러한 접근법은 (1) 행동이 단일 문서가 아닌 다수의 문서에 걸쳐 공유되는 절차적 지식이라는 사실을 간과하고, (2) 사용자가 미리 행동을 정의해야 하는 비용과 쿼리당 O(N) 연산 복잡도 때문에 실용성이 떨어진다는 한계가 있다. Gradient Atoms는 이러한 문제를 근본적으로 뒤집는다. 먼저, 파인튜닝에 사용된 N개의 훈련 샘플 각각에 대해 전체 파라미터에 대한 교차 엔트로피 손실의 그래디언트 g_i를 계산한다. 이 그래디언트는 모델이 해당 샘플을 학습할 때 이동해야 할 방향을 그대로 담고 있다. 그러나 원시 그래디언트 공간은 고곡률 방향이 과도하게 강조되는 비등방성(anisotropic) 특성을 가지고 있어, 직접적인 클러스터링이나 분해를 수행하면 의미적 구조가 가려진다. 이를 해결하기 위해 논문은 EKF‑AC(Approximate Fisher) 고유분해를 이용해 각 모듈별로 상위 k개의 고유벡터 Q^(k)와 고유값 λ^(k)를 구하고, g_i를 Q^(k)ᵀ에 투사한 뒤 λ^(k)⁻¹/² 로 사전조건화(precondition)한다. 이 과정은 그래디언트의 스케일을 균일하게 만들어, 동일한 연산 흐름을 공유하는 문서들의 그래디언트가 서로 가깝게 배치되도록 한다. 사전조건화된 그래디언트 ˆg_i는 정규화 후 희소 사전학습(dictionary learning) 알고리즘에 입력된다. 논문에서는 scikit‑learn의 MiniBatchDictionaryLearning을 사용해 K=500개의 원자(d₁…d₅₀₀)를 학습했으며, 희소성 페널티 α=0.1을 적용해 각 문서가 소수의 원자에만 할당되도록 했다. 결과적으로 각 원자는 “공유 업데이트 방향”이라는 의미를 갖게 되며, 이를 통해 동일한 작업 유형(예: 산술, 문법 교정, 예/아니오 분류 등)을 수행하는 문서들이 같은 원자에 묶인다. 원자의 품질 평가는 두 단계로 진행된다. 첫째, 원자를 활성화하는 문서 집합 S_j에 대해 원시 그래디언트 간 코사인 유사도의 평균을 구해 “coherence” 점수를 산출한다. 높은 점수는 해당 원자가 실제로 동일한 연산 흐름을 반영한다는 강력한 증거이다. 실험에서는 500개 원자 중 5개가 coherence > 0.5, 43개가 > 0.1을 기록했다. 둘째, 원자를 가중치 공간에 역투사(unproject)해 전체 파라미터 벡터 v_j를 얻고, 이를 스칼라 α와 곱해 모델 파라미터를 θ' = θ ± α·v_j 형태로 변형한다. 이렇게 만든 “steering vector”는 특정 행동을 증폭하거나 억제하는 효과를 가진다. 실험에서는 5개의 대표 원자를 선택해 알파값을 0.5, 1, 2, 5, 10으로 변동시키며 100개의 평가 프롬프트에 대해 정규표현식 기반 검출을 수행했다. 예/아니오 분류 원자(#415)는 기본 39%에서 최대 51%까지 상승하거나 0%까지 억제했으며, 코드 생성 원자(#64)는 42%→58%(+16pp) 혹은 28%(-14pp)로 변동했다. 특히 불릿 리스트 원자(#469)는 33%→94%(+61pp)까지 크게 증가시키고, 반대 방향에서는 0%까지 억제하는 양방향 선형 효과를 보였다. 거부(Refusal) 원자(#161) 역시 50%→0%로 완전 억제 가능했으며, 이는 모델이 입력 부족 상황에서 “입력을 제공해 주세요”라는 기본적인 거부 메커니즘을 별도 경로로 학습했음을 시사한다. 논문은 이러한 결과를 바탕으로 몇 가지 중요한 통찰을 제시한다. 첫째, 행동 억제는 증폭보다 일반적으로 쉬우며, 이는 하나의 공유 경로를 차단하면 해당 행동이 사라지는 반면, 증폭은 여러 경쟁 경로와의 균형을 맞춰야 하기 때문이다. 둘째, coherence와 steerability는 반드시 일치하지 않는다. 낮은 coherence를 가진 원자라도 큰 행동 변화를 일으킬 수 있음을 통해, 원자의 방향이 모델 내부에서 차지하는 비중과 기본 행동의 포화 정도가 중요한 변수임을 알 수 있다. 셋째, 원자는 작업 유형뿐 아니라 포맷(불릿 vs 번호 매기기)까지 구분한다는 점에서, LLM이 세부적인 출력 형식까지 별도 연산 흐름을 학습한다는 새로운 가설을 제공한다. 제한점으로는 전체 원자 중 대부분이 낮은 coherence를 보여 잡음 혹은 과도하게 넓은 혼합을 나타낸다. EKF‑AC 투사 차원을 6,800으로 축소하면서 일부 미세 정보가 손실될 가능성이 있다. 또한 5,000개의 인스트럭션-응답 쌍이라는 비교적 작은 데이터셋으로는 희귀 행동을 포착하기 어려우며, 정규표현식 기반의 평가가 표면적인 포맷 변화를 넘어선 의미론적 변화를 충분히 측정하지 못한다는 점이다. 결론적으로 Gradient Atoms는 훈련 단계에서 모델이 습득한 절차적 지식을 “공유 업데이트 방향”이라는 새로운 단위로 정량화하고, 이를 직접적인 모델 편집 수단으로 전환함으로써 기존 TDA와 모델 편집 연구를 효과적으로 연결한다. 향후 연구에서는 다중 원자를 조합해 복합 행동을 동시에 제어하거나, 더 큰 규모와 다양한 도메인의 데이터셋에 적용해 원자의 다양성과 품질을 향상시키는 방향이 기대된다.