EmbedOR: Provable Cluster-Preserving Visualizations with Curvature-Based Stochastic Neighbor Embeddings

Resumo Técnico do EmbedOR: Visualizações de Preservação de Agrupamentos Prováveis com Embeddings de Vizinhos Estocásticos Baseados em Curvatura

Declaração do Problema

Algoritmos de Embedding de Vizinhos Estocásticos (SNE), particularmente o t-SNE e o UMAP, são amplamente utilizados para visualizar dados de alta dimensão. No entanto, eles sofrem de limitações teóricas e práticas significativas em relação à preservação da geometria. Especificamente, esses métodos frequentemente falham em recuperar a geometria subjacente de dados de alta dimensão ruidosos ao:

1. Falhar na identificação de agrupamentos (clusters) em dados que são bem agrupáveis.
2. Introduzir fragmentação, onde componentes conectados da subvariedade de dados subjacente são esparsamente separados na visualização.

As garantias teóricas existentes para o t-SNE são restritas a pressupostos muito rigorosos sobre os dados (por exemplo, clusters esféricos e bem separados), e não existem resultados teóricos significativos para o UMAP. Além disso, estudos empíricos mostram que implementações padrão frequentemente violam esses pressupostos, levando a visualizações subótimas.

Metodologia

Os autores propõem o EmbedOR, um novo algoritmo SNE que integra a Curvatura de Ollivier-Ricci (ORC) para construir uma métrica de distância que potencializa agrupamentos. A metodologia procede em dois estágios:

1. Aprendizado de Métrica Baseado em Curvatura ($\Delta_E$):

   • O algoritmo constrói um grafo de vizinhos mais próximos a partir dos dados.
   • Ele calcula a ORC para cada aresta no grafo. A ORC quantifica a conectividade local; a curvatura positiva indica arestas dentro de uma comunidade altamente conectada, enquanto a curvatura negativa indica "gargalos" ou arestas que ligam clusters disjuntos.
   • Uma função de energia $E(\kappa; p)$ é definida com base na curvatura $\kappa$. Esta função atribui alta "energia" (penalidade) a arestas com curvatura negativa e baixa energia a arestas com curvatura positiva.
   • Uma nova métrica de distância de pares, $\Delta_E$, é calculada como a distância do caminho mais curto no grafo ponderado, onde os pesos das arestas são o produto da distância euclidiana pela função de energia baseada na curvatura. Esta métrica incentiva caminhos a atravessar arestas de curvatura positiva (dentro de clusters) e evitar arestas de curvatura negativa (entre clusters).

2. Embedding de Vizinho Estocástico Customizado:

   • O EmbedOR utiliza um framework SNE customizado que mistura elementos do t-SNE e do UMAP.
   • Afinidades de alta dimensão ($p_{ij}$): Ao contrário do t-SNE padrão que utiliza distância euclidiana, o EmbedOR usa a métrica $\Delta_E$ para calcular as afinidades.
   • Otimização: O algoritmo minimiza uma perda de "entropia cruzada difusa" (fuzzy cross-entropy) entre as afinidades de alta dimensão e as afinidades de baixa dimensão (modeladas por uma distribuição t de Student). A otimização utiliza gradiente descendente estocástico (SGD), aproveitando a eficiência computacional do UMAP (evitando a renormalização das afinidades de baixa dimensão) enquanto retém o cálculo de similaridade baseado em perplexidade do t-SNE.

Principais Contribuições

1. Extensão Teórica: Os autores provam que a métrica de distância do EmbedOR estende os resultados de consistência teórica do t-SNE (previamente limitados a dados esféricos e bem separados) para uma classe muito mais ampla de conjuntos de dados. Especificamente, eles mostram que, sob um modelo de variedade ruidosa:
   • As distâncias intra-cluster sob $\Delta_E$ convergem para a distância geodésica da variedade.
   • As distâncias inter-cluster são limitadas inferiormente por um valor que aumenta com o parâmetro de repulsão $p$.
   • Consequentemente, para um $p$ suficientemente grande, os dados satisfazem as condições necessárias para que o t-SNE produza uma visualização que preserva agrupamentos.
2. Inovação Algorítmica: A proposta do EmbedOR, que utiliza curvatura de grafo discreta para criar uma métrica de distância que enfatiza a estrutura de agrupamento subjacente sem introduzir fragmentação.
3. Ferramenta de Diagnóstico: A demonstração de que a métrica EmbedOR pode ser usada como uma ferramenta post-hoc para anotar visualizações existentes. Ao visualizar arestas com distâncias $\Delta_E$ curtas, os usuários podem identificar fragmentação (onde pontos próximos em $\Delta_E$ estão distantes na visualização), fornecendo insights sobre a geometria subjacente.

Resultados Experimentais

Os autores avaliaram o EmbedOR em conjuntos de dados sintéticos e do mundo real (incluindo MNIST, Fashion-MNIST e vários conjuntos de dados de sequenciamento de RNA de célula única) contra baselines como t-SNE, UMAP, PHATE, Isomap e Aprendizado de Autovalores Laplaciano.

• Dados Sintéticos: O EmbedOR demonstrou consistentemente um desempenho superior ou empatado como o melhor na separação de componentes conectados subjacentes, evitando a fragmentação de caminhos contínuos. Enquanto o UMAP frequentemente introduzia fragmentação e o t-SNE falhava em desenrolar estruturas como o Swiss Roll, o EmbedOR manteve a integridade topológica.
• Dados Reais (scRNA-seq): Em conjuntos de dados de trajetória de desenvolvimento e tipos celulares, o EmbedOR preservou melhor a continuidade das trajetórias e a separação dos tipos celulares em comparação com o UMAP e o t-SNE.
• Análise de Fragmentação: Análises quantitativas usando correlações de Spearman e comprimentos de aresta com pontuação z confirmaram que o EmbedOR minimamente expande as distâncias entre pontos considerados próximos pela métrica $\Delta_E$. Em contraste, o UMAP e o t-SNE mostraram expandir significativamente essas distâncias, indicando fragmentação.
• Utilidade da Métrica: Experimentos confirmaram que arestas de $\Delta_E$ curtas são significativamente mais propensas a conectar pontos dentro do mesmo cluster de verdade fundamental (ou tipo celular) do que arestas em um grafo k-NN padrão, validando a capacidade da métrica de detectar fragmentação em outros embeddings.

Significância e Alegações

O artigo afirma que o EmbedOR alcança um "equilíbrio delicado" entre destacar a estrutura de agrupamento e preservar a continuidade das variedades de dados, uma qualidade não observada nos métodos de base testados.

• Significância Teórica: O trabalho fornece a primeira justificativa teórica para o uso de métricas baseadas em curvatura para relaxar os pressupostos estritos necessários para a convergência do SNE, ampliando assim a aplicabilidade da preservação de agrupamentos provável.
• Significância Prática: O EmbedOR oferece uma alternativa robusta para visualizar dados de alta dimensão ruidosos onde manter a integridade dos componentes conectados é crítico (ex: genômica de célula única).
• Valor de Diagnóstico: Os autores enfatizam que a métrica EmbedOR serve a um propósito duplo: atua como a métrica de distância central para o algoritmo e como uma ferramenta de diagnóstico para identificar e interpretar a fragmentação em visualizações geradas por outros métodos.

Os autores permanecem modestos quanto ao escopo teórico, observando que seus resultados de convergência dependem de instâncias específicas de t-SNE e que seu algoritmo customizado mistura elementos de UMAP e t-SNE. No entanto, eles argumentam que os insights teóricos relativos à métrica $\Delta_E$ fornecem uma forte motivação para sua eficácia na prática.