Building Social World Models with Large Language Models

Social World Model (SWM) models how collective beliefs change after real-world events. We represent social beliefs with prediction-market trajectories and learn an event-conditioned transition model $s_{t+1} \sim P_\theta(s_{t+1} \mid s_t, e_t)$, where $s_t$ is the current belief state and $e_t$ is a news/event signal. This repo provides the training method and a benchmark, SWM-Bench.

Training Recipe

SWM uses a posterior-guided training recipe. A frozen hindsight LLM first identifies which candidate event best explains the observed belief shift. We then use these attribution weights to train:

• a prior attributor that predicts event relevance before seeing the future;
• a world model that predicts the belief change caused by each event.

At inference, SWM either forecasts future beliefs from real news or simulates belief shifts under hypothetical events.

Installation

git clone https://github.com/ulab-uiuc/social-world-model.git
cd social-world-model
conda create -n social-world-model python=3.10 -y && conda activate social-world-model
pip install -e .

export HF_TOKEN="..."   # dataset download / model push

Quick Start

Download SWM-Bench (price history + candidate news + posterior attributions, chronologically split into train / Polymarket-test / Kalshi-test):

python -c "from huggingface_hub import snapshot_download; \
  snapshot_download('ulab-ai/swm-bench', repo_type='dataset', local_dir='data/swm-bench')"

SWM-Bench has three parts:

• raw/ — the original Polymarket / Kalshi price series + crawled news.
• Qwen3.5-397B-attributed-data/ — the records labeled by the Qwen3.5-397B posterior attributor (our main dataset).
• Qwen3-32B-attributed-data/ — the same records labeled by Qwen3-32B.

Each record is one (history, candidate_news, target, attributions) example. attributions holds the posterior (oracle) scores; *_with_nonzero_attribution.jsonl are the splits restricted to records with at least one non-zero-score news (used for training).

Use it directly with 🤗 datasets:

from datasets import load_dataset

block = "Qwen3.5-397B-attributed-data"   # or "Qwen3-32B-attributed-data"
ds = load_dataset("ulab-ai/swm-bench", data_files={
    "train":           f"{block}/train.jsonl",                          # attributor
    "train_attr":      f"{block}/train_with_nonzero_attribution.jsonl", # world model
    "test_kalshi":     f"{block}/test_kalshi.jsonl",
    "test_polymarket": f"{block}/test_polymarket.jsonl",
})

Stage 1 — Train the attributor

The attributor is trained to reproduce the posterior's responsibility distribution over candidate news (forward-KL; one epoch is enough). It trains on train.jsonl — all records, including null events, so it learns to assign low/no score to irrelevant or news-less cases:

DATA=data/swm-bench/Qwen3.5-397B-attributed-data
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 scripts/train_attributer.py \
    --train-data-path $DATA/train.jsonl \
    --valid-data-path $DATA/valid_subset150.jsonl \
    --output-dir saves/attributer_8b --model-name Qwen/Qwen3-8B \
    --epochs 1 --max-news 30 --max-seq-length 1024 \
    --train-batch-size 2 --gradient-accumulation-steps 2 \
    --learning-rate 2e-5 --gradient-checkpointing

Stage 2 — Train the world model

The world model is trained on train_with_nonzero_attribution.jsonl (the attributed records) with a per-news, responsibility-weighted regression loss, full fine-tuning under FSDP (MODE GPUS NPROC PORT MODEL TAG SAVE EP):

# 8B, 8-GPU FSDP, 6 epochs
bash scripts/train_fc_v9odds.sh fsdp 0,1,2,3,4,5,6,7 8 29500 Qwen/Qwen3-8B wm8b saves_local 6

# 0.6B, single GPU
bash scripts/train_fc_v9odds.sh single 0 1 29501 Qwen/Qwen3-0.6B wm06b saves_local 6

Stage 3 — Joint inference (prior & posterior)

Posterior (oracle attribution, already in the test file):

DATA=data/swm-bench/Qwen3.5-397B-attributed-data
python scripts/inference_multievent_world_model.py \
    --test-data-path $DATA/test_kalshi.jsonl \
    --model-path saves_local/wm8b/final-model --model-name Qwen/Qwen3-8B \
    --output-path results/posterior_kalshi.jsonl --max-news 30

Prior (deployable) — attribute with the trained attributor, then forecast:

# (a) prior attribution: replace each record's attributions with the model's
python scripts/inference_prior_attribution.py \
    --data-path $DATA/test_kalshi.jsonl \
    --attributer-path saves/attributer_8b --model-name Qwen/Qwen3-8B \
    --output-path results/test_kalshi_prior.jsonl --max-news 30

# (b) forecast on the prior-attributed file
python scripts/inference_multievent_world_model.py \
    --test-data-path results/test_kalshi_prior.jsonl \
    --model-path saves_local/wm8b/final-model --model-name Qwen/Qwen3-8B \
    --direct-soft-routing \
    --output-path results/prior_kalshi.jsonl --max-news 30

Both inference scripts shard across GPUs with --num-shards N --shard-idx i. Each output row has pred_delta / true_delta; score MASE, MAE, directional accuracy, and correlation over the full set and the attributed subset with scripts/eval_all_vs_attr.py.

License

Apache 2.0