Customize Embeddings for Enhanced Text Analysis
OpenAI cookbook notebook that trains a matrix to customize embeddings to a task, cutting binary classification error by up to 50%.
Why it matters
Improve the relevance and accuracy of text embeddings for your specific use case. This asset helps you fine-tune embeddings to better emphasize aspects critical to your application, leading to significant improvements in tasks like classification and similarity search.
Outcomes
What it gets done
Train a custom embedding matrix using labeled text pairs (similar/dissimilar).
Generate synthetic negative examples for training if only positive data is available.
Evaluate embedding performance using cosine similarity and accuracy metrics.
Apply the optimized matrix to generate custom embeddings that enhance downstream tasks.
Install
Add it to your toolbox
Run in your project directory:
curl -fsSL https://spark.entire.vc/get/oai-customizingembeddings | bash Steps
Steps in the chain
Overview
Customizing embeddings
An OpenAI cookbook notebook that trains a matrix on labelled similar/dissimilar text pairs to customize embeddings, reducing binary classification error by up to 50%. Use when generic OpenAI embeddings don't separate a task's classes well enough and labelled or label-derivable pairs are available to train against.
What it does
Customizing Embeddings is an OpenAI cookbook notebook demonstrating how to adapt general-purpose OpenAI embeddings to a specific task. It trains on labelled pairs (text_1, text_2, label), where the label is +1 if the pair is similar and -1 if dissimilar, and produces a matrix you multiply your embeddings by to get a 'custom embedding' that better emphasizes the aspects of the text relevant to your use case. In binary classification use cases the notebook reports error-rate drops of up to 50%.
When to use - and when NOT to
Use it when off-the-shelf embeddings aren't separating your similar/dissimilar pairs cleanly enough for a classification or clustering task. The worked example uses 1,000 sentence pairs from the SNLI corpus - logically entailed pairs as positives, with synthetic negatives generated by combining sentences from different pairs. For a clustering use case, positives instead come from text pairs within the same cluster and negatives from pairs across different clusters. Good improvements have been observed with as few as ~100 training examples, though more examples perform better. It is not needed when generic embeddings already separate your classes well, and it requires labelled or label-derivable pairs to train against.
Inputs and outputs
Input is a dataset of text pairs with similarity labels, plus a chosen embedding engine and a location to cache computed embeddings. The process_input_data function is meant to be rewritten to match your data's format. The notebook proceeds step by step: load and process the input data; split into training and test sets before generating any synthetic positives or negatives (to avoid text leaking between train and test, which would inflate reported accuracy); generate synthetic negatives (or positives, for multiclass data) if the raw dataset doesn't already have both; calculate embeddings and cosine similarities, caching results so re-running the notebook doesn't re-incur embedding cost; and plot the distribution of cosine similarity for similar versus dissimilar pairs, computing accuracy as a simple threshold rule (similar if cosine similarity exceeds some cutoff X, dissimilar otherwise).
Integrations
Embeddings are normalized to length 1, so cosine similarity is equivalent to dot product in this setup; the notebook notes that L1, L2, and cosine distance all perform about the same in its experience. The final steps optimize the transformation matrix against the training data and plot before-and-after cosine-similarity distributions to show how much more cleanly the trained matrix separates similar from dissimilar pairs.
Who it's for
Developers and data scientists who have a working classification or clustering pipeline built on OpenAI embeddings - whether on the SNLI corpus-style entailment data or their own - but need better separation between classes than generic embeddings provide, and who have (or can derive) labelled similar/dissimilar text pairs to train the custom transformation on.
Source README
Customizing embeddings
This notebook demonstrates one way to customize OpenAI embeddings to a particular task.
The input is training data in the form of [text_1, text_2, label] where label is +1 if the pairs are similar and -1 if the pairs are dissimilar.
The output is a matrix that you can use to multiply your embeddings. The product of this multiplication is a 'custom embedding' that will better emphasize aspects of the text relevant to your use case. In binary classification use cases, we've seen error rates drop by as much as 50%.
In the following example, I use 1,000 sentence pairs picked from the SNLI corpus. Each pair of sentences are logically entailed (i.e., one implies the other). These pairs are our positives (label = 1). We generate synthetic negatives by combining sentences from different pairs, which are presumed to not be logically entailed (label = -1).
For a clustering use case, you can generate positives by creating pairs from texts in the same clusters and generate negatives by creating pairs from sentences in different clusters.
With other data sets, we have seen decent improvement with as little as ~100 training examples. Of course, performance will be better with more examples.
0. Imports
1. Inputs
Most inputs are here. The key things to change are where to load your datset from, where to save a cache of embeddings to, and which embedding engine you want to use.
Depending on how your data is formatted, you'll want to rewrite the process_input_data function.
2. Load and process input data
3. Split data into training test sets
Note that it's important to split data into training and test sets before generating synethetic negatives or positives. You don't want any text strings in the training data to show up in the test data. If there's contamination, the test metrics will look better than they'll actually be in production.
4. Generate synthetic negatives
This is another piece of the code that you will need to modify to match your use case.
If you have data with positives and negatives, you can skip this section.
If you have data with only positives, you can mostly keep it as is, where it generates negatives only.
If you have multiclass data, you will want to generate both positives and negatives. The positives can be pairs of text that share labels, and the negatives can be pairs of text that do not share labels.
The final output should be a dataframe with text pairs, where each pair is labeled -1 or 1.
5. Calculate embeddings and cosine similarities
Here, I create a cache to save the embeddings. This is handy so that you don't have to pay again if you want to run the code again.
6. Plot distribution of cosine similarity
Here we measure similarity of text using cosine similarity. In our experience, most distance functions (L1, L2, cosine similarity) all work about the same. Note that our embeddings are already normalized to length 1, so cosine similarity is equivalent to dot product.
The graphs show how much the overlap there is between the distribution of cosine similarities for similar and dissimilar pairs. If there is a high amount of overlap, that means there are some dissimilar pairs with greater cosine similarity than some similar pairs.
The accuracy I compute is the accuracy of a simple rule that predicts 'similar (1)' if the cosine similarity is above some threshold X and otherwise predicts 'dissimilar (0)'.
7. Optimize the matrix using the training data provided
8. Plot the before & after, showing the results of the best matrix found during training
The better the matrix is, the more cleanly it will separate the similar and dissimilar pairs.
FAQ
Common questions
Discussion
Questions & comments · 0
Sign In Sign in to leave a comment.