Что такое **agent distillation** (обучение маленького агента на траекториях большого)?

Краткий тезис

Agent distillation — это техника сжатия AI-агента, при которой маленькая модель (например, 7B) обучается имитировать поведение большого агента (например, 70B плюс инструменты). Agent|Большой агент выполняет задачи, записывая свои траектории (цепочки мыслей, вызовы функций, действия), а Agent|маленький агент учится воспроизводить их через behavior cloning (поведенческое клонирование) и DPO (Direct Preference Optimization). Результат: Agent|маленький агент работает в 5–10 раз быстрее, сохраняя 80–90% качества большого, что делает его пригодным для дешёвого инференса в production с низкой задержкой.

1. Термин: Agent distillation (дистилляция агента)

Agent distillation — это процесс переноса знаний от крупного, высококачественного агента (teacher) на более компактного агента (student). В отличие от классической дистилляции LLM, где модель учится на логарифмах вероятностей токенов (soft labels), при дистилляции агента объектом обучения являются траектории — последовательности действий, наблюдений и рассуждений, которые агент выполняет для решения задачи.

Термин агент здесь означает LLM, расширенную возможностью вызывать внешние инструменты (API, базы данных, веб-поиск) и имеющую цикл рассуждения-действия (например, ReAct). Teacher-агент обычно строится на самой мощной модели (GPT-4, Claude 3.5 Sonnet, Llama 3.1 405B) и использует полный арсенал инструментов. Student-агент — это гораздо меньшая модель (Llama 3.1 8B, GPT-4o mini, Qwen 2.5 7B), которая должна научиться действовать так же эффективно, но без затрат teacher.

2. Мотивация: зачем сжимать агента?

Производственные системы AI всё чаще используют агентов для автономного выполнения цепочек действий (бронирование, поддержка клиентов, код-ревью). Запуск большого агента дорог и медленен:

Agent distillation позволяет развернуть маленького агента на краевых устройствах (edge), в высоконагруженных системах или в приложениях, где критична задержка (чат-боты реального времени, роботизированные процессы).

3. Процесс сбора траекторий

Первый этап — генерация датасета траекторий. Траектория — это полная запись взаимодействия агента со средой: начальное наблюдение, последовательность шагов (мысль → действие → наблюдение), и финальный ответ. Большой агент запускается на наборе задач (например, 10 000 запросов из домена), и для каждой задачи сохраняется:

{
  "task": "Найди на сайте последний пост о RAG и отправь ссылку в Slack",
  "trajectory": [
    {"role": "agent", "content": "Нужно найти пост... Использую search(query='RAG')"},
    {"role": "tool", "content": "Результаты: [{'title':'Руководство по RAG','url':'...'}]"},
    {"role": "agent", "content": "Выбираю первый результат. Открываю страницу."},
    {"role": "tool", "content": "Страница загружена, найден текст поста."},
    {"role": "agent", "content": "Отправляю ссылку в Slack каналу #ml-team. Вызываю slack_send(url='...')"},
    {"role": "tool", "content": "Сообщение отправлено."},
    {"role": "agent", "content": "Готово. Сообщаю пользователю."}
  ],
  "result": "success"
}

Такие датасеты могут содержать тысячи траекторий. Для повышения разнообразия используются разные начальные состояния и вариации инструментов.

4. Метод 1: Behavioral Cloning (поведенческое клонирование)

Самый простой способ дистилляции — поведенческое клонирование (BC). Маленький агент обучается с учителем (SFT) на траекториях teacher'а, как если бы это был диалог. Для каждого шага в траектории вход — это предшествующий контекст (наблюдения и предыдущие действия), а выход — следующее действие или мысль агенту.

Формат обучения — обычный next-token prediction с кросс-энтропией. Однако есть нюанс: агент использует вызовы инструментов, которые в обучении LLM нестандартны. Поэтому архитектуру student-агента часто расширяют специальными токенами для инструментов (например, <tool_call> ) и обучают на смеси обычного текста и инструментальных вызовов.

Плюсы BC

• Простота реализации, стандартная loss.
• Быстрое обучение (несколько часов на 8 GPU).

Минусы BC

• Компаундные ошибки: маленький агент, столкнувшись с незнакомым наблюдением на раннем шаге, начинает отклоняться от траектории teacher'а, и ошибки накапливаются.
• Нет обратной связи по качеству финального результата.

5. Метод 2: Distillation с Preference Optimization (DPO)

Чтобы преодолеть компаундные ошибки, используют Direct Preference Optimization (DPO). Вместо того чтобы просто копировать действия, маленький агент учится предпочитать траектории teacher'а своим собственным — более слабым.

Процесс:

1. На том же датасете задач запускают student-агента (до обучения) и записывают его траектории.
2. Для каждой задачи создают пару:
   • chosen — траектория teacher’а (успешная, качественная);
   • rejected — траектория student’а (часто ошибочная или менее эффективная).
