adu2021 AutoMind: Adaptive Knowledgeable Agent for Automated Data Science
Core Concept
AutoMind addresses limitations in rigid LLM-driven data science workflows through three innovations: an expert knowledge base curated from Kaggle competitions and peer-reviewed papers, an agentic tree search algorithm exploring multiple solution paths, and adaptive code generation that scales complexity with task difficulty. The system achieves 41.2% human participant surpass rate on official leaderboards and outperforms prior state-of-the-art by 8% on MLE-Bench.
Architecture Overview
- Expert Knowledge Base: 3,237 Kaggle posts + papers from KDD, ICLR, NeurIPS, ICML, EMNLP; hierarchically labeled for efficient retrieval
- Agentic Knowledgeable Tree Search: Explores solution tree with three action types—drafting (create branches), improving (refine), debugging (fix failures). Stochastic heuristic guides exploration
- Self-Adaptive Code Generation: Switches between one-pass (simple plans) and stepwise (complex plans) based on solution complexity scoring. AST checking validates syntax before execution
- Dynamic Solution Space: Manages plans, code, validation metrics; greedy selection optimizes top solutions while exploring alternatives
- Scalable Model Support: GPT-4o, o3-mini, DeepSeek-V3; reduces token cost by 9.6% vs baseline
Implementation
Step 1: Expert Knowledge Base Construction
import json
from collections import defaultdict
class KnowledgeBase:
"""
Curates and indexes expert knowledge from diverse sources.
Supports hierarchical query and efficient retrieval.
"""
def __init__(self):
self.knowledge_index = defaultdict(list)
self.category_hierarchy = {
'feature_engineering': [
'categorical_encoding', 'feature_scaling', 'polynomial_features',
'interaction_features', 'domain_specific_features'
],
'model_selection': [
'tree_based', 'linear_models', 'neural_networks',
'ensemble_methods', 'domain_adapted'
],
'hyperparameter_tuning': [
'grid_search', 'random_search', 'bayesian_optimization',
'learning_rate', 'regularization'
],
'data_preprocessing': [
'missing_value_handling', 'outlier_detection', 'normalization',
'dimensionality_reduction', 'sampling_strategies'
]
}
def ingest_kaggle_posts(self, post_list):
"""
Index Kaggle competition posts and solutions.
Post format: {'id', 'title', 'content', 'competition', 'votes'}
"""
for post in post_list:
# Extract entities and techniques from post
entities = self._extract_entities(post['content'])
for entity in entities:
category = self._categorize_entity(entity)
self.knowledge_index[category].append({
'source': 'kaggle',
'post_id': post['id'],
'competition': post['competition'],
'content': post['content'],
'votes': post['votes'],
'technique': entity
})
def ingest_academic_papers(self, papers):
"""
Index papers from top-tier ML conferences.
Paper format: {'title', 'abstract', 'pdf_url', 'venue', 'year'}
"""
for paper in papers:
# Extract methodologies and benchmarks
methodologies = self._extract_methodologies(paper['abstract'])
for method in methodologies:
category = self._categorize_method(method)
self.knowledge_index[category].append({
'source': 'academic',
'paper_title': paper['title'],
'venue': paper['venue'],
'year': paper['year'],
'abstract': paper['abstract'],
'methodology': method
})
def retrieve(self, query, category=None, top_k=5):
"""
Retrieve relevant knowledge items.
Supports hierarchical category filtering.
"""
results = []
# Determine relevant categories
target_categories = [category] if category else list(self.knowledge_index.keys())
for cat in target_categories:
for item in self.knowledge_index[cat]:
# Simple relevance scoring: TF-IDF-like
relevance = self._score_relevance(query, item)
results.append((relevance, item))
# Sort by relevance and return top-k
results.sort(key=lambda x: x[0], reverse=True)
return [item for _, item in results[:top_k]]
def _extract_entities(self, text):
"""Extract technical entities from text."""
# Placeholder: would use NER or keyword extraction
return []
def _categorize_entity(self, entity):
"""Determine category for entity."""
if 'encode' in entity.lower():
return 'feature_engineering'
elif 'xgb' in entity.lower() or 'random' in entity.lower():
return 'model_selection'
return 'feature_engineering'
def _extract_methodologies(self, abstract):
"""Extract methodologies from paper abstract."""
return []
def _categorize_method(self, method):
"""Determine category for method."""
return 'feature_engineering'
def _score_relevance(self, query, item):
"""Score relevance of item to query."""
return 0.5 # Placeholder
Step 2: Agentic Tree Search Algorithm
import torch
from dataclasses import dataclass
from typing import Optional, List
@dataclass
class SolutionNode:
"""Represents a solution in the search tree."""
plan: str # High-level approach
code: str # Implementation
validation_metric: float # Accuracy/F1/RMSE
parent: Optional['SolutionNode'] = None
children: List['SolutionNode'] = None
class AgenticTreeSearch:
"""
Explores solution space using three action types:
- Draft: Create new solution branches
- Improve: Refine existing solutions
- Debug: Fix execution failures
"""
def __init__(self, knowledge_base, model_name='gpt-4o'):
self.kb = knowledge_base
self.model = model_name
self.root = None
self.best_solution = None
self.best_metric = float('-inf')
def search(self, task_description, max_iterations=10):
"""
Main search loop: iteratively explore and refine solutions.
"""
# Initialize root: generate initial solution
initial_solution = self._draft_solution(task_description)
self.root = SolutionNode(
plan=initial_solution['plan'],
code=initial_solution['code'],
validation_metric=0.0
)
current_node = self.root
for iteration in range(max_iterations):
print(f"Iteration {iteration}: Best metric={self.best_metric:.4f}")
# Choose action: draft, improve, or debug
action = self._select_action(current_node)
if action == 'draft':
# Create multiple new solution branches
for _ in range(2):
new_solution = self._draft_solution(
task_description,
parent_context=current_node
)
new_node = SolutionNode(
plan=new_solution['plan'],
code=new_solution['code'],
validation_metric=0.0,
parent=current_node
)
current_node.children.append(new_node)
elif action == 'improve':
# Refine current best solution
improved = self._improve_solution(
current_node.plan,
current_node.code
)
current_node.plan = improved['plan']
current_node.code = improved['code']
elif action == 'debug':
# Fix execution errors
if current_node.validation_metric == float('-inf'): # Error state
debugged = self._debug_solution(
current_node.code,
error_log=current_node.error_log
)
current_node.code = debugged['code']
# Evaluate current node
metric = self._execute_and_validate(current_node.code)
current_node.validation_metric = metric
# Update best solution
if metric > self.best_metric:
self.best_metric = metric
self.best_solution = current_node
print(f"New best: {metric:.4f}")
# Move to next promising node (greedy selection)
current_node = self._select_next_node()
return self.best_solution
def _draft_solution(self, task, parent_context=None):
"""Generate new solution draft using knowledge base."""
# Retrieve relevant knowledge
knowledge = self.kb.retrieve(task, top_k=3)
draft_prompt = f"""
Task: {task}
Relevant knowledge:
{self._format_knowledge(knowledge)}
Generate a solution with:
1. High-level plan (strategy, approach)
2. Python code implementation
Output JSON:
{{"plan": "...", "code": "..."}}
"""
response = self._llm_call(draft_prompt)
return self._parse_solution(response)
def _improve_solution(self, plan, code):
"""Refine solution iteratively."""
improve_prompt = f"""
Current plan: {plan}
Current code:
{code}
Suggest improvements:
1. Efficiency enhancements
2. Code quality improvements
3. Better ML practices
Provide improved plan and code.
"""
response = self._llm_call(improve_prompt)
return self._parse_solution(response)
def _debug_solution(self, code, error_log):
"""Fix execution errors."""
debug_prompt = f"""
Code:
{code}
Error:
{error_log}
Fix the error and provide corrected code.
"""
response = self._llm_call(debug_prompt)
return {'code': response}
def _select_action(self, current_node):
"""Select next action based on stochastic heuristic policy."""
# Simple policy: draft if few children, improve if no error, debug if error
if len(current_node.children) < 2:
return 'draft'
elif current_node.validation_metric == float('-inf'):
return 'debug'
else:
return 'improve'
def _select_next_node(self):
"""Greedy selection: move to best unexplored node."""
candidates = []
for node in self._get_all_nodes():
if node.validation_metric == float('-inf'): # Unevaluated
candidates.append(node)
if candidates:
return candidates[0]
# If all evaluated, pick best
return max(self._get_all_nodes(),
key=lambda n: n.validation_metric)
def _execute_and_validate(self, code):
"""Execute code and return validation metric."""
# Placeholder: would execute code and measure performance
return 0.75
def _llm_call(self, prompt):
"""Call LLM."""
return "response"
def _format_knowledge(self, knowledge):
"""Format knowledge for prompt."""
return '\n'.join([f"- {k['technique']}" for k in knowledge])
def _parse_solution(self, response):
"""Parse LLM response into solution dict."""
return {'plan': 'plan', 'code': 'code'}
def _get_all_nodes(self):
"""Collect all nodes in tree."""
def dfs(node):
if node is None:
return []
result = [node]
for child in (node.children or []):
result.extend(dfs(child))
return result
return dfs(self.root)
Step 3: Self-Adaptive Code Generation
import ast
class AdaptiveCodeGenerator:
"""
Generates code that adapts complexity to task difficulty.
Switches between one-pass (simple) and stepwise (complex).
"""
def generate_code(self, plan, task_complexity_score):
"""
Generate code with complexity matching task difficulty.
Complexity score: 0-1 (0=simple, 1=very complex)
"""
if task_complexity_score < 0.4:
return self._generate_one_pass_code(plan)
else:
return self._generate_stepwise_code(plan)
def _generate_one_pass_code(self, plan):
"""
Simple implementation: single coherent script.
For straightforward tasks (basic preprocessing, standard models).
"""
code_template = f"""
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
# Load and preprocess
X = load_features()
y = load_labels()
X_scaled = StandardScaler().fit_transform(X)
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2)
# Train model
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)
# Evaluate
accuracy = model.score(X_test, y_test)
print(f"Accuracy: {{accuracy:.4f}}")
"""
return code_template
def _generate_stepwise_code(self, plan):
"""
Complex implementation: decomposed into substeps.
For intricate tasks (feature engineering, ensemble, hyperparameter tuning).
"""
code_template = f"""
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler, PolynomialFeatures
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.model_selection import cross_val_score
# Step 1: Advanced feature engineering
def engineer_features(X):
poly = PolynomialFeatures(degree=2, include_bias=False)
X_poly = poly.fit_transform(X)
return X_poly
# Step 2: Normalize
X = load_features()
y = load_labels()
X_eng = engineer_features(X)
X_scaled = StandardScaler().fit_transform(X_eng)
# Step 3: Ensemble training with cross-validation
models = [
RandomForestClassifier(n_estimators=100),
GradientBoostingClassifier(n_estimators=100)
]
scores = []
for model in models:
cv_scores = cross_val_score(model, X_scaled, y, cv=5)
scores.append(cv_scores.mean())
print(f"{{model.__class__.__name__}}: {{cv_scores.mean():.4f}}")
print(f"Best model score: {{max(scores):.4f}}")
"""
return code_template
def validate_and_refine(self, code):
"""
Validate code syntax and refine if needed.
Uses AST parsing to catch errors before execution.
"""
try:
ast.parse(code)
return code, True
except SyntaxError as e:
print(f"Syntax error: {e}")
# Attempt automated fix
fixed_code = self._fix_syntax_error(code, e)
# Recursive validation
return self.validate_and_refine(fixed_code)
def _fix_syntax_error(self, code, error):
"""Attempt to fix common syntax errors."""
# Placeholder: would use heuristics to fix
return code
def score_plan_complexity(plan):
"""
Score plan complexity 0-1 using professional rubric.
Considers: feature engineering scope, model selection, hyperparameter tuning.
"""
complexity_factors = {
'feature_engineering': 0.3,
'model_selection': 0.3,
'hyperparameter_tuning': 0.2,
'ensemble_methods': 0.2
}
total_score = 0.0
for factor, weight in complexity_factors.items():
if factor in plan.lower():
total_score += weight
return min(total_score, 1.0)
Practical Guidance
Knowledge Base Curation:
- Kaggle: Focus on recent competitions; weight by votes/popularity
- Papers: Filter for methodology papers (not just benchmarks); include implementation details
- Organization: Hierarchical categorization enables targeted retrieval
- Update frequency: Monthly updates with new competition solutions
Search Hyperparameters:
- Max iterations: 8-12 (balance quality vs cost)
- Branching factor: 2-3 new solutions per draft action
- Metric threshold: Accept solution if metric > 0.7 of domain SOTA
Cost Optimization:
- Knowledge retrieval: TF-IDF indexes reduce LLM query overhead
- Code validation: AST parsing catches errors before execution (saves rerun cost)
- Model selection: Use cheaper model (o3-mini) for most steps; GPT-4o for complex planning
When to Use AutoMind:
- Structured tabular data problems (heterogeneous features, missing values)
- Limited time budgets (20-hour leaderboard competitions)
- Ensemble/stacking opportunities (AutoMind explores combinations)
- Transfer learning scenarios (knowledge base accelerates solution discovery)
Reference
- Agentic tree search: Combines exploration (drafting) with exploitation (improving best solutions)
- Stochastic heuristic policy: Non-deterministic action selection enables diverse solution exploration
- AST validation: Abstract Syntax Tree parsing catches errors without execution overhead