HEAL : Hierarchical Embedding Alignment Loss for Improved Retrieval and Representation Learning (2025)

Abstract: Retrieval-Augmented Generation (RAG) enhances Large Language Models (LLMs) by integrating external document retrieval to provide domain-specific or up-to-date knowledge. The effectiveness of RAG depends on the relevance of retrieved documents, which is influenced by the semantic alignment of embeddings with the domain’s specialized content. Although full fine-tuning can align language models to specific domains, it is computationally intensive and demands substantial data. This paper introduces Hierarchical Embedding Alignment Loss (HEAL), a novel method that leverages hierarchical fuzzy clustering with matrix factorization within contrastive learning to efficiently align LLM embeddings with domain-specific content. HEAL computes level/depth-wise contrastive losses and incorporates hierarchical penalties to align embeddings with the underlying relationships in label hierarchies. This approach enhances retrieval relevance and document classification, effectively reducing hallucinations in LLM outputs. In our experiments, we benchmark and evaluate HEAL across diverse domains, including Healthcare, Material Science, Cyber-security, and Applied Maths.

## introduction large language models (llms), such as @xcite , have demonstrated exceptional capabilities in natural language understanding and generation. however, llms are prone to hallucinations, generating plausible but incorrect or nonsensical content @xcite . retrieval-augmented generation (rag) frameworks @xcite mitigate this issue by integrating external knowledge through document retrieval, enhancing the factual accuracy of llm outputs. a critical component of rag systems is the embedding model used for document retrieval. standard embedding models, however, often fail to capture the hierarchical and semantic relationships within domain-specific corpora, leading to suboptimal retrieval and, consequently, increased hallucinations. this issue is particularly pronounced in domains with increased specificity such as healthcare, legal sytem, and scientific research. corpus of documents for a specialized domain inherently exhibit a high degree of semantic coherence, presenting an opportunity to align embedding models for retrieving the most contextually relevant information. hierarchical non-negative matrix factorization (hnmf) @xcite is a powerful technique for semantically categorizing documents into clusters that exhibit thematic coherence. by grouping documents into hierarchical clusters of supertopics and subtopics, hnmf provides a rich semantic categorization of the corpus, enabling a deeper understanding of document relationships. leveraging this semantic knowledge in the form of hierarchical cluster labels, we can align embedding models to preserve hierarchical information within the embedding space. this alignment enhances the embeddings to capture both coarse-grained and fine-grained document similarities, improving contextual relevance in retrieval tasks and enabling better downstream capabilities. to tackle the challenges of hallucination and suboptimal retrieval in rag systems, we introduce the hierarchical embedding alignment loss (heal), a refined extension of the hierarchical multi-label contrastive loss @xcite . heal leverages an improved hierarchical weighting scheme to align embeddings more effectively with the underlying hierarchical structure. by incorporating hierarchical label structures, heal fine-tunes embedding models to align with document clusters derived from hnmf. the method computes contrastive losses at each hierarchical level, combining them with depth-specific penalties to emphasize distinctions at higher levels of the hierarchy. ## related work contrastive learning has become a cornerstone of representation learning, particularly in computer vision and natural language processing. methods like simclr @xcite and moco @xcite have achieved state-of-the-art performance in unsupervised settings by learning representations that are invariant to data augmentations. in supervised contrastive learning, @xcite extended the contrastive loss to utilize label information, improving performance on classification tasks. similarly, the scincl framework employs neighborhood contrastive learning to capture continuous similarity among scientific documents, leveraging citation graph embeddings to sample both positive and negative examples @xcite . however, these methods generally assume flat label structures and do not exploit hierarchical relationships. hierarchical classification has been studied extensively, with approaches such as hierarchical softmax @xcite and hierarchical crossentropy loss @xcite . these methods aim to leverage hierarchical label structures to improve classification efficiency and accuracy. in the context of representation learning, @xcite introduced hierarchical semantic embedding, aligning image embeddings with wordnet hierarchies. more recent works, such as @xcite , have explored hierarchical prototypes to capture hierarchical relationships. @xcite propose a hierarchical multi-label contrastive learning framework that preserves hierarchical label relationships through hierarchypreserving losses. their method excels in scenarios with hierarchical multi-label annotations, such as biological or product classifications. in contrast, our approach focuses on enhancing information retrieval to mitigate hallucinations. rag frameworks combine retrieval models with generative models to enhance the factual accuracy of language generation @xcite . these systems rely heavily on the quality of the embeddings used for retrieval. prior work has focused on improving retrieval through better indexing and retrieval algorithms @xcite ), but less attention has been given to aligning embeddings with hierarchical document structures. ## method in this section, we propose an embedding alignment framework comprising hierarchical label extraction with hnmf, embedding alignment using heal, and retrieval with aligned embeddings as outlined in figure 1 . ## conclusion in this work, we introduced heal, a novel framework for aligning embeddings in rag systems through hierarchical fuzzy clustering and matrix factorization, integrated within a contrastive learning paradigm. heal effectively computes levelspecific contrastive losses and applies hierarchical penalties to align embeddings with domain-specific structures, enhancing both retrieval relevance and classification performance. experimental results across diverse domains -healthcare, materials science, cybersecurity, and applied mathematics -demonstrate heal's capability to significantly improve retrieval accuracy and mitigate hallucinations in llm-based systems. by bridging hierarchical semantics with contrastive alignment, heal establishes itself as a versatile and robust tool for advancing rag methodologies, enabling more precise, reliable, and domain-adaptive applications of large language models.