소형 트랜스포머로 ARC-AGI-1 44% 달성
원제: I trained a small transformer in 1.5hrs and it beats many LLMs
왜 중요한가
극소 비용과 단시간 훈련으로 다수 LLM을 능가하는 성능을 달성해 AI 샘플 효율성 연구의 새로운 가능성을 제시한다.
인도 IIT Bombay 출신 개발자 Mithil Vakde가 RTX 5090에서 1.5시간, 비용 0.67달러로 소형 트랜스포머를 처음부터 훈련시켜 ARC-AGI-1 공개 평가에서 44%를 기록했다. 다수의 대형 언어 모델을 상회하는 성능이며, 코드는 오픈소스로 공개됐다. ARC-AGI-2에서도 7%를 달성했다.
Mithil Vakde氏는 자신의 블로그에서 소형 트랜스포머를 이용해 ARC-AGI-1 공개 평가에서 44%를 달성했다고 발표했다. 훈련 시간은 약 1.5시간, 비용은 0.67달러에 불과하며, 테스트 시간 훈련(test-time training)을 수행하는 유사 모델들과 비교해 동등하거나 높은 성능을 기록했다.
주요 기술 구성은 다음과 같다. 각 입출력 쌍을 토큰 시퀀스로 변환하고, 소형 트랜스포머가 자기회귀 방식으로 학습한다. 퍼즐별 추가 임베딩으로 태스크 간 학습을 지원하며, 3D RoPE 임베딩을 사용한다. 색상 및 이면각 치환 방식으로 데이터 증강을 적용한다.
이전 모델(40% 달성)과 비교한 주요 개선 사항으로는 SwiGlu 활성화 함수, RMSNorm 적용 등 현대적 아키텍처 도입, 레이어 수를 4개에서 8개로 확장, 입력 토큰 제외 및 출력 토큰만을 손실 함수에 포함하는 지도 학습 방식 전환이 있다. 비용 절감에는 증강 횟수 대폭 감소, AdamW에서 Muon 옵티마이저로 전환, flash attention 및 flex attention 커널 활용이 기여했다.
또한 ARC-2의 ARC-1 비중복 태스크를 추가 훈련 데이터로 활용했다. 데이터 누출 방지를 위해 ARC-2 내 ARC-1 중복 773개 퍼즐을 신중하게 필터링했다. 추가 데이터 없이도 약 40%를 달성할 수 있으나 연산량이 약 2배 필요하다. 본 연구는 Lucas Beyer, Jeremy Howard, Rohan Anil 등 저명 연구자들로부터 주목받았다.