Trace: Structural Riemannian Bridge Matching for Transferable Source Localization in Information Propagation

1 Introduction

Information diffusion is a classic problem, predicting the infected nodes in the future according to the propagation pattern. Source localization, as the inverse problem of information diffusion, seeks to identify the source nodes responsible for the final diffusion observations [Yan et al., 2023]. It not only facilitates a deeper understanding of social dynamics but also enables the vital applications such as information propagation intervention and rumor source detection. In recent years, researchers have visited the advanced models including invertible graph neural network and variational autoencoder [Wang et al., 2022; Ling et al., 2022]. While achieving notable progress, source localization still presents several important issues as follows.

∗Corresponding Author: Li Sun.

Error Accumulation. Existing methods typically follow a step-by-step backward procedure [Yan et al., 2023; Huang et al., 2023] or model the round trip with forward information diffusion [Ling et al., 2022; Wang et al., 2022]. They are vulnerable to error accumulation especially in the absence of intermediate supervision. Instead, it is preferable to directly model the source and final distributions. Recent generative models achieve notable success in modeling data distributions [Song et al., 2021; Lipman et al., 2023], but they primarily focus on generating distribution from a Gaussian. So far, the transformation between arbitrary distributions, accounting for graph structure, remains under-investigated.

Transferability. Recent efforts leverage deep learning on graphs, which usually requires a large quantity of observations for effective inference [Ling et al., 2024; Dong et al., 2019]. However, acquiring observations is often laborious or expensive, limiting their widespread use in real scenarios. It calls for a transferable model that leverages the knowledge learned from previous observations to enhance the localization in new graphs. We emphasize that, given the underlying social dynamics, graph structures significantly influence information propagation, making domain adaptation essential.

Geometry. Riemannian geometry not only demonstrates superior expressiveness compared to its Euclidean counterparts [Sun et al., 2022b], but also introduces the concept of curvature to describe graph structure [Sun et al., 2023b], shedding light on quantifying the difference among social networks and enabling domain adaptation. Also, we note that recent efforts report the effectiveness of Riemannian geometry in information diffusion [Sun et al., 2024a]. However, its application to the inverse problem of source localization has not yet been explored. Nevertheless, calculating Riemannian divergence or heat kernel is computationally expensive [Sun et al., 2024b], and training models on the manifold remains challenging.

In light of the aforementioned issues, we are motivated to propose the problem of Transferable Source Localization, enabling the knowledge transfer to new graphs, and to approach this problem from a fresh geometric perspective.

In this paper, we present a novel structural Riemannian bridge matching model (termed as Trace) for transferable source localization. The key innovation is that we leverage the unified construction of Riemannian manifold to explore the transferability, and formulate a structural Schr¨odinger Bridge on the manifold along with a new learning algorithm. Specifically, to address the first issue of error accumulation, we make the first attempt, to the best of our knowledge, to introduce the Schr¨odinger Bridge [Bortoli et al., 2021] to source localization. It directly models the map between the source and final distributions via a Stochastic Differential Equation (SDE), and we establish a structural Schr¨odinger bridge with a Wiener process on the manifold. In particular, the proposed bridge is governed by the functional curvature, which is designed to encapsulate the structural information and is learnable with given structure, addressing the second issue of transferability. Note that, learning Schr¨odinger bridges on the manifold is challenging (the third issue). To fill this gap, we propose a simple yet effective algorithm, geodesic bridge matching. Concretely, we discover the optimal projection for Riemannian measures with theoretical guarantees, so that the Schr¨odinger bridge is learned through geodesic bridge matching, avoiding the calculation of Riemannian heat kernel. Finally, Trace integrates the advantages of Schr¨odinger bridge and Riemannian geometry to directly infer the sources. We emphasize that the Riemannian manifold endows our model with the transferability. With the underlying dynamics learned from previous observations, we fine-tune the functional curvature on the new graph to conduct knowledge transfer, taking the graph structure into account.

Contribution Highlights. Main contributions are four-fold:

