Relative Representations for Cognitive Graphs

Abstract

Although the latent spaces learned by distinct neural networks are not generally directly comparable, even when model architecture and training data are held fixed, recent work in machine learning [13] has shown that it is possible to use the similarities and differences among latent space vectors to derive “relative representations” with comparable representational power to their “absolute” counterparts, and which are nearly identical across models trained on similar data distributions. Apart from their intrinsic interest in revealing the underlying structure of learned latent spaces, relative representations are useful to compare representations across networks as a generic proxy for convergence, and for zero-shot model stitching [13].

In this work we examine an extension of relative representations to discrete state-space models, using Clone-Structured Cognitive Graphs (CSCGs) [16] for 2D spatial localization and navigation as a test case in which such representations may be of some practical use. Our work shows that the probability vectors computed during message passing can be used to define relative representations on CSCGs, enabling effective communication across agents trained using different random initializations and training sequences, and on only partially similar spaces. In the process, we introduce a technique for zero-shot model stitching that can be applied post hoc, without the need for using relative representations during training. This exploratory work is intended as a proof-of-concept for the application of relative representations to the study of cognitive maps in neuroscience and AI.

Appendices

Appendix A: Effect of Anchor Set Size on Reconstruction

Fig. 4.

Average cosine similarity ($\frac{u \cdot v}{\Vert u \Vert \Vert v \Vert }$) between ground-truth CSCG beliefs (messages) and their reconstructions from those of a distinct CSCG model trained on the same room and receiving the same sequence of observations, using the method in Eq. 2, plotted against number N of anchors used to define the relative representations. We begin by setting N to the dimensionality of the model’s hidden state. The average is across all 5000 messages in a test sequence.

Appendix B: Visualizing the Correspondence of Relative Representations Across Models

Fig. 5.

Example representational similarity matrix comparing relative representations of analogous message sequences (i.e. inferred from the same observation sequence) from two distinct models trained on the same environment. This differs from the (dis)similarity matrices typically used in RSA [10], as rows and columns in this case represent distinct sets of first-order representations, i.e. cell (i, j) represents the cosine similarity between $\textbf{r}^A_i$ and $\textbf{r}^B_j$. Thus the diagonal symmetry illustrates the empirical equivalence of these two sets of relative representations.

Appendix C: Comparison to LLC

Locally Linear Coordination (LLC) [22] is a method for aligning the embeddings of multiple dimensionality-reducing models so that they project to the same global coordinate system. While its aims differ somewhat from the procedure outlined in the present study, LLC is also an approach to translating multiple source embeddings to a common representational format. As noted above, there is an interesting formal resemblance between the two approaches, which we explore in this Appendix.

The LLC Representation

LLC presupposes a mixture model of experts trained on N D-dimensional input datapoints $\mathcal {X} = [ x _1, x _2, ..., x _N]$, in which each expert $m_k$ is a dimensionality reducer that produces a local embedding $z_{n_k} \in \mathbb {R}^{d_k}$ of datapoint $x_n$. The mixture weights or “responsibilities” for the model can be derived, for example, as posteriors over each expert’s having generated the data, in a probabilistic setting.

Given the local embeddings and responsibilities, LLC proposes an algorithm for discovering linear map**s $L_k \in \mathbb {R}^{d \times d_k}$ from each expert’s embedding to a common (lower-dimensional) output representation $\mathcal {Y} \in \mathbb {R}^{N \times d}$, which can then be expressed as a responsibility-weighted mixture of these projections. That is to say, leaving out bias terms for simplicity: each output image $y _n$ of datapoint $x _n$ is computed as

$$\begin{aligned} y _n = \sum _k{r_{n_k}}\big ({L_k}{z_{n_k}}\big ) \end{aligned}$$

(4)

