Aprendizado de Representações Matryoshka Explicado: O Método por Trás dos Embeddings de Texto Eficientes da OpenAI
Desenvolver modelos de aprendizado de máquina da vida real sempre envolve um trade-off entre custo e desempenho. Por exemplo, quanto maiores forem o modelo e o conjunto de dados que usamos durante o treinamento, mais capaz será, em última análise, o nosso modelo treinado. No entanto, essa maior capacidade vem com mais tempo e custos computacionais mais altos, pois o treinamento exige durações significativamente mais longas. O mesmo princípio se aplica durante a inferência, em que modelos maiores tendem a gerar representações de características maiores que exigem mais memória para armazenamento.
Como as necessidades durante o treinamento e a inferência de aprendizado de máquina podem variar significativamente de um caso de uso para outro, é essencial ter um método que nos permita trocar uma pequena parcela do desempenho de um modelo por custos reduzidos. É aqui que um método como Matryoshka Representation Learning (MRL) entra em ação. Por exemplo, o modelo text-embedding-3-small da OpenAI utiliza MRL para permitir que desenvolvedores encurtem embeddings enquanto mantêm suas propriedades centrais de representação de conceitos. Ao permitir o ajuste das dimensões dos embeddings, o MRL ajuda a encontrar o equilíbrio perfeito entre eficiência de custos e desempenho do modelo.
Este artigo explorará como o MRL funciona, sua implementação e como ele permite modelos de aprendizado de máquina escaláveis e eficientes. Vamos começar com a motivação por trás do MRL.
A Motivação por Trás do Matryoshka Representation Learning (MRL)
O trade-off entre custo e desempenho é sempre uma consideração antes de desenvolver ou usar quaisquer modelos de aprendizado de máquina.
Durante o treinamento, os parâmetros de um modelo e a quantidade de dados de treinamento influenciam diretamente seu desempenho final. Modelos e conjuntos de dados de treinamento maiores tendem a resultar em modelos mais capazes. No entanto, o custo computacional para treinar esses modelos também é mais alto.
Enquanto isso, durante a inferência, modelos maiores produzem representações de características maiores. Esse aumento nas representações de características significa que precisamos de mais memória para armazená-las.
Vamos usar BERT como exemplo no contexto de recuperação de informações. O BERT é um dos primeiros modelos de aprendizado profundo baseados em Transformer que alcançou desempenho de estado da arte em vários conjuntos de dados de benchmark em classificação de texto, Reconhecimento de Entidades Nomeadas (NER), resposta a perguntas (QnA) e muito mais.
Figura: Fluxo de trabalho da geração de embeddings usando o modelo base BERT.
Em poucas palavras, o BERT recebe uma frase ou uma palavra como entrada e a transforma em um embedding de tamanho fixo. Esse embedding captura o significado semântico da entrada original, e seu tamanho depende da variante específica do modelo usada.
O BERT tem duas variantes principais: o modelo base e o modelo large. O modelo base produz um embedding de 768 dimensões, enquanto o modelo large gera um embedding de 1024 dimensões.
Digamos que queiramos armazenar 10 milhões de embeddings em um banco de dados vetorial para um caso de uso de recuperação de informações. Se usarmos o modelo base BERT no formato FP32, precisaremos de aproximadamente 768×10M×4=30,72 GB de memória para armazenar esses embeddings. Em contraste, usar o modelo large do BERT no mesmo cenário exigiria cerca de 40,9 GB. Além disso, um tamanho maior de embedding pode desacelerar o processo de recuperação ao aumentar a complexidade computacional das buscas por similaridade. No entanto, a principal vantagem de usar embeddings maiores é sua capacidade aprimorada de recuperar informações altamente relevantes em comparação com embeddings menores.
Uma solução ideal para esses problemas em recuperação de informações é dividir todo o processo em duas partes: pré-seleção e reranqueamento. A pré-seleção refere-se à recuperação de um conjunto inicial de documentos candidatos a partir de uma vasta coleção em nosso banco de dados vetorial, como o Milvus. Em seguida, o reranqueamento pega esses candidatos pré-selecionados e os reordena para maximizar a relevância dos resultados finais.
Figura- Fluxo de trabalho de pré-seleção e reranqueamento
Figura: Fluxo de trabalho de pré-seleção e reranqueamento.
Para realizar a pré-seleção, o algoritmo calcula a similaridade entre embeddings de consulta e uma grande coleção de embeddings armazenados no banco de dados vetorial. Portanto, usar tamanhos menores de embeddings pode tornar esse cálculo mais eficiente e rápido. Por outro lado, o reranqueamento enfatiza a precisão, o que significa que nos beneficiaríamos de tamanhos maiores de embeddings.
O problema com esse método é que, uma vez que selecionamos um modelo específico, o tamanho dos embeddings ou das representações de características que ele produz é fixo. Essa falta de flexibilidade limita nossa capacidade de usar embeddings de diferentes tamanhos a partir de um único modelo.
Figura- Fluxo de trabalho de um modelo gerando um embedding de tamanho fixo..png
Figura: Fluxo de trabalho de um modelo gerando um embedding de tamanho fixo.
A abordagem Matryoshka Representation Learning (MRL) oferece uma solução intrigante para esse problema: modelos treinados com esse método podem produzir embeddings de vários tamanhos. Exploraremos o mecanismo dessa abordagem em detalhes na próxima seção.
O que é Matryoshka Representation Learning (MRL)?
Matryoshka Representation Learning (MRL) é um método para treinar redes neurais a produzir representações multiescala dentro de um único modelo. Inspirado nas bonecas Matryoshka, nas quais bonecas menores cabem dentro de maiores, o MRL permite que um modelo gere representações de tamanhos variados (do grosseiro ao fino) usando uma única passagem direta. Essa abordagem permite que modelos de aprendizado profundo compreendam melhor relações complexas e nuances nos dados. O MRL é particularmente eficaz para tarefas como busca semântica, recuperação de informações, processamento multilíngue e qualquer aplicação que exija representações matizadas de dados em diferentes níveis de abstração.
Visualização de embeddings Matryoshka com múltiplas camadas de detalhe
Figura: Visualização de embeddings Matryoshka com múltiplas camadas de detalhe
Modelos populares de embedding que adotaram a abordagem Matryoshka Representation Learning (MRL) incluem o text-embedding-3-large da OpenAI, o nomic-embed-text-v1 da Nomic e o gte-multilingual-base da Alibaba.
Como a abordagem MRL funciona
A abordagem MRL nos permite extrair representações de características de qualquer modelo de aprendizado de máquina em vários tamanhos. Por exemplo, em vez de usar as 1024 dimensões originais de um embedding, podemos utilizar as primeiras 16, 32, 64, 128 ou 256 dimensões (ou quaisquer dimensões que escolhermos). A chave para essa capacidade está na forma como o modelo é treinado ao implementar a abordagem MRL.
Figura- Treinamento da função de perda MRL e seu caso de uso durante a inferência
Figura: Treinamento da função de perda MRL e seu caso de uso durante a inferência. Fonte.
Durante o treinamento de modelos com MRL, otimizamos não apenas uma função de perda, como é típico no treinamento padrão de modelos, mas várias. Cada função de perda visa otimizar a representação de características para uma dimensão específica que definimos antecipadamente. Se definirmos cinco dimensões de características diferentes, precisaremos otimizar cinco funções de perda distintas durante o treinamento.
Digamos que queremos treinar um modelo BERT large usando a abordagem MRL. Como mencionado na seção anterior, esse modelo gera um embedding de tamanho 1024. Se definirmos as dimensões de características como 32, 64, 128, 256 e 1024, otimizaremos cinco funções de perda diferentes durante o treinamento para que o modelo aprenda a gerar embeddings otimizados nessas dimensões.
A otimização entre essas várias dimensões é direta: a MRL decompõe a função de perda geral na soma das perdas para cada dimensão individual. Usando nosso exemplo acima, isso pode ser expresso como:
O processo de otimização para um modelo treinado com MRL pode ser representado mais formalmente da seguinte forma:
Figura: Equação de otimização da abordagem MRL. Fonte.
Como você pode ver, a definição da função de perda acima é bastante geral. Portanto, a abordagem MRL pode ser aplicada a quase qualquer modelo e é totalmente independente da arquitetura do modelo. Também podemos ajustar modelos pré-treinados como BERT ou quaisquer outros modelos baseados em transformers para gerar embeddings de vários tamanhos. Continuando nosso cenário acima, agora podemos gerar embeddings de 1024 dimensões a partir do nosso modelo BERT large treinado e embeddings de 32, 64, 128 e 256 dimensões.
No geral, após treinar um modelo com MRL, as dimensões iniciais das características geradas pelo modelo treinado carregam informações mais significativas do que as dimensões posteriores. As primeiras dimensões contêm detalhes de alto nível, enquanto as dimensões posteriores se concentram em informações mais granulares, imitando como uma boneca Matryoshka é estruturada.
No entanto, isso não significa que embeddings mais curtos sejam sempre simplesmente as versões truncadas dos mais longos. Os valores de cada elemento em embeddings mais curtos podem diferir daqueles em embeddings mais longos devido ao fator de escala aplicado a cada dimensão de característica durante o treinamento. No entanto, se definirmos o fator de escala em cada dimensão como igual, os valores de cada elemento entre embeddings mais curtos e mais longos podem parecer idênticos.
Resultados Experimentais da MRL
A abordagem MRL foi avaliada em modelos de aprendizado de máquina em diferentes modalidades, incluindo texto, visão e visão-texto. Os modelos ResNet50 e ViT representam modelos baseados em visão, BERT representa modelos baseados em texto, e ALIGN representa a combinação de visão e texto. Esses modelos foram avaliados principalmente para dois casos de uso comuns: classificação e recuperação. Vamos discutir classificação primeiro.
Ao comparar o desempenho do modelo ResNet50 treinado no conjunto de dados ImageNet-1K com um ResNet50 padrão treinado de forma independente, o modelo MRL alcança uma precisão top-1 comparável em vários tamanhos de representação de características.
Para avaliar ainda mais a utilidade das representações de características para tarefas posteriores, a acurácia do 1-vizinho mais próximo (1-NN) para cada tamanho de representação de características também foi medida. A configuração para calcular o 1-NN é a seguinte: dadas 1,3K amostras de imagens no banco de dados, a tarefa é encontrar o vizinho mais próximo para cada uma das 50K consultas. Com essa configuração, a ResNet50 treinada com MRL é até 2% mais precisa em cada tamanho de representação de características em comparação com sua contraparte de características fixas.
Figura- Acurácia Top-1 e acurácia 1-NN da ResNet50 no ImageNet-1K
Figura: Acurácia Top-1 e acurácia 1-NN da ResNet50 no ImageNet-1K. Fonte.
Enquanto isso, o desempenho do modelo ViT treinado com MRL no conjunto de dados JFT-300M também é muito competitivo em todos os tamanhos de representação. Sua acurácia 1-NN é comparável à do ViT treinado com representações de características de tamanho fixo. Como mostrado na figura abaixo, o desempenho do modelo MRL também é melhor do que o de sua contraparte de tamanho fixo em representações de características menores, em parte porque características aleatórias são selecionadas do modelo de tamanho fixo para representar dimensões menores. Uma tendência semelhante é observada com o modelo ALIGN treinado usando a abordagem MRL; seu desempenho corresponde ao do modelo ALIGN treinado com representações de tamanho fixo.
Figura- Acurácia 1-NN do ViT e ALIGN no ImageNet-1K
Figura: Acurácia 1-NN do ViT e ALIGN no ImageNet-1K. Fonte.
A principal vantagem de usar a abordagem MRL é sua flexibilidade para alternar entre diferentes tamanhos de representações de características a partir do mesmo modelo. Para aproveitar totalmente essa força, testes de classificação adaptativa usando ResNet50 no ImageNet-1K também foram conduzidos.
Essa configuração envolve aprender a probabilidade máxima de softmax para determinar transições de representações de características menores para maiores. Os resultados dos testes mostram que o modelo MRL com uma representação de características de 37 dimensões é tão preciso quanto o modelo ResNet50 treinado com uma representação de características fixa de 512 dimensões, sendo apenas 0,8% menos preciso do que o modelo de 2048 dimensões.
Figura- Acurácia Top 1 do modelo ResNet50 em classificação adaptativa (esquerda) e mAP da ResNet50 em caso de uso de recuperação de imagens (direita)
Figura: Acurácia Top 1 do modelo ResNet50 em classificação adaptativa (esquerda) e mAP da ResNet50 em caso de uso de recuperação de imagens (direita). Fonte.
A seguir, vamos discutir recuperação. Como você talvez saiba, o objetivo da recuperação é encontrar um conjunto de contextos semelhantes para uma determinada consulta. Portanto, a qualidade da recuperação entre a ResNet50 treinada com MRL e representações de características de tamanho fixo usando a Precisão Média Média (mAP) foi testada. O modelo MRL mostra uma melhoria no mAP de até 3% em comparação com sua contraparte de tamanho fixo em cada tamanho de representação de características, conforme ilustrado na visualização acima.
Como podemos utilizar vários tamanhos de representação de características com MRL, também é particularmente interessante explorar como podemos acelerar o processo de recuperação mantendo a acurácia com MRL. Podemos testá-lo em um cenário chamado recuperação adaptativa.
Como mencionado anteriormente, há duas etapas principais em um processo de recuperação: pré-seleção e reranqueamento. A configuração da recuperação adaptativa envolveu o uso de uma representação de características de 16 dimensões para a pré-seleção de um conjunto de 200 candidatos, enquanto uma representação de características de 2048 dimensões é usada para o reranqueamento. Essa abordagem é comparada com outro método que usa uma representação de características de 2048 dimensões tanto para a pré-seleção quanto para o reranqueamento.
No ImageNet-1K, usar o modelo MRL pode, teoricamente, acelerar o processo de recuperação em até 128 vezes. Referimo-nos a isso como teórico porque, em aplicações do mundo real, raramente usamos um algoritmo ingênuo de vizinho mais próximo para tarefas de recuperação. Em vez disso, algoritmos de vizinho mais próximo aproximado (ANN), como FAISS, ANNOY ou HNSW, são comumente implementados para aumentar a velocidade de recuperação. Nesse caso, a configuração que usa o modelo MRL alcança uma aceleração de 14 vezes em comparação com processos de recuperação que usam o algoritmo HNSW em hardware idêntico.
Um resultado semelhante foi observado com o conjunto de dados ImageNet-4K, em que os autores do artigo usaram um modelo MRL com uma representação de características de 64 dimensões para a pré-seleção e uma representação de 2048 dimensões para o reranqueamento. Os resultados indicam uma aceleração teórica de 32 vezes e uma aceleração no mundo real de seis vezes.
A relação de compromisso entre mAP@10 vs MFLOPs/Query para Recuperação Adaptativa (AR) no ImageNet-1K (à esquerda) e ImageNet-4K (à direita).
A relação de compromisso entre mAP@10 vs MFLOPs/Query para Recuperação Adaptativa (AR) no ImageNet-1K (à esquerda) e ImageNet-4K (à direita). Fonte.
Implementação da abordagem MRL
Nesta seção, vamos dar uma olhada na implementação simples da abordagem MRL. Em particular, usaremos um modelo base MPNet treinado com a abordagem MRL no conjunto de dados NLI para gerar embeddings em diferentes tamanhos. Como você talvez já saiba, o modelo base MPNet original só consegue gerar um embedding de 768 dimensões. No entanto, este modelo MPNet treinado consegue gerar embeddings em 768, 512, 256, 128 e 64 dimensões, em que o peso de cada dimensão durante o treinamento é definido como 1.
Podemos usar este modelo com a ajuda da biblioteca SentenceTransformers com o seguinte código.
from sentence_transformers import SentenceTransformer
matryoshka_dim_short = 64
matryoshka_dim_long = 768
text = ["The weather is so nice!"]
short_embedding = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka", truncate_dim=matryoshka_dim_short).encode(text)
long_embedding = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka", truncate_dim=matryoshka_dim_long).encode(text)
print(f"Shape: {short_embedding.shape, long_embedding.shape}")
print(short_embedding[0][0:10])
print(long_embedding[0][0:10])
"""
Output:
Shape: ((1, 64), (1, 768))
[-0.33891088 0.01647538 -0.29915053 0.24952686 -0.04321517 -0.31616145
-0.12996909 -0.05221268 0.02296597 0.07074839]
[-0.33891088 0.01647538 -0.29915053 0.24952686 -0.04321517 -0.31616145
-0.12996909 -0.05221268 0.02296597 0.07074839]
"""
No código acima, usamos o modelo base MPNet treinado para gerar embeddings em dois tamanhos diferentes: 64 e 768. Como o peso ou fator de escala de cada dimensão é definido como 1, os dois embeddings têm elementos idênticos, como você pode ver nos primeiros 10 elementos dos dois embeddings acima.
Como o valor de cada elemento é idêntico, a similaridade entre os dois embeddings seria 1.
from sentence_transformers.util import cos_sim
similarities = cos_sim(short_embedding[0], long_embedding[0][:matryoshka_dim_short])
print(similarities)
# tensor([[1.]])
Se você quiser treinar seu próprio modelo com a abordagem MRL, pode dar uma olhada no repositório oficial do GitHub deste método.
Conclusão
A abordagem MRL apresenta uma solução para equilibrar a relação custo-desempenho em aprendizado de máquina. Ao permitir que qualquer modelo de aprendizado de máquina produza representações de características de tamanhos variados, o MRL oferece a flexibilidade de otimizar para velocidade ou precisão, dependendo do nosso caso de uso e dos recursos. Vimos que essa adaptabilidade é valiosa em aplicações como classificação e recuperação, onde alternar entre representações menores e maiores pode aumentar significativamente a eficiência sem sacrificar massivamente o desempenho.
Os resultados experimentais demonstram que o MRL iguala e frequentemente supera a precisão dos modelos tradicionais de tamanho fixo em vários domínios, incluindo texto, visão e tarefas multimodais. Além disso, a compatibilidade da abordagem com arquiteturas existentes como ResNet50, ViT e BERT destaca sua aplicabilidade em diferentes modalidades de modelos. Com o potencial de acelerações substanciais em cenários do mundo real e melhoria na qualidade da recuperação, o MRL representa um avanço promissor para soluções de aprendizado de máquina mais eficientes e versáteis.
Recursos Relacionados
Continue lendo
Zilliz Cloud Now Available in AWS Europe (Ireland)
Zilliz Cloud launches in AWS eu-west-1 (Ireland) — bringing low-latency vector search, EU data residency, and full GDPR-ready infrastructure to European AI teams. Now live across 30 regions on five cloud providers.
Zilliz Cloud Delivers Better Performance and Lower Costs with Arm Neoverse-based AWS Graviton
Zilliz Cloud adopts Arm-based AWS Graviton3 CPUs to cut costs, speed up AI vector search, and power billion-scale RAG and semantic search workloads.
Creating Collections in Zilliz Cloud Just Got Way Easier
We've enhanced the entire collection creation experience to bring advanced capabilities directly into the interface, making it faster and easier to build production-ready schemas without switching tools.