• Problem. We consider the geometric effect on source localization and for the first time pose the transferable source localization problem, emphasizing its transferability for practical usage, to the best of our knowledge.
• Riemannian Model. We propose a novel approach, Trace, in which the structural Riemannian Schr ¨odinger bridge is established for the initial-terminal distribution map, equipped with functional curvature to enhance knowledge transfer through domain adaptation.
• Learning Algorithm. We design a new learning algorithm for Riemannian Schr¨odinger bridges, called geodesic bridge matching, in which we prove the optimal projection holds for Riemannian measures, thereby avoiding the need for the expensive iterative procedure.
• Extensive Experiments. We evaluate the superiority of Trace on both synthetic and real datasets, analyze each proposed component with ablation study, and examine its transferability among different datasets.

2 Preliminaries and Notations

Manifold and Curvature. In Riemannian geometry, a graph structure is related to certain Riemannian manifold (M, g), a smooth manifold M attached to a Riemannian metric g. Each point on the manifold x ∈ M is associated with a tangent space TxMwhere the metric g is defined. Given two points and g, geodesic is the curve that connects them with minimal integral length. The exponential and logarithmic map are inherited from the Lie algebra for the projection between tangent space and the manifold. Riemannian measure is defined over the infinitesimal on the manifold. The notion of curvature κ describes how a smooth manifold deviates from being flat and, accordingly, it offers a geometric perspective to quantify graph structures. There exist three types of constant curvature spaces: negative curvature hyperbolic space, positive curvature hyperbolic space, and the flat Euclidean space, a special case of Riemannian geometry.

Source Localization and Generative Models. In this paper, we recast source localization as a generative task and the idea is that we leverage the distribution of final observations to generate the initial source distribution, while capturing the social dynamics in the generative model. Recently, scorebased (diffusion-based) models formulate Stochastic Differential Equations (SDEs) to generate data distribution from a simple one such as Gaussian [Song et al., 2021], while flowbased models use Ordinary Differential Equations (ODEs) to perform such probability pushforward [Lipman et al., 2023]. Though both models achieve notable success, we highlight that they fall short in transforming arbitrary distributions, i.e., source and final distributions, or require auxiliary time information to complete source localization [Huang et al., 2023]. In other words, a promising generative localizer remains largely under-explored.

Problem Statement of Transferable Source Localization. We consider a graph G(V, E) defined over node set V and edge set E ⊆ V × V. The adjacency matrix is denoted as A and graph Laplacian is given as L = I -D- 1 2 AD- 1 2 , where D is the diagonal degree matrix. For source localization, only the final snapshot is accessible and the status of each node is described in y ∈ {0, 1}|V| , where yi = 1 indicates that node i is infected, otherwise, yi = 0. The sources are collected in s ∈ {0, 1}|V| , and si = 1 iff node i is one of the sources responsible for y. Given the graph and the final snapshot, the problem of source localization aims to seek F(G, y) that is able to infer the sources s. Different from previous works, in Transferable Source Localization, we are dedicated to design F that offers the transferability among different graphs.

3 Methodology: Trace

We present a novel structural Riemannian bridge matching model (termed as Trace) for transferable source localization. The key novelty lies in that we formulate the Structural Riemannian Schr¨odinger Bridge along with a New Learning Algorithm, while exploring the advantage of Riemannian manifold for Transferability. The overall architecture is sketched in Fig. 1, and we start with the Schr¨odinger bridge in the heart of our model.

3.1 Structural Riemannian Schr¨odinger Bridge

To the best of our knowledge, we for the first time introduce the Schr¨odinger bridge to the problem of source localization. The advantage is that we can directly model the distribution map between source and final states, thereby mitigating error accumulation. In a nutshell, we establish the Riemannian Schr¨odinger bridge that is equipped with a functional curvature depicting the structural information of the graph.

Schr¨odingerBridge. It connects two arbitrary distributions with a Stochastic Differential Equation (SDE). To be specific, we consider a measure space P(C[0, 1], Rd ), which collects the continuous and finite density pt, where t ∈ [0, 1]. 1 The density at point x ∈ Rd of time t is denoted as pt(x), and d denotes the dimensionality. For a process T in measure space P, the joint distribution at t = 0, 1 is πT ∈ P(Rd × Rd ), and T|x0,x1 denotes the distribution conditioned on T 's value x0, x1 at t = 0, 1, respectively. We select a reference process W ∈ P(C[0, 1], Rd ) which is typically set as the Wiener process of Brownian motion. In particular, it satisfies the SDE of dxt = √ϵdWt, where ϵ > 0 is the volatility. Given two arbitrary distributions p0, p1 ∈ P, the (dynamic) Schr¨odinger bridge problem is formulated as the optimization,

