Generalized Schrödinger Bridge on Graphs

Resumo Técnico: Ponte de Schrödinger Generalizada em Grafos (GSBoG)

Enunciado do Problema
O transporte em grafos é um desafio fundamental em domínios que variam desde cadeias de suprimentos e infraestruturas de tráfego até cinética molecular. Embora o transporte ótimo (OT) em grafos forneça métodos fundamentados para o emparelhamento de distribuições de probabilidade de origem e destino, as abordagens existentes frequentemente resultam em acoplamentos estáticos que falham em descrever o movimento de massa dependente do tempo ou em fornecer políticas de controle executáveis. Além disso, formulações dinâmicas para transporte em grafos, como aquelas que resolvem fluxos hamiltonianos no simplex de probabilidade, sofrem de baixa escalabilidade em topologias grandes e esparsas e muitas vezes carecem de expressividade para otimizar trajetórias intermediárias sob custos dependentes do estado específicos da aplicação (ex: congestionamento ou barreiras de energia). O desafio central é construir uma evolução viável e respeitadora da topologia que controle distribuições intermediárias enquanto minimiza custos de rede e acomoda restrições operacionais, como limites de capacidade.

Metodologia
Os autores introduzem a Ponte de Schrödinger Generalizada em Grafos (GSBoG), um framework baseado em dados para aprender políticas de processos de cadeia de Markov de tempo contínuo (CTMC) executáveis e controladas em grafos arbitrários. O método estende os princípios da Ponte de Schrödinger Generalizada (GSB) de espaços de estados contínuos para domínios discretos estruturados em grafos.

1. Formulação Matemática:

   • O problema é formulado como a minimização de um funcional de custo sobre medidas de caminho de uma CTMC. O objetivo combina um custo corrente $f_t(x, p_t)$ (dependente do estado e da distribuição marginal) com um termo de divergência de Kullback-Leibler (KL) que regulariza a dinâmica controlada em direção a um processo CTMC de referência.
   • Os autores derivam condições de otimalidade discretas paralelas à formulação de estado contínuo. Eles estabelecem que as taxas de transição controladas ótimas $u^*_t(y, x)$ satisfazem uma transformação de Hopf-Cole: $u^*_t(y, x) = r_t(y, x) \exp(-V_t(y) + V_t(x))$, onde $V_t$ é uma função de potencial dependente do tempo e $r_t$ é a taxa de referência.
   • Isso leva a um sistema de equações acopladas para os potenciais de Schrödinger $(\phi_t, \hat{\phi}_t)$, análogos à equação de Fokker-Planck em configurações contínuas, mas definidos via a equação de continuidade do grafo.

2. Abordagem Algorítmica (gIPF e TD):

   • Para evitar solvers globais computacionalmente proibitivos sobre todo o grafo e horizonte temporal, a GSBoG adota uma abordagem baseada em partículas e orientada a dados usando Ajuste Proporcional Iterativo Generalizado (gIPF).
   • O método utiliza a fórmula de Dynkin para derivar análogos discretos das Equações Diferenciais Estocásticas Forward-Backward (FBSDEs) usadas na GSB de estado contínuo.
   • O objetivo de treinamento consiste em dois componentes alternados:
     • Perda de IPF: Minimiza a discrepância entre a medida de caminho aprendida e o processo de referência para satisfazer as restrições marginais de extremidade ($p_0 = \mu, p_1 = \nu$).
     • Perda de Diferença Temporal (TD): Impõe explicitamente a consistência com os custos correntes dependentes do estado e da distribuição, penalizando desvios entre os incrementos previstos pelo gerador e os incrementos reais da trajetória.
   • O algoritmo parametriza os log-potenciais usando redes neurais (ou embeddings) e os atualiza via gradiente descendente estocástico, avaliando as taxas apenas em vizinhanças locais do grafo.

Principais Contribuições

