Beyond the Unit Hypersphere: Embedding Magnitude in Contrastive Learning
Xincan Feng, Taro Watanabe
TL;DR
This work questions the default use of unit-normalized embeddings in contrastive learning, showing that embedding magnitude can encode task-relevant information, particularly for retrieval and RAG. Through a minimal 2×2 ablation and a continuous learnable normalization, the authors demonstrate that removing the unit-sphere constraint enables magnitude-based signals to improve performance, especially on hard, out-of-domain tasks, while symmetric tasks suffer if magnitude is not treated carefully. They propose the Task Symmetry Principle to guide similarity function choice: magnitude learning benefits asymmetric tasks (e.g., retrieval, QA) but can harm symmetric tasks (e.g., STS, clustering) where symmetry is essential. The results include substantial gains on RAG QA, insights into when and how magnitude helps, and practical guidance for leveraging magnitude with learnable normalization, though architectural compatibility (as in E5) can limit gains without adjustments. Overall, the paper provides both theoretical guarantees and empirical evidence that magnitude-aware contrastive learning can outperform cosine-based approaches in appropriate, asymmetric settings, with broad implications for dense retrieval and cross-modal systems.
Abstract
Cosine similarity is prevalent in contrastive learning, yet it makes an implicit assumption: embedding magnitude is noise. Prior work occasionally found dot product and cosine similarity comparable, but left unanswered WHAT information magnitude carries, WHEN it helps, and HOW to leverage it. We conduct a systematic study through a $2 \times 2$ ablation that independently controls input-side and output-side normalization across text and vision models. Our findings reveal three key insights. First, in text retrieval, output (document) magnitude strongly correlates with relevance (Cohen's $d$ up to 1.80), yielding the largest gains on reasoning-intensive tasks. Second, input and output magnitudes serve asymmetric roles: output magnitude directly scales similarity scores while input magnitude modulates training dynamics. Third, magnitude learning benefits asymmetric tasks (text retrieval, RAG) but harms symmetric tasks (STS, text-image alignment). These findings establish a task symmetry principle: the choice between cosine and dot product depends on whether the task has distinct input roles, enabling cost-free improvements by simply removing an unnecessary constraint.