The optimal process T ⋆ is referred to as the Schr¨odinger bridge [Schr¨odinger, 1932; F¨ollmer, 1988]. In other words, Schr¨odinger bridge is the process that admits the distribution p0, p1 at t = 0, 1 and is closest to the reference process in term of Kullback-Leibler divergence.

Structure-aware Bridge in Riemannian manifold. Here, we formulate the optimal process T ⋆ on the manifold with the Riemannian Wiener process. We consider a smooth manifold Mκ,d whose curvature encapsulates the structural information of graph G. Let PM(C[0, 1],Mκ,d ) denote the Riemannian measure space, and its Wiener process is written as dxt = √ϵdW M t . The Riemannian Schr¨odinger bridge is to find the minimizer of Eq. 2 in PM pinned down at p0 and p1. Different from the Euclidean version, a process on the manifold is described as dxt = f(xt, t)dt + √ϵdW M t , xt ∈ Mκ,d , and its time reversal is derived as dyt = (-f(yt, t) + ϵ∇yt log p1-t(yt)) dt + √ϵdW M t , (2) where (yt)t>0 = (x1-t)t∈[0,1] is the time reversal of xt. ∀f over Mκ,d , ∇xt f ∈ TxtMκ,d is Riemannian gradient. g(yt, t) = -f(yt, t) + ϵ∇yt log p1-t(yt) is known as drift. In fact, the optimal T ⋆ is unique and can be expressed as T ⋆ = ∫Rd×Rd W|x0,x1 dπ⋆ (x0, x1), where the joint marginal distribution π⋆ is the solution of the Entropic Optimal Transport (EOT) problem [L´eonard, 2014] as follows,

1 In this paper, we utilize t = 0 to refer to the initial of the process, while t = 1 stands for its terminal.

Π(p0, p1) is the set of transport plan over the joint P(Rd × Rd ) whose marginals are p0 and p1, and π⋆ is named as EOT plan. dg is the metric on the manifold, and H(π) is the entropy over π. However, the EOT plan π⋆ is unavailable. Hence, we are interested in learning the optimal drift g⋆ of Riemannian Schr¨odinger bridge T ⋆ .

Functional Curvature. Distinguishing from the existing Schr¨odinger bridges [Jo and Hwang, 2024; Thornton et al., 2022], the proposed Riemannian Schr¨odinger bridge T ⋆ is governed by the functional curvature κθ, a key ingredient in our design. As the curvature summarizes the graph structure, we propose to parametrize κθ with Laplacian L. Denote the tuple of L's dominant eigenvalues as l. The functional curvature is given as κθ = l⊤θ. In Trace, we learn its parameter in the self-supervised fashion, and the supervision is given by the Parallelogram Law in Riemannian geometry [Gu et al., 2019]. Concretely, given a geodesic triangle abc on the manifold, we have γabc = dg(a, m)2 + dg(b,c)2 4 - dg(a,b)2+dg(a,c)2 2 , where dg is the distance metric, and m is the midpoint of bc. Accordingly, the objective of curvature learning is derived as

where a, b, c and m in the manifold, and their encodings are introduced in Sec. 3.2.

3.2 Graph Encoding with Functional Curvature

Having the bridge between the initial and terminal distributions at hand, we specify its two ends in this part. As in Fig. 1, the initial distribution is defined as the final observation while the sources are resided in the terminal. This allows us to utilize the final observation to generate the sources in line with the setting of source localization. We first elaborate on the model space of Riemannian geometry.

Unified Construction of RiemannianManifold. Riemannian manifold offers an elegant unified construction for the constant curvature space of different curvatures. We introduce a model space Mκ,d to unify hyperbolic and hyperspherical spaces. Specifically, a d-dimensional manifold Mκ,d with curvature κ is embedded in Rd+1 space,

where xt ∈ R and xs ∈ Rd are the time and space dimension, respectively. The curvature-aware inner product is given as,

where sgn is the sign function, and the diagonal R gives the Riemannian metric in Rd+1 space. Note that, Mκ,d recovers the Lorentz model of hyperbolic space with κ < 0, and shifts to Spherical model of hyperspherical space when κ > 0.

