🤖 AI

Agentic Framework for Deep Learning workload migration via In-Context Learning

본 논문은 인컨텍스트 학습(In-Context Learning)과 실제 PyTorch 출력에서 유도된 실행 오라클(execution oracle)을 활용하여 딥러닝 모델을 PyTorch에서 JAX로 고도로 신뢰할 수 있게 자동 마이그레이션하고, 반복적인 자기 디버깅을 통해 91%의 수치적 등가성을 달야내는 자율 에이전트 프레임워크를 제안한다.

원저자: Qiyue Liang, Steven Ingram, George Vanica, Andi Gavrilescu, Newfel Harrat, Hassan Sipra, Sethuraman Sankaran

게시일 2026-06-16

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Qiyue Liang, Steven Ingram, George Vanica, Andi Gavrilescu, Newfel Harrat, Hassan Sipra, Sethuraman Sankaran

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신은 나무와 못으로 지어진 집(PyTorch)에서 살던 가족을 유리와 강철로만 만들어진 집(JAX)으로 이사시키려 한다고 상상해 보십시오. 두 집은 방의 개수도 같고 목적도 동일하지만, 집을 짓는 규칙은 완전히 다릅니다.

"나무" 집에서는 못을 박거나 나중에 수정할 수 있어 구조가 유연합니다. 반면 "유리" 집은 모든 것이 사전에 계획되어야 하고, 경직되어 있으며, 불변해야 합니다. 만약 당신이 나무 집에서 쓰던 망치질 기술을 그대로 사용하여 유리 집을 지으려 한다면, 구조물 전체가 산산조각이 날 것입니다.

이 논문은 이 이사를 자동화하기 위해 설계된 새로운 AI 건설팀( "에이전트 프레임워크")을 설명합니다. 저자들은 다음과 같은 간단한 비유를 사용하여 이 문제를 어떻게 해결했는지 설명합니다.

문제: AI 건축가가 혼란에 빠지다

보통 똑똑한 AI(대규모 언어 모델)에게 PyTorch 코드를 JAX로 번역하라고 요청하면, AI는 유리 집을 본 적이 없는 건축가처럼 행동합니다. AI는 설계도를 추측하려고 시듭로 합니다.

AI는 머릿속으로 복잡한 계산을 수행할 수 없기 때문에 수학적 계산을 틀릴 수 있습니다.
AI는 유리 집에는 방들이 연결되는 서로 다른 규칙이 필요하다는 사실을 잊어버릴 수 있습니다.
그 결과는 무엇일까요? 서류상으로는 괜찮아 보이지만, 실제로 살려고 하면 무너져 버리는 건물이 됩니다.

해결책: 3단계 안전망

저자들은 AI가 추측하는 것을 막고 엄격하고 검증된 과정을 따르도록 강제하는 시스템을 구축했습니다. 그들은 이를 "그라운디드 생성(grounded generation)" 시스템이라고 부릅니다.

1. "스타일 가이드" (인컨텍스트 러닝, In-Context Learning)

AI에게 유리 집 코드를 쓰는 법을 그냥 추측하게 하는 대신, 팀은 AI에게 엄격한 스타일 가이드를 제공했습니다.

비유: AI에게 지금까지 지어진 가장 멋진 유리 집들의 사진첩을 주며, 창문과 문이 정확히 어디에 배치되어야 하는지 보여주는 것과 같습니다.
역할: 이 "인컨텍스트 러닝(ICL)"은 닻 역할을 합니다. 이는 AI에게 "JAX에서는 코드를 이렇게 작성한다"라고 보여줍니다. 이는 AI가 환각(hallucination, 근거 없는 내용을 지어내는 것)을 일으키는 것을 방지하고 올바른 궤도에 머물게 합니다.

2. "진실 기계" (오라클, The Oracle)

이 부분이 가장 중요합니다. AI는 머릿속으로 복잡한 수학을 계산하는 데 서툽니다. 그래서 팀은 AI에게 숫자를 계산하라고 요구하지 않았습니다. 대신, 원래의 PyTorch 코드가 먼저 계산을 수행하도록 했습니다.

비유: AI가 유리 집을 짓기 전에, 팀은 나무 집을 시뮬레이터에 통과시켜 모든 방이 어떻게 작동하는지, 가구가 얼마나 무거운지, 빛이 벽에 어떻게 비치는지에 대한 "스냅샷"을 찍습니다. 그리고 이 스냅샷을 오라클(절대적인 진리의 원천)로 저장합니다.
역할: AI는 숫자를 추측하지 않습니다. AI는 이 스냅샷과 정확히 일치하도록 유리 집을 지어야 합니다. 만약 유리 집이 스냅샷과 일치하지 않는다면, 그것은 틀린 것입니다.

