LLM-AutoDP: The Framework That Lets an LLM Agent Design Its Own Training Data Pipeline

# LLM-AutoDP: Group Relative Comparison feedback mechanism
# Reconstructed from arXiv:2601.20375 Section 3.2
# The core design: LLM learns from RELATIVE performance, not absolute scores

import json
from dataclasses import dataclass
from typing import Optional

@dataclass
class DPStrategy:
    """A data processing pipeline: ordered sequence of operator groups."""
    operators: list[str]    # e.g., ["cleaning", "optimization"]
    description: str        # human-readable description for LLM prompt
    score: Optional[float] = None     # eval(strategy) result
    delta: Optional[float] = None     # score - baseline (improvement vs no processing)


class LLMAutoDPAgent:
    """
    Meta-agent that generates and refines data processing strategies.
    
    Key design: the agent itself is never fine-tuned.
    ← All learning happens through in-context prompt construction.
      This means: any capable LLM works as the meta-agent.
      No meta-agent training cost. No meta-agent memory between runs.
      The prompt IS the memory.
    """

    OPERATOR_POOL = {
        "cleaning": [
            "Apply MinHashLSH to eliminate duplicate samples",
            "Remove low-level noise such as HTML tags",
            "Maintain special character ratio within a specific range",
            "Keep token count within a specific range",
            "Keep n-gram repetition ratio within a specific range",
        ],
        "optimization": [
            "Optimize questions in samples via LLM rewrite",
            "Optimize answers in samples via LLM rewrite",
            "Optimize both questions and answers via LLM rewrite",
        ],
        "generation": [
            "Generate missing questions based on available answers",
            "Generate missing answers based on available questions",
            "Generate complete Q+A pairs from context",
        ],
        "selection": [
            "Select high-quality samples based on gradient information",
        ],
    }

    def __init__(self, llm_client, baseline_score: float):
        self.llm = llm_client
        self.baseline_score = baseline_score   # eval(no processing)
        self.history: list[list[DPStrategy]] = []

    def build_initial_prompt(self) -> str:
        """
        Round 1: no history, only operator descriptions.
        The LLM is asked to generate N diverse strategies.
        ← 'Diverse' is explicit: avoid redundant exploration in early rounds
          Start with small compositions (1-2 groups) to understand individual effects
        """
        operator_desc = json.dumps(self.OPERATOR_POOL, indent=2)
        return f"""
OBJECTIVE: Design data processing strategies to maximize fine-tuned model performance.

AVAILABLE OPERATORS:
{operator_desc}

SEARCH SPACE: Select 1-4 groups from [cleaning, optimization, generation, selection].
Order matters. Explore different combinations and orderings.

INSTRUCTIONS:
1. Generate N distinct strategies, each specifying an ordered list of operator groups
2. Start with smaller compositions (1-2 groups) to understand individual effects
3. Each strategy must differ in composition or ordering from others
4. Include a rationale for each strategy choice

Output each strategy as: {{"operators": [...], "rationale": "..."}}
"""

    def build_iterative_prompt(self, round_strategies: list[list[DPStrategy]]) -> str:
        """
        Round t+1: inject all prior round strategies and their improvement scores.
        
        ← THIS is the trick: GROUP RELATIVE COMPARISON
          All strategies from a round are shown TOGETHER, not individually.
          The LLM can compare: "strategy A (Δ=+0.15) beat strategy B (Δ=-0.03)"
          It infers: "cleaning before optimization worked; the reverse didn't"
          
          This is more informative than showing each strategy separately:
          with individual feedback, the LLM cannot identify RELATIVE differences
          that arise from ordering and composition interactions.
        """
        history_text = ""
        for round_idx, strategies in enumerate(round_strategies):
            history_text += f"\n=== ROUND {round_idx + 1} RESULTS ===\n"
            for s in strategies:
                status = "IMPROVED" if s.delta > 0 else "DEGRADED"
                history_text += (
                    f"Strategy: {s.operators}\n"
                    f"  Performance change vs baseline: {s.delta:+.4f} ({status})\n"
                    f"  Description: {s.description}\n\n"
                )

        return f"""
{self.build_initial_prompt()}

PREVIOUS RESULTS (baseline = no processing, score = {self.baseline_score:.4f}):
{history_text}

ANALYSIS REQUIRED:
- Compare strategies: which combinations/orderings performed best?
- Identify patterns: what makes high-scoring strategies different?
- Avoid configurations similar to low-scoring strategies

NEXT ROUND: Generate N new strategies that build on the above analysis.
If no further improvement is possible, output: {{"terminate": true, "best_strategy": [...]}}
"""

    def run_round(self, round_num: int) -> list[DPStrategy]:
        if round_num == 1:
            prompt = self.build_initial_prompt()
        else:
            prompt = self.build_iterative_prompt(self.history)

        response = self.llm.complete(prompt)
        strategies = self._parse_strategies(response)
        return strategies

    def converged(self, strategies: list[DPStrategy]) -> bool:
        """
        Agent-driven convergence: the LLM itself decides when to stop.
        ← Not a fixed iteration count: the agent monitors improvement rate
          If it judges no further significant improvement is likely, it terminates
          and returns its best strategy found.
        """
        return any(getattr(s, 'terminate', False) for s in strategies)