Curvature-preserving Encoding. The samples are given by the proposed curvature-preserving encoding, where we adopt the operations without coupling logarithmic and exponential maps as in [Sun et al., 2025a], and the advantage is investigated in Ablation Study. Specifically, dimension transformation is done by matrix-left-multiplication written as,

parameterized by W. ‖·‖ denotes the L2 norm. We formulate the graph convolution as the geometric midpoint on the manifold. Given a set of points {(xi,wi)}i∈Ω each attached a weight wi, the graph convolution over Ω takes the form of

where |‖a‖κ| = √-〈a, a〉κ is the modulus of the imaginary. Note that, both operations preserve the curvature.

Theorem 1 (Curvature-preserving). ∀x ∈Mκ,d1 and any w ∈ R, W ∈ Rd2×d1 , we have Convκ({(xi,wi)}i∈Ω) ∈ Mκ,d1 and Transformκ(W, x) ∈ Mκ,d2 hold for either positive or negative value of κ.

We perform dimension transformation and graph convolution to obtain the node encodings. With the curvature-preserving encoding, the structural information is delivered to the manifold where the proposed Schr¨odinger bridge is established.

3.3 Learning Schr¨odinger Bridge onManifolds

Challenges. Popular learning algorithms, such as Iterative Proportional Fitting (IPF) [Bortoli et al., 2021] and Iterative Markovian Fitting (IMF) [Shi et al., 2023], adopt an iterative procedure, which tends to result in the optimization error in each iteration [Gushchin et al., 2024]. In addition, calculating the divergence or heat kernel on the manifold is rather expensive, and the estimators are typically applied [Sun et al., 2024b]. Therefore, learning Riemannian Schr¨odinger bridges largely remains open.

Theoretical Results on Riemannian Optimal Projection. To fill this gap, we design a new learning algorithm built upon Riemannian optimal projection. The core contribution is that we discover the optimal projection for Riemannian measures.

We start with IMF algorithm in Euclidean space, a widely used successor of the IPF. Concretely, IMF first introduces the reciprocal projection projR(T) that gives the reciprocal process Tπ which is equal to a mixture of Wiener processes with given coupling π. Second, the Markovian projection projM(T) is defined by the SDE in [Shi et al., 2023]. Given p1 and p0, IMF alternates between the projections iteratively, and the sequence of {T 2n+1 = projM(T 2n ), T 2n+2 = projR(T 2n+1 )} is shown to converge to T ⋆ . Note that, the Markovian projection is problematic, which requires to solve the optimization over P(Rd × Rd ) in every odd iteration.

Instead of the Markovian projection, we propose a novel Riemannian optimal projection, generalizing the Euclidean one [Gushchin et al., 2024], and prove its optimality for Riemannian measures. To be specific, we first introduce the concept of Riemannian Schr¨odinger family S, collecting the stochastic process of Schr¨odinger bridge S pinned down at ps0 and ps1 in the Riemannian measure space PM(C[0, 1],Mκ,d ). Second, for any given π ∈ Π(p0, p1) and process T, the Riemannian reciprocal process is written as Tπ = projRM(T) = ∫M×M W M |x0,x1 dπ(x0, x1) with Riemannian Wiener process. Third, the proposed optimal projection is given as the solution to the following equation,

The intuition behind this is to find the Schr¨odinger bridge S which is closest to the Riemannian reciprocal process Tπ in terms of KL divergence.

Interestingly, we discover that the result of Riemannian optimal projection is the Riemannian Schr¨odinger bridge connecting p0, p1 ∈ PM. We show the following theorem holds for Riemannian measures. (The proof is in Appendix B.)

Theorem 2 (Optimality). Given the initial and terminal distributions p0, p1 ∈ PM, a joint distribution π ∈ Π(p0, p1), Riemannian reciprocal process Tπ is given accordingly. Riemannian optimal projection of Tπ leads to Riemannian Schr¨odinger Bridge T ∗ between the distributions p0 and p1, projSM(Tπ) = argS∈S min KL(Tπ||S) = T ∗ .

In light of the optimization in Eq. 9, we construct the Riemannian Schr¨odinger family S with tractable KL divergence. With the adjusted potential v, we consider the joint distribution πv(x0, x1) = p0(x0)πv(x1|x0), where πv(x1|x0) =

Algorithm 1 Geodesic Random Walk

Input: length N, initial value x, curvature κ, drift f, volatility σ. Output: samples {xγi }Ni=0