Crucially for what follows, with the help of a flattened (1D) index that spans the “batch” dimension N as well as the experts k, we can express this in simpler terms as $\mathcal {Y} = UL$. We define matrices $U \in \mathbb {R}^{N \times \sum _k{d_k}}$ and $L \in \mathbb {R}^{\sum _k{d_k} \times d}$ in terms of, respectively: (a) vectors $u_n$, where $u_{n_j} = r_{n_k}z^i_{n_k}$ (i.e. the jth element of $u_n$ is the ith element of k’s embedding of $x _n$ scaled its responsibility term) — and (b) re-indexed, transposed columns $l_j = l^i_k$ of the $L_k$ matrices. Intuitively, each row $u_n$ of U concatenates the experts’ responsibility-weighted embeddings $r_{n_k}z_{n_k}$ of datapoint $x _n$, while each of L’s d columns is a concatenation of the corresponding row of the projection matrices $L_k$, so that the matrix product UL returns a responsibility-weighted prediction for $y _n$ in each row (see Fig. 6).

Relationship to Our Proposal

Ignoring the motivation of dimensionality reduction which is irrelevant for present purposes, there is a precise conceptual and formal equivalence between this model and the procedure for reconstructing model B’s embeddings given those of model A described above in Sect. 5.2.

Fig. 6.

Visual schematic of the computation of a single entry of the output of (A) the projection of input $x _n$ to output $y _n$ as in the Locally Linear Coordinates (LLC) map** procedure; (B) the reconstruction of a latent embedding $\textbf{e}^B_n$ in model B’s embedding space given input $x _n$ to model A. The grou**s in brackets in (A) illustrate the concatenations of vector embeddings (scaled by responsibility terms $r_{n_k}$) in $u_n$, and of projection columns in $l_j$. $\textbf{1}_k$ in (B) denotes a row of k 1s (where k in this case denotes the number of anchors, i.e. is set to $|\mathcal {A} |$). Each entry in the column vector ${\big [\textbf{E}^B_\mathcal {A}\big ]}^T_j$ is the jth dimension of one of model B’s anchor embeddings.

Specifically, we can regard each of model A’s anchor embeddings $\textbf{e}^A_{x_k}$ as an “expert” in a fictitious mixture model, with an associated responsibility term measuring its fidelity to the input $x_i$, which in this case is given by the cosine similarity between the anchor embedding and the input embedding. Then like the rows of U, each row of $\sigma \big [\textbf{R}^A_X\big ]$, which is a relative representation $\textbf{r}^A_i = \textbf{E}^A_\mathcal {A}{\textbf{e}^A_i}$ of input i after application of the softmax, acts as a responsibility-weighted mixture of multiple “views” of the input. Similarly, since the rows of $\textbf{E}^B_\mathcal {A}$ are anchor embeddings in the output space, its columns j act precisely as do the columns of L, i.e. as columns in a projection matrix, so that $\sigma [\textbf{r}^A_i] \cdot {\textbf{E}^B_\mathcal {A}}_j$ outputs dimension j of the reconstructed target embedding $\textbf{e}^B_i$.

There is at least one important difference between LLC and our procedure: in LLC each expert uses an internal transform to generate an input-dependent embedding, which is then scaled by its responsibility term, which also depends on the input. Reconstruction via relative representations instead employs fixed stored embeddings, so that each “expert” contributes a scalar value rather than an embedding vector to the final output. However, the expression of LLC in terms of a linear index demonstrates that this makes no essential difference mathematically (conceptually, these scalar “votes” are 1D vectors; cf. Figure 6).

The point is not that these two algorithms are doing precisely the same thing (they are not, as LLC aims to align multiple embedding spaces by deriving a map** to a distinct common space, while our approach aims to recover the contents of one embedding space from another). The use of LLC to reconstruct input data $\mathcal {X}$ from its “global” embedding $\mathcal {Y}$ as in [22] is quite closely related to our procedure, however, and at this level of abstraction the approaches may be regarded as the same, with a difference in the nature of the “experts” used in the mixture model and the attendant multiple “views” of the data. The relative representation reconstruction procedure, while presumably not as expressive, may compensate to some extent for the use of scalar “embeddings” by using a large number of “experts”, and has the virtue of eschewing the need for a mixture model to assign responsibilities, or indeed for multiple intermediate embedding models, to perform such a map**.