This repository contains my experiments on generating proper message embeddings to be able to search them semantically later.
Personally, I am dissatisfied with how message search functionality is implemented with messaging apps. So I tried building my own.
My initial plan was to embed each message with a Sentence Transfomer, however, that doesn't work well in practice since there are a lot of messages that does not contain any meaningful semantics.
Instead, I opted to aggregate eight messages into a context, joined with special markers around
each message, and pulled their embeddings close to the individual messages they contain. This approach
yielded significantly better results.
Specifically, a context looks like this:
<user0>This is a message</user>
<user1>This is another message</user>
<user2>Yay another message :D</user>
<user3>I love talking</user>
<user4>Me too!</user>
<user5>Me three!</user>
<user6>I am running out of examples</user>
<user7>Shut up!</user>
Note: The <userX> tokens shown above are static, they don't contain any user related information.
That's just what I call them.
This works under the assumption is that a search query is semantically close to a message contained by a context (thus close to the context itself) which is what a sentence transformer tries to do.
This repository contains the scripts to preprocess the data, train the model and test it. Training script has Huggingface Accelerate support, so it should be able to utilize multiple GPUs, however I only worked with single-GPU setups so it is not tested well. It also has MLFlow support to log training metrics.
I tested Triplet loss, InfoNCE and CLIP losses, triplet works the best.
I started training and experimenting on an initial dataset. During my experiments the methodology changed and matured, I found out having a stride that is equal to the message context length works the best, I verified that mean pooling works best and so on.
The testing methodology is as follows: Embed all of the validation set and save it to a vector database. And measure the retrieval
performance by selecting a random message from the message group and retrieving context vectors. Check generate_embeddings_for_testing.py
and test.py for detailed info.
The two models I think worth displaying for this dataset got the scores below:
| Model Name | message_context_size | Top@1 | Top@5 | Top@8 |
|---|---|---|---|---|
| msg_emb_c5 | 5 | 86.19% | 96.79% | 98.01% |
| msg_emb_c8 | 8 | 63.00% | 83.69% | 87.73% |
Then, I went and trained several models with different context sizes with a bigger dataset. I also properly tested the effect of message context size.
| Model Name | message_context_size | Top@1 | Top@5 | Top@8 |
|---|---|---|---|---|
| context_1 | 1 | 92.48% | 94.71% | 95.06% |
| context_3 | 3 | 87.81% | 92.57% | 93.27% |
| context_5 | 5 | 78.56% | 89.06% | 90.61% |
| context_8 | 8 | 57.01% | 77.34% | 81.35% |
| context_10 | 10 | 40.02% | 61.68% | 67.06% |
From this table alone it might seem reducing the message context size gives better models but don't let it fool you. Less mcs = less information lost. In fact the problem should be trivial for mcs = 1, in that specific case user tokens act as noise.
I also tested my method agains raw embedding centroids. The results can be seen below:
| Model Name | message_context_size | Top@1 | Top@5 | Top@8 |
|---|---|---|---|---|
| context_3 | 3 | 87.81% | 92.57% | 93.27% |
| context_3_centroids | 3 | 82.98% | 90.40% | 91.50% |
| context_5 | 5 | 78.56% | 89.06% | 90.61% |
| context_5_centroids | 5 | 65.42% | 81.06% | 83.94% |
| context_8 | 8 | 57.01% | 77.34% | 81.35% |
| context_8_centroids | 8 | 42.75% | 62.12% | 67.06% |
| context_10 | 10 | 40.02% | 61.68% | 67.06% |
| context_10_centroids | 10 | 30.90% | 48.43% | 53.50% |
From the table above, it is clear that the methodology I have helps model learn to compress context information.
In order to test the effectiveness of the user tokens, I did an experiment where I fine-tuned a model without them. The results are below.
It can be seen that the results of context_8_centroids and context_8_no_user_tokens are very similar. This lead me to question
whether the model without user tokens basically end up learning to regress centroids. In order to test this idea, I set up an experiment
where I tried learning a linear transformation between the centroid embeddings to context embeddings. I used context_8_no_user_tokens
model. For the training details, please refer to centroid_vs_no_user_tokens.py. The values of the hyper-parameters used in the experiment are
the ones set default in the file, unless specified otherwise. The results are shown below:
| group embedding normalization type | initial avg. validation loss | initial avg. cosine similarity | final avg. validation loss | final avg. cosine similarity |
|---|---|---|---|---|
| none | 0.267525 | -0.004462 | 0.037941 | 0.925254 |
| l2_then_mean | 0.001387 | -0.004462 | 0.000183 | 0.933422 |
| mean_then_l2 | 0.001746 | -0.004462 | 0.000190 | 0.926644 |
It can be seen from the table that the cosine similarity between the context embedding and linearly transformed centroid embeddings is around 92.5% which
supports the claim.
Update: I ran the same test with the context_8 model, learning a linear transformation from the centroid embeddings and context embedding with user tokens and
I got 93.16% average cosine similairty, which makes the conclusion above obsolete. Perhaps without user tokens there is still something to learn but metrics-wise you
are still limited because you lack the information passing, aggregation and seperation capacity user tokens provide.