# LLM-AutoDP: Three acceleration techniques for fast strategy evaluation
# Reconstructed from arXiv:2601.20375 Section 3.3
# Without these, each round requires full fine-tuning: too slow for 10+ rounds

import numpy as np
from sklearn.linear_model import LogisticRegression
from collections import OrderedDict

# ─── TECHNIQUE 1: Distribution-Preserving Sampling (DPS) ─────────────────────
def distribution_preserving_sample(
    dataset: list[dict],
    sample_ratio: float = 0.3,
) -> list[dict]:
    """
    Sample a subset of training data while preserving class distribution.
    Fine-tune on this subset instead of the full training set.
    
    ← WHY: Each strategy evaluation requires a full fine-tuning run.
      With 5 strategies per round × 10 rounds = 50 fine-tuning runs.
      Full FT at each: prohibitive. DPS reduces FT time by using a subset.
      
    ← Key insight: distributional integrity matters more than raw size.
      A 30% sample with correct class proportions produces reliable evals.
      A 30% sample without stratification can over-represent rare classes
      and produce misleading strategy rankings.
    """
    # Stratify by domain/quality label if available; else by quantile bucket
    if 'label' in dataset[0]:
        label_groups = {}
        for item in dataset:
            label_groups.setdefault(item['label'], []).append(item)
        sample = []
        for group_items in label_groups.values():
            n_sample = max(1, int(len(group_items) * sample_ratio))
            sample.extend(np.random.choice(group_items, n_sample, replace=False))
    else:
        n_sample = int(len(dataset) * sample_ratio)
        sample = list(np.random.choice(dataset, n_sample, replace=False))

    return sample


# ─── TECHNIQUE 2: Processing Target Selection (PTS) ──────────────────────────
class ProcessingTargetSelector:
    """
    Binary classifier to identify low-quality samples.
    Only apply data processing to the low-quality subset.
    
    ← WHY: Processing ALL samples is wasteful.
      High-quality samples don't benefit from cleaning/optimization.
      Applying optimization operators to already-good samples can degrade them.
      
    ← THIS is the trick: train a cheap binary classifier F on a small labeled set.
      Apply F to the full dataset → split into D_high and D_low.
      Run expensive LLM-based operators only on D_low.
      Merge processed D_low with unmodified D_high.
      Cost reduction: proportional to the fraction of low-quality samples.
    """

    def __init__(self):
        self.classifier = LogisticRegression()
        self.fitted = False

    def fit(self, samples: list[dict], quality_labels: list[int]) -> None:
        """
        quality_labels: 0 = low quality (needs processing), 1 = high quality
        Features: perplexity proxy, length, special char ratio, n-gram repetition
        """
        features = self._extract_features(samples)
        self.classifier.fit(features, quality_labels)
        self.fitted = True

    def split(self, dataset: list[dict]) -> tuple[list[dict], list[dict]]:
        """Split dataset into (D_low, D_high) based on quality prediction."""
        if not self.fitted:
            raise ValueError("Classifier not fitted. Call fit() first.")
        features = self._extract_features(dataset)
        predictions = self.classifier.predict(features)
        d_low  = [s for s, p in zip(dataset, predictions) if p == 0]
        d_high = [s for s, p in zip(dataset, predictions) if p == 1]
        return d_low, d_high

    def _extract_features(self, samples: list[dict]) -> np.ndarray:
        """
        Cheap proxy features for quality detection:
        ← These features don't require calling a large LLM.
          They proxy for quality signals that are computationally expensive to measure directly.
        """
        features = []
        for s in samples:
            text = s.get('text', s.get('question', '') + ' ' + s.get('answer', ''))
            tokens = text.split()
            n = len(tokens)
            special_ratio = sum(1 for c in text if not c.isalnum()) / max(len(text), 1)
            uniq_ratio = len(set(tokens)) / max(n, 1)
            features.append([n, special_ratio, uniq_ratio])
        return np.array(features)


