fix: honor training_data_fraction when training shadow models

[fazelehh]

Closes #345.

Problem

construct_balanced_assignments() always assigned exactly half of the shadow population to every shadow model, so any training_data_fraction other than 0.5 was silently ignored. Worse: the cache signature records the fraction-derived data_size, while the models were actually trained on the half-population size — so with fraction ≠ 0.5 cached shadow models never matched their signature and every run retrained from scratch (the behavior reported in #345, and one of the two causes behind #432).

Fix

• construct_balanced_assignments() takes points_per_model (the same data_size the cache signature records) and uses greedy least-loaded selection with random tie-breaking.
• Every model gets exactly the requested number of points, and each point's inclusion count stays within 1 across models — the balanced IN/OUT property (Shadow model datasplit #334) now holds for any fraction, not just 0.5.
• Default stays at half the population, so existing caches for fraction 0.5 remain valid.

Tests

• test_construct_balanced_assignments_respects_points_per_model: exact per-model sizes and ≤1 inclusion-count spread for fractions 0.3 / 0.5 / 0.72 / 1.0, plus rejection of degenerate sizes.
• test_shadow_model_training_fraction_is_used_and_cached: shadow models train on int(0.75 * population) points and a re-run with the same fraction reuses the cache instead of retraining. Both fail on the old code.
• Full leakpro/tests/mia_attacks/ suite passes (49/49).

🤖 Generated with Claude Code

https://claude.ai/code/session_01U4j7ifPMD8R2Mgd7om97F2

[henrikfo]

Some changes/ additions is needed to make sure all attacks for as intended I think.

Also added some thoughts on making these leakpro changes a little bit more robust, not needed for approval though! :)

[fazelehh]

@henrikfo All four points addressed. Every one still reproduced against current main, so none had aged out. Details are in the three threads; summary here.

Branch first merged origin/main (it was cut before the get_device/seed work landed), then four commits:

Commit	Your point	What changed
89138f4	cached models and resume	construct_balanced_assignments takes prior_inclusion_counts; create_shadow_models seeds it from cached models' train_indices and generates rows only for new models
a3f5059	LiRA NaN with fraction > 0.5	Shared audit_points_with_reference_models predicate; LiRA and MS-LiRA drop unscorable points and warn with the counts and the fraction
c82101b	BASE forces 0.5	Both num_shadow_models // 2 asserts replaced with a positive-OUT-count check
d62b56e	is 1.0 allowed	gt=0.0, lt=1.0 across the eight shadow-model attack configs, plus the same guard inside the handler

Two things worth your attention before approving.

I fixed a pre-existing bug inside this PR. The BASE assert already fails on main for any odd num_shadow_models, independent of this branch. I kept the fix here because it is the same assert that blocks the feature, but say the word and I will split it into its own PR against main.

LiRA drops points, BASE raises. Deliberate asymmetry. score_samples is driven by RaMIA over range clusters, and silently dropping points there would change the clusters rather than just the sample size. Pushed back slightly on the yoqo/dts precedent in one respect: the drop is logged as a warning naming the configuration, because under balanced sampling a non-empty drop set means the fraction is wrong for the model count, so it is structural rather than incidental. A silently smaller audit set understates the result.

Not done, flagging rather than fixing: yoqo.py and dts.py still carry two verbatim copies of the filtering block. Correcting myself here, since I first wrote that switching them over was mechanical and had not checked. It is not, in general. Neither attack computes a per-point OUT statistic: yoqo uses the mask as per-model loss weights and sets them to all-ones offline, and dts uses it only for the online filter and an offline log line. So neither can hit the empty-slice NaN the predicate exists to prevent, and their online-only filtering is correct for them. A behaviour-preserving dedup is still available, about six lines, but only with the call kept inside the if self.online branch. Adopting the predicate the way the sentence above implied would make them drop points they can legitimately score. Worth a separate PR at most.

Verification: 440 passed, ruff clean, on d62b56e in leakpro_py311. Eleven new tests, including two that fail on the old code. Two failures excluded as environment issues, both confirmed identical on origin/main: four collection errors from missing transformers/diffusers extras, and test_aux_file_path, which hardcodes the repo directory name and fails in any worktree.