ByteDance Seed, Cola2 공개, 346B 토큰 학습 8개 과제 평균 60.10

ByteDance Seed가 문장을 의미 단위로 압축해 생성하는 언어 모델 Cola2의 체크포인트와 추론 코드를 Hugging Face와 GitHub에 공개했다. Cola2는 명령을 따르도록 추가 학습한 챗 비서가 아니라 베이스 모델이며, 연구자와 개발자가 내려받아 직접 돌려 볼 수 있다.

저장소에는 8개로 나뉜 safetensors 가중치 파일, 토크나이저와 생성 설정 파일, Transformers용 원격 코드 모듈이 들어 있다. 평가 과제 8종의 정규화된 JSONL 데이터와 파일 검증용 SHA256SUMS도 함께 올라왔다. 추론 구현, 명령줄 도구, 예제, 벤치마크 실행기는 GitHub 저장소에서 관리하며, 학습 코드는 이번 공개에 포함되지 않았다. 코드는 Apache License 2.0이고, 가중치·토크나이저 자산·벤치마크 데이터는 각각의 약관을 따른다.

작동 방식

Cola2는 텍스트를 의미 단위인 스팬으로 동적으로 묶고, 각 스팬을 압축된 연속 잠재 벡터 하나로 표현한다. 동적 압축 VAE, 블록 단위 인과 구조의 확산 트랜스포머(DiT), 잠재 피드백을 받는 스팬 단위 토큰 디코더로 이루어진다. DiT가 잠재 블록을 생성하면 디코더가 이를 이산 토큰으로 바꾼다.

학습은 두 단계다. 1단계에서는 동적 압축 인코더와 토큰 디코더를 함께 학습한다. 2단계에서는 토큰 쪽 감독으로 VAE를 계속 조정하면서, 약 40억 매개변수의 DiT가 분리된 잠재 분포를 플로 매칭으로 학습한다. 추론 때는 잠재 블록 제안, 토큰 디코딩, 확정된 텍스트의 재인코딩을 반복해 연속 기록과 이산 기록을 맞춘다.

성능과 한계

공개된 체크포인트는 소스 토큰 346.03B개로 학습했다. ByteDance Seed는 논문 기준 8개 과제 평균 60.10을 기록했다고 밝혔다. 잠재 생성에 무작위성이 있어 실행마다 결과가 다를 수 있고, Cola2 과제별 값은 잠재 노이즈를 달리한 평가 3회의 평균이다.

ByteDance Seed는 Cola2가 HellaSwag와 StoryCloze에서 Llama-3-8B를 넘었고, 보고된 사전 학습 토큰은 43배 넘게 적다고 밝혔다. 반면 평균은 Llama-3-8B의 66.25와 Qwen2.5-7B의 75.08보다 낮고, RACE와 MMLU 등 나머지 과제에서도 Llama-3-8B에 못 미친다. 회사는 모델마다 학습 방식, 토크나이저, 토큰 계산법이 달라 이 표가 생성 방식의 우열을 가르기보다 역량의 맥락을 보여 주는 용도라고 설명했다. 잠재 DiT가 아직 수렴하지 않은 자원 제약 체크포인트라는 점도 밝혔다.

같은 모델 크기, 데이터, 평가 조건으로 AR, LLaDA, Cola DLM, 압축 없는 Cola-MDM, 압축한 Cola2-MDM을 비교한 통제 실험도 있다. 종합 확장성은 Cola-MDM이 가장 강했고, Cola2-MDM은 수렴이 느리지만 약 150B 토큰 이후 LLaDA와의 격차를 계속 좁혔으며 약 210B 토큰에서도 포화하지 않았다. 최적의 디코더 CFG 가중치는 학습 대부분에서 -0.9 근처였다가 점차 0으로 향했다.

사용 방법

Python 3.10 이상, PyTorch 2.2 이상, Transformers 4.51 이상 5 미만, safetensors 0.4 이상이 필요하다. CUDA와 bfloat16 환경을 권장한다. 커스텀 코드를 쓰므로 모델을 불러올 때 trust_remote_code=True를 지정해야 한다.

기본 생성 설정은 논문 프리셋이다. 8단계 오일러 DiT 스케줄, DiT용 classifier-free guidance 10, 디코더 -0.9, 탐욕적 토큰 선택, max_new_tokens 32를 쓴다. 탐욕적 선택이어도 잠재 생성이 가우시안 노이즈에서 시작하므로 결과가 고정되지 않으며, 출력을 비교하려면 시드를 준 torch.Generator를 넘겨야 한다.

GitHub 저장소를 받으면 채팅과 벤치마크 스크립트를 쓸 수 있고, 200개 예제 스모크 벤치마크도 돌릴 수 있다. 논문의 시드는 공개되지 않아 참고 결과와 작은 차이가 생길 수 있다. 내려받은 파일은 hf download 후 sha256sum -c SHA256SUMS로 검증한다.