3. Обучают student’а с loss DPO, который увеличивает вероятность chosen и уменьшает вероятность rejected.

Формула DPO (упрощённо):

L_DPO = - E[ log σ( β * (log π_θ(chosen) - log π_ref(chosen) - log π_θ(rejected) + log π_ref(rejected)) ) ]

где π_θ — student модель, π_ref — reference модель (замороженная, обычно student до обучения), β — температура.

Плюсы DPO

• Учитывает финальный успех задачи, а не только пошаговое соответствие.
• Меньше компаундных ошибок — student учится не застревать в неправильных состояниях.

Минусы DPO

• Нужны траектории student’а, что удваивает затраты на генерацию.
• Чувствителен к качеству пар: если rejected не сильно хуже chosen, обучение неэффективно.

6. Сравнение методов дистилляции агента

На практике чаще всего применяют двухэтапный процесс: сначала BC (чтобы student научился общей структуре), затем DPO (чтобы отточить предпочтения и исправить типовые ошибки).

7. Оценка качества маленького агента

Для оценки используют те же метрики, что и для любого агента, но с акцентом на сравнение с teacher:

• Success Rate — доля задач, выполненных полностью (action success + correct final answer).
• Step Match Rate — процент шагов, совпадающих с teacher’ом (для диагностики).
• Latency — среднее время на задачу (должно быть в 3–10 раз меньше).
• Cost per task — стоимость инференса (в токенах).
• Faithfulness — для задач с retrieval: не выдумывает ли student факты.

Важно тестировать на задачах, не входящих в обучающий датасет, чтобы оценить обобщение.

8. Применение в production

После дистилляции маленький агент может быть развёрнут:

• На edge-устройствах (мобильные приложения, IoT) с ограниченной памятью.
• Высоконагруженные чат-боты (поддержка клиентов) — 1000+ запросов/мин.
• Мультиагентные системы, где много мелких агентов работают параллельно.
• Кеш-агенты — student используется для типовых запросов, а teacher вызывается только для сложных случаев (отказ по уверенности, confidence threshold).

Пример: компания Discord заменила GPT-4 агента для модерации на дистиллированную версию Llama 3 8B, снизив latency с 8 секунд до 1.5 секунд при сохранении 92% точности.

9. Ограничения и вызовы

• Качество teacher’а — если teacher сам ошибается, student наследует ошибки.
• Смещение (bias) — дистиллированный агент может повторять предвзятости teacher’а, а также терять способность к творческому решению нестандартных задач.
• Длительное обучение — для сложных сред (веб-навигация, робототехника) генерация траекторий может быть дорогой из-за необходимости запускать окружение.
• Недостаток разнообразия — если датасет траекторий узок, student не обобщает на невидимые сценарии.
• Отсутствие гарантий — student может придумывать новые, неоптимальные стратегии, поэтому необходим мониторинг в production.

10. Связь с другими техниками

Agent distillation тесно связана с:

• Дистилляцией LLM — снизу лежат те же принципы переноса знаний, но объект обучения другой.
• RLHF (Reinforcement Learning from Human Feedback) — DPO — более простая альтернатива RLHF для предпочтений, но в случае агента предпочтения строятся на траекториях.
• ReAct и Tool use — дистиллируемый агент должен уметь вызывать инструменты, что требует специального форматирования.
• Agent evaluation — без правильной метрики (success rate) дистилляция бессмысленна.
• Memory for agents — маленький агент может использовать более ограниченную память, поэтому дистилляция должна включать стратегии сжатия контекста (например, summarization).

Пет-проект для закрепления

Задача Обучить маленького агента (Qwen 2.5 7B) имитировать GPT-4o для задачи календарного планирования (создание событий, проверка конфликтов) с использованием календарного API.

Инструменты

• Hugging Face Transformers + TRL (SFTTrainer, DPOTrainer)
• LangChain / OpenAI API для teacher’а
• Датасет: 1000 синтетических запросов на планирование (сгенерировать с GPT-4o).
• Окружение для агента: Mock-календарь (простой Flask-сервер, эмулирующий Google Calendar API).

Шаги:

1. Создать teacher-агента (GPT-4o + календарный инструмент) и запустить на 1000 задачах, сохранить траектории (каждый шаг: мысль, action, observation).
2. Обучить student-агента (Qwen 2.5 7B) с SFT на этих траекториях (behavioral cloning) — 2 эпохи.
3. Запустить student-агента на тех же задачах (без обучения) и записать его траектории — они будут rejected.
4. Применить DPO: пары (teacher траектория, student траектория) для каждой задачи, обучить ещё 1 эпоху.
5. Оценить success rate на 200 новых задачах, измерить latency.

Ожидаемый результат
Success rate student — 85% от teacher (teacher ~98%), latency снижается с 6–8 секунд до 1–2 секунд. Стоимость одного запроса уменьшается в 10 раз.

Связь с другими вопросами

Навигация

• Предыдущий: 583
• Следующий: 585
• Индекс: 00. Индекс разборов