AnglE

AnglE is an open-source Python library from the paper "AnglE-optimized Text Embeddings" (ACL 2024). It lets users train sentence embedding models and run inference with a variety of transformer-based embedding models. It is distributed as the angle-emb package under the MIT License, and the latest listed release is v0.6.0.

What It Is

AnglE is a sentence-embedding training and inference framework. Its core idea, per the paper, is angle optimization in a complex space, which addresses the vanishing-gradient problem caused by the saturation zones of the cosine function in embedding objectives. The library supports BERT-based models (BERT, RoBERTa, ModernBERT), LLM-based models (LLaMA, Mistral, Qwen) and bidirectional LLMs via BiLLM.

Training and Loss Options

Users prepare a Hugging Face dataset in one of three formats: paired texts with similarity scores, query-positive pairs, or query-positive-negative triples. Training runs through the angle-trainer CLI or a Python fit API, on a single GPU or multiple GPUs (including FSDP with accelerate). Supported losses are AnglE, contrastive, CoSENT and Espresso (2DMSE). Options include LoRA for LLMs, prompt templates, column renaming for older datasets, and teacher or self-distillation to reduce catastrophic forgetting. Trained models can be converted to sentence-transformers format.

Inference and Pretrained Models

AnglE.from_pretrained loads official models, such as WhereIsAI/UAE-Large-V1 for general English, UAE-Code-Large-V1 for code similarity and the pubmed-angle models for medical similarity. It can also load third-party transformer embedding models. Prompts can be applied for retrieval tasks. Espresso/Matryoshka models support layer truncation and reduced embedding sizes, and dynamic batching is available through the batched library.

Research Context

The README reports that the UAE-Large-V1 model and mixedbread's mxbai-embed-large-v1, both trained with AnglE, reached top scores on the MTEB Leaderboard at the time of those entries. The paper reports that AnglE outperforms state-of-the-art STS models that ignore the cosine saturation zone.