Neural Bandit Based Optimal LLM Selection for a Pipeline of Tasks

2508.09958v1 cs.CL, cs.LG 2025-08-15
Авторы:

Baran Atalar, Eddie Zhang, Carlee Joe-Wong

Резюме на русском

## Контекст В последние годы высокая популярность больших языковых моделей (LLMs) вызвана их универсальностью и применением во многих областях. Однако, при выполнении сложных задач, одна LLM может не обеспечить достаточного качества вывода. Чтобы улучшить результаты, задачи предлагаются разбивать на меньшие подзадачи, которые затем могут быть выполнены разными LLMs, каждая из которых оптимизирована для конкретной подзадачи. Например, в области медицины могут быть выделены LLMs для сводки медицинских записей, проверки итогов, и извлечения диагноза. Большинство существующих методов LLM-селекции или маршрутизации не учитывают зависимость выхода одного этапа от другого, что влияет на качество и стоимость решения задачи. Мы предлагаем алгоритм, основанный на контекстной бандитной структуре, чтобы решить эту проблему, учитывая уникальные связи между задачами и подзадачами. ## Метод Мы предлагаем алгоритм, основанный на моделировании LLM-поведения с помощью контекстно-нейронаучных моделей и бандитной структуры. Этот алгоритм обучает нейронные сети, которые предсказывают успешность каждой LLM на каждой подзадаче. Модели обучаются онлайн, что позволяет им оптимизировать выбор LLMs даже в условиях отсутствия исторических данных о их производительности. Ключевой идеей является моделирование зависимостей между подзадачами, чтобы предоставить улучшенные выборки LLMs в зависимости от результатов предыдущих этапов. Мы также используем нейронные модели, которые могут обучаться в реальном времени, чтобы учитывать изменения поведения LLMs в зависимости от контекста. ## Результаты Мы проверили наш алгоритм на двух наборах данных: телекоммуникационных запросов и медицинских диагнозах. Наши эксперименты показали, что онлайн-обучение моделей позволяет создавать более точные предсказания успешности LLMs на каждой подзадаче. Мы сравнили нашу модель с другими LLM-селекционными алгоритмами и показали, что у нас есть значительные выигрыши в качестве решений и снижении стоимости. Благодаря нейронной модели мы можем эффективно учитывать связи между подзадачами и выбирать LLMs, которые создают меньшую зависимость в качестве решений. ## Значимость Наш алгоритм может быть применен для различных приложений, которые требуют последовательного использования LLMs в различных подзадачах. Он намного эффективнее существующих методов, так как учитывает связи между этапами и уменьшает уровень ошибок в выборе LLMs. Это может привести к более эффективному использованию ресурсов, снижению затрат и улучшению общего качества решений. Мы также планируем расширить исследования, включая применение нашего подхода

Abstract

With the increasing popularity of large language models (LLMs) for a variety of tasks, there has been a growing interest in strategies that can predict which out of a set of LLMs will yield a successful answer at low cost. This problem promises to become more and more relevant as providers like Microsoft allow users to easily create custom LLM "assistants" specialized to particular types of queries. However, some tasks (i.e., queries) may be too specialized and difficult for a single LLM to handle alone. These applications often benefit from breaking down the task into smaller subtasks, each of which can then be executed by a LLM expected to perform well on that specific subtask. For example, in extracting a diagnosis from medical records, one can first select an LLM to summarize the record, select another to validate the summary, and then select another, possibly different, LLM to extract the diagnosis from the summarized record. Unlike existing LLM selection or routing algorithms, this setting requires that we select a sequence of LLMs, with the output of each LLM feeding into the next and potentially influencing its success. Thus, unlike single LLM selection, the quality of each subtask's output directly affects the inputs, and hence the cost and success rate, of downstream LLMs, creating complex performance dependencies that must be learned and accounted for during selection. We propose a neural contextual bandit-based algorithm that trains neural networks that model LLM success on each subtask in an online manner, thus learning to guide the LLM selections for the different subtasks, even in the absence of historical LLM performance data. Experiments on telecommunications question answering and medical diagnosis prediction datasets illustrate the effectiveness of our proposed approach compared to other LLM selection algorithms.

Ссылки и действия