3. "자기 수정 검사관" (반복적 디버깅, Iterative Debugging)

AI가 유리 집의 초안을 만듭니다. 그러면 로봇 검사관이 "진실 기계" 스냅샷과 대조하여 이를 점검합니다.

비유: 만약 검사관이 창문의 크기가 틀리거나 문이 열리지 않는 것을 발견하면, 단순히 "실패"라고 말하지 않습니다. 대신 설계도를 AI 건축가에게 돌려주며 이렇게 말합니다. "보세요, 빛이 이 벽에 비치는 모습이 사진과 다릅니다. 수정하세요."
역할: AI는 다시 시도하고, 오류를 읽고, 코드를 수정합니다. 유리 집이 스냅샷과 완벽하게 일치할 때까지 이 과정을 반복합니다.

결과: "아마도"에서 "완벽함"으로

팀은 이 시스템을 두 가지 유형의 작업에 테스트했습니다:

단순 수학 연산 (레벨 1): 벽돌 한 장을 옮기는 것과 같은 작업.
복잡한 신경 모듈 (레벨 2): 복잡한 배선이 있는 집의 한 구역 전체를 짓는 것과 같은 작업.

이 시스템 없이 (베이스라인):

복잡한 모듈에 대해 AI가 수학을 맞춘 비율은 단 **9%**였습니다. AI는 대부분 추측만 하고 있었습니다.

지침만 있을 때 (오라클 없음):

AI는 규칙을 따르는 능력은 좋아졌지만, 여로 여전히 수학적 계산은 **73%**의 확률로 틀렸습니다. 겉모습은 그럴듯하지만 제대로 작동하지 않는 집을 짓고 있었던 것입니다.

전체 시스템 적용 시 (논문의 방법):

단순 작업: 100% 성공.
복잡한 작업: 91% 성공.

이 시스템은 SAM(컴퓨터가 이미지를 "보는" 것을 돕는 모델), T5(텍스트 번역기), Code Whisper(코딩 어시스턴트)와 같은 유명한 모델들을 PyTorch에서 JAX로 높은 정확도로 성공적으로 이식했습니다.

핵심 요약

이 논문은 복잡한 코드를 "번역"하라고 AI에게 요청하고 결과가 좋기를 바라는 것만으로는 부족하다는 것을 증명합니다. 반드시 다음 세 가지를 제공해야 합니다:

어떻게 제대로 하는지에 대한 예시 (스타일 가이드).
결과물이 어떤 모습이어야 하는지에 대한 증거 (오라클).
그 증거를 바탕으로 실수를 바로잡을 수 있는 기회 (검사관).

이 세 가지를 결려 함으로써, 그들은 인간이 모든 코드 한 줄 한 줄을 일일이 확인할 필요 없이 딥러닝 모델을 서로 다른 프레임워크 간에 안정적으로 이동시키는 신뢰할 수 있는 자동화 방법을 만들어냈습니다.

기술 요약: 인컨텍스트 러닝을 통한 딥러닝 워크로드 이전을 위한 에이전트 프레임워크

1. 문제 정의

PyTorch의 이가적(eager), 객체 지향적, 가변 상태 패러다임에서 JAX의 함수형, 무상태, JIT 컴파일(XLA) 패러다임으로 딥러닝 모델을 마이그레이션하는 것은 중대한 엔지니어링 병목 현상입니다. 이러한 전환은 Google의 TPU와 같이 행렬 곱셈 속도 및 비용 효율성이 뛰어난 하드웨어 특화 최적화를 활용하기 위해 필수적입니다. 그러나 다음과 같은 근본적인 차이로 인해 수동 마이그레이션은 오류를 범하기 쉽습니다:

실행 모델: PyTorch의 이가적 실행(eager execution) 대 JAX의 Just-In-Time(JIT) 컴파일.
상태 관리: PyTorch의 가변 객체 상태 대 JAX의 함수형, 무상태 flax.linen 추상화.
제어 흐름: 표준 Python 루프 대 jax.lax.scan과 같은 컴파일러 친화적 구조.
수치적 불일치: 기본 텐서 레이아웃(NCHW 대 NHWC), 정밀도(fp32 대 bfloat16), 그리고 어텐션 레이어의 패딩 동작 차이.

