Hi! Thanks for releasing the code for the free-embedding experiments, I've been replicating and extending them over the past week and the methodology has honestly been a pleasure to work with. Wanted to flag one packaging issue: the script as released crashes on every training run when you invoke it with its own documented example command.
Repro (clean clone, jax[cpu]==0.6.2):
python code/free_embedding_experiment.py --d=4 --k=2 --enable_critical_n_search=11 \
--results_output_path='d4_k2.json' --device=cpu
Every probe fails right away with
TypeError: optimize_embeddings() got an unexpected keyword argument 'device'
and the search then degenerates quietly: the gallop concludes "N is between None and 1" in under a second and every result row is an error record, so it looks like it ran when nothing actually trained.
Root cause: main() folds every non-None absl flag into base_params (around line 775), and run_experiment_base() forwards all of them to optimize_embeddings(**params_for_opt) after popping only "q" (around line 519). The signature of optimize_embeddings() doesn't accept device (that one is consumed separately via experiment_data["device_context"]), so the first extra kwarg raises.
A minimal fix that worked for me:
allowed = set(inspect.signature(optimize_embeddings).parameters) - {"experiment_data", "q"}
params_for_opt = {k: v for k, v in params_for_opt.items() if k in allowed}(or just pop "device" explicitly and restrict the flag harvest to the experiment's own flags rather than all of FLAGS.)
For what it's worth, once it's driven with a corrected parameter set the code reproduces the paper's published critical-n values in my hands, d=16 gives exactly 79, so this is purely a release-packaging thing and not a science thing. One small note for fellow replicators while I'm here: the paper text describes incrementing n by 1 until failure, while the released code actually does galloping plus binary search plus a small local sweep. Might be worth a line in the README.
Happy to open a PR with the fix if that's helpful. Thanks again for putting the code and data out there!
Hi! Thanks for releasing the code for the free-embedding experiments, I've been replicating and extending them over the past week and the methodology has honestly been a pleasure to work with. Wanted to flag one packaging issue: the script as released crashes on every training run when you invoke it with its own documented example command.
Repro (clean clone,
jax[cpu]==0.6.2):Every probe fails right away with
and the search then degenerates quietly: the gallop concludes "N is between None and 1" in under a second and every result row is an error record, so it looks like it ran when nothing actually trained.
Root cause: main() folds every non-None absl flag into base_params (around line 775), and run_experiment_base() forwards all of them to optimize_embeddings(**params_for_opt) after popping only "q" (around line 519). The signature of optimize_embeddings() doesn't accept device (that one is consumed separately via experiment_data["device_context"]), so the first extra kwarg raises.
A minimal fix that worked for me:
(or just pop "device" explicitly and restrict the flag harvest to the experiment's own flags rather than all of FLAGS.)
For what it's worth, once it's driven with a corrected parameter set the code reproduces the paper's published critical-n values in my hands, d=16 gives exactly 79, so this is purely a release-packaging thing and not a science thing. One small note for fellow replicators while I'm here: the paper text describes incrementing n by 1 until failure, while the released code actually does galloping plus binary search plus a small local sweep. Might be worth a line in the README.
Happy to open a PR with the fix if that's helpful. Thanks again for putting the code and data out there!