Skip to content

⚡ Bolt: vectorize BasicEstimator.predict and pre-calculate norms#19

Open
guesswh0 wants to merge 1 commit into
masterfrom
bolt-optimize-predict-10186505443181998047
Open

⚡ Bolt: vectorize BasicEstimator.predict and pre-calculate norms#19
guesswh0 wants to merge 1 commit into
masterfrom
bolt-optimize-predict-10186505443181998047

Conversation

@guesswh0
Copy link
Copy Markdown
Owner

💡 What

Vectorized the BasicEstimator.predict method using NumPy operations and the squared Euclidean distance expansion formula. Additionally, updated the fit method to pre-calculate and store the squared norms of the fitted embeddings.

🎯 Why

The original implementation used a Python loop to iterate over each input embedding, performing a np.linalg.norm calculation against all fitted embeddings in each iteration. This was inefficient for large batches of input embeddings or many fitted samples.

📊 Impact

Expected performance improvement of approximately 12x for batch predictions.

  • Before: ~0.245 seconds for 500 input vs 2000 fitted embeddings.
  • After: ~0.020 seconds for the same workload.

🔬 Measurement

Verified using benchmark_predict.py and extra/verify_optimization.py. The latter ensures numerical consistency (within rtol=1e-5) between the iterative and vectorized versions, and validates backward compatibility for older models.

⚡ Bolt Journal Entry

Added a learning entry to .jules/bolt.md regarding the use of np.maximum(..., 0) to maintain stability when using the expansion formula.


PR created automatically by Jules for task 10186505443181998047 started by @guesswh0

Optimized the `BasicEstimator` by vectorizing the distance calculation
in `predict` and pre-calculating squared norms in `fit`.

Key improvements:
- Switched from an iterative loop over input embeddings to a fully
  vectorized approach using the squared Euclidean distance expansion
  formula: ||a-b||^2 = ||a||^2 + ||b||^2 - 2<a, b>.
- Pre-calculate squared norms of fitted embeddings in `fit` to avoid
  redundant O(N*D) calculations during each prediction.
- Added backward compatibility for models loaded without pre-calculated
  norms.
- Handled empty input scenarios gracefully.

Impact:
Measured ~12x speedup for 500 input embeddings against 2000 fitted
embeddings (from ~0.245s down to ~0.02s).

Co-authored-by: guesswh0 <10531675+guesswh0@users.noreply.github.com>
@google-labs-jules
Copy link
Copy Markdown
Contributor

👋 Jules, reporting for duty! I'm here to lend a hand with this pull request.

When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down.

I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job!

For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with @jules. You can find this option in the Pull Request section of your global Jules UI settings. You can always switch back!

New to Jules? Learn more at jules.google/docs.


For security, I will only act on instructions from the user who triggered this task.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant