\section{Introduction}

Recurrent neural networks have been used to perform spatial navigation tasks and the subsequent study of their internal representations has yielded dynamics and structures that are strikingly biological. Metric \citep{Cueva2018EmergenceLocalization, Banino2018Vector-basedAgents} and non-metric \citep{Recanatesi2019SignaturesManifolds} representations mimicking grid \citep{Fyhn2004SpatialCortex} and place cells \citep{O'Keefe1978} respectively form once the recurrent network has learned a predictive task in the context of a complex environment. \cite{Cueva2020Emergence} demonstrates not only the emergence of characteristic neural representations, but also hallmarks of head direction system cells when training a recurrent network on a simple angular velocity integration task.
Biologically, non-metric representations are associated with landmark spatial memory, in which place cells within the mammalian hippocampus fire when the associated organism is present in a corresponding place field. Extrafield firing of place cells occurs when these neurons spike outside of these contiguous place field regions. Here we show that recurrent neural networks (RNNs) not only form corresponding attractor landscapes, but also produce representations with internal dynamics that closely resemble those found experimentally in the hippocampus when performing goal-directed behaviour.

Research in neuroscience such as that of \cite{Johnson2007NeuralPoint}, shows that spatial representations in mice in the CA3 region of the hippocampus frequently fire nonlocally. \cite{Griffin2007SpatialTask} show that a far higher proportion of hippocampal neurons in the CA1 region in rats performing an episodic task in a T-shaped maze encode the phase of the task rather than spatial information (in this case trajectory direction). \cite{Ainge2007HippocampalPoints} show CA1 place cells encode destination location at the start position of a maze. \cite{Lee2006GradualLocations} demonstrate that place fields of CA1 neurons gradually drift toward reward locations throughout reward training on a T-shaped maze.

In this work we show that a recurrent neural network learning a choice-reward based task using reinforcement learning, in conjunction with predictive sensory learning in a T-shaped maze produces an internal representation with consistent extrafield firing associated with consequential decision points.
In addition we find that the network's representation, once trained, follows a forward sweeping pattern as identified by \citet{Johnson2007NeuralPoint}. We then show that a higher proportion of units in the trained network show strong selectivity for the encoding or choice phase of the task than the proportion showing selectivity for spatial topology. Importantly, these properties only emerge during predictive learning, where task learning is much faster compared to traditional deep Q learning.

\section{Method}

\begin{figure}[h]
    \centering
    \begin{minipage}{0.5\textwidth}
        \centering
        \includegraphics[width=0.7\textwidth]{final_figures/wall_prediction.png}
    \end{minipage}\hfill
    \begin{minipage}{0.5\textwidth}
        \centering
        \includegraphics[width=0.68\textwidth]{final_figures/maze.png}
    \end{minipage}
\caption{Left, the wall observation and cue received by the network at each timestep. Right, the entangled predictive task the LSTM network is pre-trained on in order to generate a non-metric map of the maze environment.}
\label{fig:CueMaze}
\end{figure}
We use a form of the cued-choice maze used by \citet{Johnson2007NeuralPoint} which has a central T structure with returning arms, shown in Figure \ref{fig:CueMaze}. All walls of the maze are tiled with distinct RGB colours which are generated at random and remain fixed throughout.
An agent is initially learning to predict the next sensory stimulus given its movement. This combination of unsupervised learning and exploration has been shown previously to produce place cell-like encoding of the agent's position \citep{Recanatesi2019SignaturesManifolds}. Next, rewards at four possible locations are introduced and the agent is tasked with associating a cue with the rewarding trajectory. The agent has four vision sensors, one in each cardinal direction, reading the wall RGB colours they intersect.
The cue tone is played to the agent as it passes the halfway point of the central maze stem. A low frequency cue indicates that the agent will turn left at the top of the maze stem and a high frequency cue indicates a right turn. These cue tones take the form of a high or low valued scalar perturbed with normally distributed noise if at a cue point, with a zero value given at all other locations.
These four RGB colours as well as the cue frequency at the current location make up the total input received by the agent.

The agent is controlled by a recurrent neural network comprised of a 380 unit Long-Short term memory \citep{Hochreiter1997LongMemory} (LSTM) network with a single layered readout for the prediction of RGB values.
We first pre-train the network by tasking it with predicting the subsequent observation of wall colours from the currently observable wall colours given its trajectory through the maze. The agent's starting location is at the bottom of the central stem of the T maze and a trajectory of left or right at the top of the central stem is chosen pseudorandomly, depicted with red and blue arrows respectively in Figure \ref{fig:CueMaze} and corresponding to the low (red trajectory) or high (blue trajectory) cue tone value given halfway up the stem. As in the experiments by \citet{Johnson2007NeuralPoint}, during pre-training the agent does not choose any of its actions and is only learning to predict the sequence of wall colors it encounters. In a given pre-training iteration, we collect all observations as the agent traverses the maze until it returns to the start location at the bottom of the central stem and finally train the LSTM on the entire collected trajectory. The network is trained with a mean-squared error loss of predicted and target wall colours (Eq. 1), with model parameters optimised using Adam \citep{Kingma2015Adam:Optimization} and a learning rate of 0.001.
\begin{equation} \label{eq:loss_rgb}
loss_{rgb} = \frac{1}{n}\sum_{i=1}^{n}(y_{rgb} - (W_{rgb}h_t + b_{rgb}))^2
\end{equation}
To solve this task, the network has to maintain the cue tone played in its internal memory for several time steps in order to predict subsequent wall colours from the top of the central stem. In our model, this is achieved through the network forming a non-metric representation (attractor landscape) of the maze environment, as also demonstrated by \cite{XU2020Implementing}. Similarly, behavioural experiments typically have a comparable familiarisation phase with the environment before reward-based tasks are introduced \citep{Johnson2007NeuralPoint,Griffin2007SpatialTask}.

\begin{figure}[h]
\begin{center}
\includegraphics[width=0.8\textwidth]{final_figures/RLschematicsvg.png}
\end{center}

\caption{For the joint task to be learned by the LSTM network, we introduce secondary cue points, where the same cue tone as that played at the primary cue point will be repeated if and only if the agent has proceeded in turning in the direction corresponding to the cue tone frequency given at the primary cue location. The agent is free to choose the next action to be taken when traversing the maze at either the choice point at the top of the stem of the maze or at the secondary cue locations. There are two potential reward sites on both returning arms, with the reward sites being active if the agent is on the returning arm corresponding to the cue tone frequency.}
\label{fig:Mazeschematic}
\end{figure}

Once the LSTM has formed an internal representation of the maze, the agent is tasked with navigating towards potential reward sites whose location is indicated by the cue signal: a low frequency cue indicates active reward sites on the left return arm and a high frequency cue indicates active reward sites on the right return arm - the cue tone and corresponding side of active reward sites are together chosen randomly at each iteration with a secondary cue given if the agent has turned correctly. In this phase there are three choice points wherein the agent is able to choose its next action and is constrained to follow the forward maze direction elsewhere: at the top of the maze stem and at the two secondary choice points (Figure \ref{fig:Mazeschematic}), with initially random movement at these points during reward training. There are 5 steps between the cue and choice points and 7 steps from the choice point to the first reward site on either return arm. The inclusion of the secondary cues as additional choice points was motivated by the experimental set up used by \cite{Johnson2007NeuralPoint}, to compare the network activity at these points to experimental data. These secondary points also give the agent the opportunity to backtrack on its decision made at the primary choice point in light of further environmental observation (the presentation or lack thereof of the secondary cue), and make learning more efficient in our model. This may explain how it speeds up training the animals in the same task.

We additionally introduce a new single layered readout for the LSTM network which predicts state-action values associated with the four cardinal directions in relation to the agent's current position and direction. At each timestep, this ensemble receives the agent's environment observation and the agent follows an epsilon-greedy policy (starting with fully random movement at choice points and a decaying epsilon thereafter) for choosing optimal actions of those available at each of the three choice points.

The recurrent network controlling the agent is trained on a weighted combined loss of a reinforcement learning (RL) task loss and the previously described predictive wall colour loss:
\begin{equation} \label{eq:total_loss}
loss_{combined} = |Q(s, a) - (r + \gamma \cdot Q'(s', \argmax_{a'} Q(s', a'))| + \lambda \cdot loss_{rgb}
\end{equation}

The first component of this loss is the difference between predicted and observed state-action values which are represented by Q-values \citep{Watkins1992Q-learning}, which are a prediction of future global reward:
\begin{equation} \label{eq:q}
Q(s, a) = W_{Q}h_t + b_{Q}
\end{equation}

We use double-Q learning \citep{VanHasselt2016DeepQ-Learning} to train the agent on the task, updating the target Q value predictor ($Q'$ - a LSTM with same number of units) every 15 training iterations. Double-Q learning allows for optimal performance on the reward task in drastically fewer agent maze traversals and network training iterations than with standard DQN \citep{mnih2013playing} based Q-learning which suffers from overestimation of Q-values. We settle on a discount factor ($\gamma$) of 0.8 as values higher than this regularly cause the network to converge on solutions wherein the agent does not take the most direct path to reward locations, with backtracking at secondary choice points. The second loss component is the sensory prediction task which we used to pre-train the network ($\lambda = 0.5$). After the network has been pre-trained (optimising to minimise Eq. \ref{eq:loss_rgb}), the network achieves perfect performance on the predictive task. This loss component is included when training the network on the reward task so that the spatial map (non-metric attractor landscape) of the maze environment formed during pre-training is maintained throughout Q-learning. This ensures the map is not overwritten as would happen when Q-learning is performed alone, and leads to faster task learning (see results). We optimise the network for this joint task using Adam and a learning rate of 0.0005, which we find improves the rate of convergence with optimal task performance, as opposed to higher learning rates which still converge but with backtracking often inherent in task solutions.

In contrast to much of the previous work on spatial representations in recurrent networks, we do not give the network any indication of the agent's location or movement. This makes the task considerably more difficult due to the unpredictable movement possible at choice points during the reward task. The network is coerced into storing the current movement direction of the agent in its representation, in addition to storing the cue frequency. As such, a network of Gated Recurrent Units \citep{Cho2014LearningTranslation} (GRUs) or vanilla RNN units was unable to perform well in either the pre-training or joint RL task due to these prevalent long term dependencies (18 steps between cue and final reward).

To analyse the representations formed by the network, we train a further single layered fully connected network with a softmax layer (shown in green in Figure \ref{fig:Mazeschematic}) to predict the agent's next location using the activity of the LSTM. There is no backpropagation of gradients between this predictor and the LSTM network, and the predictor is trained at the end of reward training. The plots in Figures \ref{fig:presweep} and \ref{fig:postsweep} are distributions indicating the probability of agent location inferred from LSTM activity by the predictor. This is used in place of the decoding algorithm used by \cite{Johnson2007NeuralPoint} to predict the neuronally inferred maze location of rats when performing a cue based task.

\section{Results}
The agent learns the sensory prediction task to a high degree of recall and after around a thousand training iterations (combined loss with pre-training in Figure \ref{fig:performace}), the agent was able to achieve perfect performance on the reward task when the LSTM network had 380 or more units (Fig. \ref{fig:performace}, right). We trained the reinforcement learning (Eq. \ref{eq:total_loss}) portion of the task in an epsilon greedy manner, with a steadily decaying epsilon to ensure that the agent would choose the rewarding path consistently once actions were chosen at choice points completely by the network. Notably, the agent did not turn at either of the secondary choice points once training had completed - only at the primary choice point.

We attempted to run the reinforcement learning task alone in a maze with no sensory input except the reward cue. In this scenario the network is not able to learn the task due to a lack of self-localisation and is unable to perform the task based on step counting between the cue and choice point. In addition, the reward based reinforcement learning task was attempted using Q-learning alone with a loss function that did not include the wall colour prediction error, both with and without pre-training (shown in Fig. \ref{fig:performace}, left). In both cases we find that the reward task is not learnable with the same higher rate of epsilon decay we use for the combined loss function with pre-training, as the network quickly forgets the attractor landscape of the maze formed during pre-training, which we maintain through the combined loss (Eq. \ref{eq:total_loss}). We also find the network can solve the reward task using the combined loss without pre-training, albeit in around 3 times the number of maze traversals as with the use of the spatial map formed in the pre-trained case (Fig. \ref{fig:performace}, left).
\begin{figure}[h]
\begin{center}
\includegraphics[width=0.94\textwidth]{final_figures/both_smaller.png}
\end{center}
\caption{Left: Success rate (proportion of direct traversals to reward locations) of each set of training paradigms on the reward task, averaged over 10 initial conditions and random wall colours using optimal rate of epsilon decay for each paradigm, each shown with a 95\% confidence interval. Attractor landscape formed during pre-training alongside combined loss allows network to achieve perfect performance on reward task in relatively few maze traversals. Q-learning alone without pre-training also achieves perfect performance in more than twice the number of maze traversals. Q-learning alone with pre-training takes far more maze traversals to converge (and is less likely to be optimal) due to the non-random initial state of network and inability to utilise the spatial map formed. Combined training without pre-training also takes relatively many maze traversals to converge due to a relatively difficult joint task with no biased initial state. Right: Pre-trained network optimised with combined loss converges at similar rates with different network sizes above 380 units.}
\label{fig:performace}
\end{figure}
\subsection{Extrafield place cell firing}
First, we investigate the representation learned by the the network during these two stages of training. Pre-training causes the formation of discrete attractors that resemble place cells in the hippocampus. Individual units in the network generally have well isolated place fields, which together cover the whole maze and therefore allow reliable decoding of agent location. In addition to an increase in activity in a particular unit when the agent moves across its respective place field, we also observe substantial extrafield firing of these units. This activity occurs mainly at the primary cue location and at the first choice point after pre-training. After training on the reward task, in addition to the place fields, the network also has units with extrafield activity at the secondary choice points (Fig. \ref{fig:extrafield}E).

In the top row of Figure 3(A-D) we show activity in 4 reward trained LSTM units obtained through the collection of unit activity from a full left sided trajectory from the maze start point returning to the start point with cues presented, together with a full right sided trajectory. We show all activity from this activity collection in the top row of Figure 3(A-D) and proceed to outline the maze areas for each unit with activity higher than 30\% of the peak activity of that particular unit (mirroring the experimental threshold used by \cite{Johnson2007NeuralPoint}), denoting them as place fields corresponding to these LSTM units.
In experiments, rodents seem to pause at high consequence decision points \citep{Johnson2007NeuralPoint} with alternating head movement behaviour signifying vicarious trial and error (VTE) \citep{Muenzinger1938VicariousEfficiency, Hu1995AFunction}.
In the activity plots in the bottom row of Figure 3(A-D), we simulate this using our reward trained model by running the agent from the start position at the bottom of the maze stem, then pausing it at the top of the stem, with a left cue presented halfway up. We show activity above 60\% of unit peak activity (identified with the previously collected aggregated activity) shown in addition to the previously identified place fields.

The network representation seems to sample both return arms, with surprisingly high extrafield activity in the shown LSTM units when the agent is paused at the maze choice point, a location for which these units do not usually have corresponding activity (Fig. \ref{fig:extrafield}A-D). We define nonlocal firing as unit activity above 60\% of peak averaged unit activity when running the agent along the central stem (bottom row Fig. \ref{fig:extrafield}A-D) to observe only the most poignant extrafield behaviour.

\begin{figure}[h]
\begin{center}
\includegraphics[width=0.89\textwidth]{final_figures/extrafieldsvg.png}
\end{center}
\caption{\textbf{A}-\textbf{D}) \textit{Top row}: Activity maps showing well isolated place fields of four LSTM units (acting as place cells) indicated in dotted regions after the reward task. Place fields determined by contiguous locality with average activity exceeding 30\% peak unit activity during a single left trajectory followed by a right trajectory. \textit{Bottom row}: LSTM unit activity exceeding 60\% of previously averaged peak unit activity for the given neuron when agent run from bottom of maze stem to top of stem and given a low frequency (left) cue tone halfway up the stem, then stationary at choice point with LSTM network repeatedly receiving observation from choice point for timesteps thereafter (shown in addition to previously determined unit place fields in dotted regions). \textbf{A}, \textbf{B}) Strong extrafield firing contiguously from cue to choice point. \textbf{C}, \textbf{D}) High extrafield firing at choice point while agent is paused at top of stem. \textbf{E}) Place fields (determined from average activity on both trajectories) of four LSTM units outlined in dotted areas after reward based task. High levels of consistent extrafield firing at primary and secondary cue points in 56\% of LSTM units.}
\label{fig:extrafield}
\end{figure}

\subsection{Forward moving representation}
The internal dynamics of the LSTM network has an inherently forward looking representation of the maze once pre-trained in a predictive manner. As depicted in Fig. \ref{fig:presweep}, whilst the agent is stationary, the dynamics of the LSTM network moves forward through the maze, incorporating the trajectory modulation of the cue played halfway up the maze stem.
The forward movement of the representation is also notable for having an inconsistent velocity, where the LSTM inferred agent location jumps \citep{Hasselmo2009ACells} from the top of the maze to lower down the arm (timestep 16 in Fig. \ref{fig:presweep}).
\begin{figure}[h!]
\begin{center}
\includegraphics[width=0.69\textwidth]{final_figures/Pre_sweep.png}
\end{center}
\caption{LSTM inferred agent position after pre-training on maze. The agent is run from the start at timestep 1 to timestep 4 where it receives a low frequency cue (indicating a left turn). At timestep 9 the agent is stopped at the top of the maze stem and the LSTM is given the environment observation from this location for the remainder of the shown timesteps. The inferred position then moves left according to the cue with the position seeming to jump abruptly between timesteps 15 and 16. The inferred position then moves back to the starting position at timestep 26. We observe an analogous inferred forward moving representation on the right side of the maze with a high frequency cue.}
\label{fig:presweep}
\end{figure}

\begin{figure}[h!]
\begin{center}
\includegraphics[width=0.79\textwidth]{final_figures/Post_sweep.png}
\end{center}
\caption{LSTM representation after reward training. As previously, we run the agent from the start position to the top of the stem of the maze at timestep 9 with a low frequency (left) cue tone at timestep 4. Again, the agent is stopped at this position with the LSTM network receiving the environment observation from this position for the remainder of the shown timesteps. As with the network purely trained on the predictive task, the representation moves in the direction corresponding to the frequency of the given cue tone. Then between timesteps 14 and 15, the inferred position jumps from the return arm with active reward sites to the alternate arm, with the inferred position moving from this position to the start location fairly consistently. Then the inferred position jumps again at timestep 32 to the rewarding return arm and moves constantly to the start position.}
\label{fig:postsweep}
\end{figure}

In stark contrast to the dynamics of the LSTM network after predictive pre-training, following training on the reward task the forward representation of the LSTM is still looking ahead of the agent but is now displaying sweeping behaviour (Fig. \ref{fig:postsweep}) which is identified experimentally in rats by \cite{Johnson2007NeuralPoint} when performing cue based tasks. When the agent is stationary at the choice point, we observe the representation moving ahead of the agent - first in the direction corresponding to the cue given at the first cue point and then abruptly down the opposing arm of the maze towards the starting location, thereafter the representation moves down the correct arm (corresponding to the cue) and becomes stationary at the maze start location. This path switching behaviour is reliably observed in networks trained on the combined loss (Eq. \ref{eq:total_loss}) with and without pre-training, with differing numbers of units and initial conditions as long as the reward task is solved without backtracking at secondary cue locations. The network lacks a sweeping or forward moving representation when trained on the reward task with Q-learning alone, regardless of pre-training. Thus pre-training does not contribute to sweeping or path switching behaviour.

\begin{figure}[h]
\begin{center}
\includegraphics[width=0.89\textwidth]{final_figures/umapsvg.png}
\end{center}
\caption{UMAP manifold of LSTM network dynamics of complete left trajectory (dark blue) and complete right trajectory (red) shown along with manifold of dynamics when agent run from start location to choice point with left cue (light blue) and right cue (pink) given at cue point and agent paused in place at the top of the maze stem. A few timesteps after the agent is paused, the dynamics of the left cue paused agent (light blue) switches manifold path abruptly from running alongside the complete left trajectory path (blue) and joins the right trajectory path (red), following this for many timesteps before ultimately resulting at the same manifold end position as the complete left trajectory manifold path (blue). This is analogous for the right cue paths (red and pink).}
\label{fig:umap}
\end{figure}
We further investigate the network representation using Uniform Manifold Approximation and Projection (UMAP) \citep{McInnes2018UMAP:Projection}. Figure \ref{fig:umap} shows generally connected manifolds, with closer inspection revealing the dynamics which leads to the sweeping arm behaviour in Figure \ref{fig:postsweep} when the agent is stationary at the primary choice point. Zeroing visual input while the agent is paused at the choice point gives comparable representation dynamics to that observed in Figures \ref{fig:postsweep} and \ref{fig:umap}.

Separately, we find that place fields of particular LSTM units drift forwards from their original firing positions after pre-training, towards the reward locations on the return arms throughout reward training, as shown experimentally in CA1 neurons in \cite{Lee2006GradualLocations}. We observe this behaviour in 50 out of 380 network units (13\%), with final resting locations of place fields at reward locations (seen in Appendix Figure \ref{fig:drift}). This is possibly explained by the gradient of Q values (prediction of predicted reward) spreading backwards from reward locations \citep{Hasselmo2005ABehavior} and becoming stronger throughout training.

\subsection{Selectivity of neuronal units}
In addition to a forward sweeping representation, this trained network also exhibits neural selectivity that closely matches hippocampal circuits. \cite{Griffin2007SpatialTask} reported that after reward learning, hippocampal neurons were more strongly selective for the encoding or choice phase of a task rather than the direction of the organism's trajectory. We garner the preference of selectivity of each neuronal unit in our network using a discrimination index used by \cite{Griffin2007SpatialTask} for the turn direction selectivity ($DI_{\text{turn}}$) and the phase selectivity ($DI_{\text{phase}}$):

