A Multi-agent AI System for Deep Learning Model Migration from TensorFlow to JAX

1. 문제 정의 (Problem Statement)

• 배경: 구글은 수천 개의 머신러닝 모델을 운영하며, 기존 표준 프레임워크인 TensorFlow 에서 TPU 최적화 및 현대적 아키텍처에 적합한 JAX 로 전환 중입니다.
• 도전 과제:
  • 아키텍처 차이: TensorFlow 는 상태 유무 (stateful) 와 정적 실행 그래프, 객체 지향적 Keras API 를 기반으로 하는 반면, JAX 는 상태 없음 (stateless), 함수형 프로그래밍, 동적 실행을 지향합니다.
  • 복잡성: 커스텀 연산 (custom ops), 복잡한 의존성, 엄격한 내부 코딩 스타일 가이드, 그리고 수천 줄에 달하는 코드베이스를 일관성 있게 변환하는 것은 인간 전문가에게도 수백 년의 엔지니어링 노력이 필요한 과업입니다.
  • 일반 LLM 의 한계: 단순한 프롬프트 기반의 단일 에이전트는 긴 컨텍스트, 복잡한 의존성, 그리고 구체적인 도메인 지식이 필요한 작업에서 할루시네이션 (허위 정보 생성) 이나 스타일 불일치, 빌드 불가 코드 생성 등의 문제를 겪습니다.

2. 방법론 (Methodology)

저자들은 TensorFlow 에서 JAX 로의 마이그레이션을 위해 Planner(계획자), Orchestrator(조정자), Coder(코더) 로 구성된 다중 에이전트 시스템을 설계했습니다.

2.1. 계층적 플레이북 (Hierarchical Playbooks)

에이전트에게 필요한 지식과 스타일을 주입하기 위해 마크다운 형식의 '플레이북'을 계층 구조로 구성하여 시스템 프롬프트에 포함시켰습니다.

• General Instruction: 저장소 구조, 빌드 도구 등 공통 지침.
• Style Playbook: 구글 내부 코딩 스타일 규칙 (예: nn.compact 사용 등).
• Task-Specific Playbook: TensorFlow 와 JAX 의 API 차이점 (기본 파라미터 등) 에 대한 지침.
• Client-Specific Playbook: 특정 팀 (예: YouTube) 의 요구사항과 선호 라이브러리를 포함. 골든 예시 (Golden Examples) 로부터 자동 생성된 Few-shot 예시를 활용하여 도메인 특화 지식을 제공합니다.

2.2. 다중 에이전트 아키텍처

• Planner (계획자):
  • 마이그레이션 대상 모델의 전체 의존성을 정적 분석 도구 (Kythe) 를 통해 파악합니다.
  • 복잡한 작업을 빌드 가능하고 검증 가능한 작은 단위 (Step) 로 분해하여 JSON 형태의 마이그레이션 계획을 수립합니다.
  • 각 단계가 독립적으로 실행 가능하도록 범위를 제한합니다.
• Orchestrator (조정자):
  • Planner 의 계획을 Coder 가 처리할 수 있는 논리적 덩어리 (Sub-step) 로 그룹화합니다.
  • Coder 의 컨텍스트 윈도우를 관리하고, 실패 시 재시도, 건너뛰기, 중단 등의 전략적 결정을 내립니다.
  • 진행 상황을 메모리 뱅크 (파일 기반 공유 저장소) 에 저장하여 장기 작업 중 단절 시 복구 가능하게 합니다.
• Coder (코더):
  • ReAct 패턴을 따르는 에이전트로, 파일 읽기/쓰기, 검색, 빌드, 단위 테스트 실행 등의 도구를 사용합니다.
  • 할당된 Sub-step 에 대해 코드를 생성하고, 빌드 및 테스트를 통과할 때까지 자기 수정 (Self-correction) 을 수행합니다.
  • 각 단계 완료 후 변경 사항 요약 (Markdown) 을 생성하여 다음 단계의 컨텍스트로 전달합니다.

