A single-backbone, multi-role framework for LLM-based tool-use agents.
RAMTL replaces the multi-LLM overhead of conventional multi-agent pipelines with a single frozen backbone augmented by five lightweight LoRA role adapters and a learned Gumbel-Softmax router. The result is a parameter-efficient agent that dynamically switches between specialised roles (planner, retriever, executor, verifier, recovery) at each reasoning step.
Paper: See
paper/ramtl_paper_full.texfor the full manuscript.
┌────────────────────────────────────────────────────────────────┐
│ RAMTL Pipeline │
│ │
│ User Query │
│ │ │
│ ▼ │
│ ┌──────────────────────────────────────────────────────┐ │
│ │ Frozen 4-bit QLoRA Backbone (shared) │ │
│ │ (Mistral-7B / Gemma-2-2B / Gemma-2-9B / Phi-3-Mini)│ │
│ └───────────────────┬──────────────────────────────────┘ │
│ │ hidden state (last token) │
│ ▼ │
│ ┌──────────────────────────────────────────────────────┐ │
│ │ Gumbel-Softmax MLP Router │ │
│ │ │ │
│ │ Inputs: │ │
│ │ • backbone hidden state (H-dim) │ │
│ │ • schema summary (H-dim, mean-pooled) │ │
│ │ • previous action embed (64-dim) │ │
│ │ • failure flags (4 binary features) │ │
│ │ │ │
│ │ MLP: [H+H+64+4] → 512 → LN → GELU → 256 → LN → │ │
│ │ GELU → K=5 logits → Gumbel Top-2 │ │
│ └───────────────────┬──────────────────────────────────┘ │
│ │ top-2 role selection │
│ ▼ │
│ ┌──────────────────────────────────────────────────────┐ │
│ │ LoRA Role Adapter Bank (K = 5) │ │
│ │ │ │
│ │ ┌──────────┐ ┌──────────┐ ┌──────────┐ │ │
│ │ │ Planner │ │Retriever │ │ Executor │ │ │
│ │ │ LoRA │ │ LoRA │ │ LoRA │ │ │
│ │ │ r=16 │ │ r=16 │ │ r=16 │ │ │
│ │ └──────────┘ └──────────┘ └──────────┘ │ │
│ │ ┌──────────┐ ┌──────────┐ │ │
│ │ │ Verifier │ │ Recovery │ α=32, dropout=0.05 │ │
│ │ │ LoRA │ │ LoRA │ targets: q_proj,v_proj│ │
│ │ └──────────┘ └──────────┘ │ │
│ └───────────────────┬──────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌──────────────────────────────────────────────────────┐ │
│ │ Tool Execution Environment │ │
│ │ (ToolBench / API-Bank / StableToolBench) │ │
│ └──────────────────────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ Observation → next step │
│ (up to 15 steps with retry & fallback) │
└────────────────────────────────────────────────────────────────┘
| Role | Adapter | Responsibility |
|---|---|---|
| Planner | LoRA #0 | Decompose tasks, select tool sequences, produce final answers |
| Retriever | LoRA #1 | Search API schemas, compress tool descriptions, match tools to subtasks |
| Executor | LoRA #2 | Format valid JSON tool calls with correct arguments |
| Verifier | LoRA #3 | Check outputs for correctness, detect hallucinations |
| Recovery | LoRA #4 | Handle API failures, timeouts, schema mismatches; retry or backtrack |
RAMTL uses a three-stage training pipeline:
- Stage 1 — Joint SFT (870 steps): Supervised fine-tuning of all 5 LoRA adapters + router on tool-use trajectories from ToolBench, API-Bank, and StableToolBench.
- Stage 2 — Stability SFT: Fine-tuning with injected failure modes (timeouts, API unavailability, schema drift, malformed outputs, stale caches) at 30% injection rate.
- Stage 3 — DPO: Direct Preference Optimisation (β=0.1) on preference pairs to refine role specialisation.
ramtl-release/
├── README.md
├── LICENSE # Apache 2.0
├── pyproject.toml # Dependencies & build config
├── .gitignore
│
├── src/
│ ├── __init__.py
│ ├── train_sft.py # Entry point: SFT training (Stages 1 & 2)
│ ├── train_dpo.py # Entry point: DPO training (Stage 3)
│ ├── run_eval.py # Entry point: Evaluation
│ │
│ ├── models/
│ │ ├── backbone.py # Backbone loading (4-bit QLoRA, 4 models)
│ │ ├── role_adapters.py # LoRA adapter bank (5 roles)
│ │ ├── router.py # Gumbel-Softmax MLP router
│ │ └── tool_state.py # ToolSpec, Trajectory, ExecutionRecord
│ │
│ ├── training/
│ │ ├── trainer_sft.py # SFT trainer with router loss
│ │ └── trainer_dpo.py # DPO trainer with router entropy callback
│ │
│ ├── inference/
│ │ └── adaptive_inference.py # RAMTL adaptive decoding loop
│ │
│ ├── eval/
│ │ └── evaluate.py # Metrics: TCR, VCR, SoPR, SoWR, etc.
│ │
│ ├── data/
│ │ └── trajectory_builder.py # Raw data → Trajectory conversion
│ │
│ ├── environments/
│ │ ├── env_toolbench.py # ToolBench wrapper
│ │ ├── env_apibank.py # API-Bank wrapper
│ │ └── env_stabletoolbench.py # StableToolBench wrapper (+ failure injection)
│ │
│ └── utils/
│ ├── config.py # OmegaConf config helpers
│ ├── io.py # JSONL, atomic I/O, checksums
│ ├── logging.py # Logging setup
│ └── seed.py # Reproducibility seed utilities
│
├── configs/
│ ├── exp/ # Per-backbone experiment configs
│ │ ├── ramtl_mistral7b_sft.yaml
│ │ ├── ramtl_mistral7b_dpo.yaml
│ │ ├── ramtl_mistral7b_eval.yaml
│ │ ├── ramtl_gemma2_2b_sft.yaml
│ │ ├── ramtl_gemma2_2b_dpo.yaml
│ │ ├── ramtl_gemma2_2b_eval.yaml
│ │ ├── ramtl_gemma2_9b_sft.yaml
│ │ ├── ramtl_gemma2_9b_dpo.yaml
│ │ ├── ramtl_gemma2_9b_eval.yaml
│ │ ├── ramtl_phi3mini_sft.yaml
│ │ ├── ramtl_phi3mini_dpo.yaml
│ │ └── ramtl_phi3mini_eval.yaml
│ ├── model/ # Backbone configs
│ ├── lora/ # LoRA defaults
│ ├── train/ # Training defaults (sft, dpo)
│ ├── data/ # Data paths
│ ├── eval/ # Evaluation defaults
│ └── inference/ # Inference defaults
│
├── paper/
│ ├── ramtl_paper_full.tex # Full NeurIPS paper
│ └── tables_publication.tex # Publication-ready tables
│
├── data/ # Data directory (not tracked; see below)
│ └── .gitkeep
│
└── outputs/ # Training/eval outputs (not tracked)
└── .gitkeep
- Python ≥ 3.11
- CUDA-capable GPU (≥ 16 GB VRAM for Gemma-2-2B; ≥ 24 GB for Mistral-7B; ≥ 48 GB for Gemma-2-9B)
# Clone the repository
git clone git@github.com:zhengtaoyao/RAMTL.git
cd RAMTL
# Create conda environment
conda create -n ramtl python=3.11 -y
conda activate ramtl
# Install PyTorch (adjust for your CUDA version)
pip install torch>=2.2.0 --index-url https://download.pytorch.org/whl/cu121
# Install RAMTL and all dependencies
pip install -e .RAMTL is evaluated on three tool-use benchmarks. Download and place data as follows:
- Source: https://github.com/OpenBMB/ToolBench
- Setup: Follow the ToolBench instructions to download the dataset and tool API definitions.
- Place files at:
data/raw/toolbench/ ├── tools/ # API JSON files (one per tool) ├── train.jsonl # Training trajectories └── eval.jsonl # Evaluation trajectories
- Source: https://github.com/AlibabaResearch/DAMO-ConvAI/tree/main/api-bank
- Description: 73 API tools, 314 annotated dialogues (eval), 753 annotated API calls.
- Place files at:
data/raw/apibank/ ├── api_data.json # Tool/API specifications ├── train_dialogs.jsonl # Training dialogues └── eval_dialogs.jsonl # Evaluation dialogues
- Source: https://github.com/zhichengg/StableToolBench
- Description: Extends ToolBench with a virtual API server and cached responses for reproducible evaluation. Introduces Solvable Pass Rate (SoPR) and Solvable Win Rate (SoWR).
- Place files at:
data/raw/stabletoolbench/ ├── tools/ # Tool definitions (same format as ToolBench) ├── train.jsonl # Training data └── eval.jsonl # Evaluation data
| Backbone | HuggingFace Model ID | Hidden Size | VRAM (4-bit) |
|---|---|---|---|
| Mistral-7B-Instruct-v0.2 | mistralai/Mistral-7B-Instruct-v0.2 |
4096 | ~6 GB |
| Gemma-2-2B-it | google/gemma-2-2b-it |
2304 | ~2 GB |
| Gemma-2-9B-it | google/gemma-2-9b-it |
3584 | ~8 GB |
| Phi-3-mini-4k-instruct | microsoft/Phi-3-mini-4k-instruct |
3072 | ~3 GB |
# Mistral-7B
python -m src.train_sft --config configs/exp/ramtl_mistral7b_sft.yaml
# Gemma-2-2B
python -m src.train_sft --config configs/exp/ramtl_gemma2_2b_sft.yaml
# Gemma-2-9B
python -m src.train_sft --config configs/exp/ramtl_gemma2_9b_sft.yaml
# Phi-3-Mini
python -m src.train_sft --config configs/exp/ramtl_phi3mini_sft.yamlOverride any config parameter via CLI:
python -m src.train_sft --config configs/exp/ramtl_mistral7b_sft.yaml \
train.learning_rate=1e-4 \
train.num_epochs=2 \
lora.r=32python -m src.train_dpo --config configs/exp/ramtl_mistral7b_dpo.yamlNote: DPO requires a completed SFT checkpoint. Update
sft_checkpointin the DPO config to point to your trained model.
# Full evaluation (all 3 benchmarks × 3 seeds)
python -m src.run_eval --config configs/exp/ramtl_mistral7b_eval.yaml
# Evaluate on a single benchmark with limited samples
python -m src.run_eval --config configs/exp/ramtl_mistral7b_eval.yaml \
eval.benchmarks='[toolbench]' \
eval.max_samples=50| Metric | Description |
|---|---|
| Task Completion Rate (TCR) | Fraction of tasks with a non-empty, correct final answer |
| Valid Call Rate (VCR) | Fraction of tool calls that executed successfully |
| Exact Argument Accuracy | Fraction of calls with exactly correct arguments |
| Recovery Success Rate | Fraction of error tasks eventually resolved |
| SoPR | StableToolBench Solvable Pass Rate |
| SoWR | StableToolBench Solvable Win Rate |
| Router Entropy | Mean entropy of router decisions (role specialisation measure) |
| Pass Rate | ToolBench/ToolEval-compatible pass rate |
| Win Rate | ToolBench/ToolEval-compatible win rate (vs reference) |
| Parameter | Default | Description |
|---|---|---|
lora.r |
16 | LoRA rank for each adapter |
lora.lora_alpha |
32 | LoRA scaling factor |
router.routing_mode |
top2 |
single / top2 / soft |
router.gumbel_tau |
1.0 | Gumbel-Softmax temperature |
router.entropy_coeff |
1e-3 | Entropy regularisation coefficient |
train.lambda_router |
0.3 | Router loss weight during SFT |
train.stage2_inject_failure_rate |
0.3 | Failure injection rate for stability training |
dpo.beta |
0.1 | DPO β parameter |
inference.max_steps |
15 | Maximum reasoning steps per task |
inference.max_retries |
3 | Max retries on tool call failure |
RAMTL adds only ~25.2M trainable parameters per backbone:
| Component | Parameters |
|---|---|
| 5 × LoRA adapters (r=16, q_proj + v_proj) | ~24.4M |
| MLP Router | ~0.8M |
| Total | ~25.2M |
The backbone remains fully frozen (4-bit NF4 quantisation with double quantisation).
@article{ramtl2025,
title={RAMTL: Role-Adaptive Multi-Tool Learning for Single-Backbone Tool-Use Agents},
author={Yao, Zhengtao},
year={2025}
}This project is licensed under the Apache License 2.0. See LICENSE for details.
