Three upgrades to the DynamicSimilaritySampler infrastructure:
1. **GPU kNN** (`dss_knn_device="cuda"`, default):
Moves the per-seed similarity matmul to the GPU. At 25K train items
this cuts per-epoch sampler generation from 17s to 1.6s — a 10.8x
speedup. Negligible VRAM (100MB for the [N, 1024] embedding tensor).
2. **LSH index** (`src/datasets/lsh_index.py`, opt-in via `dss_use_lsh=True`):
Random-projection cosine-LSH with H tables of B bits each. When enabled,
the sampler narrows the candidate pool per seed via hash-bucket lookup
before exact refinement. At 25K it's a wash (pool already fits in VRAM)
but provides a scaling path for 100K+ where the N² similarity matrix
would stop fitting. Default off.
3. **Embedding cache** (`src/datasets/embedding_cache.py`, `dss_cache_dir` config):
Disk-backed cache for drone query embeddings, keyed by epoch. Skips
re-embedding on --resume and lets ablations replay from a snapshot.
Atomic writes via `.tmp` → `.replace`.
Measured on 25K train entries, 1024-dim random embeddings:
CPU kNN: 17.44s
GPU kNN: 1.62s (10.8x)
GPU + LSH: 1.42s (LSH candidate pool 0.05% of N)
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
d98d853455