• Primeira GSB Baseada em Dados em Grafos: O artigo apresenta a primeira formulação da Ponte de Schrödinger Generalizada em grafos arbitrários de maneira baseada em dados, distinta dos trabalhos recentes de difusão discreta que focam em geração de grafos.
• Aprendizado de Política Escalável: Ao aprender políticas ao nível de trajetória via amostragem baseada em partículas, em vez de resolver fluxos expandidos no tempo globais, a GSBoG alcança escalabilidade em grafos grandes e esparsos onde os métodos tradicionais de OT dinâmico falham devido a restrições de memória.
• Dinâmica Consciente de Custos: O framework otimiza explicitamente o comportamento intermediário sob custos correntes gerais (ex: penalidades de congestionamento), garantindo que as políticas resultantes sejam não apenas topologicamente viáveis, mas também operacionalmente eficientes.
• Derivação Teórica: Os autores fornecem uma derivação rigorosa das condições de otimalidade discretas e um objetivo de aprendizado tratável baseado no gerador da CTMC, unindo a teoria de controle estocástico e o transporte em grafos discretos.

Resultados Experimentais
Os autores avaliam a GSBoG em três topologias de grafos do mundo real desafiadoras:

1. Rede de Cadeia de Suprimentos (9.559 nós): A GSBoG aprendeu com sucesso políticas de roteamento que corresponderam às distribuições de oferta/demanda de extremidade, enquanto reduziu significativamente o congestionamento intermediário e a ocupação de pico de nós em comparação com baselines como Attraction Flow e Dynamic W1 flow. Notavelmente, a GSBoG foi o único método que manteve os fluxos de arestas induzidos dentro dos limites de capacidade sem restrições rígidas explícitas, enquanto solvers globais (GrSB) falharam em completar devido à exaustão de memória.
2. Tarefa de Atribuição Balanceada: Em grafos bipartidos com tamanhos variados ($n=6$ a $20$), a GSBoG alcançou uma sobreposição quase perfeita com a solução de OT estático ótimo (100% de precisão para $n \le 10$), enquanto aprendia um acoplamento suave que minimizava os custos de atribuição.
3. Dinâmica Molecular Discretizada (Dobramento de Chignolin): Aplicada a um Modelo de Estado de Markov (MSM) de dobramento de proteínas, a GSBoG direcionou transições de eventos raros de um estado desenovelado para um estado dobrado. Ela alcançou uma taxa de dobramento de 99,7% com uma barreira de energia significativamente reduzida (1,80 $k_B T$) em comparação com dinâmicas não controladas ou vieses guiados por committor, e melhorou ainda mais para 1,39 $k_B T$ ao incorporar custos correntes de energia livre.

Análise de Escalabilidade
Experimentos demonstram que a GSBoG escala linearmente com o fator de ramificação local, em vez do número total de nós. Em grafos com até $10^6$ nós, a GSBoG permaneceu viável com baixo uso de memória (1,4 GB) e tempo de execução (1 hora), enquanto os métodos de OT Dinâmico tornaram-se inviáveis (N/A) ou exigiram significativamente mais recursos (ex: 22 GB de memória, 4,6 horas).

Significância e Alegações
O artigo afirma que a GSBoG abre novos caminhos para o transporte dinâmico consciente de custos em grafos gerais. Ao ir além de acoplamentos estáticos e solvers globais, o framework permite o aprendizado de políticas executáveis e respeitadoras da topologia que otimizam tanto o emparelhamento de extremidades quanto os custos de trajetória intermediários. Os autores posicionam este trabalho como um passo para integrar a dinâmica moderna baseada em aprendizado com os princípios de transporte ótimo para sistemas de rede complexos, particularmente em cenários onde restrições operacionais (como congestionamento ou barreiras de energia) são críticas. O trabalho é modesto em suas limitações, observando que atualmente assume uma topologia de grafo fixa e totalmente conhecida, e identifica extensões para grafos variantes no tempo ou parcialmente observados como trabalhos futuros.