대규모 언어 모델(LLM)이 일반적인 코드 생성에는 능숙함을 보여주었으나, 엄격한 프레임워크 번역에는 어려움을 겪습니다. 신뢰할 수 있는 검증 없이는, LLM은 고차원 텐서 산술, 특히 고차원 텐서 연산에 관한 수학적 연산과 의미론적 로직을 환각(hallucination)하여, 코드는 컴파일되지만 수치적으로 동일한 결과를 생성하지 못하는 경우가 발생합니다.

2. 방법론

저자들은 PyTorch 모듈을 JAX로 번역하기 위해 설계된 완전 자율형 에이전트 파이프라인을 제안합니다. 이 시스템은 인컨텍스트 러닝(ICL)과 실행 기반 오라클(execution-based oracle)을 결합한 4단계 워크플로(그림 1 참조)를 사용하여 프레임워크 번역을 "근거 기반 생성 문제(grounded generation problem)"로 취급합니다.

1단계: 인컨텍스트 러닝(ICL)을 통한 앵커링(Anchoring)

"중간에서 길을 잃는(Lost in the Middle)" 현상을 완화하고 광범위한 검색 방식(RAG 등)으로부터 발생하는 노이즈를 줄이기 위해, 시스템은 엄격하게 제어된 ICL 전략을 채택합니다. 복잡한 코드 조각을 대량으로 검색하는 대신, 컨텍스트를 다음을 포함하는 작고 밀도 높은 참조 세트로 제한합니다:

소스 PyTorch 모듈.
관용적인(idiomatic) JAX 번역본.
관련 테스트 케이스.
이 예시들은 구조적 앵커(structural anchors) 역할을 하여, LLM이 컨텍스트 창을 압도하지 않으면서도 올바른 구문 형식과 기능적 패턴을 채택하도록 유도합니다.

2단계: 동적 실행 오라클(Dynamic Execution Oracle)

LLM이 고차원 텐서 산술을 신뢰성 있게 추론할 수 없다는 점을 인식하여, 시스템은 변경 불가능한 실행 오라클을 구축합니다.

소스 PyTorch 모듈을 무작위 입력 텐서로 실행합니다.
그 결과 발생하는 동적 계산 상태(초기화된 가중치(state_dict), 입력 텐서, 출력 활성화 값 포함)를 고정된 아티팩트(.pkl 파일)로 직렬화합니다.
이 오라클은 정적 분석이 놓칠 수 있는 동작(예: 암묵적 패딩 방법, 동적 형상 변화)을 포착하여, 타겟 JAX 코드에 대한 절대적인 수학적 참값(ground truth)을 제공합니다.

3단계: 에이전트 번역 및 테스트 생성

에이전트는 ICL 앵커에 의해 제약된 상태에서 PyTorch 소스를 flax.linen 추상체로 번역합니다. 결정적으로, 에이전트는 2단계 오라클을 직접 기반으로 하는 커스텀 테스트 하네스를 생성하는 과업을 수행합니다. 평가는 다음과 같은 계층적 3단계 과정을 따릅니다:

컴파일 체크: JAX 모듈이 런타임 오류 없이 초기화되고 트레이싱되는지 확인합니다.
형상(Shape) 정확성: 레이아웃 시프트(예: NCHW에서 NHWC로)를 처리하며 차원의 일관성을 확인합니다.
수치적 등가성: 오라클 가중치를 역직렬화하여 JAX 모델에 매핑하고, 엄격한 절대 오차 한계인 $10^{-7}$ 을 사용하여 numpy.allclose로 수치적 등가성을 확인합니다. 이를 "실버(Silver)" 테스트 케이스라고 부릅니다.

4단계: 반복적 자기 디버깅(Iterative Self-Debugging)

생성된 코드가 오라클 조건 테스트를 통과하지 못하면, 시스템은 자기 수정 루프에 진입합니다. 스택 트레이스, 표준 에러, 정밀한 수치 차이를 포함한 실행 피드백이 에이전트의 컨텍스트로 다시 전달됩니다. ICL 템플릿과 이러한 객관적 피드백에 의해 유도되어, LLM은 엄격한 수학적 제약을 통과할 때까지 코드를 반복적으로 개선합니다.

3. 주요 기여

본 논문은 세 가지 주요 기여를 명시합니다:

