Unraveling the Localized Latents: Learning Stratified Manifold Structures in LLM Embedding Space with Sparse Mixture-of-Experts

2025-02-19

AI Summary

Real-world data in large language model (LLM) embedding spaces exhibit complex, heterogeneous local geometry that cannot be adequately captured by a single smooth manifold. Method: We propose and empirically validate the “semantics-driven hierarchical manifold structure” hypothesis—asserting that inputs of varying perplexity and domain specificity reside on local manifolds of differing intrinsic dimensions, collectively forming a hierarchical geometric structure in the embedding space. To operationalize this, we design an interpretable, sparse Mixture-of-Experts (MoE) framework integrating sparse dictionary learning, attention-based soft gating, manifold intrinsic dimension estimation, and expert assignment entropy/distance analysis. Contributions/Results: (1) First empirical confirmation of hierarchical manifold structure in LLM embeddings; (2) Semantic-consistent, interpretable submanifold clustering; (3) Expert assignment statistics uncover intrinsic hierarchical regularities in input data, establishing a novel paradigm for embedding space modeling.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Natural Language Processing: (Large) Language Models
• Search and Optimization: Learning to Search

Application Category

• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

However, real-world data often exhibit complex local structures that can be challenging for single-model approaches with a smooth global manifold in the embedding space to unravel. In this work, we conjecture that in the latent space of these large language models, the embeddings live in a local manifold structure with different dimensions depending on the perplexities and domains of the input data, commonly referred to as a Stratified Manifold structure, which in combination form a structured space known as a Stratified Space. To investigate the validity of this structural claim, we propose an analysis framework based on a Mixture-of-Experts (MoE) model where each expert is implemented with a simple dictionary learning algorithm at varying sparsity levels. By incorporating an attention-based soft-gating network, we verify that our model learns specialized sub-manifolds for an ensemble of input data sources, reflecting the semantic stratification in LLM embedding space. We further analyze the intrinsic dimensions of these stratified sub-manifolds and present extensive statistics on expert assignments, gating entropy, and inter-expert distances. Our experimental results demonstrate that our method not only validates the claim of a stratified manifold structure in the LLM embedding space, but also provides interpretable clusters that align with the intrinsic semantic variations of the input data.

Problem

• Analyzing stratified manifold structures in LLM embedding space
• Using Mixture-of-Experts to model local data complexities
• Validating semantic stratification through specialized sub-manifolds learning

Innovation

• Mixture-of-Experts with dictionary learning
• Attention-based soft-gating network
• Stratified manifold structure analysis