Автоматизируем машинное обучение с помощью ИИ-агентов

Решая соревнования на Kaggle начинаешь замечать паттерн. Baseline сделать просто: загрузить данные, запустить CatBoost или LightGBM, получить baseline метрику. Это занимает полчаса. Но чтобы попасть в топ решений, нужно перепробовать десятки вариантов препроцессинга, сотни комбинаций фичей и тысячи наборов гиперпараметров.

Существующие AutoML системы не сильно помогают. Они работают по фиксированному сценарию: пробуют предопределенный набор алгоритмов, выбирают лучший по метрике и возвращают результат. AutoGluon обучает несколько моделей и делает многоуровневый ансамбль, но каждый запуск начинается с нуля. TPOT генерирует pipeline через генетический алгоритм, но не учится на ошибках предыдущих запусков.

Главная проблема в том, что эти системы не рассуждают. Они не анализируют почему конкретный подход сработал или провалился. Они не адаптируются к специфике задачи. Они не накапливают опыт между запусками. Каждая новая задача для них как первая.

Человек работает иначе. Если дата-саентист видит несбалансированные классы, он сразу знает что нужна стратификация и подбор порога. Если видел похожую задачу раньше, применяет то, что сработало тогда. Если первая попытка провалилась, анализирует почему и пробует другой подход.

С появлением языковых моделей появилась возможность создать систему, которая работает ближе к человеку. LLM умеют анализировать данные, рассуждать о выборе методов и учиться на примерах. Но одна модель недостаточна. Она может пропустить очевидную ошибку или зациклиться на неправильном подходе. Нужна архитектура, которая позволит системе проверять саму себя и накапливать опыт.

Идея с двумя агентами

Первая версия была простой: один агент получает данные, обучает модель и возвращает предсказания. Проблема выявилась быстро. LLM иногда пропускает проверки данных или забывает обработать пропуски.

Решение пришло из reinforcement learning. В Actor-Critic методах один агент действует, второй оценивает эти действия. Почему бы не применить этот подход к AutoML?

Actor получает данные и набор инструментов для анализа, обработки и обучения моделей. Он исследует датасет, решает какие шаги нужны и генерирует решение. Critic смотрит на результат со стороны и проверяет, все ли сделано правильно. Если Critic находит проблемы, Actor получает обратную связь и пытается снова.

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

На схеме видно как работает система. Actor имеет доступ к пяти специализированным MCP серверам с инструментами для работы с данными и моделями. Critic работает без инструментов, только анализирует отчеты Actor. Между ними происходит итеративный обмен: решение от Actor идет на оценку к Critic, обратная связь возвращается обратно. После каждой итерации опыт сохраняется в память, откуда извлекается при работе с похожими задачами.

Инструменты для агента

LLM сам по себе может рассуждать, но чтобы он мог работать с данными, нужны инструменты. Я разделил их на несколько категорий: предпросмотр данных, статистический анализ, обработка и обучение моделей.

Например, инструмент для предпросмотра возвращает структурированную информацию:

{
    "shape": (150, 5),
    "columns": ["sepal_length", "sepal_width", "petal_length", "petal_width", "species"],
    "dtypes": {"sepal_length": "Float64", "species": "String"},
    "sample": [{"sepal_length": 5.1, "species": "setosa"}, ...]
}

Агент видит размерность, типы колонок и первые строки. Этого достаточно для принятия решений о дальнейших шагах.

Важный момент для обработки данных: агент должен применять одинаковые трансформации к train и test. Если на обучающей выборке категориальный признак закодировали как {"red": 0, "blue": 1}, на тестовой нужно использовать ту же кодировку. Для этого маппинги сохраняются в JSON файлы:

# Сохраняем маппинг для декодирования предсказаний
mapping_path = Path(output_dir) / f"{column}_mapping.json"
with open(mapping_path, "w") as f:
    json.dump(mapping, f)

Это критично для категориальной классификации. Если целевая переменная была закодирована в числа, модель вернет числа, а нужны исходные метки классов.

Каждый инструмент для обучения возвращает три вещи: путь к модели, путь к предсказаниям и метрики на train. Пути генерируются с timestamp и UUID, поэтому агент может экспериментировать с несколькими алгоритмами одновременно без конфликтов.

Протокол Model Context Protocol

Но что произойдет, если потребуется добавить больше инструментов, например, для специализированного домена? Когда количество инструментов превысит десяток, возникнут проблемы с их поддержкой, управлением и масштабированием.

Model Context Protocol, и в частности фреймворк FastMCP, решает эту задачу. MCP позволяет упаковать инструменты в отдельные серверы, которые агент вызывает по мере необходимости.

Я создал пять MCP серверов: file_operations для работы с файлами, data_preview для предпросмотра CSV, data_analysis для статистики, data_processing для трансформаций и machine_learning для обучения моделей и ансамблирования их предсказаний.

Оценка решений

Actor генерирует решение в виде структурированного отчета с четырьмя секциями: анализ данных, предобработка, обучение моделей и результаты. Critic получает этот отчет и анализирует каждую секцию отдельно.

Вместо одного LLM-судьи я использовал четыре специализированных. Первый проверяет качество анализа данных: изучил ли агент распределения, проверил ли пропуски, обнаружил ли аномалии. Второй смотрит на предобработку: правильно ли обработаны пропуски, корректно ли закодированы категории, нет ли утечки данных. Третий оценивает выбор модели и гиперпараметров. Четвертый анализирует результаты и общую методологию.

judges = [
    LLMJudge(rubric="Evaluate data_analysis: Is exploration thorough?"),
    LLMJudge(rubric="Evaluate preprocessing: Are steps appropriate?"),
    LLMJudge(rubric="Evaluate model_training: Is selection justified?"),
    LLMJudge(rubric="Evaluate results: Are metrics calculated correctly?"),
]