ICL을 통한 구조적 앵커링: 정교하게 큐레이션된 참조 번역 및 테스트 세트가 생성 과정을 효과적으로 앵커링하여, 광범기한 검색이나 지시문 전용 방식에 비해 환각 위험을 크게 낮춘다는 것을 입증했습니다.
동적 실행 오라클: 소스 PyTorch 구현을 참값으로 사용하여 LLM의 산술적 한계를 우회하는 방법을 도입했습니다. 이는 코드 평가를 이론적 추정이 아닌 실제 계산을 통한 실행 기반 평가로 전환하여 수학적 충실도를 보장합니다.
에이전트 워크플로를 위한 절제 연구(Ablation Guidance): 지시문, ICL, 실행 조건부 자기 디버깅의 영향을 분리하여 입증함으로써, 향-에이전트 코드 번역 시스템을 위한 청사진을 제공하는 경험적 증거를 제시했습니다.

4. 결과

시스템은 세 가지 난이도 수준으로 분류된 홀드아웃 세트에서 평가되었습니다:

레벨 1: 기초적인 수학적 연산 (9개 태스크).
레벨 2: 신경망 모듈 및 아키텍처 추상화 (MLP, Transformer, GAN 등을 포함한 11개 태스크).
레벨 3: 전체 저장소(Repository) 마이그레이션 (SAM, T5, Code Whisper, FinGPT를 포함한 10개 저장소).

정량적 성능 (표 1):

베이스라인 (지시문 전용): 레벨 2 모듈에서 수치적 등가성 9% 달성.
지시문 + 자기 디버깅: 레벨 2에서 수치적 등가성 27%로 개선되었으나, 형상 일치는 55%에 머묨.
전체 파이프라인 (ICL + 오라클 + 자기 디버깅): 레벨 1에서 100%, 레벨 2에서 **91%**의 수치적 등가성 달성.

저장소 마이그레이션 (표 2):
10개의 레벨 3 저장소 전반에 걸쳐, 전체 파이프라인은 코드 변환에서 100% 완결성을 달성했습니다. 수치적 등가성은 모델 복잡도에 따라 다양했습니다:

높은 충실도 (93%–100%): Code Whisper, Multimodal Transformer, 시계열 예측, StyleGAN2, T5, Two-tower 모델.
중간 충실도: FinGPT (86%), SAM2 (80%), DETR (57%).
가독성: 인간 평가자는 변환된 파일의 가독성을 5점 리커트 척도 기준으로 3.75에서 4.9 사이로 평가했습니다.

실패 분석:
논문은 베이스라인의 세 가지 주요 실패 모드를 식별했습니다: 구문적 환각(래퍼 클래스 누락), API 시그니처 불일치(잘못된 kwargs), 상태 매핑 불일치(가중치 전송 중 KeyError 발생). 전체 파이프라인의 반복적 디버깅 루프는 에이전트가 오라클에 대조하여 구조 및 가중치 매핑 로직을 해결하도록 강제함으로써 이를 해결합니다.

5. 의의 및 주장

저자들은 자신들의 근거 기반 에이전트 구동 파이프라인이 매우 신뢰할 수 있고 확장 가능한 청사진을 제공한다고 주장합니다. 본 연구의 의의는 다음과 같습니다:

수동 테스트 제거: 시스템은 실행 오라클에 근거한 "실버" 테스트 케이스를 생성함으로써 수동 테스트 생성의 필요성을 제거합니다.
LLM의 한계 극려: LLM의 역할을 수학적 추론(오라클이 담당)에서 구조적 번역 및 디버깅으로 분리함으로써, 순수 LLM 접근 방식이 실패하는 지점에서 높은 충실도의 결과를 달성합니다.
확장성: 복잡한 최신 모델(예: SAM, T5)을 높은 수치적 등가성으로 성공적으로 마이그레이션함으로써 실무 엔지니어링 과제에 대한 유용성을 입증했습니다.

한계점:
저자들은 두 가지 현재의 한계점을 솔직하게 언급했습니다:

완전한 구조적 의존성 관리가 필요한 완전 자동 저장소 수준의 마이그레이션으로 아직 확장되지 않았습니다.
ICL 컨텍스트에 사용되는 예시의 변화가 미치는 구체적인 영향에 대해 체계적인 탐구가 이루어지지 않았습니다.

이러한 한계에도 불구하고, 본 논문은 구조적 앵커링과 동적 실행 기반 피드백을 결합하는 것이 프레임워크 번역 병목 현상을 효과적으로 해결하여, 최소한의 수동 개입으로 최적화된 JAX 환경으로 딥러닝 워크로드를 마이그레이션할 수 있음을 결론짓습니다.