CLI crashes on every run: optimize_embeddings() gets an unexpected 'device' kwarg

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:

```python
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!

Start in code/free_embedding_experiment.py, reviewing main() around line 775 and run_experiment_base() around line 519, then inspect optimize_embeddings() and its accepted parameters. Run the documented command with jax[cpu]==0.6.2 and verify that training proceeds without the unexpected device keyword error and produces non-error result rows.

Written by the indexing model from the issue text.