늦깎이 공대생의 인공지능 연구실
TinyAgent: Edge에서의 함수 호출 본문
자연어(예: 영어)를 통해 명령을 실행하는 LLM의 능력 덕분에 적절한 도구들(ToolFormer, Gorilla)을 조율하여 사용자 쿼리를 완료할 수 있는 에이전트 시스템을 구현할 수 있게 되었습니다. 이는 GPT-4o 또는 Gemeni-1.5 모델과 같은 최근의 멀티모달 연구와 함께 AI 에이전트의 가능성의 영역을 확장시켰습니다. 이는 매우 흥미로운 일이지만, 이러한 모델의 규모가 크고 계산 요구사항이 많기 때문에 클라우드에서 연산을 수행해야 하는 경우가 많습니다. 이로 인하여 광범위한 도입에는 몇 가지 문제가 발생할 수 있습니다. 무엇보다도 비디오, 오디오 또는 텍스트 문서와 같은 데이터를 타사 클라우드 공급업체에 업로드 하게되면 개인정보 보호문제가 발생할 수 있습니다. 둘째로, 클라우드/Wi-Fi 연결이 필요하지만, 항상 연결이 안정적이지만은 않을 수 있습니다. 예를 들어, 현실 세계에 배치된 로봇의 연결 상태가 항상 안정적이지 않을 수 있습니다. 또한, 대량의 데이터를 클라우드에 전송하고 응답을 기다리면 응답 지연 시간이 길어져 문제 해결에 걸리는 시간이 길어질 수 있기 때문에 지연 시간도 문제가 될 수 있습니다. 이러한 문제는 LLM 모델을 Edge에서 Local로 배포하면 해결할 수 있습니다.
하지만 GPT-4o나 Gemini-1.5와 같은 최신 LLM은 Local 배포에 사용하기에는 용량이 너무 큽니다. 그 원인 중 하나는 모델 크기의 상당 부분이 특수한 다운스트림 애플리케이션에는 필요하지 않을 수 있는 일반적인 정보를 파라메트릭 메모리에 저장하게 된다는 점입니다. 예를 들어, 이러한 모델에서 역사적 사건이나 잘 알려진 인물과 같은 일반적인 사실에 대한 질문을 하면 프롬프트에 추가 컨텍스트가 없어도 파라메트릭 메모리를 사용하여 결과를 생성할 수 있습니다. 그러나 이러한 학습 데이터의 파라메트릭 메모리에 대한 암묵적 암기는 맥락 내 학습 및 복잡한 추론과 같은 LLM의 "새로운" 현상과 관련이 있으며, 이는 모델 규모를 확장하는 원동력이 되었습니다.
그러나, 여기서 흥미로운 궁금증이 생갑니다.
파라메트릭 메모리가 훨씬 적은 소규모 언어 모델이 대규모 언에 모델의 놀라운 능력을 모방할 수 있을까?
만약 이러한 목표를 달성한다면 에이전트 시스템의 계산 공간을 크게 줄여 효율적이면서 개인정보를 보호하는 Edge 배포가 가능해집니다. 이 연구는 일반적인 범용 지식을 불러올 필요가 없는 전문화된 고품질 데이터를 학습함으로써 소규모 언어 모델에서도 이러한 작업이 가능하다는 것을 보여줍니다.
이러한 시스템은 특히 AI 에이전트의 역할이 자연어로 된 사용자 쿼리를 이해하고 ChatGPT 유형의 질문 답변으로 응답하는 대신 사용자의 명령을 수행하기 위ㅜ해 적절한 도구와 API셋을 조율하는 의미 체계인 Semantic System에 유용할 수 있습니다. 예를 들어, Siri와 유사한 애플리케이션에서 사용자가 언어 모델에 특정 ㅊ참석자가 포함된 캘린더 초대를 만들도록 요청할 수 있습니다. 캘린터 항목 생성을 위한 사전 정의된 스크립트가 이미 존재하는 경우 LLM은 올바른 입력 인수(예: 참석자의 이메일 주소, 이벤트 제목, 시간)로 이 스크립트를 호출하는 방법을 학습하기만 하면 됩니다. 이 과정에는 Wikipedia와 같은 자료롤부터 세간의 지식을 기억/암기하는 것이 아니라 올바른 함수를 호출하고 이를 올바르게 조율하는 추론과 학습이 필요합니다.
우리의 목표는 Edge에서 안전하고 비공개적으로 배포할 수 있는 복잡한 추론이 가능한 소규모 언어 모델(SLM)을 개발하는 것입니다. 지금부터 이를 위해 추진하고 있는 연구 방향에 대해 설명하겠습니다. 먼저, 에이전트 시스템의 핵심 구성 요소인 함수 호출을 정확하게 수행할 수 있도록 소규모 오픈소스 모델을 지원하는 방법에 대해 살펴봅니다. 공개된 소형 모델은 함수 호출 성능이 매우 낮은 것으로 나타납니다. 우리는 전문화된 Mac 어시스턴트 에이전트를 구동 애플리케이션으로 사용하여 함수 호출을 위한 고품질 데이터를 체계적으로 큐레이션함으로써 이 문제를 해결하는 방법에 대해 알아보니다. 그런 다음 이렇게 선별된 고품질 데이터셋에서 모델을 미세조정하면 SLM이 GPT-4-Turbo의 함수 호출 성능을 훨씬 능가할 수 있음을 보여드리겠습니다. 그 다음 새로운 Tool RAG 방법을 통해 이를 더욱 개선하고 효율적으로 만들 수 있다는 것을 보여드리겠습니다. 마지막으로 실시간 응답을 통해 최종 모델을 Edge에 효율적으로 배포할 수 있는 방법을 보여드리겠습니다.
Macbook M3 Pro에서 로컬로 배포된 Whisper-v3와 실행 가능한 TinyAgent-1B 데모. 이 프레임워크는 오픈 소스로 제공되며 https://github.com/SqueezeAILab/TinyAgent 에서 확인할 수 있습니다.
LLM에 함수 호출 학습시키기
위에서 살펴보았듯이, 우리의 주요 관심 분야는 AI 에이전트가 사용자 쿼리를 일련의 함수 호출로 변환하여 작업을 완료하는 애플리케이션입니다. 이러한 애플리케이션에서는 함수(또는 API)가 대부분 미리 정의되어 있고 이미 사용 가능하므로 모델이 함수 정의 자체를 작성할 필요가 없습니다. 따라서 모델이 해야 할 일은 (i) 호출할 함수, (ii) 해당 입력 인수, (iii) 함수 호출 간에 필요한 상호 의존성을 기반으로 이러한 함수를 호출하는 올바른 순서(즉, 함수 오케스트레이션)를 결정하는 것입니다.
첫 번째 질문은 함수 호출을 수행하기 위해 SLM을 효과적으로 갖추는 방법을 찾는 것입니다. GPT-4와 같은 대형 모델은 함수 호출을 수행할 수 있지만 오픈 소스 모델에서는 어떻게 이를 구현할 수 있을까요? 우리 그룹의 최신 프레임워크인 LLMCompiler는 입력 인수 및 종속성과 함께 호출해야 하는 함수 집합을 포함하는 Function Calling Plan을 출력하도록 LLM에 지시함으로써 이를 가능하게 합니다(그림 1의 예시 참조). Function Calling Plan이 생성되면 이를 파싱하고 종속성을 기반으로 각 함수를 호출할 수 있습니다.
여기서 가장 중요한 부분은 모델이 올바른 구문과 종속성으로 Function Calling Plan을 생성하도록 가르치는 것입니다. 원래의 LLMCompiler 논문에서는 프롬프트에 충분한 지침이 제공될 때 계획을 생성할 수 있는 복잡한 추론 기능을 갖춘 LLaMA-2 70B와 같은 대형 모델만 고려했습니다. 하지만 더 작은 모델에서도 동일한 방식으로 올바른 Function Calling Plan을 출력하도록 프롬프트를 제공할 수 있을까요? 안타깝게도 실험 결과, TinyLLaMA-1.1B와 같은 기존의 소형 모델(또는 더 큰 Wizard-2-7B 모델)은 올바른 계획을 출력할 수 없는 것으로 나타났습니다. 오류는 잘못된 함수 집합 사용, 잘못된 이름, 잘못된 종속성, 일관되지 않은 구문 등과 같은 다양한 문제에서 비롯되었습니다.
이러한 소규모 모델은 일반적인 데이터셋에서 학습되었고 주로 모델의 배경 지식과 일반적인 추론 또는 기본 명령어 추종 능력을 테스트하는 일반적인 벤치마크에서 좋은 정확도를 달성하는 것을 목표로 했기 때문에 이는 어느 정도 예상할 수 있는 결과입니다. 이 문제를 해결하기 위해 함수 호출 및 계획을 위해 특별히 선별된 고품질 데이터셋에서 이러한 모델을 미세 조정하면 목표 작업에 대한 소규모 언어 모델의 정확도를 개선하여 잠재적으로 대규모 모델보다 더 나은 성능을 낼 수 있는지 살펴봤습니다. 먼저 이러한 데이터 세트를 생성한 방법을 살펴본 다음 미세 조정 접근 방식에 대해 논의합니다.
데이터셋 생성
작동 애플리케이션으로 위와 같이 사용자의 일상적인 작업을 해결하는 Apple의 Macbook용 로컬 에이전트 시스템을 예로 들어 보겠습니다. 특히 이 Agent는 Mac의 다양한 애플리케이션과 상호 작용할 수 있는 16가지 기능을 갖추고 있습니다
- Email: 새 이메일을 작성하거나 이메일에 회신/전달하기
- Contacts: 연락처 데이터베이스에서 전화번호 또는 이메일 주소 가져오기
- SMS: 연착처에서 문자 메시지 보내기
- Calendar: 제목, 시간, 참석자 등의 세부 정보가 포함된 캘린더 이벤트 생성
- Notes: 다양한 폴더의 노트에 콘텐츠를 만들거나, 열거나, 추가하기
- Reminder: 다양한 활동 및 작업에 대한 미리 알림 설정
- File management: 다양한 파일 경로에서 문서 열기, 읽기, 요약하기
- Zoom meetings: Zoom 미팅 예약 및 구성
이러한 각 기능/도구마다 사전 정의된 Apple 스크립트가 존재하며, 모델이 해야 할 일은 사전 정의된 API를 활용하여 주어진 작업을 수행하는 데 적합한 Function Calling Plan을 결정하기만 하면 됩니다. 하지만 앞서 설명한 것처럼 소규모 언어 모델은 기존 함수 호출 기능이 수준 이하이므로 평가 및 학습을 위한 데이터가 필요합니다.
다양한 Function Calling Plan으로 수작업으로 데이터를 생성하는 것은 어렵고 확장성이 떨어집니다. 하지만 GPT-4-Turbo와 같은 LLM을 사용하여 합성 데이터를 큐레이션할 수 있습니다. 이러한 접근 방식은 유능한 LLM에 주어진 샘플 예제 또는 템플릿 세트와 유사한 데이터를 생성하도록 지시하는 일반적인 방법이 되어가고 있습니다(LLM2LLM 및 Self-Instruct 참조). 본 작업에서도 비슷한 접근 방식을 사용했지만, 위의 예제처럼 일반적인 사용자 쿼리를 템플릿으로 제공하는 대신 다양한 함수 집합을 제공하고 해당 함수가 작업을 수행하는 데 필요한 실제 사용자 쿼리를 관련 Function Calling Plan 및 입력 인수와 함께 생성하도록 지시했습니다. 생성된 데이터의 유효성을 검증하기 위해 Function Calling Plan에 건전성 검사를 통합하여 실행 가능한 그래프를 형성하는지, 함수 이름과 입력 인자 유형이 올바른지 확인했습니다. 이 접근 방식을 통해 8만 개의 학습 데이터, 1천 개의 유효성 검사 데이터, 1천 개의 테스트 데이터를 생성했으며 총 비용은 약 500달러에 불과했습니다.
함수 호출 추론 개선을 위한 파인 튜닝
데이터셋을 확보했으니 이제 기존 SLM을 미세 조정하여 함수 호출 기능을 향상시킬 수 있습니다. 먼저 두 가지 기본 소형 모델인 TinyLlama-1.1B(instruct-32k 버전)와 Wizard-2-7B를 사용해 보겠습니다. 이러한 모델을 미세 조정하려면 먼저 성능을 평가할 수 있는 기준을 정의해야 합니다. 우리의 목표는 이러한 모델이 올바른 계획을 정확하게 생성하는 것이며, 여기에는 올바른 함수 집합을 선택하는 것뿐만 아니라 올바른 순서로 올바르게 오케스트레이션하는 것도 포함됩니다. 따라서 두 기준이 모두 충족되면 1을, 그렇지 않으면 0을 할당하는 성공률 기준을 정의합니다. 모델이 올바른 함수 호출 집합을 선택했는지 확인하는 것은 간단합니다. 이러한 함수의 오케스트레이션이 올바른지 추가로 확인하기 위해 위의 그림과 같이 의존성을 기반으로 함수 호출의 방향 비순환 그래프(Directed Acyclic Graph)를 구성합니다. 여기서 각 노드는 함수 호출을 나타내고 노드 A에서 B까지의 연결 변수는 상호 의존성을 나타냅니다(즉, 함수 A가 실행된 후에만 함수 B를 실행할 수 있음). 그런 다음 이 DAG가 기준 데이터 계획과 동일한지 비교하여 종속성의 정확성을 검증합니다.
평가 지표를 정의한 후, LoRA를 적용하여 8만 개의 학습 예제에 대해 7e-5의 학습률을 사용하여 3회의 epoch 동안 모델을 미세 조정하고 검증 성능을 기반으로 최상의 체크포인트를 선택했습니다. 미세 조정을 위해 기준값 함수(즉, 기준값 계획에 사용된 함수)에 대한 설명뿐만 아니라 관련 없는 다른 함수도 네거티브 샘플로 포함시켰습니다. 네거티브 샘플은 주어진 쿼리에 적합한 도구를 선택하는 방법을 모델에 가르치는 데 특히 효과적이어서 학습 후 성능을 향상시키는 데 도움이 된다는 사실을 발견했습니다. 또한 쿼리가 어떻게 함수 호출 계획으로 변환되는지 보여주는 몇 가지 상황별 예시도 포함되어 있습니다. 이러한 컨텍스트 내 예제는 학습 데이터 세트의 데이터에서 사용자 쿼리를 기반으로 검색 증강 생성(RAG) 프로세스를 통해 선택됩니다.
위의 설정을 사용하여 TinyLlama-1.1B/Wizard-2-7B 모델을 미세 조정했습니다. 미세 조정 후 1.1B 모델의 성공률은 12.71%에서 78.89%로, 7B 모델의 성능은 41.25%에서 83.09%로 향상되어 GPT-4-Turbo보다 약 4% 더 높았습니다.
Tool RAG를 사용한 효율적인 추론
GPT와 같은 비공개 소스 모델이 배포되는 GPU에 비해 사용 가능한 연산 및 메모리 리소스가 제한된 Macbook에 TinyAgent 모델을 로컬로 배포하는 것이 우리의 주요 목표입니다. 짧은 지연 시간으로 효율적인 성능을 달성하려면 모델 크기가 작아야 할 뿐만 아니라 입력 프롬프트가 가능한 한 간결해야 합니다. 후자는 시퀀스 길이에 대한 관심의 복잡성이 4제곱으로 증가하기 때문에 지연 시간과 계산 리소스 소비에 중요한 영향을 미칩니다.
앞서 설명한 미세 조정된 TinyAgent 모델은 프롬프트에 사용 가능한 모든 도구에 대한 설명으로 미세 조정되었습니다. 하지만 이는 매우 비효율적입니다. 사용자 쿼리를 기반으로 관련 도구에 대한 설명만 포함하면 프롬프트 크기를 크게 줄일 수 있습니다. 예를 들어, 위 그림에 표시된 예제에서 사용자가 두 사람과 함께 캘린더 초대를 만들 것을 요청하는 경우를 생각해 보겠습니다. 이 경우 LLM은 프롬프트에 이메일 주소를 가져오고 캘린더 이벤트를 생성하는 기능만 필요합니다.
이러한 특성을 활용하려면 사용자의 명령을 수행하는 데 필요한 함수를 결정해야 하는데, 이를 검색 증강 생성(Retrieval Augmented Generation)의 작동 방식과 유사하기 때문에 Tool RAG라고 부릅니다. 하지만 여기에는 중요한 미묘한 차이가 있습니다. 사용자 쿼리의 임베딩을 계산하고 이를 사용하여 관련 도구를 검색하는 기본 RAG 방법을 사용하면 성능이 매우 저하됩니다. 사용자 쿼리를 완료하려면 여러 개의 보조 도구를 사용해야 하는 경우가 많은데, 보조 도구의 임베딩이 사용자 쿼리와 유사하지 않은 경우 단순한 RAG 방법으로는 이를 놓칠 수 있기 때문입니다. 예를 들어 위 그림에 표시된 예는 사용자 쿼리가 캘린더 초대를 생성하는 것에 대한 질문인데도 get_email_address 함수를 호출해야 합니다.
이 문제는 어떤 도구가 필요한지 분류하는 것으로 해결할 수 있습니다. 이를 위해 아래 그림과 같이 16-way 분류를 수행하기 위해 학습 데이터에 대해 DeBERTa-v3-small 모델을 미세 조정했습니다. 사용자 쿼리를 이 모델에 입력으로 제공한 다음, 마지막에 CLS 토큰을 768x16 크기의 간단한 Fully Connected Layer를 통과시켜 16차원 벡터(전체 도구의 개수)로 변환합니다. 이 레이어의 출력은 Sigmoid를 통과하여 각 도구의 선택 확률을 생성합니다. 추론하는 동안 50% 이상의 확률을 가진 도구를 선택하고, 그럴 경우 프롬프트에 해당 도구에 대한 설명을 포함합니다. 평균적으로 3.97개의 도구만 검색되어 0.998의 Recall값을 보이는 반면, 기본 RAG에서는 상위 6개 도구를 사용하여 0.968의 도구 Recall을 달성해야 합니다.
도구 RAG를 통합한 후 모델 성능을 평가해보았습니다. 그 결과는 아래 표 1에 나와 있으며, 여기에는 단순 RAG 시스템과 미세 조정된 DeBERTa 접근 방식의 성능을 비교한 결과가 나와 있습니다. 표에서 볼 수 있듯이, DeBERTa 기반 도구 RAG 방식은 거의 완벽한 Recall 성능을 달성하고 기준 정확도를 개선하는 동시에 프롬프트 크기를 최대 2배까지 줄였습니다.
표1: 기본 RAG 설정과 RAG 설정이 없는 경우의 TinyAgent 성능 비교
| Tool RAG Method | Tool Recall | Prompt Size (Tokens) | TinyAgent 1.1B Success Rate (%) | TinyAgent 7B Success Rate (%) |
| No RAG (all tools in the prompt) | 1 | 2762 | 78.89 | 83.09 |
| Basic RAG | 0.949 (top 3) | 1674 | 74.88 | 78.50 |
| Fine-tuned DeBERTa-v3-small (Ours) | 0.998 (tools with >50% prob) | 1397 | 80.06 | 84.95 |
정량화를 통한 빠른 Edge 배포
모델 파라미터를 불러오면 사용 가능한 메모리의 상당 부분을 소모할 수 있기 때문에 일반 소비자용 MacBook과 같은 Edge에 모델을 배포하는 것은 O(1B) 파라미터의 작은 모델에서도 여전히 어려울 수 있습니다. 이러한 문제에 대한 해결책은 양자화를 통해 모델을 더 낮은 비트 정밀도로 저장할 수 있습니다. 양자화는 스토리지 요구 사항과 모델 공간을 줄일 뿐만 아니라 모델 가중치를 메모리에 로드하는 데 필요한 시간과 리소스를 줄여 전반적인 추론 지연 시간도 줄여줍니다(양자화에 대한 자세한 내용은 여기를 참조하세요).
보다 효율적인 모델 배포를 위해 모델을 4 비트, 그룹 크기 32로 양자화했으며, 이는 양자화 인식 학습을 통해 llama.cpp 프레임워크에서 지원됩니다. 표 2에서 볼 수 있듯이 4비트 모델을 사용하면 지연 시간이 30% 개선되고 모델 크기가 4배 감소합니다. 또한 시뮬레이션 양자화를 통한 추가적인 미세 조정으로 인해 약간의 정확도 향상도 확인할 수 있습니다.
표2: 정량화 전후의 지연 시간, 크기, 성공률. 지연 시간은 Function Calling Planner의 End-to-End 지연 시간으로, 프롬프트 처리 시간 및 생성을 포함함
| Model | Weight Precision | Latency (seconds) | Model Size (GB) | Success Rate (%) |
| GPT-3.5 | Unknown | 3.2 | Unknown | 65.04 |
| GPT-4-Turbo | Unknown | 3.9 | Unknown | 79.08 |
| TinyAgent-1.1B | 16 | 3.9 | 2.2 | 80.06 |
| TinyAgent-1.1B | 4 | 2.9 | 0.68 | 80.35 |
| TinyAgent-7B | 16 | 19.5 | 14.5 | 84.95 |
| TinyAgent-7B | 4 | 13.1 | 4.37 | 85.14 |
종합
여기서는 Macbook Pro M3에 배포된 최종 TinyAgent-1.1B 모델의 데모를 보여드리며, 실제로 Mac에 다운로드하여 설치하여 테스트해 볼 수 있습니다. 컴퓨터에서 모든 모델 추론을 로컬로 실행할 뿐만 아니라 오디오를 통해 명령을 제공할 수도 있습니다. 저희는 whisper.cpp 프레임워크를 사용하여 로컬로 배포된 OpenAI의 Whisper-v3 모델을 사용하여 오디오도 로컬에서 처리합니다. 가장 놀라웠던 점은 1.1B 모델의 정확도가 GPT-4-Turbo의 정확도를 능가하며, 로컬 및 디바이스에 비공개로 배포된 상태에서 현저하게 빠르다는 점입니다.
요약하자면, TinyAgent를 소개하면서 작은 언어 모델을 훈련하고 이를 사용하여 사용자 쿼리를 처리하는 시맨틱 시스템을 구동하는 것이 실제로 가능하다는 것을 보여주었습니다. 특히, Mac용 Siri와 유사한 어시스턴트를 구동 애플리케이션으로 고려했습니다. 이를 구현하기 위한 핵심 요소는 (i) LLMCompiler 프레임워크를 통해 함수 호출을 수행할 수 있도록 상용 SLM을 학습시키고, (ii) 주어진 작업에 맞는 고품질의 함수 호출 데이터를 큐레이션하며, (iii) 생성된 데이터에 대해 상용 모델을 미세 조정하고, (iv) ToolRAG라는 방법을 통해 사용자 쿼리에 따라 필요한 도구만 검색하고 추론 자원 소비를 줄이기 위해 모델 배포를 양자화하여 프롬프트 크기를 최적화함으로써 효율적인 배포를 가능하게 하는 것이 핵심입니다. 이러한 단계를 거친 후 최종 모델은 TinyAgent1.1.B와 7B 모델의 경우 80.06%와 84.95%를 달성하여 이 작업에서 GPT-4-Turbo의 성공률인 79.08%를 뛰어넘었습니다.
참고자료: https://bair.berkeley.edu/blog/2024/05/29/tiny-agent/
TinyAgent: Function Calling at the Edge
The BAIR Blog
bair.berkeley.edu
