Adds image_features parameter to predict for pre-computed embeddings

[NetZissou]

Disclaimer: This PR was developed with assistance from Claude Opus 4.6 (1M context). The author has reviewed all code changes and test additions. CI has been executed successfully in the forked repo. Opening this PR to request review from the package maintainers for further feedback and iteration.

Summary

This PR adds an optional image_features parameter to predict() on TreeOfLifeClassifier and CustomLabelsClassifier (CustomLabelsBinningClassifier inherits this through CustomLabelsClassifier). When provided, the method skips image encoding and computes classification directly from pre-computed embeddings.

Embedding validation

The method validates input embeddings before classification:

• Verifies tensor is 2D (N, embedding_dim)
• Checks that embedding_dim matches the model's expected dimension (model.visual.output_dim)
• Normalizes the embedding vector via L2 norm only if not already normalized

Test plan

New tests in TestPredictFromEmbeddings:

• Results from embeddings match results from images exactly (all fields except file_name)
• Species, family, and multi-image predictions
• Unnormalized features auto-normalized with correct classifications
• CustomLabelsClassifier and CustomLabelsBinningClassifier
• Error cases: no inputs, wrong tensor dim, wrong embedding dim, image-embedding length mismatch

Closes #167

[hlapp]

Thanks @NetZissou. The part of the implementation approach that I don't like here is that now creating probabilities is taking place redundantly in two different functions. This also creates more code noise than I think should be needed in the predict() method.

Instead, shouldn't the clean way to handle this in the predict() method be to see whether image_features are already provided. If they are, apply basic checks like correct dimensions etc. If they are not, create them (like the are being created now from images). Then proceed with creating probabilities from image_features.

[hlapp]

@NetZissou just FYI, it might be advisable to rebase on main to bring in the changes from #179. It's well possible your changes so far are not in conflict at all, but some stuff did get moved around.

[NetZissou]

@egrace479 @hlapp

@NetZissou, @hlapp and I discussed the implementation this morning. It would make the most sense to expect (require) the list of image identifiers (e.g., list of filename strings) that correspond to the embeddings to be passed as well. These can then be used for the image keys. The embeddings would go through the size checks previously discussed and you have the sanity check that the list lengths match. We just want to make sure it's well-documented then.

Re-implemented based on the comment above. Ready for re-review.

[hlapp]

Thanks @NetZissou. Some thoughts and comments:

• Since this is trying to optimize performance for the case that someone has the image embeddings already, I think we can expect at least by default that they are already normalized, rather than re-normalizing them (for testing whether they are normalized) each time. So I think that should be conditional to a parameter that by default is off?

• The _resolve_image_features() function doesn't really resolve image features, given that it returns probabilities, not resolved image features. So instead it should probably better be named something like _create_probabilities_for_imgs_or_img_feats() (or maybe even just _create_probabilities()?). And given the number of parameters and None being allowed in several places, all invocations should arguably use named parameter syntax.

• In this implementation (as per _resolve_image_features()), there's actually really nothing directly in common between creating probabilities from image features and creating them from images: there is no code before or after the conditional that would run either way. I guess the shared part is indirect: creating probabilities from images invokes create_batched_probabilities_for_images(), whereas creating them from image features invokes create_probabilities(). Of course, indirectly the former invokes create_probabilities_for_images() which then invokes create_probabilities(), too. It seems what this really means is the following:

  • If I have images and no image features, encode the images to obtain image features. This should be done batched for performance optimization.
  • Once I have image features, create probabilities. This need not be batched.

  In your current implementation, creating probabilities from image features will be batched if starting from images, and the batching will be bypassed if image embeddings are already in hand. Is this reasonable?

[NetZissou]

Since this is trying to optimize performance for the case that someone has the image embeddings already, I think we can expect at least by default that they are already normalized, rather than re-normalizing them (for testing whether they are normalized) each time. So I think that should be conditional to a parameter that by default is off?

I think it's worth adding an explicit param normalize_features: bool. By turning it off, user can assure that "I promise input vectors are already normalized". Regarding to the norm-check performance decrease, I think the main goal of this PR is to reduce the unnecessary image encoding (ViT forward pass).

N: number of entries
D: vector dimension
C: number of classes in the text embeddings

Dimensions: 
image_features: (N, D)
txt_embeddings: (D, C)
logits/probs: (N, C)

norm check: N x D
logits matmul: N x C x D

(norm check) / matmul = 1/C

The norm check costs about 1/C of the logits matmul. For BioCLIP 2, norm check is about 10^5 cheaper than matmul. And they both are dwarfed by image encoding. I think we can afford to keep this cheap norm sanity check that warns or raises instead of silently producing wrong probs. Setting normalize_features = False bypass the norm check.

---

The _resolve_image_features() function doesn't really resolve image features, given that it returns probabilities, not resolved image features. So instead it should probably better be named something like _create_probabilities_for_imgs_or_img_feats() (or maybe even just _create_probabilities()?). And given the number of parameters and None being allowed in several places, all invocations should arguably use named parameter syntax.

Agreed. I prefer _create_probabilities_for_imgs_or_img_feats() because _create_probabilities() name would confuse with the existing create_probabilities().

---

In your current implementation, creating probabilities from image features will be batched if starting from images, and the batching will be bypassed if image embeddings are already in hand. Is this reasonable?

You're right, the image features path should be batched!

Current implementation:

• if only images, no image features provided, the image encoding is batched create_batched_probabilities_for_images
• if images, and image features are both provided, the logits/probs calculation image_features @ txt_features is not batched, which can cause OOM errors for large N.

[hlapp]

The norm check costs about 1/C of the logits matmul. For BioCLIP 2, norm check is about 10^5 cheaper than matmul. And they both are dwarfed by image encoding. I think we can afford to keep this cheap norm sanity check that warns or raises instead of silently producing wrong probs.

Probably. But your argument assumes that norm check and matmul use the same hardware acceleration, which may or may not be true, and it's also just a matter of future-proofing principle.

[NetZissou]

@hlapp thanks, this makes sense. In _validate_image_features, I'll drop unit norm check torch.allclose entirely, still keeping the tensor dimension checks to be defensive, good for error catching.

Introduce new parameter normalize_features:

• True: unconditionally F.normalizes the input
• False (default): skip normalization entirely

[hlapp]

@NetZissou were you still going to implement the change as per above? Just wanted to make I didn't misunderstand your response/comment.

[NetZissou]

@hlapp yes that's still in the work