Simultaneous Multi-Objective Alignment Across Verifiable and Non-verifiable Rewards

This repository provides data, code, and models for the paper: Simultaneous Multi-objective Alignment Across Verifiable and Non-verifiable Rewards (arXiv:2510.01167).

Content

1. Set Up
2. Domains and Data
3. Process Reward Model Training
4. MAH-DPO Training
5. PRM-Guided Decoding
6. Evaluation
7. Trained Model Checkpoints
8. Citation

Set Up

This repository is developed under the MosaicML LLM Foundry environment. We recommend using the official Docker image:

docker pull mosaicml/llm-foundry:2.7.0_cu128-latest

Then clone the code:

git clone https://github.com/pearls-lab/multiobj-align.git
cd multiobj-align

Alternatively, install the packages:

pip install -r requirements.txt

Our experiments used 8×H100 80GB GPUs.

Notes:

• mah_dpo/requirements.txt contains packages specifically for MAH-DPO training

Domains and Data

Our approach is implemented across three diverse domains, each with multiple objectives:

• Mathematics: Mathematical problem-solving with step-by-step reasoning using the MATH dataset, optimizing for accuracy and engagement.
• Human Values: Open-domain conversations using the UltraFeedback dataset, optimizing for helpfulness, honesty, and truthfulness.
• Socratic Mind (AI Tutoring): Multi-turn AI tutoring dialogues for teaching programming collected from Socratic Mind, optimizing for accuracy and engagement.

Process Reward Model Training

We train Process Reward Models (PRMs) across all domains.

cd prm_train

# Math accuracy (Value model with hindsight relabeling)
bash scripts/train_math_value.sh

# Math engagement (PRM with binary classification)
bash scripts/train_math_engagement.sh

# Human values (Bradley-Terry RM)
bash scripts/train_human_values.sh

# Socratic Mind engagement (PRM with binary classification)
bash scripts/train_sm_engagement.sh

# Socratic Mind accuracy (PRM with binary classification)
bash scripts/train_sm_accuracy.sh

# Unified model (all 7 dimensions)
bash scripts/train_unified_model.sh

Note: For Human Values Bradley–Terry RM training, please follow the installation instructions from the RLHFlow/RLHF-Reward-Modeling repository before running the script.

MAH-DPO Training

Download Training Data

The training data for MAH-DPO is available on Google Drive. Download the data and place it in the mah_dpo/data/ directory:

1. Download the data from Google Drive
2. Extract the downloaded files
3. Place all files in the mah_dpo/data/ directory

MAH-DPO Training

First, perform supervised fine-tuning on base model. Then, train models with objective-specific heads for multi-objective alignment:

cd mah_dpo
python3 -m venv env
source env/bin/activate
pip install -r requirements.txt

# Mathematics (2 Heads)
bash scripts/train_sft_math.sh
bash scripts/train_mahdpo_math.sh

# Human Values (3 Heads) 
bash scripts/train_sft_human_values.sh
bash scripts/train_mahdpo_human_values.sh

# Socratic Mind (2 Heads)
bash scripts/train_sft_sm.sh
bash scripts/train_mahdpo_sm.sh

Inference

Use domain-specific inference scripts multiple configurations (base model, head 0, ensemble, SFT baseline). Edit the script to run the desired configuration or allocate different GPUs for different splits:

# Mathematics
  bash scripts/inference_math.sh

# Human Values 
  bash scripts/inference_human_values.sh

# Socratic Mind
  ## Step 1: Generate assistant responses (edit scripts/inference_sm.sh to select configuration)
  bash scripts/inference_sm.sh

  ## Step 2: Convert to message format (example with ensemble output)
  python convert_inference_to_messages.py \
    --input results/sm_ensemble_split0.jsonl \
    --output results/sm_ensemble_formatted.jsonl

  ## Step 3: Generate student responses (for evaluation)
  python generate_next_student_response_from_messages.py \
    --input results/sm_ensemble_formatted.jsonl \
    --output results/sm_ensemble_with_student.jsonl