• 1: Set step size γ ← T/N.
• 2: for i = 0 to N -1 do
• 3: ¯zi+1 ∼ N(0, Id), zi+1 ← Transportκxγi (¯zi+1);
• 4: wi+1 ← γf(iγ, xγi ) + √γσ(iγ, xγi )zi+1;
• 5: xγi+1 ←Expκxγi (wi+1);
• 6: end for

exp(〈x0,x1〉κ/ϵ)v(x1) ∫M exp(〈x0,x1〉κ/ϵ)v(x1)dx1 . Then, the process S|x0,x1 = W M |x0,x1 is the Schr¨odinger bridge between p0 and pv(x1) = ∫M πv(x0, x1)dx0 on the manifold. The corresponding KL divergence exists closed-form expression regarding Riemannian reciprocal process Tπ, detailed in Appendix B. Note that, the advantage is that we no longer need the iterative procedure but conduct only one step of the proposed projection.

Algorithm. Consequently, it enables a new objective function for learning Riemannian Schr¨odinger bridges. Riemannian Wiener process (also known as Brownian motion) can be simulated by geodesic random walk [Bortoli et al., 2022], whose procedure is described in Algo. 1. Transport, Exp and Log denote the parallel transport, exponential map and logarithmic map, respectively. With the time reversal process described in the SDE, dxt = Logxt (x1) 1-t dt + √ϵdW M t , we derive the objective function as follows.

Theorem 3 (Objective Function). Given p0, p1 ∈ PM and Riemannian Wiener process, we have T ⋆ described in dxt = gθ(xt, t)dt+√ϵdW M t is the Riemannian Schr¨odinger bridge connecting p0 and p1, if the drift gθ is the minimzer of the following objective,

where U[0,1] denotes the uniform distribution over [0, 1].

The strength lies in that Eq. 10 offers a simple yet effective mean squared error objective, avoiding the expensive Riemannian divergence or heat kernel. Next, we develop the geometric understanding with the notion of geodesic.

Geodesic. Given x0, x1 in the geodesically complete manifold (Mκ , g), the geodesics xt = Expκx0 (tLogκx0 (x1)) is the shortest curve connecting them under g.

The intuition behind is that the initial and terminal distributions approach each other along the shortest path. The remaining issue is on the form of gθ. We parameterize the adjusted potential v with a mixture of wrapped Gaussian on the manifoldNM(x|µ,Σ). Accordingly, the drift takes the form of gθ(xt, t) = ϵ∇xt log (NM(xt|0, ϵ(1- t)Id)v(xt, t)), where v(xt, t) is given by the aforementioned mixture. Note that, g is easy to compute with the closed form and diagonal covariance. In summary, the overall training procedure is given in Algo. 2. Finally, Trace integrates the advantages

Algorithm 2 Training Trace.

Input: Graph G, The set of the sources and final infection pairs. Output: Learned parameters of the transferable localizer.

• 1: Perform normalization in light of transferability;
• 2: repeat
• 3: Sample triangles over G for the functional curvature;
• 4: Sample batch {x0, x1} ∼ π and {tn} ∼ U[0,1];
• 5: Sample batch {xt} ∼W[x0,x1] by Algo. 1;
• 6: Compute the overall loss J = αLmanifold+(1-α)Lbridge;
• 7: Update parameters via the gradient descent method;
• 8: until converged

of Schr¨odinger bridge and Riemannian geometry to directly infer the sources without any iterative procedure.

ConnectionwithRecent Advances. In Riemannian geometry, [Thornton et al., 2022] and [Jo and Hwang, 2024] follow the iterative procedure in spirit of IPF and IMF, respectively, while we conduct the one step projection. [Chen and Lipman, 2024] discusses the pushforward from a Gaussian but we are interested in a map of arbitrary distributions. In Euclidean space, [Liu et al., 2023] conducts such map without randomness, while we consider the randomness in the social dynamic process. LightSB-M [Gushchin et al., 2024] can be regarded as a special case of ours, given that Logxt (x1) converges to x1 -xt in the limit at zero curvature.

On Transferability. It is the Riemannian manifold that endows Trace with transferability, as it offers a unified construction for different curvatures, corresponding to diverse graph structures. In the transfer setting, we leverage the pretrained parameters of Schr¨odinger bridge to capture underlying social dynamics. Its functional curvature is fine-tuned on the new graph, so that we adapt the knowledge according to the structural information encapsulated in the curvature.