\begin{equation} \label{eq:1}
DI_{\text{turn}} = \frac{FR_{\text{right}} - FR_{\text{left}}}{FR_{\text{right}} + FR_{\text{left}}}\quad
DI_{\text{phase}} = \frac{FR_{\text{cue}} - FR_{\text{choice}}}{FR_{\text{cue}} + FR_{\text{choice}}}
\end{equation}

where $FR_{\text{right}}$ for a particular LSTM unit is the mean firing rate from the cue point on the central stem to the choice point at the top of the stem on trajectories where the agent turns right at the choice point. Similarly $FR_{\text{left}}$ is the mean stem firing rate when the agent turns left. $FR_{\text{cue}}$ is the firing rate at the cue (encoding) point averaged over both left and right trajectories and similarly $FR_{\text{choice}}$ is the averaged firing rate at the choice (sampling) point.

\begin{figure}[h]
\begin{center}

\includegraphics[width=0.8\textwidth]{final_figures/selectivesvg.png}
\end{center}
\caption{Histograms showing LSTM unit discrimination index for turn direction selectivity ($DI_{\text{turn}}$) vs task phase selectivity ($DI_{\text{phase}}$). A highly negative selectivity index for turn direction indicates a neuronal unit which exhibits high levels of selectivity (uniquely high network activity) for a leftward trajectory and a highly positive selectivity index indicates selectivity for a rightward trajectory. A negative selectivity for task phase indicates a neuron which is highly selective for the choice (retrieval) phase of the goal based task whereas a positive index indicates a neuron which is highly selective for the cue (encoding) phase of the task.}
\label{fig:selectivity}
\end{figure}

The firing areas used for selectivity measurement are insets in Figure \ref{fig:selectivity}. We use the stem above the cue point to assess turn direction selectivity, and the cue/choice points to assess encoding and sampling ($DI_{\text{phase}}$). Figure \ref{fig:selectivity} shows a higher proportion of LSTM units are strongly task selective rather than turn selective, with significantly more units having large absolute $DI_{\text{phase}}$ indices than $DI_{\text{turn}}$ indices.

In addition, the reward trained network is found to have a disproportionately high number of units (163 out of 380 LSTM units) with place fields at the start location of the maze.
Moreover, we find evidence of conditional destination encoding in these units which were heavily differentiated in their firing with respect to particular rewarding locations, as shown experimentally in CA1 hippocampal place cells \citep{Ainge2007HippocampalPoints, Wood2000HippocampalLocation, Ferbinteanu2003ProspectiveHippocampus}. 59.5\% of units with a place field at the maze start location fired uniquely at this point for rewarding locations on a particular return arm.

\section{Discussion}

In this work we show that networks trained with a combined predictive and goal-based objective exhibit functional dynamics and selectivity behaviour coinciding with that of hippocampal neurons. We demonstrate that extrafield firing activity of network units emerge when a simulated agent, which is trained on a goal based reward task in a T-shaped maze, pauses at decision points - suggesting intrinsic dynamics are encoding the future trajectory of the agent. This mirrors experimental results in hippocampal place cells in rats \citep{Johnson2007NeuralPoint, Frank2000TrajectoryCortex}. At the same time, we find that networks using this combined objective, following pre-training only on a sensory prediction task, can learn the correct goal-directed behaviour much faster than an equivalent network with only a Q learning objective.

Previous work shows that metric neural representations of environments form when an RNN is optimised to predict agent position from agent velocity \citep{Cueva2018EmergenceLocalization, Banino2018Vector-basedAgents} and non-metric representations form when an RNN is trained to predict future sensory events given direction of movement \citep{Recanatesi2019SignaturesManifolds}. When training our model we do not provide the LSTM network with any explicit information about location or direction, it only receives sensory information. This is similar to the purely contextual input received by the model pre-trained by \cite{XU2020Implementing} where no velocity input is given, however, the network used by these authors is still trained on position and landmark prediction in a supervised way.

Instead, our training paradigm forces the LSTM to maintain an implicit notion of movement within its internal state in relation to environmental observations. This, in conjunction with the consideration that model-free RL methods such as Q-learning perform poorly on tasks in dynamic environments such as ours \citep{Dolan2013GoalsBrain}, and the long term dependency on the delayed cue in perspective of the choice location, makes the task outlined in Figure \ref{fig:Mazeschematic} particularly challenging.

Training on a sensory predictive task causes the formation of a non-metric place cell like representation in the activations of network units, similarly to \cite{Recanatesi2019SignaturesManifolds}. These units demonstrate nonlocal extrafield firing (Appendix Figure \ref{fig:extrafield-pre}) and after reward training (Figure \ref{fig:extrafield}).
\cite{Johnson2007NeuralPoint} find that this extrafield firing is particularly striking at consequential decision points where rats usually pause in order to sample previously seen trajectories. We observe that cue or choice point extrafield activity is evident in most LSTM units after training on the reward task. This is likely due to the increased precedence these points have in the agent reaching reward locations. Together the trained LSTM network units form a representation which sweeps along the paths available to the agent, first down the reward path and then the other, as shown in Figure \ref{fig:postsweep} and demonstrated in rats in \cite{Johnson2007NeuralPoint}.

 Although hippocampal place cells are critical for spatial memory \citep{Nakazawa2002RequirementRecall, Florian2004HippocampalMice, Sandi2003RapidTraining, Redish1998TheMaze,Miller2020HumanMemories}, it is currently unclear by what mechanism an ensemble of place cells contributes to a representation of goal-directed behaviour \citep{morris1990does}. Our model and training paradigm is in keeping with the hypothesis that the hippocampus is involved in maintaining a conjunctive representation of cognitive maps and sensory information \citep{whittingtonTolmanEichenbaumMachineUnifying2019}. We show that this paradigm can be extended with predictive learning of Q-values of anticipated future reward, and show that the resulting representation is well suited for learning actions leading from a cue to a reward. Importantly, this representation emerges solely from sampling sensory inputs and predicted rewards, while reinforcement learning itself remains model-free and is initially random. The surprising similarity of the task-dependent activity in our simulations and experimentally recorded neural activity in similar tasks suggests that the model may replicate central aspects of learning and planning in the hippocampus. Our trained model could improve understanding of hippocampal function by testing hypotheses regarding previously unobserved dynamics inexpensively. This could be performed on maze environments such as this work, or more open arena settings once the model is retrained.


\section{Introduction}
\subfile{chapters/introduction}

\section{Preliminaries}\label{sec:preliminaries}
\subfile{chapters/preliminaries}

\section{Partial OT for Graph Matching}\label{sec:method}
\subfile{chapters/method}

\section{Heterogeneous Graph Matching}\label{sec:ppgmh}
\subfile{chapters/ppgmh}

\section{Experiments}\label{sec:experiment}
\subfile{chapters/experiment}

\section{Conclusion}\label{sec:conclusion}

In this paper, we propose a novel partial Gromov-Wasserstein learning framework for partially matching both homogeneous graphs and heterogeneous graphs.
Our framework models the source graph and the target graph as discrete distributions and transports a fraction of the probability mass to match nodes that are relatively similar.
Incorporating the heterogeneous embedding learning, our framework can match general heterogeneous graphs.
Empirical results demonstrate that our framework can improve the graph alignment quality.
In the future, we plan to extend our framework to multi-graph matching, graph partitioning, graph compression, and network retrieval.

\section*{Ethics Statement}

This work proposes a framework for partially matching graphs, which has many applications including aligning protein-protein networks from different species, linking user accounts in social platforms, aligning entities in knowledge graphs, to name a few.
Our framework can improve the matching quality hence benefit these applications.
We admit that negative aspects may exist.
Our framework basically does not raise ethical concerns.
But when it is maliciously used, our framework can be used to determine the identity of anonymous users in a social platform and cause privacy leakage problems.

\bibliographystyle{plainnat}

\subsection{Proposed Model}
We simultaneously learn the transport map $\mathbf{T}$ indicating the correpsondence between graphs and the node embeddings $\mathbf{Z}^s$ and $\mathbf{Z}^t$ that are parameterized by $\mbox{\boldmath$\theta$\unboldmath}^s$ and $\mbox{\boldmath$\theta$\unboldmath}^t$ respectively, which leads to the following optimization objective
\begin{equation}\label{objective}
	\begin{aligned}
	\min_{\mathbf{T},\mbox{\boldmath$\theta$\unboldmath}^s,\mbox{\boldmath$\theta$\unboldmath}^t}&\Big\langle \mathbf{L}\Big(\mathbf{C}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{C}^t(\mbox{\boldmath$\theta$\unboldmath}^t),\mathbf{T}\Big),\mathbf{T}\Big\rangle\\
	+\alpha\Big\langle &\mathbf{K}\Big(\mathbf{Z}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{Z}^t(\mbox{\boldmath$\theta$\unboldmath}^t)\Big),\mathbf{T}\Big\rangle+\alpha_1 R\Big(\mathbf{Z}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{Z}^t(\mbox{\boldmath$\theta$\unboldmath}^t)\Big),
	\end{aligned}
