social-data-attribution-w2/docs/SOC28_HANDOFF.md

# SOC-28 Influence Heatmap Handoff

This document covers all the data, code, figures, and findings from the SOC-27/28/30/31 influence visualization work. Everything described here is available on GitHub, HuggingFace, and Linear.

## Where to find everything

| Resource | Location |
|----------|----------|
| Code + artifacts | [PR #74](https://github.com/eilab-gt/social-data-attribution/pull/74), branch `worktree-trackstar-visuals` |
| Data + figures (download) | [HuggingFace: HCAI-Lab/dolma3-influence-heatmaps](https://huggingface.co/datasets/HCAI-Lab/dolma3-influence-heatmaps) |
| Linear tickets | SOC-27, SOC-28, SOC-30, SOC-31 (all Done) |
| Raw score matrices (368 GB) | PACE ICE: `/storage/ice-shared/cs7634/staff/TDA/trackstar/scores_full/base/20260326T163642Z_1102443/` |
| Source training shards (41 GB) | PACE ICE: `/storage/ice-shared/cs7634/staff/TDA/trackstar/shards_10k/sample_10000_docs/` |

## What was done

We took the raw per-document influence scores from SOC-156 (5.68M docs scored against 4 benchmarks using TrackStar/Bergson on OLMo3-7B Base) and produced:

1. **Bin-level aggregation** (SOC-27): collapsed 5.68M per-doc scores into 576 bins (24 topics x 24 formats from WebOrganizer) per benchmark
2. **Influence heatmaps** (SOC-28): 47 publication-quality figures showing influence patterns
3. **Top-bin ranking tables** (SOC-30): ranked bins per benchmark + contrastive table
4. **Correctness stratification** (SOC-31): separate aggregation for queries the model answered correctly vs incorrectly

## Corpus and model

| Property | Value |
|----------|-------|
| Model | `allenai/Olmo-3-1025-7B` (OLMo3 7B Base) |
| Corpus | 5,678,621 docs from stratified 10K docs/bin sample |
| Source | `HCAI-Lab/dolma3-6t-sample-10000-docs` |
| Bins | 576 (24 topics x 24 formats), 559 fully filled at 10K docs |
| Benchmarks | GSM8K (1,319 queries), SocialIQA (10,000), MMLU-SS (3,077), MMLU-STEM (3,018) |
| Scoring method | TrackStar Mode A dot-product scoring with mixed preconditioner |
| Run ID | `20260326T163642Z_1102443` |

## Aggregation method

The primary metric is **per-query median influence**. For each bin:

1. For each training shard: group docs by bin, sum scores across docs per query
2. After all shards: divide by doc_count to get per-query bin means
3. Report median across queries

This avoids a problem where averaging across all queries (`mean(axis=1)`) makes benchmarks with different query counts incomparable. Raw per-element scores are ~0.0034 for all benchmarks, but mean-across-queries creates a 6.5x magnitude difference between GSM8K (1.3K queries) and SocialIQA (10K queries). The per-query median keeps all benchmarks within 8% of each other.

## Data files

### Aggregated bin scores (`artifacts/influence_bin_scores/`)

| File | Description |
|------|-------------|
| `queries_*_bin_scores_perquery.csv` | Per-query aggregation (primary, 4 files) |
| `queries_*_bin_scores.csv` | Legacy mean-across-queries aggregation (4 files) |

Perquery CSV columns: `topic_label`, `format_label`, `median_influence`, `mean_influence`, `p25_influence`, `p75_influence`, `std_influence`, `median_abs_influence`, `mean_abs_influence`, `doc_count`

### Correct/incorrect split (`artifacts/influence_bin_scores_split/`)

| File | Description |
|------|-------------|
| `queries_*_bin_scores_correct.csv` | Aggregated over correct-only queries (4 files) |
| `queries_*_bin_scores_incorrect.csv` | Aggregated over incorrect-only queries (4 files) |

### Proponent examples (`artifacts/proponent_examples/`)

| File | Description |
|------|-------------|
| `proponents_*.csv` | Top-3 most influential training docs for 10 queries per benchmark |

Columns: `query_id`, `query_text`, `is_correct`, `rank`, `score`, `doc_id`, `doc_snippet`

These show the actual training document text that most influenced each query. Balanced mix of correct and incorrect queries.

### Tables (`artifacts/paper_figures/table_*.csv`)

| File | Description |
|------|-------------|
| `table_top_bins_*.csv` | Top 20 bins per benchmark by signed median influence |
| `table_contrastive_socialiqa_vs_gsm8k.csv` | Top 20 bins by SocialIQA - GSM8K difference |
| `table_correctness_diff_*.csv` | Top 20 bins by correct - incorrect difference |

## Figures

All figures are in `artifacts/paper_figures/` as PNGs. Each heatmap has two versions:

- **Value-ordered** (default): rows/columns sorted by influence magnitude. Highlights strongest signals.
- **Canonical-ordered** (`_canonical` suffix): fixed alphabetical order. Same layout across all benchmarks for direct comparison.

### Priority figures for the paper

| Priority | Figure | File |
|----------|--------|------|
| 1 | **Contrastive difference (SocialIQA - GSM8K)** | `fig_influence_diff_socialiqa_vs_gsm8k.png` |
| 2 | **Paired topic bars** | `fig_influence_topic_paired_socialiqa_vs_gsm8k.png` |
| 3 | Signed heatmap SocialIQA | `fig_influence_signed_socialiqa.png` |
| 4 | Signed heatmap GSM8K | `fig_influence_signed_gsm8k.png` |
| 5 | Correct vs incorrect (SocialIQA) | `fig_influence_diff_socialiqa__correct_vs_socialiqa__incorrect.png` |

### Full figure inventory (47 PNGs)

**Per-benchmark (4 benchmarks x 2 orderings x 2 types = ~24 heatmaps):**
- `fig_influence_abs_*.png` / `fig_influence_abs_*_canonical.png`
- `fig_influence_signed_*.png` / `fig_influence_signed_*_canonical.png`

**Contrastive (SocialIQA vs GSM8K):**
- `fig_influence_diff_socialiqa_vs_gsm8k.png` (+ canonical)
- `fig_influence_compare_abs_socialiqa_vs_gsm8k.png` (+ canonical)
- `fig_influence_compare_signed_socialiqa_vs_gsm8k.png` (+ canonical)
- `fig_influence_topic_paired_socialiqa_vs_gsm8k.png`

**Correctness stratification (4 benchmarks):**
- `fig_influence_diff_*__correct_vs_*__incorrect.png` (+ canonical)

**Supplementary:**
- `fig_influence_topic_*.png` (topic marginal bars, 4 benchmarks)
- `fig_influence_format_*.png` (format marginal bars, 4 benchmarks)
- `fig_influence_radar_abs.png` / `fig_influence_radar_signed.png`
- `fig_influence_hist_*.png` (per-benchmark + overlay)
- `fig_influence_facets_gsm8k.png`
- `fig_sample_topic_doc_count.png` / `fig_sample_bin_doc_count.png` (sample verification)

## Key findings

### Contrastive pattern (SocialIQA vs GSM8K)

The contrastive signal is driven primarily by GSM8K's negative side: Documentation and Legal Notices formats in Industrial, Health, and Politics topics have strong negative influence on math performance. SocialIQA shows near-zero or mildly positive influence from these same bins.

Top contrastive bins (SocialIQA - GSM8K difference):

| Bin | SocialIQA | GSM8K | Difference |
|-----|-----------|-------|------------|
| Industrial / Documentation | +0.000030 | -0.000413 | +0.000443 |
| Health / Documentation | +0.000013 | -0.000411 | +0.000424 |
| Industrial / Legal Notices | +0.000010 | -0.000268 | +0.000277 |

### Correctness stratification signal strength

| Benchmark | Max diff / std ratio | Assessment |
|-----------|---------------------|------------|
| GSM8K | 2.04 | Strong |
| MMLU-SS | 0.75 | Strong |
| MMLU-STEM | 0.59 | Strong |
| SocialIQA | 0.47 | Moderate |

GSM8K shows the clearest correctness stratification. SocialIQA is weaker, suggesting social reasoning draws from more diffuse training data.

### Proponent examples

The top proponent training docs for GSM8K math questions are not math content. They include cooking instructions, social media posts, and server logs. This suggests influence is driven by structural/formatting patterns (Q&A format, numbered lists) rather than topic content.

## Code modules

### Aggregation (runs on PACE ICE)

| Module | Entry point | Purpose |
|--------|-------------|---------|
| `src/data_attribution/attribution/trackstar/bin_aggregate.py` | `data-attribution-trackstar-bin-aggregate` | Base pooled aggregation |
| `src/data_attribution/attribution/trackstar/bin_aggregate_perquery.py` | `data-attribution-trackstar-bin-aggregate-perquery` | Per-query aggregation (comparable) |
| `src/data_attribution/attribution/trackstar/bin_aggregate_split.py` | `data-attribution-trackstar-bin-aggregate-split` | Correct/incorrect split |
| `src/data_attribution/attribution/trackstar/proponent_examples.py` | `data-attribution-trackstar-proponent-examples` | Top-K doc text extraction |

### Visualization (runs locally)

| Module | Purpose |
|--------|---------|
| `src/dolma/distribution_report/influence_loader.py` | Load CSVs, auto-detect perquery format, normalize labels |
| `src/dolma/distribution_report/influence_figures.py` | Absolute + signed 24x24 heatmaps |
| `src/dolma/distribution_report/influence_comparison.py` | Side-by-side + contrastive difference heatmaps |
| `src/dolma/distribution_report/influence_marginals.py` | Topic/format marginal bar charts |
| `src/dolma/distribution_report/influence_radar.py` | 24-axis radar fingerprint chart |
| `src/dolma/distribution_report/influence_facets.py` | Format-conditioned topic bars |
| `src/dolma/distribution_report/influence_histograms.py` | Score distribution histograms |
| `src/dolma/distribution_report/influence_tables.py` | Top-bin ranking + correctness diff tables |
| `src/dolma/distribution_report/influence_runner.py` | Orchestrates all influence figure generation |

### SLURM batch scripts

| Script | Purpose |
|--------|---------|
| `scripts/slurm/attribution/trackstar_bin_aggregate.sbatch` | Base aggregation (CPU, 8GB, 6h) |
| `scripts/slurm/attribution/trackstar_bin_aggregate_perquery.sbatch` | Per-query aggregation (CPU, 8GB, 6h) |
| `scripts/slurm/attribution/trackstar_bin_aggregate_split.sbatch` | Correct/incorrect split (CPU, 16GB, 12h) |
| `scripts/slurm/attribution/trackstar_proponent_examples.sbatch` | Proponent extraction (CPU, 4GB, 1h) |

## How to regenerate figures

From the repo root on the `worktree-trackstar-visuals` branch:

```bash
PYTHONPATH=src python -m dolma.distribution_report.cli \
  --eda-dir artifacts/dolma_eda \
  --output-dir artifacts/paper_figures \
  --influence-dir artifacts/influence_bin_scores \
  --influence-split-dir artifacts/influence_bin_scores_split \
  --format all \
  --dummy
```

The `--dummy` flag generates sampling comparison figures with placeholder data (the real sampling manifests are separate). Remove it and provide `--representative-manifest` and `--stratified-manifest` if those are available.

## How to re-run aggregation

If new score data is produced or the sample changes:

```bash
# On PACE ICE, from the worktree
sbatch --export=SCORES_DIR=<scores_path>,SHARD_DIR=<shards_path>,MANIFEST=<manifest_path>,OUTPUT_DIR=<output_path> \
  scripts/slurm/attribution/trackstar_bin_aggregate_perquery.sbatch
```

Transfer the output CSVs locally and regenerate figures.

## Open items

- Publication formatting: figures use Plotly defaults. May need font size adjustment for camera-ready COLM submission dimensions.
- Appendix tables: current top-bin tables show top 20. The `top_bins_table()` function accepts a `top_k` parameter for longer lists.
- Proponent examples: current selection is 10 queries x 3 docs per benchmark. Can be expanded with `--max-queries` and `--max-rank` flags.