# ─── TECHNIQUE 3: Cache-and-Reuse Mechanism (CRM) ────────────────────────────
class CacheAndReuse:
    """
    Cache intermediate processed datasets for operator prefix reuse.
    When a new strategy shares a prefix with a cached strategy, reuse the cached output.
    
    ← WHY: Many strategies share operator prefixes.
      E.g., [Cleaning → Optimization] and [Cleaning → Generation] both start with Cleaning.
      Without CRM: Cleaning is applied twice to the full dataset.
      With CRM: Cleaning result is cached after round 1; round 2 starts from the cache.
      
    ← THIS is the trick: treat the strategy as a prefix tree.
      The cache key is the operator sequence up to a given step.
      New strategies look up the longest matching cached prefix and continue from there.
    """

    def __init__(self):
        # OrderedDict: key = tuple of operators applied so far, value = processed dataset
        self.cache: OrderedDict[tuple, list[dict]] = OrderedDict()

    def get_cached_prefix(
        self,
        strategy: list[str],
        raw_data: list[dict],
    ) -> tuple[list[dict], list[str]]:
        """
        Find the longest cached prefix of this strategy.
        Returns: (cached_data, remaining_operators_to_apply)
        """
        for end_idx in range(len(strategy), 0, -1):
            prefix = tuple(strategy[:end_idx])
            if prefix in self.cache:
                remaining = strategy[end_idx:]
                return self.cache[prefix], remaining
        # No prefix match: start from raw data
        return raw_data, strategy

    def store(self, strategy: list[str], processed_data: list[dict]) -> None:
        """Cache the processed result after applying this strategy."""
        for end_idx in range(1, len(strategy) + 1):
            prefix = tuple(strategy[:end_idx])
            if prefix not in self.cache:
                # ← Store intermediate results at every prefix, not just the final output
                #   This maximizes reuse opportunities for future strategies
                self.cache[prefix] = processed_data

Dataset: D_train = 50,000 medical Q+A pairs (raw, noisy)
Model: 7B LLM, pre-trained on general text
Baseline (no processing): eval score = 0.72 on held-out medical QA
Meta-agent: GPT-4 class LLM
Privacy constraint: no human may access D_train directly

Agent generates 5 diverse strategies:
  S1: [Cleaning]                      → removes duplicates, HTML, length filters
  S2: [Optimization]                  → rewrites Q+A pairs with LLM
  S3: [Cleaning → Optimization]       → clean first, then rewrite
  S4: [Selection]                     → gradient-based quality selection
  S5: [Generation → Cleaning]         → augment first, then clean

DPS applied: D_sample = 15,000 samples (30% of D_train, stratified)
PTS applied: D_low = 12,000 samples, D_high = 38,000 samples

Each strategy evaluated:
  FT 7B LLM on DPS(PTS(D_train, strategy_k)), eval on D_val
  
Round 1 results:
  S1: score = 0.79, Δ = +0.07
  S2: score = 0.74, Δ = +0.02  (optimization alone: marginal)
  S3: score = 0.83, Δ = +0.11  (← best this round)
  S4: score = 0.76, Δ = +0.04
  S5: score = 0.77, Δ = +0.05

CRM cache after Round 1:
  ("cleaning",) → processed D_sample (13,500 clean samples)
  ("optimization",) → processed D_sample (14,200 optimized samples)
  ("cleaning", "optimization") → 12,800 processed samples
  ("selection",) → 11,200 selected samples

Prompt injected with: all 5 strategies + their Δ scores (group comparison)
Agent observes: [Cleaning → Optimization] (+0.11) > [Cleaning] (+0.07) > [Optimization] (+0.02)
Agent infers: cleaning before optimization is the correct order; optimization alone insufficient

Agent generates 5 new strategies:
  S6: [Cleaning → Optimization → Selection]
  S7: [Selection → Cleaning → Optimization]  (different order)
  S8: [Cleaning → Generation → Optimization]
  S9: [Cleaning → Optimization → Generation]
  S10: [Selection → Optimization]

CRM reuse:
  S6 shares prefix ("cleaning", "optimization") with S3 → reuse S3's cache → only run Selection
  S8 shares prefix ("cleaning",) with S1 → reuse S1's cache → run Generation → Optimization
  S9 shares prefix ("cleaning", "optimization") with S3 → reuse S3's cache → run Generation

Round 2 results:
  S6: score = 0.86, Δ = +0.14  (← best overall)
  S7: score = 0.81, Δ = +0.09  (order matters: Selection before Cleaning is worse)
  S8: score = 0.84, Δ = +0.12
  S9: score = 0.82, Δ = +0.10
  S10: score = 0.78, Δ = +0.06

Agent observes: [Cleaning → Optimization → Selection] is the best seen.
Agent generates 5 variations with different selection thresholds and cleaning parameters.
All Δ scores < 0.01 improvement over S6.
Agent terminates: "No further improvement possible. Best strategy: [Cleaning → Optimization → Selection]"

Final strategy applied to full D_train (50,000 samples):
  Step 1 (Cleaning): 50,000 → 44,200 samples (dedup, noise removal)
  Step 2 (Optimization): LLM rewrites 44,200 Q+A pairs for quality
  Step 3 (Selection): gradient filter → 38,000 highest-quality samples

Final model trained on 38,000 processed samples:
  Eval score: 0.88 (vs baseline: 0.72)
  Improvement: +22% over unprocessed
  Win rate vs AutoML baselines: ~65% (documented)
  
Total search time with acceleration:
  Naive (no DPS/PTS/CRM): ~10x longer (3 rounds × 5 strategies × full FT each)
  With DPS + PTS + CRM: practical search time (documented 10x speedup)
  Human expert access to D_train: zero (privacy constraint satisfied)