Local AI World Model Part 2 - Deep NN을 사용하여 이미지를 플레이 가능한 캐릭터로 변환하는 방법 (프롬프트 전환

지난번에 이 subReddit에 제 작업을 공유했을 때, 작은 데모 하나만으로도 너무 많은 분들이 관심을 가져주셔서 거의 눈물을 흘릴 뻔했습니다. 지난 2개월 동안 새로운 모델을 훈련해 왔는데, 이번에는 실제 텍스트 가이드를 사용했습니다. 그래서 이전 모델에 대해 조금 말씀드리겠습니다. 일반적인 비디오 모델은 너무 크고 실시간으로 소비자 하드웨어에서 실행되도록 설계되지 않았습니다. 정적인 클립은 생성할 수 있지만, 심지어 빠르더라도 실시간 비디오는 아직 로컬 환경에서 완전히 해결된 것은 아닙니다. 많은 월드 모델 데모들이 나왔지만, 이들은 거대한 데이터센터 GPU에서 실행되거나, 아니면 WAN이나 LTX 같은 인기 있는 모델을 자가회귀적(Autoregressive) 방식으로 작동하도록 증류한 것들입니다 (참고로 이것 역시 로컬 환경에서는 실시간이 아닙니다). 위에 있는 비디오는 RTX 5090에서 단지 30%의 활용률로 작동하는 것입니다. 이 1B 모델의 최대 FPS는 rtx5090에서 50-60이지만, 저는 강제로 소프트웨어 스로틀링을 걸어 12fps로 설정했습니다. 그리고 일부 테스트에 따르면 이는 이 모델이 다른 RTX 카드인 40, 30 시리즈에서도 작동할 수 있음을 의미합니다 (곧 시도해 볼 예정입니다). 저는 MacBook을 가지고 있고 아직 MLX로 포팅하지 않았지만, 벤치마크를 해보니 M5 MacBook에서 30fps로 실행됩니다. 아키텍처에 대해 말씀드리자면, 이 모델은 순수 트랜스포머(transformer)이며 블록 인과 마스크(block causal mask)와 함께 작동합니다. 이는 훈련 시 과거 프레임이 미래 프레임을 볼 수 없게 하여 마치 LLM처럼 학습하게 한다는 의미입니다. 제가 훈련하는 데 사용한 또 다른 중요한 방법은 '확산 강제(diffusion forcing)'라고 불립니다. 이는 일반적인 비디오 모델 훈련과 달리, 훈련 시 각 프레임을 독립적으로 노이즈 처리하여 모델이 노이즈가 있는 과거에도 편안함을 느끼도록 학습하게 한다는 의미입니다. 이 모델은 28개의 블록, 20개의 헤드를 가지고 있으며 약 ~960M의 파라미터를 가집니다. 추론 시(At inference) 프레임당 2-5단계의 확산(diffusion)을 실행하고, 한 프레임의 노이즈 제거가 완료되면 이를 KV 캐시(KV cache)에 추가합니다. 이는 LLM의 디코딩 단계와 유사합니다. LLM과 가장 큰 차이점은 모든 KV 컨텍스트, 즉 모든 과거 컨텍스트를 유지하지 않고 슬라이딩 윈도우를 사용한다는 것입니다. 따라서 실제로 남아있는 것은 과거 80프레임 분량의 컨텍스트뿐입니다. 이전 모델은 MMDiT였는데, 이는 텍스트에 대한 교차 어텐션(cross attention)이 없었다는 것을 의미합니다.

이것은 월드 모델(world model)에서 좋지 않은데, 왜냐하면 이전 프레임의 kv와 텍스트 kv가 소프트맥스(softmax) 내에서 문자 그대로 경쟁하기 때문에 신뢰할 수 있는 텍스트 가이던스를 절대 할 수 없기 때문입니다. 지난 모델 역시 텍스트-비디오로 학습되지 않았으므로 어차피 무용지물입니다. 현재 모델은 다시 텍스트 교차 어텐션(cross attention)을 사용하며, 저는 많은 텍스트-비디오 사전 학습(pretraining)을 진행했습니다. 제가 라이브로 제공하는 키보드 액션(예: WASD를 이용한 adaln 추가 항)을 받아들이고, 가장 재미있는 부분은 텍스트 프롬프트 전환입니다. '사막에 연못을 추가해라', '빨간 후디 입혀라', '환경을 얼음으로 바꿔라'와 같은 것입니다. 모델이 많은 텍스트-비디오 정렬(alignment)로 학습되었기 때문에 실제로 프롬프트를 따를 수 있게 되었습니다. 일관성이나 품질 개선 등 여전히 많은 한계가 있다는 것을 알고 있지만, 진심으로 올해 말까지 RTX GPU나 새로운 MacBook을 가진 누구나 시도해 볼 수 있는 결과물을 출시할 수 있기를 바랍니다. 이 초기 이미지를 특별히 선택한 이유는 지난번 subReddit 게시물에서도 동일하게 사용했기 때문입니다. PS: 지난 게시물에서 많은 분들이 저와 자금 지원에 대해 물어보셨는데, 저는 Bangalore를 기반으로 하는 대학생(부분적으로 휴학)이며 학생 인큐베이터의 지원을 받고 있습니다. 혼자서 작업하며 팀이나 실제 회사는 없습니다. 위의 모델은 8x H100 SXM에서 약 3~4주 동안 학습되었습니다. 제가 만드는 모든 모델은 데이터센터 제출용이 아니라 명시적으로 로컬 추론(local inference)용입니다. /u/lucidml_lover [링크] [댓글]