[Paper Review] Octopus v2: On-device language model for super agent

Wei Chen, Zhiyuan Li|arXiv (Cornell University)|Apr 2, 2024

Multi-Agent Systems and Negotiation4 citations

TL;DR

Octopus v2 introduces a 2B-parameter on-device language model that achieves GPT-4-level accuracy in function calling while reducing latency by 35x and context length by 95% through fine-tuning with functional tokens. The method enables production-ready, low-latency AI agents on edge devices like smartphones, outperforming both GPT-4 and Llama-7B with RAG in speed and efficiency.

ABSTRACT

Language models have shown effectiveness in a variety of software applications, particularly in tasks related to automatic workflow. These models possess the crucial ability to call functions, which is essential in creating AI agents. Despite the high performance of large-scale language models in cloud environments, they are often associated with concerns over privacy and cost. Current on-device models for function calling face issues with latency and accuracy. Our research presents a new method that empowers an on-device model with 2 billion parameters to surpass the performance of GPT-4 in both accuracy and latency, and decrease the context length by 95\%. When compared to Llama-7B with a RAG-based function calling mechanism, our method enhances latency by 35-fold. This method reduces the latency to levels deemed suitable for deployment across a variety of edge devices in production environments, aligning with the performance requisites for real-world applications.

Motivation & Objective

• To enable high-accuracy, low-latency function calling on edge devices using a compact on-device language model.
• To address privacy and cost concerns associated with cloud-based LLMs like GPT-4 by deploying models locally on smartphones and edge devices.
• To reduce context length and inference latency for on-device AI agents, making them viable for real-world, battery-constrained environments.
• To develop a scalable method for fine-tuning small models with functional tokens to match or exceed the performance of large cloud-based models.
• To support efficient deployment across diverse applications via LoRA adapters, enabling application-specific function calling without full retraining.

Proposed method

• The model is fine-tuned from Gemma-2B using functional tokens—special tokens like <nexa_0> to <nexa_N-1> that represent specific APIs or functions.
• A weighted cross-entropy loss function is applied, assigning higher weights to special functional tokens to address class imbalance during training and improve convergence.
• Full model fine-tuning and LoRA (Low-Rank Adaptation) are compared, with LoRA enabling efficient, application-specific adaptation while maintaining high accuracy.
• The framework supports parallel and nested function calls by generating 4K data points per API to maintain single-call performance levels.
• Functional tokens are integrated into both the tokenizer and the model’s attention heads, allowing the model to learn direct mappings between natural language queries and function calls.
• The approach reduces context length by over 95% by replacing lengthy function descriptions with compact tokens, significantly improving inference speed and efficiency.

Experimental results

Research questions

• RQ1 Can a 2B-parameter on-device language model achieve function calling performance comparable to GPT-4 in accuracy and latency?
• RQ2 How can functional tokens be effectively integrated into a language model to improve function calling accuracy and reduce context length?
• RQ3 What is the impact of weighted loss functions on convergence and performance when training with rare functional tokens?
• RQ4 Can LoRA-based fine-tuning maintain high accuracy while enabling efficient, application-specific deployment of on-device AI agents?
• RQ5 To what extent can on-device models reduce inference latency and context length compared to RAG-augmented cloud models?

Key findings

• Octopus v2 achieves GPT-4-level accuracy in function calling while reducing latency by 35-fold compared to Llama-7B with RAG.
• The model reduces context length by over 95% during inference, enabling significantly more function calls per battery charge on mobile devices.
• For iPhone deployment, the model supports 37 times more function calls than a comparable 7B model due to reduced context length and energy efficiency.
• LoRA fine-tuning maintains high accuracy with only a minor performance drop, enabling efficient, application-specific adaptation without full retraining.
• The use of weighted cross-entropy loss improves convergence during training, especially for rare functional tokens, without affecting final inference performance.
• The model enables production-ready, low-latency AI agents on edge devices, overcoming previous limitations of on-device models in speed and accuracy.