MAXS: meta-adaptive exploration with LLM agents

This repository contains the official implementation of MAXS, a sophisticated decoding algorithm for Large Language Model Agents. MAXS enhances complex reasoning capabilities by integrating tool-use (code execution and web search) with a lookahead planning mechanism. It is particularly well-suited for challenging, multi-step tasks such as advanced mathematics and science problems.

The core of MAXS is a path selection strategy that evaluates potential reasoning steps not just on their immediate quality, but on their projected future success. This is achieved by combining multiple reward signals:

1. Advantage Score: The improvement a step offers over the previous state.
2. Step-Level Variance: The consistency of the reasoning path, penalizing erratic jumps in confidence.
3. Slope-Level Variance: The directional smoothness of the path, rewarding consistent trends.

These signals are normalized and combined to guide the model toward more robust and logical solutions.

✨ Key Features

• Lookahead Planning: Simulates future reasoning steps (foresight rollouts) to evaluate the long-term viability of each potential action.
• Integrated Tool Use: Seamlessly executes Python code (<step type="code">) and performs web searches (<step type="search">) to ground reasoning in external knowledge and computation.
• Advanced Path Scoring: Employs a novel, multi-faceted reward function inspired by control theory and mathematical analysis to select the most promising reasoning paths.
• Adaptive Decoding Strategy: Automatically switches to a faster, greedy "fast-forward" mode when the algorithm is confident in its path, optimizing for both performance and efficiency.
• Robust Output Parsing: Includes a "rescue" mechanism to handle malformed model outputs, ensuring the reasoning chain remains intact.
• Multi-Dataset Support: Comes with data loaders for various complex reasoning benchmarks, including MathVista, OlympiadBench, TheoremQA, DeepMindMath, and EMMA.

⚙️ Setup and Installation

1. Clone the Repository or Download the zip

   # git clone [REPOSITORY_URL]
   cd maxs-decoding

2. Install Dependencies

   We recommend using a virtual environment (e.g., conda or venv). The primary dependency is vLLM for high-throughput inference.

   conda create -n maxs-decoding python=3.10
   conda activate maxs-decoding
   pip install -r requirements.txt

   Note: vLLM may have specific installation requirements depending on your CUDA version. Please refer to the official vLLM documentation for details.

3. Download Models

   Download the LLM you wish to use (e.g., Qwen2.5-VL-7B-Instruct) and place it in a local directory.

   # Example using huggingface-cli
   huggingface-cli download Qwen/Qwen2.5-VL-7B-Instruct --local-dir ./models/Qwen2.5-VL-7B-Instruct

4. Prepare Datasets

   Download the full datasets and place them in the ./dataset directory following the structure expected by the code. For example:

   ./dataset/
   └── olympiadbench/
       └── OE_MM_maths_zh_CEE.parquet
   └── mathvista/
       └── mathvista.parquet
   ...
   

🚀 Running the Algorithm

Execute the main script maxs.py with the desired configuration. The following command runs MAXS on the OlympiadBench dataset.

python maxs.py \
    --model_path ./models/Qwen2.5-VL-7B-Instruct \
    --datasets olympiadbench \
    --data_path ./dataset/olympiadbench/OE_MM_maths_zh_CEE.parquet \
    --output_dir ./results/olympiadbench \
    --file_name "olympiad_qwen2.5" \
    --gpus 1 \
    --step_beam_size 1 \
    --num_rollout 4 \
    --num_foresight 4 \
    --record_process True

Command-Line Arguments

• --model_path: Path to the local directory containing the LLM.
• --gpus: Number of GPUs to use for inference.
• --datasets: The dataset to process (e.g., olympiadbench, mathvista, theoremqa).
• --data_path: Path to the specific dataset file.
• --output_dir: Directory to save the final results.
• --file_name: Basename for the output JSON file.
• --traj_path: Directory to save detailed trajectory logs (if record_process is True).
• --step_beam_size: Number of candidate paths to maintain at each step (currently supports 1).
• --num_rollout: Number of candidate next-steps to generate at each planning stage.
• --num_foresight: Number of future steps to simulate in each foresight rollout.
• --record_process: If True, saves detailed logs of the entire reasoning process.
• --seed: Random seed for reproducibility.

📂 Output Structure

The program generates two main types of output files:

1. Results File: ({output_dir}/{file_name}.json)

   • A JSON list containing the final results for each problem.
   • Each entry includes the id, question, ground_truth answer, and the final response generated by the model.

2. Trajectory File: (``{output_dir}/{traj_path}/TRAJ_INFO-{file_name}.json`)

   • This file is generated if --record_process is True.
   • It contains highly detailed, step-by-step logs of the entire decoding process for each problem.
   • This includes every candidate generated at each step, their foresight rollouts, and their raw and normalized scores, making it invaluable for analysis and debugging.

Some of the MathVista experiment results are placed in ./result folder.

⚖️ License

This project is licensed under the MIT License. See the LICENSE file for details.