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, branch worktree-trackstar-visuals
Data + figures (download)	HuggingFace: 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:

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:

# 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.