PRM-Guided Decoding with Continuing Hidden State

Our novel test-time optimization method that maintains hidden-state continuity during guided generation.

cd prm_guided_decoding

# Human Values with ORM Guidance
  ## Base model + ORM guidance
  bash scripts/guide_human_values_base.sh
  ## MAH-DPO model + ORM guidance
  bash scripts/guide_human_values_mahdpo.sh

# Math with Value Model Guidance
  ## Base model + Value model guidance
  bash scripts/guide_math_value_base.sh
  ## MAH-DPO model + Value model guidance
  bash scripts/guide_math_value_mahdpo.sh

# Math with Engagement PRM Guidance
  ## Base model + Engagement PRM guidance
  bash scripts/guide_math_eng_base.sh
  ## MAH-DPO model + Engagement PRM guidance
  bash scripts/guide_math_eng_mahdpo.sh

# Socratic Mind with Candidate Selection
  ## GPT-4o workflow: generate candidates with GPT-4o → select best with PRM → generate student response
  bash scripts/guide_sm.sh
  ## MAH-DPO workflow: generate candidates with base model → select best with PRM → generate student response
  bash scripts/guide_sm_base.sh
  ## MAH-DPO workflow: generate candidates with MAH-DPO → select best with PRM → generate student response
  bash scripts/guide_sm_mahdpo.sh

Unified Model Guidance (All Dimensions)

# Math domain with unified model
bash scripts/guide_math_unified.sh

# Human values domain with unified model
bash scripts/guide_hv_unified.sh

# Socratic mind domain with candidates generation → select best with unified model → generate student response for later evaluation
bash scripts/guide_sm_unified.sh

Evaluation

All evaluation scripts are in the eval/ directory.

cd eval

# Mathematics 
  ## Mathematical Accuracy
    ### Preprocess inference output
    python 01_math500_data_process.py \
      -i results/math_output.jsonl \
      -o results/math_output_processed.json

    ### Evaluate correctness
    python 02_math500_accuracy_eval.py \
      -i results/math_output_processed.json \
      -o results/math_accuracy_eval.json

  ## Mathematical Engagement
  python math500_eng_eval.py \
    results/math_output.jsonl \
    -o results/math_engagement_eval.json

# Human Values
  python hv_eval.py \
    --input_file results/hv_output.jsonl \
    --output_file results/hv_eval.json \
    --model_type dpo \
    --devices cuda:0 cuda:1 cuda:2 \
    --batch_size 4

# Socratic Mind (requires simulated student responses)
  ## Socratic Mind Accuracy
  python reward_model_accuracy_evaluator_stu.py \
    --input results/sm_with_student.jsonl \
    --out results/sm_accuracy_eval.jsonl

  ## Socratic Mind Engagement
  python reward_model_engagement_evaluator_stu.py \
    --input results/sm_with_student.jsonl \
    --out results/sm_engagement_eval.jsonl

Trained Model Checkpoints

We provide trained models for all components. All models are available on HuggingFace Hub and can be loaded directly via transformers.

Mathematics	Human Values	Socratic Mind	Unified
SFT	SFT	SFT	Unified PRM
MAH-DPO	MAH-DPO	MAH-DPO
Accuracy PRM	Helpfulness PRM	Engagement PRM
Accuracy Value Model	Honesty PRM	Accuracy PRM
Engagement PRM	Truthfulness PRM

Citation

If you found our work useful please consider citing it:

@article{shen2025simultaneous,
  title={Simultaneous Multi-objective Alignment Across Verifiable and Non-verifiable Rewards},
  author={Shen, Yiran and Xia, Yu and Chang, Jonathan and Ammanabrolu, Prithviraj},
  journal={arXiv preprint arXiv:2510.01167},
  year={2025},
  url={https://arxiv.org/abs/2510.01167}
}