\end{equation}
subject to $\mathbf{T}\in\Pi^b(\mbox{\boldmath$\mu$\unboldmath}^s,\mbox{\boldmath$\mu$\unboldmath}^t)$,
where we explain the three terms as follows
\begin{enumerate}
	\item The GW-distance $\Big\langle \mathbf{L}\Big(\mathbf{C}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{C}^t(\mbox{\boldmath$\theta$\unboldmath}^t),\mathbf{T}\Big),\mathbf{T}\Big\rangle$ is the cost for pairwise assignment.
	$\mathbf{C}^s(\mbox{\boldmath$\theta$\unboldmath}^s)$ and $\mathbf{C}^t(\mbox{\boldmath$\theta$\unboldmath}^t)$ are defined as follows
	\begin{equation}\label{eq:C definition}
		\mathbf{C}^z(\mbox{\boldmath$\theta$\unboldmath}^z)=(1-\alpha)\mathbf{G}^z+\alpha \mathbf{K}\Big(\mathbf{Z}^z(\mbox{\boldmath$\theta$\unboldmath}^z),\mathbf{Z}^z(\mbox{\boldmath$\theta$\unboldmath}^z)\Big),
	\end{equation}
	for $z=s,t$,
	where $\mathbf{G}^z=[g(w_{ij}^z)]\in\mathbb{R}^{|\mathcal{V}^z|\times|\mathcal{V}^z|}$ represents the dissimilarity derived according to the graph weight matrix $\mathbf{W}^z$, $\mathbf{K}\Big(\mathbf{Z}^z(\mbox{\boldmath$\theta$\unboldmath}^z),\mathbf{Z}^z(\mbox{\boldmath$\theta$\unboldmath}^z)\Big)=\big[ k\big(\mathbf{z}_i(\mbox{\boldmath$\theta$\unboldmath}^z),\mathbf{z}_j(\mbox{\boldmath$\theta$\unboldmath}^z)\big)\big]\in\mathbb{R}^{|\mathcal{V}^z|\times|\mathcal{V}^z|}$ with each element $k\big(\mathbf{z}_i(\mbox{\boldmath$\theta$\unboldmath}^z),\mathbf{z}_j(\mbox{\boldmath$\theta$\unboldmath}^z)\big)$ measuring the difference between node embeddings $\mathbf{z}_i$ and $\mathbf{z}_j$ within the same graph.
	$g(\cdot)$ and $k(\cdot,\cdot)$ can take forms as in Sec. \ref{sec:experiment}.
	$\alpha$ is a hyperparameter characterizing the contribution of embedding-based dissimilarity to $\mathbf{C}^z(\mbox{\boldmath$\theta$\unboldmath}^z)$.
	\item
	The Wasserstein distance $\Big\langle\mathbf{K}\Big(\mathbf{Z}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{Z}^t(\mbox{\boldmath$\theta$\unboldmath}^t)\Big),\mathbf{T}\Big\rangle=\sum_{i,i'}k\Big(\mathbf{z}_i(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{z}_{i'}(\mbox{\boldmath$\theta$\unboldmath}^t)\Big)T_{ii'}$ measures the cost induced by unary assignment.
	This term is also based on the difference between node embeddings,
	and its contribution is controlled by the same hyperparameter{} $\alpha$, following \citep{xu2019gromov}.
	\item $R\Big(\mathbf{Z}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{Z}^t(\mbox{\boldmath$\theta$\unboldmath}^t)\Big)$ regularizes the embedding learning and is formulated in Sec. \ref{sec:homo embed} and \ref{sec:hete embed} for homogeneous and heterogeneous graphs, respectively.
\end{enumerate}

Although (\ref{objective}) is a complicated non-convex{} optimization problem, we can solve it effectively by alternately updating the transport map and the embedding, as we shall see in the next two subsections.

\subsection{Learning Optimal Transport}

Given $\mbox{\boldmath$\theta$\unboldmath}^s_m$ and $\mbox{\boldmath$\theta$\unboldmath}^t_m$, we solve the following sub-problem{}
\begin{equation}\label{eq:ot subproblem}
	\min_{\mathbf{T}}\Big\langle \mathbf{L}\Big(\mathbf{C}^s(\mbox{\boldmath$\theta$\unboldmath}^s_m),\mathbf{C}^t(\mbox{\boldmath$\theta$\unboldmath}^t_m),\mathbf{T}\Big)+\alpha \mathbf{K}\Big(\mathbf{Z}^s(\mbox{\boldmath$\theta$\unboldmath}^s_m),\mathbf{Z}^t(\mbox{\boldmath$\theta$\unboldmath}^t_m)\Big),\mathbf{T}\Big\rangle,
\end{equation}

subject to $\mathbf{T}\in\Pi^b(\mbox{\boldmath$\mu$\unboldmath}^s,\mbox{\boldmath$\mu$\unboldmath}^t)$.
Following \citep{xie2018a,xu2019gromov}, we solve (\ref{eq:ot subproblem}) iteratively with a proximal point method.
Specifically, in the $n$-th iteration, we update $\mathbf{T}$ by solving
\begin{equation}\label{proximal subproblem}
	\begin{aligned}
	\min_{\mathbf{T}\in\Pi^b(\mbox{\boldmath$\mu$\unboldmath}^s,\mbox{\boldmath$\mu$\unboldmath}^t)}&A_n(\mathbf{T}):=\Big\langle \mathbf{L}\Big(\mathbf{C}^s(\mbox{\boldmath$\theta$\unboldmath}^s_m),\mathbf{C}^t(\mbox{\boldmath$\theta$\unboldmath}^t_m),\mathbf{T}\Big)\\
	+\alpha \mathbf{K}\Big(&\mathbf{Z}^s(\mbox{\boldmath$\theta$\unboldmath}^s_m),\mathbf{Z}^t(\mbox{\boldmath$\theta$\unboldmath}^t_m)\Big),\mathbf{T}\Big\rangle+ \gamma \operatorname{KL}(\mathbf{T}\|\mathbf{T}_n),
	\end{aligned}
\end{equation}
which is solved via projected gradient descent, where both the gradient step and the projection are computed according to the KL divergence \citep{peyre2016gromov}.
More specifically, iterations of projected gradient descent with stepsize{} $\tau$ are given by
\begin{equation}\label{eq:projected gradient descent}
	\mathbf{T}_{l+1,n}=\operatorname{Proj}_{\mathbf{T}\in\Pi^b(\mbox{\boldmath$\mu$\unboldmath}^s,\mbox{\boldmath$\mu$\unboldmath}^t)}^{\operatorname{KL}}\Big(\mathbf{T}_{l,n}\odot e^{-\tau\nabla A_n(\mathbf{T}_{l,n})}\Big),
\end{equation}
where the projection according to the KL divergence onto a convex set $\mathcal{C}$ is defined as $\operatorname{Proj}_{\mathbf{T}\in\mathcal{C}}^{\operatorname{KL}}(\mathbf{A})=\mathop{\rm argmin}_{\mathbf{T}'\in\mathcal{C}}\operatorname{KL}(\mathbf{T}'\|\mathbf{A})$.

Problem (\ref{eq:projected gradient descent}) can be solved via periodic projections \citep{benamou2015iterative}.
Let $\mathcal{C}_1=\{\mathbf{T}|\mathbf{T}\mathbf{1}\le\mbox{\boldmath$\mu$\unboldmath}^s\}$, $\mathcal{C}_2=\{\mathbf{T}|\mathbf{T}^\top\mathbf{1}\le\mbox{\boldmath$\mu$\unboldmath}^t\}$, and $\mathcal{C}_3=\{\mathbf{T}|\mathbf{1}^\top\mathbf{T}\mathbf{1}=b\}$.
The KL projection onto these three sets have closed form
\begin{equation}
	\begin{aligned}
		\operatorname{Proj}_{\mathcal{C}_1}^{\operatorname{KL}}(\mbox{\boldmath$\gamma$\unboldmath})&=\operatorname{Diag}\Big(\min\{\frac{\mbox{\boldmath$\mu$\unboldmath}^s}{\mbox{\boldmath$\gamma$\unboldmath}\mathbf{1}},\mathbf{1}\}\Big)\mbox{\boldmath$\gamma$\unboldmath}\\
		\operatorname{Proj}_{\mathcal{C}_2}^{\operatorname{KL}}(\mbox{\boldmath$\gamma$\unboldmath})&=\operatorname{Diag}\Big(\min\{\frac{\mbox{\boldmath$\mu$\unboldmath}^s}{\mbox{\boldmath$\gamma$\unboldmath}^\top\mathbf{1}},\mathbf{1}\}\Big)\mbox{\boldmath$\gamma$\unboldmath}\\
		\operatorname{Proj}_{\mathcal{C}_3}^{\operatorname{KL}}(\mbox{\boldmath$\gamma$\unboldmath})&=\mbox{\boldmath$\gamma$\unboldmath}\frac{b}{\mathbf{1}^\top\mbox{\boldmath$\gamma$\unboldmath}\mathbf{1}},
	\end{aligned}
\end{equation}
where the division and minimum is element-wise.
Extending the indexing of the sets by $3$-periodicity so that they satisfy $\mathcal{C}_{k+3}=\mathcal{C}_k$ for all $k\in\mathbb{N}$, one computes
\begin{equation}
	\mbox{\boldmath$\gamma$\unboldmath}^{(k)}=\operatorname{Proj}_{\mathcal{C}_k}^{\operatorname{KL}}(\mbox{\boldmath$\gamma$\unboldmath}^{k-1}),\forall k>0
\end{equation}
where $\mbox{\boldmath$\gamma$\unboldmath}^{(0)}=\mathbf{T}_{l,n}\odot e^{-\tau\nabla A_n(\mathbf{T}_{l,n})}$.
The following proposition states the convergence of periodic projections.
\begin{proposition}[Proposition 5 of \cite{benamou2015iterative}]
	$\mbox{\boldmath$\gamma$\unboldmath}^{(k)}$ converges towards the unique solution of (\ref{eq:projected gradient descent}).
\end{proposition}

Hence, we decompose (\ref{eq:ot subproblem}) into a sequence of sub-problems{}, each of which is then solved via the projected gradient descent.
The convergence is guaranteed as follows.
\begin{proposition}
	Every limit point generated by the above procedure, i.e., $\lim_{n\rightarrow\infty} \mathbf{T}_n$, is a stationary point of Eq. (\ref{eq:ot subproblem}).
\end{proposition}
The proof is given in the appendix.

We use the learned $\hat{\mathbf{T}}$ to construct an extended transport plan $\tilde{\mathbf{T}}=[\hat{\mathbf{T}},\mathbf{1}-\hat{\mathbf{T}}\mathbf{1}]\in\mathbb{R}^{|\mathcal{V}^s|\times(|\mathcal{V}^t|+1)}$.
This can be thought of adding a \emph{dummy} point to $\mathcal{G}^t$, which is in the same spirit as \citep{chapel2020partial}.
By choosing the largest $\tilde{T}_{ij}$ for each i, we find the correspondence.
Nodes matched to the dummy node are considered to have no counterparts.
When $b=1$, the partial graph matching problem is reduced to full graph matching.
We can control the number of matched node pairs by varying $b$, as is depicted in Figure \ref{fig:matching result w.r.t. b}.

\begin{figure}[ht]
	\centering
	\begin{subfigure}[b]{.32\textwidth}
		\centering

		\includegraphics[width=\textwidth]{imgs/twonodes.pdf}
		\caption{Two two-node graphs.}
		\label{fig:two-node graph}
	\end{subfigure}
	\hfill
	\begin{subfigure}[b]{.32\textwidth}
		\centering

		\includegraphics[width=\textwidth]{imgs/T1.pdf}
		\caption{The 1\textsuperscript{st} row of $\tilde{\mathbf{T}}$ vs. $b$.}
		\label{fig:T1}
	\end{subfigure}
	\hfil
	\begin{subfigure}[b]{.32\textwidth}
		\centering

		\includegraphics[width=\textwidth]{imgs/T2.pdf}
		\caption{The 2\textsuperscript{nd} row of $\tilde{\mathbf{T}}$ vs. $b$.}
		\label{fig:T2}
	\end{subfigure}
	\caption{Maching result of a pair of two-node graphs for varying $b$ when we only consider the Wasserstein distance with cost matrix $\mathbf{K}\Big(\mathbf{Z}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{Z}^t(\mbox{\boldmath$\theta$\unboldmath}^t)\Big)=\big[\begin{smallmatrix}
			2 & -1\\
			1 & 2
		\end{smallmatrix}\big]$.
		The dotted line denotes the ground truth{} correspondence, i.e., only node s1 in the source graph has its counterpart.
		By definition, $\mbox{\boldmath$\mu$\unboldmath}^s=\mbox{\boldmath$\mu$\unboldmath}^t=[0.5, 0.5]^\top$.
		For each node in the source graph, we select its matched node according to the calculated extended transport plan $\tilde{\mathbf{T}}$, which we demonstrate in Figure \ref{fig:T1} and Figure \ref{fig:T2}.
		If $0.75<b\le1$, s1 is matched to t2 and s2 is matched to t1.
		If $0.25<b<0.75$, only s1 is matched to t2 while s2 is matched to the dummy node in the target graph.
		If $0<b<0.25$, both s1 and s2 are matched to the dummy node.
		Therefore, we can correctly recover the correspondence of nodes with $0.25<b<0.75$.}
	\label{fig:matching result w.r.t. b}
\end{figure}

\paragraph{Complexity for learning optimal transport.}
One of the most computationally expensive steps in calculating the optimal transport given stored embeddings, is evaluating $\mathbf{L}\Big(\mathbf{C}^s(\mbox{\boldmath$\theta$\unboldmath}^s_m),\mathbf{C}^t(\mbox{\boldmath$\theta$\unboldmath}^t_m),\mathbf{T}_{l,n}\Big)$ while calculating $\nabla A_n(\mathbf{T}_{l,n})$.
Fortunately, $\mathbf{L}\Big(\mathbf{C}^s(\mbox{\boldmath$\theta$\unboldmath}^s_m),\mathbf{C}^t(\mbox{\boldmath$\theta$\unboldmath}^t_m),\mathbf{T}_{l,n}\Big)$ can be computed efficiently when common loss functions like square loss or KL divergence are adopted for $\ell(\cdot,\cdot)$ \citep{peyre2016gromov,xu2019gromov}.
To be more specific, $\ell (a,b)$ can be written as $\ell(a,b)=d_1(a)+d_2(b)-h_1(a)h_2(b)$.
Then the loss matrix $\mathbf{L}$ can be obtained by performing element-wise operations as $\mathbf{L}(\mathbf{C}^s(\mbox{\boldmath$\theta$\unboldmath}_m^s),\mathbf{C}^t(\mbox{\boldmath$\theta$\unboldmath}_m^t),\mathbf{T}_{l,n})=d_1\Big(\mathbf{C}^s(\mbox{\boldmath$\theta$\unboldmath}_m^s)\Big)\mbox{\boldmath$\mu$\unboldmath}^s\mathbf{1}^{\top}_{|\mathcal{V}^t|}+\mathbf{1}_{|\mathcal{V}^s|}\mbox{\boldmath$\mu$\unboldmath}^{t\top} d_2\Big(\mathbf{C}^t(\mbox{\boldmath$\theta$\unboldmath}_m^t)\Big)-h_1\Big(\mathbf{C}^s(\mbox{\boldmath$\theta$\unboldmath}_m^s)\Big)\mathbf{T}_{l,n} h_2\Big((\mathbf{C}^t(\mbox{\boldmath$\theta$\unboldmath}_m^t)\Big)$.
Because $\mathbf{T}$ tends to be sparse quickly \citep{peyre2016gromov,xu2019gromov}, $\mathbf{L}(\mathbf{C}^s,\mathbf{C}^t,\mathbf{T})$ involves computational complexity $\mathcal{O}(V^3)$, where $V=\max\{|\mathcal{V}^s|,|\mathcal{V}^t|\}$.
Given the embedding vectors, the computational cost for $\mathbf{K}\Big(\mathbf{Z}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{Z}^t(\mbox{\boldmath$\theta$\unboldmath}^t)\Big)$ is $\mathcal{O}(V^2d)$, where $d$ is the dimension of embeddings.
In addition, if we run projected gradient descent for $N$ iterations, the complexity is $\mathcal{O}(NV^3)$.
Therefore, $\mathcal{O}(V^2d+NV^3)$ computational cost is involved in learning the optimal transport.

\subsection{Updating Embeddings}\label{sec:homo embed}
With slight abuse of notation, let $\ell(\cdot,\cdot)$ in Eq. (\ref{eq:gw distance}) be an element-wise loss function.
Then $R\Big(\mathbf{Z}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{Z}^t(\mbox{\boldmath$\theta$\unboldmath}^t)\Big)$ for homogeneous graphs is given by
\begin{equation}\label{eq:homo regularizer}
	\sum_{z=s,t}\ell\bigg(\mathbf{K}\Big(\mathbf{Z}^z(\mbox{\boldmath$\theta$\unboldmath}^z),\mathbf{Z}^z(\mbox{\boldmath$\theta$\unboldmath}^z)\Big),\mathbf{G}^k\bigg)+\underbrace{\ell\bigg(\mathbf{K}\Big(\mathbf{Z}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{Z}^t(\mbox{\boldmath$\theta$\unboldmath}^t)\Big),\mathbf{G}^{st}\bigg)}_{\text{optional}}.
\end{equation}

Note that if some ground truth{} correspondences across $\mathcal{G}^s$ and $\mathcal{G}^t$ are available,
we can calculate a distance matrix between the two graphs, denoted as $\mathbf{G}^{st}\in\mathbb{R}^{|\mathcal{V}^s|\times|\mathcal{V}^t|}$,
and require the distance between the embeddings to be close to $\mathbf{G}^{st}$,
as shown in Eq. (\ref{eq:homo regularizer}).

Given the learned transport plan $\hat\mathbf{T}_m$, we update the embeddings as follows
\begin{equation}\label{eq:homo embedding}
	\begin{aligned}
		\min_{\mbox{\boldmath$\theta$\unboldmath}^s,\mbox{\boldmath$\theta$\unboldmath}^t}&\Big\langle \mathbf{L}\Big(\mathbf{C}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{C}^t(\mbox{\boldmath$\theta$\unboldmath}^t),\hat\mathbf{T}_m\Big),\hat\mathbf{T}_m\Big\rangle\\
		+\alpha&\Big\langle \mathbf{K}\Big(\mathbf{Z}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{Z}^t(\mbox{\boldmath$\theta$\unboldmath}^t)\Big),\hat\mathbf{T}_m\Big\rangle+\alpha_1 R\Big(\mathbf{Z}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{Z}^t(\mbox{\boldmath$\theta$\unboldmath}^t)\Big),
	\end{aligned}
\end{equation}
which can be solved effectively by SGD.
Note that $\mbox{\boldmath$\theta$\unboldmath}^s$ and $\mbox{\boldmath$\theta$\unboldmath}^t$ are initialized randomly and $\mathbf{Z}^s$ and $\mathbf{Z}^t$ are thus unreliable at early iterations.
We adjust $\alpha$ as follows:
with the total number of rounds as $M$, in the $m$-th iteration,
we set $\alpha=\frac{m}{M}$.
\text{PPGM}{} is summarized in Algorithm \ref{alg:ppgm}.
Both the vertex embedding and the optimal transport learning can be conducted on GPUs.

\begin{algorithm}[ht]
	\caption{Partial optimal transport-based Partial Graph Matching (\text{PPGM}{})\label{alg:ppgm}}
	\begin{algorithmic}[1]
		\State {\textbf{Input:}} $\alpha_1$, total rounds $M$, graphs $\mathcal{G}^s$ and $\mathcal{G}^t$, total transport mass $b$.
		\State {\textbf{Output:}} Correspondence set $\mathcal{C}$.
		\State Initialize $\mbox{\boldmath$\theta$\unboldmath}^s$ and $\mbox{\boldmath$\theta$\unboldmath}^t$ randomly, $\hat{\mathbf{T}}_0=\mbox{\boldmath$\mu$\unboldmath}^s\mbox{\boldmath$\mu$\unboldmath}^{t\top}$.
		\For{$m=0,\dots,M-1$}
		\State $\alpha=\frac{m}{M}$
		\State Calculate $\hat\mathbf{T}_{m+1}$ by solving  (\ref{eq:ot subproblem}).
		\State Update embedding via solving (\ref{eq:homo embedding}) or (\ref{eq:hete embedding})
		\State Store the embedding
		\EndFor
		\State $\hat{\mathbf{T}}=\hat\mathbf{T}_M$
		\State $\tilde{\mathbf{T}}=[\hat{\mathbf{T}},\mathbf{1}-\hat{\mathbf{T}}\mathbf{1}]$
		\State Initialize correspondence set $\mathcal{C}=\emptyset$
		\For{$u_i\in\mathcal{V}^s$}
		\State $j=\arg\max_j \tilde{T}_{ij}$.
		\If{$j\neq(|\mathcal{V}^t|+1)$}
		\State $\mathcal{C}=\mathcal{C}\cup\{(u_i,v_j)\}$
		\EndIf
		\EndFor
	\end{algorithmic}
\end{algorithm}

\subsection{Heterogeneous Network Embedding}\label{learning embedding vectors}

In this subsection, we elaborate how we obtain the embeddings $\{f(u;\mbox{\boldmath$\theta$\unboldmath}^z)|u\in\mathcal{V}^z,z=s,t\}$, which consists of three parts: (1) sampling heterogeneous neighbours; (2) encoding contents; (3) aggregating heterogeneous neighbours.

\paragraph{Sampling heterogeneous neighbours.}
We adopt a \textbf{r}andom \textbf{w}alk  with \textbf{r}estart  strategy $\text{RWR}(p,\eta,\iota,n)$ to determine the set of neighbours $\{\mathcal{N}_u^z|u\in\mathcal{V}^z,z=s,t\}$ \citep{tong2006fast,zhang2019heterogeneous}.
It contains two steps:
\begin{enumerate}
	\item For each graph $\mathcal{G}^z$, $z=s,t$, we start a random walk from vertex $u\in\mathcal{V}^z$ and iteratively travels to a randomly selected adjacent vertex.
	With probability $p$, the walk returns to the starting vertex.
	\item After sampling $\iota$ random walk paths (each has length $\eta$), for each vertex type $r$, we choose the $n$ most frequent vertices as the $r$-type neighbours of $u$, denoted as $\mathcal{N}_{r,u}^z$.
	Different types of neighbours are then grouped together to form $\mathcal{N}_u^z$, i.e., $\mathcal{N}_u^z=\cup_{r} \mathcal{N}_{r,u}^z$.
\end{enumerate}
Such procedure makes use of the local topological structure of the graph and is more flexible than meta-path based approaches \citep{sun2011pathsim,shi2019heterogeneous}.

\paragraph{Encoding contents.} We use a multilayer perceptron $\MLPh r$ to encode contents for the $r$-type vertices. Specifically, we denote the feature representation of vertex $u\in\mathcal{V}^z$ as $\mathbf{x}_u\in\mathbb{R}^{D}$. The content is then encoded using $\text{MLP}: \mathbb{R}^{D}\rightarrow\mathbb{R}^d$ where $d$ is the dimension of an embedding vector, i.e.,
\begin{equation*}
	f_1(u)=\text{MLP}(\mathbf{x}_u;\mbox{\boldmath$\theta$\unboldmath}_{\MLPh r}^z),
\end{equation*}
for vertex $u$ of type $r$ and $z=s,t$.

\begin{table*}[ht]
	\centering
	\begin{tabular}{c|ccc|ccc|ccc}
		\hline
		\multirow{2}{*}{$\rho$} & \multicolumn{3}{c|}{0.7} & \multicolumn{3}{c|}{0.5} & \multicolumn{3}{c}{0.3} \\ \cline{2-10}
		& recall    & precision   & F1   & recall    & precision   & F1   & recall    & precision   & F1   \\ \hline
		SuperGlue               &  $.40\pm.10$        &    $.36\pm.12$    &  $.38\pm.11$    &      $.40\pm.10$    &    $.33\pm.08$    &    $.36\pm.09$  &   $.23\pm.01$       &   $.19\pm.01$     &   $.21\pm.01$   \\
		ZAC                    &   $.69\pm.03$       &  $.62\pm.01$  &     $.65\pm.02$ &   $.35\pm.07$ &  $.30\pm.08$ &  $.32\pm.07$  &  $.20\pm.02$    &  $.17\pm.02$  & $.19\pm.02$ \\
		BB-GM                   &  $.25\pm.03$  &  $.20\pm.03$   &   $.23\pm.03$ &  $.22\pm.04$  &  $.16\pm.03$  &  $.18\pm.03$ &  $.19\pm.03$  &  $.12\pm.01$ &  $.14\pm.01$ \\
		\hline
		\text{PPGM}{}                &   $\mathbf{.83\pm.02}$  & $\mathbf{.73\pm.01}$  &  $\mathbf{.78\pm.02}$ &  $\mathbf{.46\pm.02}$  &  $\mathbf{.42\pm.01}$  & $\mathbf{.44\pm.01}$  & $\mathbf{.38\pm.02}$  &  $\mathbf{.35\pm.03}$  & $\mathbf{.36\pm.02}$ \\ \hline
	\end{tabular}
	\caption{The performance of \text{PPGM}{} and state-of-the-art methods on K-NN graph datasets{} with varying overlap ratio $\rho$.\label{tb:knn}}
\end{table*}

\paragraph{Aggregating Neighbours.} Aggregating content vectors $\{f_1(u)|u\in\mathcal{V}^z,z=s,t\}$ includes two steps:
\begin{enumerate}
	\item \emph{Same type neighbours aggregation.} We employ Bi-LSTM networks to aggregate content vectors for $\mathcal{N}_{r,u}^k$, i.e.,
	\begin{equation*}
	\begin{aligned}
		f_2^r(u)=\frac{1}{|\mathcal{N}_{r,u}^z|}&\sum_{u'\in\mathcal{N}_{r,u}^z}\Big[\overset{\rightarrow}{\text{LSTM}}\big(f_1(u');\mbox{\boldmath$\theta$\unboldmath}^z_{\text{LSTM-} r}\big)\\
		 &\bigoplus \overset{\leftarrow}{\text{LSTM}}\big(f_1(u');\mbox{\boldmath$\theta$\unboldmath}^z_{\text{LSTM-} r}\big)\Big],
	\end{aligned}
	\end{equation*}
	for vertices $u'$ of type $r$, where $\text{LSTM}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d/2}$ and $\bigoplus$ means vector concatenation.
	\item \emph{Types combination.} Different types of neighbours make different contributions to the final embeddings. To combine the vectors generated by the previous step, we adopt an attention mechanism \citep{velivckovic2017graph}. The output embedding vector is thus defined as
	\begin{equation*}\label{eq:feature map}
		\begin{aligned}
		f(u)&=\sum_{f_i\in\mathcal{F}(u)}a^{u,i}f_i,\text{ with }\\
		a^{u,i}&=\frac{\exp(\text{LeakyReLU}\Big(\mathbf{h}_u^\top[f_i\bigoplus f_1(u)]\Big))}{\sum_{f_i\in\mathcal{F}(u)}\exp(\text{LeakyReLU}\Big(\mathbf{h}_u^\top[f_i\bigoplus f_1(u)]\Big))}
		\end{aligned}
	\end{equation*}
	where $\mathcal{F}(u)=\Big\{f_1(u)\cup\Big(f_2^r(u),\forall r\Big)\Big\}$, attention parameter $a^{u,i}$ indicates the importance of component vectors and $\mathbf{h}_u$ is the parameter to be learned.
\end{enumerate}

Leveraging the \emph{negative sampling} technique \citep{mikolov2013distributed,cai2018comprehensive,zhang2019heterogeneous}, $\log \mathbb{P}\big(\mathcal{N}_u^z|f(u;\mbox{\boldmath$\theta$\unboldmath}^z)\big)$ in Sec. \ref{sec:hete graph prelim} is defined as follows
\begin{equation*}\label{negative sampling}
	\begin{aligned}
		\log \mathbb{P}\Big(\mathcal{N}_u^z|&f(u;\mbox{\boldmath$\theta$\unboldmath}^z)\Big)=\sum_{u_p\in\mathcal{N}_{r,u}^z}\log\sigma\Big(f(u;\mbox{\boldmath$\theta$\unboldmath}^z)^\top f(u_p;\mbox{\boldmath$\theta$\unboldmath}^z)\Big)\\
		&+\sum_{u_n\in\mathcal{V}^z\backslash\mathcal{N}_u^z}\log\sigma\Big(-f(u;\mbox{\boldmath$\theta$\unboldmath}^z)^\top f(u_n;\mbox{\boldmath$\theta$\unboldmath}^z)\Big),
	\end{aligned}
\end{equation*}
where $z=s,t$, $r$ is the type of $u$, $\sigma(\cdot)$ is the sigmoid function, and $\mbox{\boldmath$\theta$\unboldmath}^z=\Big[\mbox{\boldmath$\theta$\unboldmath}_{\MLPh r}^z\bigoplus \mbox{\boldmath$\theta$\unboldmath}^z_{\text{LSTM-} r}|\forall r\Big] \bigoplus \Big[\mathbf{h}_u|u\in\mathcal{N}^z\Big]$.

\begin{table*}[ht]
	\centering
	\begin{tabular}{c|ccc|ccc|ccc}
		\hline
		\multirow{2}{*}{$\rho$} & \multicolumn{3}{c|}{0.7} & \multicolumn{3}{c|}{0.5} & \multicolumn{3}{c}{0.3} \\ \cline{2-10}
		& recall    & precision   & F1   & recall    & precision   & F1   & recall    & precision   & F1   \\ \hline
		SuperGlue               &  $.34\pm.06$        &    $.31\pm.05$    &  $.32\pm.05$    &      $.29\pm.03$    &    $.25\pm.04$    &    $.27\pm.04$  &   $.26\pm.02$       &   $.19\pm.01$     &   $.22\pm.01$   \\
		ZAC                    &   $.37\pm.03$       &  $.34\pm.04$  &     $.35\pm.04$ &   $.34\pm.01$ &  $.30\pm.01$ &  $.32\pm.01$  &  $.24\pm.06$    &  $.21\pm.06$  & $.23\pm.06$ \\
		BB-GM                   &  $.31\pm.07$  &  $.26\pm.06$   &   $.28\pm.06$ &  $.25\pm.03$  &  $.19\pm.02$  &  $.21\pm.03$ &  $.18\pm.01$  &  $.12\pm.01$ &  $.14\pm.01$ \\
		\hline
		\text{PPGM}{}                &   $\mathbf{.64\pm.09}$  & $\mathbf{.54\pm.08}$  &  $\mathbf{.58\pm.08}$ &  $\mathbf{.55\pm.06}$  &  $\mathbf{.43\pm.06}$  & $\mathbf{.48\pm.06}$  & $\mathbf{.48\pm.06}$  &  $\mathbf{.35\pm.05}$  & $\mathbf{.40\pm.05}$ \\ \hline
	\end{tabular}
	\caption{The performance of \text{PPGM}{} and state-of-the-art methods on BA graph datasets{} with varying overlap ratio $\rho$.\label{tb:ba}}
\end{table*}

\subsection{Increasing Inter-type Separability}\label{sec:hete embed}
We explicitly use the type information to increase the separability of inter-type embedding vectors.
Let $\mathcal{H}^K$ be the Reproducing Kernel Hilbert Space induced by feature map $f(\cdot)$, with $K$ being the kernel.
Then the optimal classifier for different types of nodes takes the form $h(\mathbf{z})=\sum_{i=1}^{n}\beta_i K(\mathbf{z}_i,\mathbf{z})$ by the Representer Theorem \citep{belkin2006manifold},
where $\mbox{\boldmath$\beta$\unboldmath}=[\beta_1,\dots,\beta_n]^\top$ is the coefficients vector.
To sum it up, $R\Big(\mathbf{Z}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{Z}^t(\mbox{\boldmath$\theta$\unboldmath}^t)\Big)$ for heterogeneous graph matching is defined as
\begin{equation*}
	\begin{aligned}
		R\Big(\mathbf{Z}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{Z}^t(\mbox{\boldmath$\theta$\unboldmath}^t)\Big)&=\sum_{z=s,t}\bigg(-\log \mathbb{P}\Big(\mathcal{N}_u^z|f(u;\mbox{\boldmath$\theta$\unboldmath}^z)\Big)\\
		&+\zeta\sum_{u\in\mathcal{V}^z}\Big(r_u-h^z\big(\mathbf{z}(\mbox{\boldmath$\theta$\unboldmath}^z)\big)\Big)^2\bigg).
	\end{aligned}
\end{equation*}
Therefore, given the optimal transport $\hat\mathbf{T}_m$ obtained in (\ref{eq:ot subproblem}), we update the embeddings by solving the following optimization problem
\begin{equation}\label{eq:hete embedding}
	\begin{aligned}
		\min_{\mbox{\boldmath$\theta$\unboldmath}^s,\mbox{\boldmath$\theta$\unboldmath}^t,\mbox{\boldmath$\beta$\unboldmath}^s,\mbox{\boldmath$\beta$\unboldmath}^t}&
		\Big\langle \mathbf{L}\Big(\mathbf{C}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{C}^t(\mbox{\boldmath$\theta$\unboldmath}^t),\hat\mathbf{T}_m\Big),\hat\mathbf{T}_m \Big\rangle\\
		+\alpha\Big\langle& \mathbf{K}\Big(\mathbf{Z}^s(\mbox{\boldmath$\theta$\unboldmath}^s),\mathbf{Z}^t(\mbox{\boldmath$\theta$\unboldmath}^t)\Big),\hat\mathbf{T}_m\Big\rangle\\
		+\alpha_1\sum_{z=s,t}&\bigg(-\log \mathbb{P}\Big(\mathcal{N}_u^z|f(u;\mbox{\boldmath$\theta$\unboldmath}^z)\Big)\\
		&+\zeta\sum_{u\in\mathcal{V}^z}\Big(r_u-h^z\big(\mathbf{z}(\mbox{\boldmath$\theta$\unboldmath}^z)\big)\Big)^2\bigg),
	\end{aligned}
\end{equation}
where $r_u$ is the type of $u$ and the last line is the proposed regularizer promoting the separability of inter-type embeddings.

Our method learns node embeddings of the two graphs on the same manifold.
By explicitly using the type information of vertices, we distribute embeddings of different types of vertices to different regions of the manifold.
As is shown in Figure \ref{fig:embed}, the distances between embedding vectors indicate the correspondence across the two graphs.

\begin{figure}[ht]
	\begin{subfigure}{0.23\textwidth}
		\centering

		\includegraphics[width=\linewidth]{imgs/gwl_nc.pdf}
		\caption{Alignment of GWL}
		\label{fig:gwl_nc}
	\end{subfigure}
	\hfill
	\begin{subfigure}{0.23\textwidth}
		\centering

		\includegraphics[width=\linewidth]{imgs/hgm_nc.pdf}
		\caption{Alignment of ours}
		\label{fig:hgm_nc}
	\end{subfigure}
	\hfil
	\begin{subfigure}{0.23\textwidth}
		\centering

		\includegraphics[width=\linewidth]{imgs/gwl_em.pdf}
		\caption{Embeddings of GWL}
		\label{fig:gwl_embed}
	\end{subfigure}
	\hfill
	\begin{subfigure}{0.23\textwidth}
		\centering

		\includegraphics[width=\linewidth]{imgs/hgm_em.pdf}
		\caption{Embeddings of ours}
		\label{fig:hgm_embed}
	\end{subfigure}
	\caption{Comparison of node alignments and embedding vectors between our method (with $b=1$) and GWL \citep{xu2019gromov}, a Gromov-Wasserstein learning algorithm that disregards the type information of nodes, on a toy example. We denote correct and wrong alignments by black dashed lines and red dotted lines, respectively.
		Our method successfully aligns all nodes while GWL only aligns two pairs of nodes correctly. Figure \ref{fig:gwl_embed} and Figure \ref{fig:hgm_embed} demonstrate that the distances between embedding vectors directly indicate the node correspondence.
		In GWL, embedding vectors of different types are in the same region and cannot be distinguished, which worsens the alignment.}
	\label{fig:embed}
\end{figure}

Note that while the last term in Eq. (\ref{eq:hete embedding}) increases the separability of inter-class embeddings,
it may reduce the distances between intra-class embedding vectors.
Therefore $\zeta$ is usually a small number and the last term only regularizes the node embedding for heterogenous graphs.
Problem (\ref{eq:hete embedding}) can also be solved efficiently by SGD.
\subsection{Synthetic Datasets{}}

\begin{figure}[!ht]
	\begin{center}
		\includegraphics[width=0.6\linewidth]{../imgs/synth.pdf}
	\end{center}\caption{The performance of tested algorithms on synthetic dataset{}. The left column compares the node correctness of HGM (ours) with WAVE and SANA. The matching performance of HGM in environments with larger scales of noise is shown in the right column. It is also shown that the objective values (cyan lines) of HGM correspond to the node correctness well.}\label{fig synth}
\end{figure}

We verify the efficacy of out HGM method on two kinds of synthetic datasets{}, K-NN graph dataset{} \cite{wiki:knn} and BA graph dataset{} \cite{barabasi2016network}. The baselines include WAVE and SANA\footnote{WAVE, MAGNA++, and SANA are extended to heterogeneous graphs by \cite{gu2018homogeneous}. SANA is shown to have better empirical performance than MAGNA++.}. The two datasets{} share the same source graph that has two types of vertices, $\mathcal{V}_a$ and $\mathcal{V}_b$. An edge connects vertex $u$ and $v$ if the distance between $u$ and $v$ is among the $k$-th smallest distances from $u$ to other objects in the graph. To mimic the interactions betweent the two types of vertices, $\rho\%|\mathcal{V}_a||\mathcal{V}_b|$ edges are additionally added into the source graph. For the K-NN graph dataset{}, the target graph is constructed by adding $q\%$ noisy edges to the source graph. For the BA dataset{}, we simulate the target graph by adding $q\%|\mathcal{V}^s|$ more vertices and corresponding edges.

We set $|\mathcal{V}^s|=50$ and $|\mathcal{V}_a|=|\mathcal{V}_b|$. $\rho=20$. Each tested algorithm is run for $5$ times.
The \emph{node correctness} is chosen for the evaluation metric: Given the learned correspondence set $\mathcal{C}$ and the ground truth{} alignment $\mathcal{C}_{real}$, we calculate the percent of correctly aligned nodes $NC=\frac{|\mathcal{C}\cap\mathcal{C}_{real}|}{|\mathcal{C}_{real}|}\times 100 \%$. In addition, we report the final corresponding objective values of our HGM algorithm.

Figure \ref{fig synth} visualizes the performance of our HGM method and baseline methods. WAVE and SANA have poor performance even when the percent of noisy nodes and edges is as small as $5\%$ while HGM can align almost all vertices correctly when $q\%=5\%$, as the left column shows. HGM is much more robust to noise (the green lines in the right column). The right column also shows that the correspondence between the objective (\ref{objective}) (cyan lines) and the matching quality (green lines).

\subsection{Feature Matching}
Feature matching, an essential part of a wide variety of computer vision applications such as object recognition, camera callibration, and image retrieval, is the task of estimating correspondences between two images of the same scene/object.

We test HGM, WAVE and SANA on a subset of the celebA dataset{} \cite{liu2015faceattributes}, which contains $51$ face images. Each face has $68$ keypoints (\emph{landmarks}) which we aim to align. Vertices are divided into $5$ families: points on the cheek boundary, eyebrows, eyes, nose, and mouth. The positions of these points are determined using the open-source code of \cite{bulat2017far}, which exhibits superior accuracies. Given an image, an edge connects keypoints $u$ and $v$ in this image if the distance between $u$ and $v$ is among the $k$-th smallest distances from $u$ to other keypoints. The edge weight is calculated based on the distance. An appropriate $k$ ensures there exist edges connecting different types of landmarks. One image is randomly selected to generate the source graph. Similar to the first experiment, we add $q\%$ noisy edges into the target graph. We align the chosen graph to the rest of the generated graphs and test the algorithms for $5$ times. We set $k=7$ and report the node correctness results in Table \ref{tab celebA result}.

\begin{table}[!ht]
	\begin{center}\caption{Node correctness results on celebA}\label{tab celebA result}
		\begin{tabular}{ c c c c }
			\hline
			Algorithms & WAVE & SANA & HGM \\
			\hline
			$q\%=2\%$ & $0.237\pm0.057$ & $0.289\pm0.022$ & $\mathbf{0.421\pm0.132}$ \\
			\hline
			$q\%=3\%$ & $0.235\pm0.096$ & $0.245\pm0.029$ & $\mathbf{0.417\pm0.044}$\\
			\hline
			$q\%=4\%$ & $0.234\pm0.103$ & $0.245\pm0.051$ & $\mathbf{0.412\pm 0.059}$  \\
			\hline
			\end{tabular}
			\end{center}
\end{table}

Due to the limited number of pages, more experimental results are shown in the supplementary material.

\subsection{Homogeneous Synthetic Datasets{}}

We first verify the efficacy of our framework on two kinds of commonly used synthetic datasets{}, K-NN graph datasets{} \citep{wiki:knn} and BA graph datasets{} \citep{barabasi2016network}.
For both kinds of graphs, we sample the feature content vectors from Gaussian distributions $N(\mathbf{0}, \mathbf{1})$ and we slightly disturb the feature vectors in the target graph, that is,
if $u_i\in\mathcal{V}^s$ and $u_{i'}\in\mathcal{V}^t$ are counterparts,
the feature vector of $u_{i'}$ is given by $\mathbf{x}^t_{i'}=\mathbf{x}^s_i+0.1\mbox{\boldmath$\sigma$\unboldmath}_i$, where $\mbox{\boldmath$\sigma$\unboldmath}_i\sim N(\mathbf{0},\mathbf{1})$.
For the K-NN graphs, an edge connects vertex $u$ and $v$ if the distance between $u$ and $v$ is among the $k$ smallest distances from $u$ to other nodes in the graph.
We add new nodes and edges to the BA graphs following the \emph{growth} procedure of the BA model, i.e.,
each new node is connected to $m$ existing nodes with a probability that is proportional to the number of links that the existing nodes already have.
We set the weight of each edge connecting two nodes in one graph as the cosine similarity, i.e., $w_{ij}^z=\frac{1}{2}+\frac{\langle \mathbf{x}_i^z, \mathbf{x}_j^z\rangle}{2\|\mathbf{x}_i\|\|\mathbf{x}_j\|}$ for $z=s,t$.
The datasets{} have $k=3$, $m=3$, and $\#\{\text{ground truth{} matching}\}=50$ when we generate the graphs.
We report the performance of \text{PPGM}{} for varying overlap ratios in Table \ref{tb:knn} and Table \ref{tb:ba}.
Formally, the overlap ratio is defined as $\rho=\frac{\#\{\text{ground truth{} matching}\}}{\min\{|\mathcal{V}^s|,|\mathcal{V}^t|\}}$.
We set $\alpha_1=1\times10^{-2}$ and the total transport mass $b=\rho$.
As is shown, \text{PPGM}{} consistently outperforms state-of-the-art methods on both types of graphs.
With the overlap ratio decreasing, the performance of all methods deteriorates.
Such phenomenon is also observed in \cite{wang2020zero}.

\subsection{Heterogeneous Real-world{} Graphs}\label{sec:hete exp}
\begin{table}[ht]
	\centering
	\begin{tabular}{ccc}
		\hline
		\multicolumn{1}{c}{Data}    & \multicolumn{1}{c}{Node} & \multicolumn{1}{c}{Edge} \\ \hline
		\multirow{3}{*}{Source Graph} & \# author: 286                 & \# author-paper: 618            \\
		& \# paper: 286                  & \# paper-paper: 133              \\
		& \# venue: 18                    & \# paper-venue: 286              \\
		\hline
		\multirow{3}{*}{Target Graph} & \# author: 286                    & \# author-paper: 618              \\
		& \# paper: 286                    & \# paper-paper: 123               \\
		& \# venue: 18                     & \# paper-venue: 286              \\
		\hline
	\end{tabular}
	\caption{Datasets used in the heterogeneous graph matching experiment.}\label{tb:hete exp data}
\end{table}

We extract two academic social networks from the public Aminer data\footnote{\url{https://www.aminer.org/data}} and preprocess data following \cite{zhang2019heterogeneous}.
The main statistics of the graphs are summarized in Table \ref{tb:hete exp data}.
The overlap ratio for the three types of nodes $\{\text{author},\text{paper},\text{venue}\}$ is 0.7, 0.7, and 1 respectively.
We set $\alpha_1=1\times10^{-2}$, $\zeta=1\times10^{-4}$ and $b=0.7$.
\text{PPGM}{} excels all baseline methods by $77\%$ in terms of F-measure.
\begin{table}[ht]
	\centering
	\begin{tabular}{c|ccc}
		\hline
		method & recall    & precision   & F1    \\ \hline
		SuperGlue               &  $.019\pm.005$        &    $.026\pm.008$    &  $.022\pm.006$    \\
		ZAC                    &   $.006\pm.001$       &  $.005\pm.001$  &     $.005\pm.001$ \\
		BB-GM                   &  $.001\pm.001$  &  $.005\pm.005$   &   $.002\pm.002$  \\
		\hline
		\text{PPGM}{}                &   $\mathbf{.030\pm.007}$  & $\mathbf{.058\pm.011}$  &  $\mathbf{.039\pm.006}$  \\ \hline
	\end{tabular}
	\caption{The performance of \text{PPGM}{} and state-of-the-art methods on academic social networks.\label{tb:aca}}
\end{table}

Due to the limited space, more experimental results are provided in the supplementary material.

\subsection{Notation}\label{notation}

We use bold lowercase symbols, bold uppercase letters, uppercase calligraphic fonts, and Greek letters to denote vectors, matrices, spaces (sets), and measures, respectively.
For example, $\alpha=\sum_i \delta(\mathbf{x}_i)$ where $\delta(\mathbf{x})$ is the Dirac at position $\mathbf{x}$, is a discrete measure supported on $\mathbf{x}_1,\dots,\mathbf{x}_n\in \mathcal{X}$.
$\pi(\alpha,\beta)$ is a coupling measure between measures $\alpha$ and $\beta$, that is, for any subset $\mathcal{A}\subset \mathcal{X}$, $\pi(\mathcal{A}\times\mathcal{Y})=\alpha(\mathcal{A})$, and for any $\mathcal{B}\subset\mathcal{Y}$, $\pi(\mathcal{X}\times\mathcal{B})=\beta(\mathcal{B})$.
$\Pi(\alpha,\beta)$ is the set of coupling measures.
$\Sigma^n$ is a probability symplex with $n$ bins, namely the set of probability vectors in $\mathbb{R}^n_{+}$.
We denote the cardinality of set $\mathcal{A}$ by $|\mathcal{A}|$.
$(\cdot)^\top$ is the transpose of a vector.
For two matrices $\mathbf{A}$ and $\mathbf{B}$ that are of the same size, $\langle \mathbf{A},\mathbf{B}\rangle=\operatorname{trace}(\mathbf{A}^\top\mathbf{B})$ is the Frobenius dot-product.

\subsection{Optimal Transport}
Let $\mathcal{X}=\{\mathbf{x}_i\}_{i=1}^m$ and $\mathcal{Y}=\{\mathbf{y}_j\}_{j=1}^n$ be two sample spaces.
We assume two discrete probability distributions over $\mathcal{X}$ and $\mathcal{Y}$ respectively, i.e.,
\begin{equation*}
	\mathbf{p}=\sum_{i=1}^{m}p_i\delta(\mathbf{x}_i) \text{ and }\mathbf{q}=\sum_{j=1}^{n}q_j\delta(\mathbf{y}_j),\text{ s.t. }\mathbf{p}\in\Sigma^m,\mathbf{q}\in\Sigma^n.
\end{equation*}
Optimal transport (OT) addresses the problem of optimally transporting $\mathbf{p}$ toward $\mathbf{q}$, given a cost $K_{ij}$ measuring the distance between $\mathbf{x}_i$ and $\mathbf{y}_j$.
The $p$-Wasserstein distance between $\mathbf{p}$ and $\mathbf{q}$ is given by
\begin{equation*}
	W_p^p(\mathbf{p},\mathbf{q})=\min_{\mathbf{T}\in\Pi(\mathbf{p},\mathbf{q})}\langle \mathbf{K}^p,\mathbf{T} \rangle,
\end{equation*}
where $\mathbf{K}^p=[K_{ij}^p]$ and $\mathbf{T}=[T_{ij}]$.
Gromov-Wasserstein (GW) distance is a generalization of Wasserstein distance \citep{memoli2011gromov,patterson2019hausdorff}.
Endowing the spaces $\mathcal{X}$ and $\mathcal{Y}$ with metrics (distance functions) $d_\mathcal{X}$ and $d_\mathcal{Y}$,
the GW distance is defined as
\begin{equation*}
	GW_p^p(\mathbf{p},\mathbf{q})=\min_{\mathbf{T}\in\Pi(\mathbf{p},\mathbf{q})}\sum_{i,j=1}^{m}\sum_{i',j'=1}^{n}D_{ii'jj'}^pT_{ii'}T_{jj'},
\end{equation*}
where $D_{ii'jj'}=|d_\mathcal{X}(\mathbf{x}_i,\mathbf{x}_j)-d_\mathcal{Y}(\mathbf{y}_{i'},\mathbf{y}_{j'})|$.

The partial OT problem focuses on transporting only a fraction $0\le b\le 1$ of the mass with minimum transportation costs \citep{figalli2010optimal,caffarelli2010free,chapel2020partial}, that is, the set of admissible couplings is given by
\begin{equation*}
	\Pi^b(\mathbf{p},\mathbf{q})=\Big\{\mathbf{T}\Big|\mathbf{T}\in\mathbb{R}_+^{m\times n},\mathbf{T}\mathbf{1}\le\mathbf{p},\mathbf{T}^\top\mathbf{1}\le\mathbf{q},\mathbf{1}^\top\mathbf{T}\mathbf{1}=b\Big\}.
\end{equation*}
and the partial Wasserstein distance and the partial GW distance are defined as follows
\begin{equation*}
	\begin{aligned}
		PW_p^p(\mathbf{p},\mathbf{q})&=\min_{\mathbf{T}\in\Pi^b(\mathbf{p},\mathbf{q})}\langle \mathbf{K}^p,\mathbf{T} \rangle,\\
		PGW_p^p(\mathbf{p},\mathbf{q})&=\min_{\mathbf{T}\in\Pi^b(\mathbf{p},\mathbf{q})}\sum_{i,j=1}^{m}\sum_{i',j'=1}^{n}D_{ii'jj'}^pT_{ii'}T_{jj'}.
	\end{aligned}
\end{equation*}

\subsection{GW Distance for Graph Matching}

Following \citep{titouan2019optimal,barbe2020graph}, we consider undirected and attributed graphs as tuples of the form $\mathcal{G}=(\mathcal{V},\mathcal{E},\mathbf{X},\mathbf{W},f)$, where $\mathcal{V}$ and $\mathcal{E}$ are the sets of vertices\footnote{In this paper, we use words "vertex" and "node" interchangeably.} and edges of the graph, respectively.
Node features are summarized in a $|\mathcal{V}|\times D$ matrix $\mathbf{X}$.
$\mathbf{W}\in\mathbb{R}^{|\mathcal{V}|\times|\mathcal{V}|}$ assigns each observed edge $(u_i,u_j)\in\mathcal{E}$ for $u_i,u_j\in\mathcal{V}$ a weight $w_{ij}$.
$f:\mathbf{X}\rightarrow\mathbf{Z}$ associates each vertex $u_i$ with some representation vector $\mathbf{z}_i$.
A cost matrix $\mathbf{C}=[c_{ij}]\in\mathbb{R}^{|\mathcal{V}|\times|\mathcal{V}|}$ is further defined with each $c_{ij}$ characterizing the dissimilarity between vertices $i$ and $j$.

Given two graphs, denoted as source graph $\mathcal{G}^s$ and target graph $\mathcal{G}^t$, we assume that $|\mathcal{V}^s|\le|\mathcal{V}^t|$, without loss of generality.
Following the definition of \emph{measure network} \citep{chowdhury2019gromov,xu2019gromov}, we assign two distributions $\mbox{\boldmath$\mu$\unboldmath}^s=[\mu^s_i]$ and $\mbox{\boldmath$\mu$\unboldmath}^t=[\mu^t_i]$ to the two graphs, where
\begin{equation}
	\mu^z_i=\frac{\sum_j w_{ij}^z}{\sum_{ij}w_{ij}^z}, \text{ for }z=s,t.
\end{equation}
The GW distance between $\mathcal{G}^s$ and $\mathcal{G}^t$ is then defined as
\begin{equation}\label{eq:gw distance}
	\begin{aligned}
	GW(\mbox{\boldmath$\mu$\unboldmath}^s,\mbox{\boldmath$\mu$\unboldmath}^t)&=\min_{\mathbf{T}\in\Pi(\mbox{\boldmath$\mu$\unboldmath}^s,\mbox{\boldmath$\mu$\unboldmath}^t)}\sum_{i,j,i',j'}\ell(c_{ij}^s,c_{i'j'}^t)T_{ii'}T_{jj'}\\
	&=\min_{\mathbf{T}\in\Pi(\mbox{\boldmath$\mu$\unboldmath}^s,\mbox{\boldmath$\mu$\unboldmath}^t)}\langle \mathbf{L}(\mathbf{C}^s,\mathbf{C}^t,\mathbf{T}),\mathbf{T}\rangle,
	\end{aligned}
\end{equation}
where $\ell(c_{ij}^s,c_{i'j'}^t)$ measures the distance between two scalars $c_{ij}^s$ and $c_{i'j'}^t$.
Accordingly, $\mathbf{L}(\mathbf{C}^s,\mathbf{C}^t,\mathbf{T})=[l_{jj'}]\in\mathbb{R}^{|\mathcal{V}^s|\times|\mathcal{V}^t|}$ and each $l_{jj'}=\sum_{i,i'}\ell(c_{ij}^s,c_{i'j'}^t)T_{ii'}$.
The optimal transport map indicates the correpsondence of vertices across the two graphs.

\subsection{Heterogeneous Graphs}\label{sec:hete graph prelim}

A heterogeneous graph contains multiple types of nodes and edges.
Random walk based methods are often employed to capture the structural properties of heterogeneous graphs \citep{cai2018comprehensive,zhang2019heterogeneous,shi2019heterogeneous}.
To form a random walk path, they recursively select one adjacent vertex of the current vertex at random.
When the number of vertices within the path reaches	a pre-set number called \emph{walk length} $\eta$, the sampling procedure terminates.
After repeating the above procedure $\iota$ times to gather random walk paths, one can capture the \emph{proximities} among vertices carried by the graph and construct the neighbourhood $\mathcal{N}_u$ of vertex $u$.
Neural networks like LSTM are then applied to generate node embeddings which
preserve the vertex \emph{proximities}  \citep{cai2018comprehensive,zhang2019heterogeneous}.
Specifically, one can learn the embedding vector of each node by optimizing the following objective
\begin{equation*}
	\max_f\sum_{u\in\mathcal{V}}\log\mathbb{P} \Big(\mathcal{N}_u|f(u)\Big),
\end{equation*}
where $\mathbb{P}(\cdot)$ denotes the probability an event happens, $f:\mathcal{V}\rightarrow\mathbb{R}^d$ is the neural network to be trained.