4 Experiment

4.1 Experimental Settings

Datasets andBaselines. The experiments are conducted on both synthetic and real datasets. The synthetic datasets consist of Jazz, Network Science, Cora-ML, and Facebook [Ling et al., 2022]. We follow [Ling et al., 2022] for diffusion simulation, and the ratio of the sources is set as 10%. The four real datasets include Android, Christianity, Douban, and Twitter of real information cascades [Huang et al., 2023]. We follow [Huang et al., 2023] to define the sources and final observations, and the time ratio of the sources is 5%. The 7 baselines are Netsleuth [Prakash et al., 2012], LPSI [Wang et al., 2017], DDMSL [Yan et al., 2023], TGASI [Hou et al., 2023], SLVAE [Ling et al., 2022], IVGD [Wang et al., 2022] and GCNSI [Dong et al., 2019]. Transferability in source localization is under-explored, and our Trace is proposed to bridge this gap. (Datasets and baselines are detailed in Appendix D.)

Evaluation Metrics. We evaluate the comparison methods with three popular metrics, i.e., AUC, PR and F1 score [Ling et al., 2022; Hou et al., 2023]. Each case undergoes 10 independent runs for statistical robustness. We report the mean value with the standard deviation.

Model Configuration and Reproducibility. In light of transferability, Trace performs a graph normalization process. Concretely, the input features are given by the eigenvectors of K largest eigenvalues in graph Laplacian, so that the feature dimension of different datasets is normalized to the predefined K. Correspondingly, the functional curvature κθ is fed by the tuple of K largest eigenvalues. Then, we use the exponential map with κθ to obtain the initial encoding on the manifold.

4.2 Empirical Results &Discussion

Main Results on Real and Synthetic Datasets. The performance on synthetic and real datasets are collected in Table 1 and Table 2, respectively. On synthetic datasets, we generally obtain the best performance except two cases of runner up. On the real datasets, infection observations are fed into the localization methods and the infection simulation process is disabled. IVGD runs into Out-Of-Memory (OOM) error on Facebook dataset. Comparing to the synthetic datasets, all models have experienced the significant decline on real datasets given the complex dynamics in reality. The proposed Trace consistently achieves the best results in terms of AUC, PR and F1 score in Table 2. In addition, to evaluate the efficiency, we show the training time in each epoch in Fig. 3. Trace exhibits satisfactory efficiency compared to the baselines. The following part demonstrates how the proposed components contribute to Trace.

AblationStudy. We study the effectiveness of (a) structureawareness of Schr¨odinger bridge, (b) fully Riemannian operations, and (c) the new learning algorithm. In Trace, the Schr¨odinger bridge is governed by κθ that describes the graph structure. Accordingly, w/oFCurv variant gives a counterpart regardless of structural information. Concretely, we conduct logarithmic map on the encoding and replace our design with the usual Euclidean Schr¨odinger bridge [Thornton et al., 2022]. As for Riemannian operations, w/oFRiem variant leverages an additional tangent space with logarithmic and exponential maps. To evaluate the proposed learning algorithm, w/oMPara variant utilizes a usual parameterized drift of [Thornton et al., 2022] and employs the traditional IMF procedure. The results are summarized in Table 3, and the main findings are listed as follows: 1) Trace consistently outperforms w/oFCurv, and it suggests the importance of structural information in modeling the distribution map for source localization. 2) Comparing to w/oFRiem, we receive performance gain with the proposed operations. 3) w/oMPara exhibits inferior accuracy. It is due to the optimization error in IMF procedure, and a similar result has been reported in Euclidean space [Gushchin et al., 2024]. More importantly, the proposed algorithm demonstrates impressive efficiency, e.g., its training time is less than 1/12 of IMF on Douban as in Fig. 2.

On Geometry. Trace is presented as a Riemannian model, while existing methods perform localization in the Euclidean geometry. This part evaluates the impact of representation space with geometric ablation. Hyperbolic and hyperspherical Trace are created by setting the functional curvature as -1 and 1, respectively. Without loss of generality, the Euclidean counterpart is instantiated with corresponding operations. The comparison results are given in Fig. 4. It shows that the geometry plays a significant role in source localization, as evidenced by the performance across different geometries. Note that, the learnable curvature potentially enhances accuracy through improved geometric matching and offers the advantage for transferability, as in the next part.

