Search recall optimization - what appropriate loss function to use?

Question

I am studying machine learning and wanted to work on a project of my own so that I have better chances after graduating college. I'm studying the application of ML to improve searches using a toy dataset. Here's a couple of scenarios:

1. Document retrieval (search) system: We have a (source) document with us and we're trying to find a matching document in a database. The source document has text and image attributes - for simplicity let's say a title and a single image. Each search result will also be a document - with a title and at most one image.
2. A search engine: We have a query comprised of both text and an image (like google image search allows text to be added to the query as well). Each search result will be a website with text and image attributes (for simplicity, webpage title and at most one image)

More generally, I have a search system - whatever we're trying to search for has text and an image associated with it. Each search result will also have its own attributes. Let's denote each source document (or in case of search engine, a query) with $D$ and each search result with $R$.

The metric that makes the most natural sense to me is top $k$ search recall (not sure what the more widely used name for this is): suppose we have a sample of source documents or queries $D_1,\ldots,D_n$, and let the search results for $D_i$ be $\{R_{i1},R_{i2},\ldots,R_{im_i}\}$. If any of $R_{i1},\ldots,R_{ik}$ is a match for $D_i$, then we can say the search worked well for $D_i$. So the top $k$ search recall is the proportion of source docs/queries for which at least one match was found in the top $k$ search results.

Labels

Just to clarify, I'm looking for a search result scoring/ranking model that uses source doc/query attributes and attributes of search results.

Now to tweak my search system, let's say I also have a dataset with source document - search result pairs with labels - whether the search result is a match for the source document. i.e. something like

D1   R11   No Match
D1   R12   No Match
D1   R13   Match
.
.
D1   R1100 No Match
D2   R21   No Match
D2   R22   Match
D2   R23   Match
.
.
D2   R2100 No Match

(as you can imagine it would be a highly imbalanced dataset)

My current approach to tweak the system is like this:

Features and modelling

I can form features linking a source doc/query to a search result. e.g. some sort of embedding similarity between titles, or jaccard similarity of titles, or embedding similarity of their images, etc. There will be many null values too since some results may not have images, or some source docs/queries may not have images, or a result may not have a clear title, etc.

I'm using an XGB classifier - training, validating, testing it using the labeled dataset - the usual routine. I'm using the predict_proba method of the classifier to output the probability of a doc-result pair being a match. I use that probability as a score. Then for each $D_i$, the results with the top $k$ scores are shortlisted. And then the top $k$ search recall is measured for the $D_1,\ldots,D_n$ sample.

Problems

I'm not completely satisfied with the approach since I don't think I'm using the optimal loss function for the classifier. Right now it's just the standard XGBClassifier loss function. The classes are highly imbalanced as I said before. Artificially balancing the classes leads to a model that can better classify a matching or non-matching result, BUT it actually reduces top $k$ search recall compared to a model trained on unbalanced dataset for some value of $k$.

What else have I tried: So I'm using jaccard text similarity and image similarity as two of the features.

• One modification I tried was using all features as separate "scores", and then calculated ranking for each of those "scores". So now I have different feature ranks - I used reciprocal rank fusion to combine them into a single rank. RRF top $k$ does not do as well as the classifier top $k$
• Another modification: I calculated classifier top $k-5$. Then I clubbed those results together with jaccard text similarity top $3$ and image similarity top $2$. This method does better than the classifier top $k$

The last point clearly indicates that my way of combining various features into a model is suboptimal. Otherwise the model itself should have learned to accommodate this top $3$ and top $2$ into its top $k$.

My question is, what is the recommended loss function for simultaneously maximizing the top $k$ search recall and minimizing $k$? What is the best way to combine all this feature info into the optimal top $k$?