Anthropic은 신뢰할 수 있고, 해석 가능하며, 제어 가능한 AI 시스템을 만들기 위해 연구하는 AI 안전 및 연구 기업입니다.
이 글과 위의 라운드테이블 영상에서, Anthropic 연구진은 과학적 진전과 엔지니어링 발전이 얼마나 밀접하게 맞물려 있는지를 돌아보며, 해석가능성 연구를 훨씬 더 큰 AI 모델로 확장하는 과정에서 부딪힌 기술적 난제들을 이야기합니다.
지난해 10월, Anthropic 해석가능성 팀은 소규모 트랜스포머 모델에 사전 학습(dictionary learning) 기법을 적용한 논문 Towards Monosemanticity를 발표했습니다. 올해 5월에는 동일한 기법을 수 자릿수(orders of magnitude) 더 큰 모델에 적용한 Scaling Monosemanticity를 공개했습니다. 이 연구에서 우리는 Claude 3 Sonnet 내부에서 수천만 개의 "피처(feature)"—의미 개념에 대응하는 뉴런 조합—를 발견했으며, 이는 AI 모델의 내부 작동 원리를 이해하는 데 있어 중요한 진전이었습니다.
이러한 연구를 계속 전진시키려면, 더 많은 엔지니어가 필요합니다.
초기 논문(예: Frameworks, Toy Models of Superposition)만 읽어본 분이라면 의아할 수도 있습니다. 당시에는 엔지니어링 비중이 크지 않았으니까요. 하지만 최신 연구를 살펴보면, 우리가 마주한 엔지니어링 과제의 규모가 얼마나 큰지 바로 체감하실 수 있을 것입니다.
아래에서는 최신 연구에서 실제로 해결해야 했던 기술적 엔지니어링 문제 두 가지를 소개합니다. 이 사례들을 통해 현재 우리 엔지니어들이 어떤 종류의 문제에 매달리고 있는지 보여드리고, 엔지니어링이 AI 해석가능성, 나아가 AI 안전성 연구의 가장 큰 병목 중 하나가 될 것이라고 판단하는 이유를 설명하고자 합니다.
엔지니어라면 바로 이 글이 여러분을 위한 것입니다. 아래에서 소개하는 엔지니어링 문제들에 공감하고 흥미를 느끼신다면, Research Engineer 포지션에 꼭 지원해 주세요.
우리가 "피처"를 분석하는 데 사용하는 희소 오토인코더(Sparse Autoencoder)는 트랜스포머의 활성화 값으로 학습됩니다. 이때 활성화 값이 순서에 의존하는 허위 패턴을 학습하지 않도록 반드시 셔플해야 합니다. 처음에는 학습 데이터가 GPU 하나에 모두 올라갔기 때문에 셔플은 아무 문제가 되지 않았습니다. 하지만 메모리에 담을 수 없을 만큼 데이터가 커지면서 상황이 달라졌습니다. 카드 한 벌을 섞는 건 쉽지만, 창고 가득한 카드를 섞어야 한다고 상상해 보세요—완전히 다른 차원의 문제입니다.
이 시점에서 페타바이트 규모까지 확장 가능한 분산 셔플을 구현할 수도 있었습니다. 하지만 우리는 빠르게 구현할 수 있되, 확장성은 다소 떨어지는 방식을 택했습니다. 셔플을 K개의 작업으로 나누고, 각 작업이 셔플된 출력 데이터의 1/K을 담당하도록 했습니다. 순열(permutation)을 생성한 뒤, 각 작업이 전체 학습 데이터를 스트리밍 방식으로 읽어 자기 몫의 출력을 기록하는 구조였습니다. 덕분에 규모를 더 키울 수 있었지만, 단점은 명확했습니다. 모든 작업이 전체 학습 데이터를 읽어야 한다는 것이었죠. 처음에는 몇 시간이 걸렸고, 나중에는 며칠이 걸렸습니다. Towards Monosemanticity 작업 당시에는 학습 데이터가 100TB(각 1KB인 데이터 포인트 1,000억 개)에 달했고, 셔플은 심각한 골칫거리가 되어 있었습니다.
확장 가능한 분산 셔플 자체는 새롭거나 최첨단의 문제가 아닙니다. 하지만 과학적 진전을 위해 빠르게 해결해야 했던 수많은 엔지니어링 문제 중 하나였습니다.
이 경우, 유용한 블로그 포스트를 발견하고 거기서 소개된 접근법을 다중 패스 방식으로 확장했습니다. 한 패스에서 N개의 작업을 실행합니다. 각 작업은 데이터셋의 1/N을 읽어 셔플한 뒤, 각각 1/NK 크기인 K개의 파일로 기록합니다. 각 작업에서 쓰인 첫 번째 파일들을 모으면 최종 셔플 결과의 처음 1/K에 해당하지만, 아직 그 안에서의 셔플은 남아 있습니다. 두 번째 파일도 마찬가지이고, 이후도 동일합니다. 한 패스를 거치면, 전체 데이터 셔플 한 번이 N개의 셔플로 분해되며, 각 셔플의 크기는 K배 줄어듭니다. 이 시점에서 셔플이 단일 머신의 메모리에 들어간다면 바로 처리하면 됩니다. 들어가지 않는다면, 패스를 한 번 더 돌리면 됩니다.
각 작업이 메모리에 100GB를 유지할 수 있고, 1GB짜리 파일 100개를 기록한다고 가정해 봅시다. 패스 한 번마다 필요한 셔플 크기가 100배씩 줄어듭니다. 1패스로 100GB, 2패스로 10TB, 3패스로 1PB, 4패스로 100PB까지 셔플할 수 있는 셈입니다.
이 방식을 도입한 이후, 셔플은 더 이상 고민거리가 아닙니다. 빠르게 끝나고, 문제도 발생하지 않습니다. 물론 더 나은 접근법이나 더 빠른 구현이 존재하겠지만, 이 방법으로 병목이 해소됐고, 덕분에 다음 문제에 집중할 수 있게 되었습니다.
또 다른 엔지니어링 과제는 피처 시각화에 필요한 기반 데이터를 생성하는 작업이었습니다. 피처 시각화를 통해 사용자는 개별 피처에서 가장 강하게 활성화되는 특정 토큰과 기타 정보를 확인할 수 있습니다(Towards Monosemanticity 논문의 Feature Browser를 이 링크에서 확인하세요).
각 피처에 대해 다양한 활성화 수준에서 이를 활성화하는 데이터셋 예시를 찾아, 전체 분포를 탐색해야 합니다. 수백만 개의 피처에 대해 이를 효율적으로 수행하는 것은 그 자체로 흥미로운 분산 시스템 문제입니다. 처음에는 단일 작업으로 처리했지만, 곧 그 한계를 넘어섰습니다. 아래는 현재 접근 방식의 개략적인 구조입니다.
시각화용 데이터셋은 1억 토큰이며, 수백만 개의 피처를 처리해야 합니다. 먼저 데이터셋과 피처를 여러 부분으로 "샤딩(shard)"합니다. 각 작업은 자기 몫의 데이터셋을 순회하면서, 담당하는 피처들에 대해 가장 높은 활성화 값을 보이는 상위 K개의 토큰과, 해당 피처를 활성화하는 토큰 중 무작위로 선택한 10*K개를 추적합니다(트랜스포머 활성화 값은 이미 S3에 캐싱해 두었으므로, 다시 계산할 필요가 없습니다).
다음으로, 피처 단위로 샤딩하여 이전 패스의 결과를 집계합니다. 이를 통해 전체 데이터셋에 걸쳐 각 피처를 가장 강하게 활성화하는 토큰과 무작위 활성화 토큰 집합을 얻게 됩니다. 이것이 피처 시각화에서 보여줄 예시들입니다.
각 예시에 대해, 해당 피처가 주변 토큰에서 어떻게 반응하는지도 계산해야 합니다. 처음에는 피처 단위로 샤딩하는 방식을 사용했습니다. 각 작업이 담당 피처의 예시에 해당하는 트랜스포머 활성화 값을 로드하는 구조였죠. 문제는 이 예시들이 데이터셋 전체에 무작위로 흩어져 있어서, 작업에 필요한 데이터만 골라 읽을 방법이 마땅치 않다는 점이었습니다.
이를 개선하기 위해, 데이터셋 단위로 샤딩하는 패스를 추가했습니다. 이 구조에서는 각 작업이 트랜스포머 활성화 값의 일부를 담당하고, 각 피처 그룹에 필요한 활성화 값을 별도 파일로 저장합니다. 그런 다음 피처 단위 패스를 실행하면, 필요한 데이터만 간편하게 접근할 수 있습니다. 피처가 주변 토큰에서 얼마나 반응하는지 계산한 뒤, 관련 데이터를 프론트엔드 웹사이트에서 읽고 표시할 수 있는 형식으로 기록합니다.
희소 오토인코더 연구를 시작했을 때, 이 접근법이 성공할지 확신할 수 없었습니다. 실험을 거듭할수록 연구에 대한 확신이 커졌고, 그에 따라 더 큰 실험을 돌릴 수 있도록 인프라에도 더 많이 투자하게 되었습니다. 이 과정이 가장 최근 논문인 Scaling Monosemanticity까지 이어졌습니다.
이러한 과정은 해석가능성 팀에서 이루어지는 엔지니어링 업무의 성격을 잘 보여줍니다. 우리에게 연구와 엔지니어링은 분리할 수 없는 하나입니다. 팀원들은 연구와 엔지니어링 사이를 수시로 오가며, 현재 시스템에서 가능한 한 더 큰 규모를 쥐어짜내 새 실험을 돌린 뒤 다시 연구로 돌아갑니다. 많은 연구 아이디어가 실패로 끝나기 때문에, 성과가 확인되기 전까지는 인프라에 과도하게 투자하지 않습니다.
연구는 팀의 노력이며, 아이디어를 구상하는 것만큼 구현하는 것도 중요합니다. 가설만 세우는 것이 아니라, 검증하고, 만들고, 반복하고, 확장합니다.
이런 이유로, 우리는 다양한 영역을 유연하게 넘나들 수 있는 제너럴리스트 엔지니어를 특히 찾고 있습니다. 파이프라인 구축이든, ML 실험 수행이든, GPU 활용 최적화든 가리지 않고 뛰어들 수 있는 분이면 좋겠습니다.
이런 역량을 갖추고 AI 안전성에 열정을 가진 엔지니어라면, 꼭 지원해 주세요. Research Engineer 포지션 상세 내용과, 해석가능성 팀의 다른 채용 공고는 채용 페이지에서 확인하실 수 있습니다.