Figure 3: Time

Table 4: Performance in the transfer setting with different pretraining datasets in terms of F1 score (%).

Transferability and Impact of Pre-training Dataset. We demonstrate the transferability of the proposed model. Tables 1 and 2 collect the results of training and testing models on the same dataset. Here, Trace is first pre-trained on one dataset, and then the model is fine-tuned and tested on another. We do experiments on different pre-training datasets for comprehensive evaluation. F1 scores are reported in Table 4, which shows the effectiveness of transferring the pretrained model to new graphs. Additionally, Trace exhibits robustness over different pre-training datasets. We argue that the Schr¨odinger bridge on the manifold captures the underlying social dynamics, while the learnable curvature adapts knowledge according to the graph structure.

Case Study & Visualization. We conduct a case study on Karate dataset, and visualize the inference results in Fig. 5, where the inferred sources are marked in red color. Compared to TGASI in Fig. 5(c), Trace shows a more precise alignment with the actual source nodes in Fig. 5(a).

5 RelatedWork

Source Localization. Early practices leverage statisticmetrics, e.g., Jordan centrality, to infer the sources [Zhu et al., 2017]. Nowadays, the deep localization models can be roughly categorized into two groups. The first group primarily focuses on static or temporal characteristics of the sources [Cheng et al., 2024; Hou et al., 2023; Wang et al., 2023; Hou et al., 2024], while the second group models the backward dynamics along with the forward information diffusion [Wang et al., 2022; Ling et al., 2022; Huang et al., 2023; Yan et al., 2023; Zhang et al., 2024]. We are in line with the latter, and the Schr¨odinger bridge presents the duality via forward and backward SDEs. Recently, [Wang et al.,

2024] jointly exploits co-related networks to improve localization accuracy, while [Ling et al., 2024] studies a different problem of source localization from one network to another. Distinguishing from prior Euclidean models, our Riemannian Trace is endowed with the transferability.

Generative Models & Schr¨odinger Bridge. Diffusionand flow-based models emerge as popular generative models, described as SDEs [Song et al., 2021] and probability ODEs [Lipman et al., 2023], respectively. Both of them conduct data generation from a simple distribution, e.g., Gaussian, while we are devoted to bridge the source and final distributions. We notice that rectified flow [Liu et al., 2023] builds such distribution map without randomness. However, the randomness is often nested in the information propagation process. Hence, we consider the Schr¨odinger bridge that models the map between two arbitrary distributions [Shi et al., 2023; Gushchin et al., 2024; Liu et al., 2024]. It is noteworthy that structural information has not yet been encapsulated into the Schr ¨odinger bridge, and we offer an alternative to fill this gap.

Riemannian Graph Learning. Compared to the traditional Euclidean space, Riemannian manifolds show superiority in graph clustering [Sun et al., 2024c; Sun et al., 2024b; Sun et al., 2023b], node classification [Sun et al., 2022b; Sun et al., 2024d; Sun et al., 2024e; Sun et al., 2025a; Fu et al., 2023; Li et al., 2022; Wei et al., 2024], graph structural learning [Sun et al., 2023a] as well as modeling graph dynamics [Sun et al., 2023c; Sun et al., 2022a; Sun et al., 2021]. Recently, ODEs and SDEs are extended to Riemannian manifolds for graph generation [Sun et al., 2025b; Sun et al., 2024a; Fu et al., 2024]. We notice that [Jo and Hwang, 2024] build a Riemannian Schr¨odinger bridge. Besides the different designs on curvature, another key contribution is that we design a new learning algorithm more efficient than previous ones.

6 Conclusion

This paper presents the Riemannian source localizer named Trace with the transferability to different graphs. Specifically, the structural Schr¨odinger bridge governed by functional curvature is established to model the direct transformation between source and final distributions. Furthermore, we design a new learning algorithm that conducts the efficient geodesic bridge matching with the provable optimal projection for Riemannian measures. Extensive experiments show the effectiveness and transferability of the proposed model.

Acknowledgments

This work is supported in part by National Natural Science Foundation of China under Grant 62202164, and the Fundamental Research Funds for the Central Universities 2025MS024. Philip S. Yu is supported in part by NSF under grants III-2106758, and POSE-2346158. The corresponding author is Li Sun.