SAFO

This repository contains the implementation of SAFO, built upon LLaMA-Factory.

Table of Contents

• Overview
• Installation
• Datasets
• Dataset Preparation
• Training
• Project Structure
• Acknowledgements

Overview

SAFO introduces novel training strategies for large language models to improve both utility and fairness in social simulation tasks. Our implementation extends LLaMA-Factory with custom loss computation methods.

Installation

Please follow the official LLaMA-Factory installation guide to set up the environment For more details, refer to the LLaMA-Factory documentation.

Datasets

We evaluate SAFO on three large-scale social survey datasets that span diverse countries and cultural contexts. These datasets encompass a wide range of demographic attributes (e.g., age, gender, education, income) and survey topics (e.g., social attitudes, political views, economic perceptions), providing a comprehensive testbed for assessing both utility and fairness in social simulation.

Dataset Preparation

Due to licensing restrictions, we provide sample data and processing scripts instead of the full datasets.

1. Download Raw Data

• GSS: Download from GSS Data Explorer
• CGSS: Download from CGSS Official Website
• ESS: Download from ESS Data Portal

2. Data Format

The processed data should follow this JSON format:

[
  {
    "instruction": "The survey ...",
    "input": "",
    "output": "B. Better",
    "respondent_id": 2692,
    "question_index": 4,
    "question": "compared to your parents when they were the age you are now, do you think your own standard of living is now better, about the same, worse than theirs was?",
    "answer": "Better",
    "background": {
      "age": "45-54",
      "sex": "male",
      "race": "other",
      "marital": "separated",
      "degree": "high school",
      "income": "$25,000 or more",
      "wrkstat": "working full time",
      "res16": "in a small city or town (under 50,000)",
      "region": "east north central"
    },
    "group_label": "45-54|male|other|separated|high school|$25,000 or more|working full time|in a small city or town (under 50,000)|east north central"
  }
]

Field Descriptions

Field	Description
instruction	Complete prompt including demographic background and survey question
input	Additional input context (typically empty)
output	Ground truth answer from the survey respondent
respondent_id	Unique identifier for the survey respondent
question_index	Index of the survey question
question	The survey question text
answer	The selected answer option
background	Dictionary containing all demographic attributes
group_label	Concatenated string of all demographic attributes for group identification

4. Register the Dataset

Add your dataset configuration to data/dataset_info.json:

{
  "gss": {
    "file_name": "gss_train.json",
    "formatting": "alpaca",
    "columns": {
      "prompt": "instruction",
      "query": "input",
      "response": "output"
    }
  },
  "cgss": {
    "file_name": "cgss_train.json",
    "formatting": "alpaca",
    "columns": {
      "prompt": "instruction",
      "query": "input",
      "response": "output"
    }
  },
  "ess": {
    "file_name": "ess_train.json",
    "formatting": "alpaca",
    "columns": {
      "prompt": "instruction",
      "query": "input",
      "response": "output"
    }
  }
}

5. Train/Test Split

To evaluate generalization under realistic test-time novelty, we follow the unified evaluation protocol where the held-out test set may include novel questions and/or novel demographic compositions compared to training.

Training

We provide training scripts for different datasets in examples/train_lora/.

Training on GSS

bash examples/train_lora/gss.sh

Training on CGSS

bash examples/train_lora/cgss.sh

Training on ESS

bash examples/train_lora/ess.sh

Project Structure

SAFO/
├── data/
│   ├── dataset_info.json          # Dataset configuration
│   ├── gss_train.json             # GSS training data
│   ├── cgss_train.json            # CGSS training data
│   └── ess_train.json             # ESS training data
├── examples/
│   └── train_lora/
│       ├── cgss.sh                # CGSS training script
│       ├── ess.sh                 # ESS training script
│       └── gss.sh                 # GSS training script
├── src/
│   └── llamafactory/
│       └── train/
│           ├── data/              # Modified data processing
│           └── sft/
│               └── trainer.py     # Core algorithm (compute_loss)
├── scripts/
│   └── process_data.py            # Data processing script
├── README.md
└── requirements.txt

Key Modifications

Compared to the original LLaMA-Factory repository, we have modified:

1. src/llamafactory/train/data/: Custom data loading and preprocessing logic for survey data with demographic attributes
2. src/llamafactory/train/sft/: Modified SFT training pipeline
3. src/llamafactory/train/sft/trainer.py: Core SAFO algorithm implementation in compute_loss method

Acknowledgements

This project is built upon LLaMA-Factory. We thank the authors for their excellent work.

License

This project is licensed under the Apache License 2.0 - see the LICENSE file for details.