뉴런 활성화 그래프를 통한 타겟 지향적 사전 학습 데이터 선택

LLM의 내부 뉴런 활성화 패턴을 분석하여 특정 타겟 작업에 가장 적합한 학습 데이터를 선별하는 새로운 프레임워크를 제안한다. 기존의 블랙박스 방식과 달리 모델의 '기능적 중추'를 직접 활용함으로써 학습 효율을 극대화하고 데이터 선택 과정의 투명성을 확보한다.

기존의 데이터 선택 방식은 텍스트를 고차원 벡터로 변환하는 Embedding에 의존한다. 하지만 Embedding은 텍스트의 표면적인 의미는 잘 포착할지라도, 실제 모델 내부에서 해당 데이터를 처리할 때 어떤 연산 경로를 거치는지(즉, 어떤 능력을 요구하는지)는 명확히 보여주지 못하는 블랙박스라는 한계가 있다.

이 논문은 모델 내부의 '뉴런 활성화'라는 기초 개념에 집중한다. 특정 수학 문제를 풀 때 모델 내부의 특정 뉴런들이 강하게 반응한다면, 그 뉴런들을 활성화시키는 다른 데이터들 역시 수학적 추론 능력을 기르는 데 도움이 될 것이라는 원리이다. 이를 위해 각 레이어에서 영향력이 큰 뉴런들을 선별하여 Neuron-Activated Graph(NAG)를 구성한다.

결과적으로 NAG는 모델이 특정 지식을 습득하기 위해 사용하는 실제 '연산 경로'를 데이터 선택의 기준으로 삼는다. 이는 단순한 텍스트 유사도를 넘어 모델의 내부 동작 구조와 데이터의 특성을 직접 연결함으로써, 훨씬 더 정밀하고 해석 가능한 데이터 필터링을 가능하게 한다.

전체 프로세스는 뉴런 영향력 측정, NAG 구축, 데이터 랭킹의 세 단계로 구성된다.

먼저 뉴런 영향력(Neuron Impact)을 측정하기 위해 특정 뉴런을 비활성화했을 때 레이어 출력의 변화량을 계산한다. 입력값 h_in과 가중치 행렬 W의 k번째 열 W:,k가 주어질 때, ||h_in^T W:,k||_2를 계산하여 해당 뉴런이 출력값의 크기에 기여하는 정도를 수치화한다. 이 값이 클수록 해당 입력 처리에 핵심적인 뉴런임을 의미한다.

다음으로 각 레이어별로 영향력이 가장 높은 상위 K개의 뉴런 인덱스를 추출하여 Neuron-Activated Graph(NAG)를 생성한다. 타겟 데이터셋 D_target에 대해 각 데이터의 NAG를 집계하여 '타겟 프로필'을 만든다. 후보 데이터 c가 들어오면 해당 데이터의 NAG와 타겟 프로필 간의 교집합 크기를 합집합 크기로 나누는 Dice-style 유사도를 계산하여 점수를 부여한다.

최종적으로 이 유사도 점수를 기준으로 후보 데이터를 내림차순 정렬하고, 설정된 비율 r_f만큼 상위 데이터를 추출하여 사전 학습에 사용한다. 이 과정은 별도의 모델 학습이 필요 없으며 기존 모델의 추론 신호만으로 수행된다.

HellaSwag 벤치마크에서 기존 SOTA 방식인 BETR 대비 5.3% 높은 정확도를 기록했으며, 무작위 샘플링 대비 평균 4.9%의 성능 향상을 달성했다. 특히 일반적인 품질 필터인 FineWeb-Edu가 놓치는 타겟 특화 신호를 효과적으로 포착함을 확인했다.

다중 타겟(Multi-target) 설정에서도 NAG는 강력한 성능을 보였다. 6개의 벤치마크를 동시에 타겟팅했을 때, 기존 방식들이 성능 저하를 겪는 것과 달리 NAG는 무작위 샘플링 대비 3.1%, FineWeb-Edu 대비 0.6% 더 높은 성능을 유지했다.

분석 실험에서는 NAG가 선택한 0.12%의 뉴런만 제거해도 모델 성능이 60.6%에서 37.1%로 23.5%p 폭락하는 결과를 얻었다. 반면 무작위로 동일한 수의 뉴런을 제거했을 때는 성능 변화가 거의 없었으며, 이는 NAG가 모델의 실제 '기능적 중추'를 정확히 찾아내고 있음을 시사한다.

NAG는 트랜스포머 아키텍처의 FFN UP 프로젝션 레이어에서 추출된 뉴런 신호가 가장 효과적임을 발견했다. 이는 확장 레이어(up_proj)가 고차원 잠재 공간에서 작동하여 작업 특화 신호를 더 잘 분리하기 때문으로 분석된다. 반면 잔차 연결에 가까운 레이어들은 정보가 압축되어 있어 변별력이 낮았다.

또한 단일 레이어가 아닌 전체 레이어의 신호를 통합하는 것이 필수적이다. 마지막 레이어의 정보만 사용할 경우 무작위 샘플링보다 성능이 낮아지는 경우가 발생했는데, 이는 작업 관련 신호가 모델 전체 층에 걸쳐 '계산 궤적' 형태로 분포되어 있음을 의미한다.

구현 측면에서 NAG 추출은 H100 GPU 기준 150B 토큰당 192 GPU-hour가 소요되며, 이는 모델 전체 학습 비용에 비해 매우 저렴한 수준이다. 한 번 추출된 뉴런 특징은 타겟 작업이 바뀌어도 재사용이 가능하여 확장성이 뛰어나다.

본 연구는 1.2B 규모의 모델과 30B 토큰 학습 환경을 주된 실험 대상으로 삼았으며, 더 거대한 모델 규모(7B 이상)와 수조 단위 토큰에서의 일반화 가능성에 대해서는 추가적인 검증이 필요하다. 또한 다중 타겟 설정에서 단순한 혼합 방식을 사용했으므로 더 정교한 데이터 믹싱 전략과의 결합 연구가 요구된다.