WorldTUI 360M 모델, 터미널 시뮬레이션 학습

TL;DR

WorldTUI는 360M 파라미터 Qwen2.5-Coder 모델을 터미널 에뮬레이터로 동작하도록 미세 조정했습니다. 각 키 입력이 모델에 공급되고, 모델은 화면 그리드 업데이트를 예측합니다. 결과 데모는 완전히 브라우저에서 실행되며(208MB int4 모델), 파일 생성, vim, tmux 실행, 스네이크 게임까지 가능하고, 명령 간 간단한 상태를 기억합니다.

모델이 터미널을 표현하는 방법

결론: 터미널은 문자 그리드로 인코딩되며, 모델은 선형 토큰 스트림 대신 행 단위 델타 연산을 예측합니다.

  • 입력은 화면 크기, 내용, 커서 위치를 설명하는 스냅샷 토큰 블록으로 시작합니다(예: BOS SCREEN_START 32 96 …).
  • 각 키 입력은 원시 바이트('l', \r, ESC [ D 등)로 전송되며 INPUT_START/INPUT_END 토큰으로 감싸집니다.
  • 모델은 OUTPUT_START/OUTPUT_END 안에 ROW_0 17 'l', CURSOR 0 18 1 같은 연산 집합을 출력하여 덮어쓸 셀과 최종 커서 상태를 나타냅니다.
  • 특수 WAIT 토큰은 경과 시간을 인코딩하여 스네이크 게임 같은 자율 업데이트를 가능하게 합니다.
  • 어휘는 333개 토큰으로 축소되었습니다: 256개 바이트 값 토큰, 64개 행 식별자, 그리고 몇 가지 제어 토큰(예: COPY, WAIT). 이 컴팩트한 어휘 덕분에 모델이 브라우저 친화적인 크기로 들어갈 수 있습니다.

행 복사 연산으로 효율적인 화면 업데이트

결론: COPY 연산을 도입하여 스크롤 및 줄 조작의 토큰 비용을 행당 수천 개에서 몇 개로 줄였습니다.

  • 터미널이 스크롤될 때, 모델은 새 화면의 모든 문자를 출력하는 대신 ROW_i COPY j 명령을 내보내 이전 행을 새 위치로 복사합니다.
  • 이로 인해 일반적인 3,000토큰 전체 화면 다시 그리기가 영향을 받는 행당 약 3토큰으로 줄어듭니다.
  • 동일한 메커니즘으로 화면 지우기, vim에서 줄 삽입/삭제, man 페이지 업데이트를 처리합니다.

학습 데이터 생성

결론: 저자는 무작위 환경의 실제 Docker 셸을 기록하여 336k 합성 터미널 세션(약 2B 토큰)을 생성했으며, 키 입력과 화면 간 완벽한 정렬을 보장했습니다.

  1. 무작위 사용자 이름, 호스트 이름, 프롬프트, 터미널 크기.
  2. 개발자, 시스템 관리자 등 페르소나를 반영한 무작위 파일 시스템 레이아웃.
  3. 가중치가 적용된 활동 선택(vim 편집, python REPL, tmux 등).
  4. 명령을 현실적으로 유지하기 위한 상태 추적(예: 기존 파일 열기).
  5. 1.2% 오타율의 인간형 타이핑.
  6. 세션의 35%는 애플리케이션 중간에 시작하여 모델이 임의 상태에서 계속하도록 학습.

최종 데이터셋은 32×96, 24×80, 40×120 화면 크기를 포함합니다.

미세 조정 및 자기 수정

결론: DAgger 스타일 자기 수정 예제를 추가하면 메모리 일관성이 크게 향상되어 PresetBench 점수가 24/47에서 38/47(int4)로 올라갔습니다.

  • 초기 지도 미세 조정 후, 모델을 기록된 세션에서 실행하여 자체 실수를 수집했습니다.
  • 각 실수에 대해 모델의 화면을 실제 정답으로 되돌릴 수정 델타를 구성하여 학습 예제로 사용했습니다.
  • 수정은 실제 화면에 이미 나타난 토큰만 사용할 수 있어 환각을 방지합니다.
  • 샘플링된 세션의 약 1/3에서 유용한 수정이 얻어졌으며, 이 과정을 세 번 반복했습니다.

훈련(수정 라운드 포함)은 RTX 4090에서 약 75 GPU 시간이 소요되었습니다.

벤치마크

결론: WorldTUI는 명령 출력 정확도와 간단한 파일 메모리 작업 모두에서 프롬프트 기반 기본 모델을 능가합니다.

모델 ShellBench (106개 명령) 파일 읽기 (13개)
Qwen2.5-Coder-0.5B (프롬프트) 52 2
Qwen3-0.6B (프롬프트) 60 3
WorldTUI 80 9

가장 큰 향상은 메모리 테스트에서 나타나며, WorldTUI는 13개 중 9개 읽기 시나리오를 정확히 재현하는 반면 기본 모델은 2~3개에 그칩니다.

속도 최적화

결론: 병렬 행 디코딩과 추측 검증으로 RTX 4090에서 대형 프레임 지연 시간을 6.7초에서 0.87초로 줄였습니다.

  • 병렬 행: 각 행은 배치에서 독립적으로 디코딩되며 KV 캐시를 공유하므로 디코딩 비용은 모든 행의 합(약 3,000단계)이 아닌 가장 긴 행(약 100단계)과 같습니다.
  • 추측 디코딩: 병렬 패스는 모든 행에 대한 초안을 생성하고, 동일한 모델로 두 번째 검증 패스가 초안을 확인하고 첫 번째 불일치를 수정하여 일관성을 유지합니다(예: 올바른 vim 창 배치).
  • Mac M4에서 단일 문자 지연 시간은 약 60ms이고, 전체 man 페이지는 순차적으로 2.5초, 병렬 초안으로 0.9초가 걸립니다.

브라우저용 모델 크기 축소

결론: 임베딩 테이블을 333토큰 어휘로 줄이고 int4 양자화를 적용하여 모델을 208MB로 축소하여 WebGPU 추론을 가능하게 했습니다.

  • 원래 Qwen2.5-Coder-0.5B: 494M 파라미터, 990MB.
  • 사용되지 않는 토큰 임베딩 제거 후: 약 360M 파라미터, 716MB.
  • Int4 양자화: 208MB 다운로드 크기.
  • 추론은 사용자 정의 WGSL 셰이더를 통해 완전히 클라이언트 측에서 실행되며, 활성화는 32비트 부동 소수점으로 유지됩니다.

증류 실험

결론: 레이어를 제거하면 다음 토큰 정확도가 높게 유지되어도 메모리 일관성이 손상되어, 메모리가 특정 더 깊은 표현에 의존함을 나타냅니다.

  • 레이어 7만 제거하면 토큰 정확도에 거의 영향을 미치지 않았지만(99.3% → 98.9%), 정확한 메모리 읽기 수가 64/64에서 0/64로 감소했습니다.
  • 레이어 1, 7, 15, 22를 유지하는 16레이어 학생 모델은 메모리에서 교사와 일치했지만(64/64), ShellBench 성능이 떨어졌습니다(정확한 명령 65 vs 78).
  • 이러한 트레이드오프로 인해 추가 증류 시도는 중단되었습니다.

시연된 기능

결론: WorldTUI는 대화형 전체 화면 애플리케이션(vim, tmux, snake)을 실행하고 세션 간 간단한 상태를 유지할 수 있지만, 복잡한 파일 시스템 조작에서는 여전히 환각을 일으킵니다.

  • 메모리: 별도의 tmux 창에서 생성, 편집, 읽은 파일이 대부분의 경우 올바르게 유지됩니다.
  • Vim: 모드 전환(Esc), 삽입, 저장, 종료를 지원합니다.
  • Tmux: 창 생성, 탐색, 환경 변수 확장을 처리합니다.
  • Python REPL: 간단한 표현식을 받아들이고, 사소한 오타를 허용합니다(pythoon이 Python을 실행).
  • Snake: 계산 집약적임에도 불구하고 원활하게 실행됩니다.

한계

결론: 모델에는 실제 파일 시스템이나 외부 도구가 없어 환각과 가끔 경로 추적 오류가 발생합니다.

  • Docker 컨테이너 내부에는 인터넷 접근이 없어 lynx 같은 도구는 작동하지 않습니다.
  • 복잡한 파일 작업(예: 기존 파일에 추가)은 종종 잘못된 출력을 생성합니다.
  • 중첩 디렉터리 탐색 시 현재 경로를 잃을 수 있습니다.
  • Firefox WebGPU 성능은 Chrome보다 약 20배 느립니다. 원인은 알 수 없습니다.

커뮤니티 반응 (선택된 댓글)

"WorldTUI는 지능적인 출력을 예측하는 모델과 지능 자체의 차이라는 질문을 탐구하기에 아주 좋은 방법입니다." – OtherShrezzing

"매우 멋집니다. 재미있는 프로젝트입니다. 좋은 기존 모델을 미세 조정하는 방법을 이해하면 탐구할 넓은 분야가 있다고 느껴집니다." – oilkillsbirds

"시도해 봤는데 데모 GIF만큼 잘 작동하지 않습니다. 폴더가 없는데 이미 존재한다고 말하고, ls 후에는 폴더가 사라져 있습니다." – donperignon

"다음 단계는 인터넷인 척하는 모델을 훈련하는 것입니다. (이미지 없이 훈련/테스트하기 위해 그 터미널에서 lynx를 실행할 수 있습니다)" – amelius

"파일 시스템용 도구 호출 API, 내장 busybox, 그리고 '여기서 busybox가 할 일을 출력'이라는 단일 출력 토큰이 있다면 작은 모델이 얼마나 많은 일을 할 수 있을지 궁금합니다." – Lerc

시사점

  • WorldTUI는 완벽하게 정렬된 키 입력-화면 데이터로 훈련할 때 적당한 360M 파라미터 모델이 일관된 대화형 터미널 세계를 학습할 수 있음을 증명합니다.
  • DAgger 스타일 예제를 통한 자기 수정은 지속적인 메모리 향상에 핵심입니다.
  • 병렬 행 디코딩과 추측 검증은 소비자 GPU 및 브라우저에서 거의 실시간 지연 시간을 달성하는 데 필수적입니다.
  • 향후 작업은 더 풍부한 상태 표현(실제 파일 시스템 API, 도구 호출, 확산 기반 화면 업데이트)을 탐구하여 환각을 줄이고 적용 범위를 확장할 수 있습니다.

사용해 보기: 데모는 https://simedw.com/projects/worldtui 에서 사용할 수 있습니다(208MB 다운로드, Chrome/Edge 권장).

Sources

관련