2.3. 평가 및 검증 메커니즘

• 체크리스트 기반 LLM 저지 (Judge):
  • 기존 TensorFlow 소스 코드를 참조하지 않고, 마이그레이션된 JAX 코드만 보고 "골든 체크리스트" (모델 아키텍처, 손실 함수, 메트릭 등 필수 요소 포함) 의 충족 여부를 평가합니다.
  • 프레임워크에 구애받지 않는 기능적 논리 (예: "사용자 임베딩이 L2 정규화되었는가?") 를 검증합니다.
• 수치적 동등성 테스트: 빌드 및 테스트 단계에서 Coder 가 자체적으로 생성하여 실행하지만, 최종 품질 지표로는 체크리스트 점수를 주로 사용합니다.

3. 주요 기여 (Key Contributions)

1. 정적 분석과 AI 지시 결합: 매우 복잡한 코드 컴포넌트에 대해 정적 분석 (Kythe) 을 통해 의존성을 파악하고, 이를 기반으로 AI 가 신뢰할 수 있는 마이그레이션 계획을 수립하도록 했습니다.
2. 품질 지표 및 AI 저지 도입: 테스트가 없거나 엄격한 스타일이 요구되는 환경에서도 개발 속도를 가속화하기 위해 체크리스트 기반의 AI 저지 시스템을 도입했습니다.
3. 실제 대규모 환경에서의 검증: 상업적 실사용 사례 (YouTube 모델) 를 통해 시스템의 효과성을 입증하고, AI 가 자신의 개발 워크플로우를 지원하여 선순환을 만드는 것을 시연했습니다.

4. 결과 (Results)

• 성능 비교:
  • 다중 에이전트 + 클라이언트 특화 플레이북 (multi_agent_yt_specific) 구성이 가장 우수한 성능을 보였습니다.
  • 단일 에이전트 (Single Agent) 는 복잡한 모델에서 아키텍처는 맞췄으나, 손실 함수, 메트릭, 핵심 로직 누락 등 중요한 요소를 빠뜨리는 경우가 많았습니다.
  • 복잡한 모델일수록 다중 에이전트 시스템의 우월성이 두드러졌습니다 (단순 모델에서는 차이가 미미함).
• 시간 단축:
  • 수동 마이그레이션 대비 6.4 배 ~ 8 배의 속도 향상 (Speedup) 을 달성했습니다.
  • 특히 중간 규모의 복잡한 모델에서 가장 큰 효율성을 보였습니다.
  • 작은 모델의 경우 1.7 배의 속도 향상만 있었으나, 이는 수동 작업의 고정 비용과 AI 생성 코드 검토 시간에 제한받았기 때문입니다.
• 품질: 생성된 코드는 빌드가 가능하고, 체크리스트 기반 평가에서 높은 점수를 받았으며, 인간 전문가의 검토를 통해 최종 품질이 보장되었습니다.

5. 의의 및 결론 (Significance & Conclusion)

• AI 에 의한 AI 개발: 이 시스템은 AI 기술이 AI 모델 자체의 개발 및 마이그레이션 워크플로우를 가속화하는 '선순환 (Virtuous Circle)'을 보여줍니다.
• 확장성: TensorFlow-JAX 마이그레이션에 국한되지 않고, 다른 프레임워크 전환 및 일반적인 코드 변환 작업에도 적용 가능한 범용적인 접근법입니다.
• 미래 과제: 향후 마이그레이션 후 발생하는 디버깅 (학습 불안정성, 추론 지연 등) 을 자동으로 해결하는 '디버깅 에이전트' 개발을 통해 엔드투엔드 자동화를 완성할 계획입니다.

이 연구는 대규모 머신러닝 인프라의 현대화 과정에서 발생하는 기술적 부채를 해결하기 위해, 단순한 코드 생성을 넘어 계획, 조정, 검증이 통합된 다중 에이전트 시스템의 필요성과 효과를 입증한 중요한 사례입니다.