Каждый судья возвращает оценку от 0 до 1 с обоснованием. Средняя оценка сравнивается с порогом принятия (обычно 0.75). Если выше, решение принимается. Иначе Critic формирует обратную связь из комментариев всех судей и передает Actor для следующей итерации.

Этот подход работает стабильнее одного судьи. Один LLM может быть слишком строгим или пропустить очевидную ошибку. Четыре специализированных судьи сглаживают субъективность.

Изолированное рабочее пространство

Когда агент работает с файлами, нельзя давать ему доступ ко всей файловой системе. Нужна изоляция. Я создал выделенную директорию для каждой сессии в ~/.scald/actor/ с тремя поддиректориями: data для копий исходных данных, output для промежуточных файлов и workspace для моделей и предсказаний.

Исходные CSV копируются в data. Все инструменты работают только внутри этих директорий. Агент не может случайно перезаписать важные файлы или прочитать чужие данные.

После завершения работы все артефакты копируются в директорию сессии с timestamp, а workspace очищается. Можно спустя время открыть эту директорию и понять что именно делал агент: какие модели обучал (загрузить их из .pkl), какие метрики получил, какие шаги выполнил.

Память и обучение

После каждой итерации система сохраняет опыт. Отчет Actor и оценка Critic записываются в векторную базу ChromaDB. Когда агент получает новую задачу, система ищет похожие прошлые решения по семантическому сходству через embedding модель Jina.

self.mm.save(
    actor_solution=actor_solution,
    critic_evaluation=critic_evaluation,
    task_type=task_type,
    iteration=iteration,
)

# Поиск похожих
actor_memory, critic_memory = self.mm.retrieve(
    actor_report=actor_solution.report,
    task_type=task_type,
    top_k=5,
)

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

Интересно, что полезны не только успешные решения. Когда Critic говорит "вы забыли обработать пропуски", это ценная информация для будущих задач. Семантический поиск находит и такие случаи.

Основной цикл

Когда все компоненты готовы, остается собрать их вместе. Цикл работает до достижения максимального числа итераций или принятия решения Critic. На каждой итерации Actor решает задачу с учетом обратной связи, Critic оценивает решение, опыт сохраняется в память, извлекается релевантный контекст для следующей попытки.

Интересно наблюдать как Actor учится на обратной связи. Первая итерация обычно простая: базовая предобработка и одна модель. Critic находит проблемы: "вы не проверили баланс классов" или "не сделали feature engineering". Вторая итерация уже аккуратнее: агент добавляет недостающие шаги, пробует несколько моделей, делает ансамбль.

Проблемы при разработке

Кодирование целевой переменной

Первая версия ломалась на категориальной классификации. Агент кодировал target в числа, обучал модель, но забывал декодировать предсказания обратно. На выходе получались числа вместо меток классов.

Решение потребовало явных инструкций в системный промпт:

If you encode target column, you MUST DECODE predictions before returning.
Use decode_categorical_label with the mapping path from encoding step.

Уникальные имена файлов

Когда агент экспериментирует с несколькими моделями, файлы перезаписывают друг друга. Я пытался делегировать это агенту через prompt, но LLM не всегда генерирует уникальные имена. Правильное решение оказалось в том, чтобы делать это на уровне инструментов с timestamp и UUID.

Что получилось в итоге

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

Я протестировал систему на задачах из OpenML. На датасете christine она показала F1-score 0.743, обогнав Random Forest (0.713) на 4% и превзойдя AutoGluon с FLAML, которые вообще не справились с этой задачей. На cnae-9 результат составил 0.980 против 0.945 у лучшего конкурента FLAML, что на 3.5% лучше.

Были и провалы. На датасете Australian система показала 0.836, проиграв AutoGluon (0.860) и другим baseline методам. Интересно, что на blood-transfusion результат 0.756 оказался лучше Random Forest (0.712) и AutoGluon (0.734), но хуже FLAML (0.767).

Стоимость запуска варьируется от $0.14 до $3.43 в зависимости от сложности задачи и количества итераций. Время работы непредсказуемо: от минуты до получаса.

На самом деле, ценность результата не в том, что метрика лучше или хуже того или иного AutoML-фреймворка. Ценность в том, что современные LLM позволяют автоматизировать рутину более интеллектуально, а модульность MCP открывает путь к созданию экосистемы специализированных агентов. Это решает изначальную проблему "фиксированного сценария" классического AutoML: теперь систему можно адаптировать к любой задаче, просто подключив нужных агентов, сохраняя при этом единый цикл итеративного улучшения и накопления опыта.

Ограничения понятны. Система хороша для табличных данных с алгоритмами градиентного бустинга. Для deep learning или временных рядов нужны другие инструменты. Качество сильно зависит от размера базовой LLM.

Использование

Для тех, кто хочет попробовать, установка простая:

pip install scald

А использовать можно либо из терминала через CLI:

scald --train data/train.csv --test data/test.csv --target price --task-type regression

Либо через Python API:

from scald import Scald

scald = Scald(max_iterations=5)
predictions = await scald.run(
    train="data/train.csv",  # на вход можно подать как .csv, так и датафрейм
    test="data/test.csv",
    target="target_column",
    task_type="classification",
)

Для работы необходим API-ключ от провайдера, совместимого с OpenAI (например, OpenRouter). Также потребуется ключ от Jina для эмбеддингов в системе памяти (при регистрации сервис предоставляет большое количество бесплатных токенов).

Весь код упакован в библиотеку и доступен на GitHub