summaryrefslogtreecommitdiff
path: root/Group-Entropy-Equalization/README.md
diff options
context:
space:
mode:
Diffstat (limited to 'Group-Entropy-Equalization/README.md')
-rw-r--r--Group-Entropy-Equalization/README.md29
1 files changed, 29 insertions, 0 deletions
diff --git a/Group-Entropy-Equalization/README.md b/Group-Entropy-Equalization/README.md
index 33bd020..0ea2010 100644
--- a/Group-Entropy-Equalization/README.md
+++ b/Group-Entropy-Equalization/README.md
@@ -57,6 +57,35 @@ Checkpoints are saved under `checkpoints/<model>/<run_name>/`.
---
+### Group-wise Entropy Equalization (GEE)
+
+GEE balances sensitive groups by:
+- Group mass parity (push group probability mass toward target pi)
+- Group entropy equalization (normalize and equalize per-group entropy)
+- Optional anchors to keep global token-entropy and sensitive-union mass close to baseline
+
+Default groups file: `groups/gender.json`.
+
+Run on Colab (example):
+
+```bash
+!python train.py \
+ --model_name Qwen2.5-1.5B \
+ --model_path Qwen/Qwen2.5-1.5B \
+ --train_data dataset/1shot_rlvr/pi1_r1280.parquet \
+ --effective_batch 4 --micro_batch_size 1 \
+ --temperature 0.5 --learning_rate 2e-5 --sample_temp 0.5 \
+ --max_steps 15 --log_steps 1 --save_steps 5 \
+ --run_name colab_gee15 --wandb_project one-shot-em \
+ --no_deepspeed --mixed_precision no \
+ --gee_enable --gee_groups_path groups/gender.json \
+ --gee_alpha 1.0 --gee_beta 0.3 --gee_lambda 0.0 --gee_gamma 0.0 --gee_tau 1e-3 --gee_top_m 50
+```
+
+You can customize groups and target proportions in the JSON.
+
+---
+
### Reproducing One-shot EM Training (SOTA)
```bash