Embedding Alignment

Embedding Alignment

August 2025

Why?

Given the increasing use of vendored vector store solutions and dependencies on AI for mission critical systems, it is important to determine ways to ensure systems stay reliable when third party APIs degrade. Embedding alignment can be an approach to re-map embeddings from one vendored system to another in the hope that systems stay reliable at the cost of minimal degradation at inference time. This approach requires only maintaining a linear mapping between embeddings rather than duplicating vector stores for multiple vendors which can be expensive in terms of ownership, processes and infrastructure.

Method

Embedding alignment can be framed as a Procrustes problem. The general approach referenced in Creating Embeddings of Heterogeneous Relational Datasets for Data Integration Tasks and Word Translation Without Parallel Data

The high level procedure is to determine anchors by constructing pair-wise embeddings and then finding an appropriate (linear) mapping, typically using Procurstes Analysis. In practise, we’re interested in the general class of Procrustes Problem which devolves into solving for (linear) homomorphism since there is no particular reasons metric spaces of arbitrary embeddings reside in the same metric space/same number of dimensions.

As an approprimation, we can greedily estimate a ‘good enough’ solution using linear regression.

 1import numpy as np
 2from model2vec import StaticModel
 3from sklearn.linear_model import Ridge
 4from sklearn.multioutput import MultiOutputRegressor
 5
 6model_8m = StaticModel.from_pretrained("minishlab/potion-base-8M")
 7model_4m = StaticModel.from_pretrained("minishlab/potion-base-4M")
 8
 9# define anchor corpus
10corpus = [
11    "It's dangerous to go alone! Take this.",
12    "All your base are belong to us",
13    "The cake is a lie.",
14    "Hey you, you're finally awake",
15    "You must construct additional pylons!"
16]
17
18X = model_4m.encode(corpus)
19Y = model_8m.encode(corpus)
20
21# determine linear transformation matrix
22regr = MultiOutputRegressor(Ridge(random_state=123, fit_intercept=False)).fit(X, Y)
23transformation_matrix = np.vstack([x.coef_ for x in regr.estimators_])
24print(np.allclose(regr.predict(X), X.dot(transformation_matrix.T)))