LaMDAgent: An Autonomous Framework for Post-Training Pipeline Optimization via LLM Agents

Introduction: The Post-Training Bottleneck

While pre-training provides Large Language Models (LLMs) with foundational knowledge, the subsequent post-training phase—comprising Supervised Fine-Tuning (SFT), preference learning, and model merging—is what transforms these models into useful, task-specific assistants. Currently, designing an effective post-training pipeline is often described as a form of "alchemy." Researchers must manually experiment with data ordering, hyperparameter tuning, and merging strategies, a process that is not only labor-intensive but also prone to missing optimal configurations within a vast search space.

To address this, we introduce LaMDAgent (Language Model Developing Agent), an autonomous framework designed to navigate the end-to-end post-training process. By treating various optimization techniques as a unified action space, LaMDAgent discovers high-performing training recipes that outperform manual baselines and reveal non-obvious strategies for multi-skill acquisition.

The LaMDAgent Framework

LaMDAgent treats post-training as an iterative search problem managed by an LLM agent (specifically GPT-4o in our experiments). The framework operates through a continuous four-step loop:

1. Action Enumeration: The system identifies all possible actions, such as performing SFT on a specific dataset with defined hyperparameters or merging existing models using techniques like TIES-Merging.
2. Action Selection: The agent selects the most promising next step based on its "memory" of previous trials and specific exploration directives.
3. Model Evaluation: The chosen action is executed, and the resulting model is evaluated against target benchmarks.
4. Memory Update: The agent records the results, summarizes the findings, and updates its internal strategy for the subsequent iteration.

By unifying heterogeneous methods—such as training on new data and mathematically merging weights—into a single agentic workflow, the framework can discover complex, multi-step pipelines that human researchers might overlook.

Experimental Methodology

The researchers evaluated LaMDAgent across two primary domains using the Gemma 2 model family:

• Multi-Skill Teaching: This experiment focused on balancing mathematical reasoning (GSM8k), commonsense reasoning (CommonsenseQA), and reading comprehension (TriviaQA). Generalization was tested using out-of-distribution sets such as GSM-Symbolic and Natural Questions.
• Tool Usage: This setting aimed to enhance tool-utilization capabilities using the Agent-FLAN, ToolACE, and WizardLM datasets, with performance measured via AceBench and MT-Bench for general instruction-following.

Key Insights and Discoveries

The "Math First" Strategy

One of the most significant findings was the agent’s discovery of a specific curriculum: learning mathematical skills before integrating general skill data was more effective for maintaining a balance across all capabilities. This suggests that certain foundational reasoning skills may serve as a better scaffold for subsequent learning, a strategy that is not always intuitive to human designers.

Mitigating the "Alignment Tax"

Standard fine-tuning on diverse data often results in a performance trade-off known as the "alignment tax," where specialized gains come at the cost of general capabilities. LaMDAgent found that specific sequences of training followed by strategic model merging could preserve general instruction-following while significantly boosting specialized performance.

Scaling Dynamics

The study revealed a divergence between data and model scaling. While effective pipelines discovered on small data subsets could be successfully scaled to larger datasets, the same was not true for model size. A pipeline optimized for a 2B parameter model was not necessarily the best for a 9B model, suggesting that the "optimal recipe" is highly dependent on the underlying architecture's capacity.

Agentic vs. Random Exploration

Crucially, the LLM agent significantly outperformed random action selection. This demonstrates that the agent's ability to reason about past failures and successes is a viable and superior method for navigating the complex optimization landscape of modern LLMs.

Results and Performance

LaMDAgent demonstrated substantial improvements over standard baselines:

• Reasoning Gains: In multi-skill experiments, the agent-discovered pipeline improved math-related task accuracy by 3.7 points over a standard fully fine-tuned baseline.
• Tool-Use Excellence: In the tool-use domain, the framework boosted accuracy on AceBench by 9.0 points compared to the base Gemma2 2B Instruct model, without degrading general instruction-following scores.
• Operational Efficiency: The framework identified that while training is superior for acquiring entirely new skills, model merging is a powerful tool for combining those skills without the interference often seen in joint fine-tuning.

Limitations and Future Work

Despite its success, LaMDAgent faces several challenges. The current experiments were limited to the Gemma 2 architecture and English-language datasets. Furthermore, the lack of transferability between model scales (e.g., 2B to 9B) means that exploration must currently be re-run for different model sizes, which is computationally expensive.

Future research will focus on expanding the action space to include Preference Learning (DPO/PPO) and investigating methods to make discovered pipelines more robust across varying model scales and architectures.