\chapter{Supplementary: Learning Grid Cells by Predictive Coding}\label{app:grid}

\section{Algorithms}\label{appendix:alg}
Below is the training algorithm for a sparse, non-negative PCN given spatial inputs $\bfp$. We obtain the grid cells shown in the main text directly by taking the converged latent activities $\bfg$ after training.
\begin{algorithm}
\caption{Learning latent representations of space with a PCN}\label{alg:pc}

\begin{algorithmic}[1]
    \LComment{Training}
    \While{$\bfW$ not converged}
    \State Initialize $\bfg$ randomly;
    
    \State Input: $\bfp$
    \While{$\bfg$ not converged} 
        \State $\bfg \leftarrow \texttt{ReLU}(\bfg_t + \Delta \bfg_t$) (Eq.~\ref{eq:pc-inf})
    \EndWhile
    \State Update $\bfW$ (Eqs.~\ref{eq:pc-learn})
    \EndWhile
\end{algorithmic}
\end{algorithm}

Below is the training algorithm for tPCN in path integration tasks. The testing performance and grid cells shown in the main text are obtained by performing a forward pass through the model after training, given an unseen trajectory $\{\bfv_t, \bfp_t\}$.
\begin{algorithm}
\caption{Path integration with tPCN}\label{alg:tpc}
\begin{multicols*}{2}
\begin{algorithmic}[1]
    \LComment{Training}
    \While{$\bfW$, $\bfF$, $\bfB$ not converged}
    \State Initialize $\hat{\bfg}_{0}$ randomly or from $\bfp_0$ via a PCN;
    \For{$t=1,...,T$}
        \State Input: $\bfp_t$, $\hat{\bfg}_{t-1}$ and optionally $\bfv_t$
        \State Initialize $\bfg_t=f(\bfW\hat{\bfg}_{t-1})$ 
        \For{$k=1,...,K$} 
            \State $\bfg_t \leftarrow \bfg_t + \Delta \bfg_t$ (Eq.~\ref{eq:tpc_inf})
        \EndFor
        \State Update $\bfW$, $\bfF$, $\bfB$ (Eqs.~\ref{eq:tpc_learn})
        \State $\hat{\bfg}_t \leftarrow \bfg_t$
    \EndFor
    \EndWhile
    \\
    \LComment{Testing}
    \State Initialize $\bfg_{0}$ randomly or from $\bfp_0$ via a PCN;
    \For{$t=1,...,T$}
        \State Input: $\bfg_{t-1}$ and optionally $\bfv_t$
        \State Obtain $\bfg_t$, $\hat{\bfp}_{t}$ via Eq. \ref{eq:tpc-forward}
    \EndFor
\end{algorithmic}
\end{multicols*}
\end{algorithm}

\section{Formulation of principal component analysis (PCA) within the context of spatial learning}\label{app:pca}
Here we provide an intuition for the PCA formulation presented in Eq.~\ref{eq:PCA} as a constrained least square problem, by relating it to the more well-known ``maximum variance'' formulation. Recall that the goal of PCA is to find a lower-dimensional projection of the data ($\mathbf{P}$ in our case), such that the projected data has the maximized variance. Suppose the $i$th place cell has its activity across $N_x$ locations $\bfp(i)\in \mathbb{R}^{N_x}$ and we (tentatively) want to find a $1$-dimensional projection of it through a projection vector $\bfg(1)\in\mathbb{R}^{N_x}$. The scalar, $1$-dimensional projection is performed by $\bfg(1)^\top \bfp(i)$. For all $i=1,...,N_p$ place cells, assume the mean of their activities across $N_x$ locations is $0$ (without loss of generality), the variance of the projected dataset is then:
\begin{equation}
    \frac{1}{N_p} \sum_{i=1}^{N_p} (\bfg(1)^\top \bfp(i))^2 = \bfg(1)^\top \left(\frac{1}{N_p} \sum_{i=1}^{N_p} \bfp(i) \bfp(i)^{\top}\right)\bfg(1) = \bfg(1)^\top \bfS_p \bfg(1)
\end{equation}
where $\bfS_p$ is the $N_x\times N_x$ covariance of matrix of the place cells, which represents ``how similar are two positions, meausred as the similarity between their place cell representations or codes?''. To achieve the maximum-variance objective of PCA, we want to maximize $\bfg(1)^\top \bfS_p \bfg(1)$ with respect to $\bfg(1)$. An obvious but trivial solution is $\bfg(1)=\infty$. Therefore, to ensure our solution is sensible, we should constrain the norm $\Vert \bfg(1) \Vert_2$ to some finite value, which, without loss of generality, can simply be $1$. The optimization problem is then turned into a constrained one:
\begin{equation}
    \max_{\bfg(1)} \bfg(1)^\top \bfS_p \bfg(1) \quad \tn{s.t.} \quad \bfg(1)^\top \bfg(1) = 1
\end{equation}

We can solve this problem using a Lagrangian multiplier $\lambda_1$ and reformulate it as:
\begin{equation}
    \max_{\bfg(1)} \bfg(1)^\top \bfS_p \bfg(1) + \lambda_1 (1-\bfg(1)^\top \bfg(1))
\end{equation}

By setting the derivative of the above objective with respect to $\bfg(1)$ to $0$ we get
\begin{equation}
    \bfS \bfg(1) = \lambda_1 \bfg(1)
\end{equation}
which says that $\bfg(1)$ must be the eigenvector of $\bfS_p$. If we now multiply both sides of the equation by $\bfg(1)^\top$, we see that the variance of the projected data becomes:
\begin{equation}
    \bfg(1)^\top\bfS \bfg(1) = \lambda_1
\end{equation}
which means that if we want to obtain the maximum variance of projected data $\bfg(1)^\top\bfS \bfg(1)$, $\lambda_1$ has to be the largest eigenvalue. In other words, $\bfg(1)$ has to be the first eigenvector of $\bfS_p$.

To relate this view of PCA to the least square formulation in Eq.~\ref{eq:PCA}, we first rewrite the Frobenius norm to the sum of vector $l2$ norms:
\begin{equation}
    \mathcal{F}_{PCA} = \sum_{i=1}^{N_p} \Vert \bfp(i) - \mathbf{G} \mathbf{m}(i)\Vert_2^2\quad\tn{s.t.}\quad \mathbf{G}^\top \mathbf{G} = \bfI
\end{equation}
where $\mathbf{G}$ is an $N_x \times N_g$ matrix, whose columns are $\bfg(1),...,\bfg(N_g)$ and $\mathbf{m}(i)$ is a $N_g$ dimensional column vector. Taking derivative of $\mathcal{F}_{PCA}$ with respect to $\mathbf{m}(i)$ an set it to $0$ we get:
\begin{equation}
    \mathbf{G}^\top \bfp(i) - \mathbf{G}^\top\mathbf{G} \mathbf{m}(i) = 0 \Rightarrow \mathbf{m}(i) = \mathbf{G}^\top \bfp(i)
\end{equation}
with the constraint that $\mathbf{G}^\top \mathbf{G} = \bfI$. We can thus rewrite the objective as:
\begin{equation}
\begin{split}
    \mathcal{F}_{PCA} &= \sum_{i=1}^{N_p} \Vert \bfp(i) - \mathbf{G}\mathbf{G}^\top \bfp(i)\Vert_2^2
    \\&=\sum_{i=1}^{N_p} \bfp(i)^\top \bfp(i) - 2\bfp(i)^\top \mathbf{G}\mathbf{G}^\top \bfp(i) + \bfp(i)^\top \mathbf{G}\mathbf{G}^\top\mathbf{G}\mathbf{G}^\top \bfp(i)
    \\&=\sum_{i=1}^{N_p} \bfp(i)^\top \bfp(i) - \bfp(i)^\top \mathbf{G}\mathbf{G}^\top \bfp(i)
\end{split}
\end{equation}

If we minimize the above expression with respect to $\bfG$, it  becomes equivalent to maximize:
\begin{equation}
    \begin{split}
        \sum_{i=1}^{N_p} \bfp(i)^\top \mathbf{G}\mathbf{G}^\top \bfp(i) &= tr(\bfP^\top \mathbf{G}\mathbf{G}^\top \bfP) \\& = tr(\bfG^\top \bfP\bfP^\top \bfG) \\&= tr(\bfG^\top \bfS_p \bfG) \\&= \sum_{j=1}^{N_g}\bfg(j)^\top \bfS_p \bfg(j)
    \end{split}
\end{equation}
where $tr$ denotes the trace of a matrix. To maximize for the case $N_g=1$, we have shown above that $\bfg(1)$ has to be the first eigenvector of $\bfS_g$. Since $\bfS_g$ is a positive semi-definite matrix, we know that for any $j$,  $\bfg(j)^\top \bfS_p \bfg(j) \geq 0$, and thus to maximize the sum  $\sum_{j=1}^{N_g}\bfg(j)^\top \bfS_p \bfg(j)$, $\bfg(j),j=1,...,N_g$ has to be the first $N_g$ eigenvectors of $\bfS_g$ so that the sum is the sum of the biggest $N_g$ eigenvalues. We have now formally linked the least-square formulation of PCA to the maximum variance interpretation. Bringing it back to place/grid cells, we have established that grid cells are, in essence, the top-$N_g$ eigenvectors of the spatial covariance matrix $\bfS_p$ \citep{dordek2016extracting,sorscher2023unified}.

\section{Derivations of learning dynamics}\label{appendix:bptt}

Here we derive the recurrent weight update rules for $\bfW^{\tn{RNN}}$ (Equation \ref{eq:bptt-learning-rule}) and $\bfW^{\tn{tPCN}}$ (Equation \ref{eq:tpc-learning-rule}). For RNN, we assume that the weights are updated at each time step and therefore $\bfW^{\tn{RNN}}$ is updated following the chain rule:
\begin{equation}
    \Delta \bfW^{\tn{RNN}} = -\frac{d \mathcal{F}_{\tn{RNN}, t}}{d \bfW} = -\frac{d \mathcal{F}_{\tn{RNN}, t}}{d \bfg_t} \frac{d \bfg_t}{d \bfW}
\end{equation}

We first look at the term $\frac{d \bfg_t}{d \bfW}$, which, following the rule of partial derivatives, can be written as:
\begin{equation}
\begin{split}
    \frac{d \bfg_t}{d \bfW} &= \frac{\partial \bfg_t}{\partial \bfW} + \frac{\partial \bfg_t}{\partial \bfg_{t-1}}\frac{d \bfg_{t-1}}{d \bfW} \\ &= \frac{\partial \bfg_t}{\partial \bfW} + \frac{\partial \bfg_t}{\partial \bfg_{t-1}}\left(\frac{\partial \bfg_{t-1}}{\partial \bfW} + \frac{\partial \bfg_{t-1}}{\partial \bfg_{t-2}}\frac{d \bfg_{t-2}}{d \bfW} \right) \\ &= ... 
    \\ &= \sum_{k=1}^t \frac{\partial \bfg_t}{\partial \bfg_k} \frac{\partial \bfg_k}{\partial \bfW}
\end{split}
\end{equation}
due to the recursive and implicit dependency of $\bfg_t$ on $\bfg_{t-1}$ and $\bfg_{t-1}$ on $\bfW^{\tn{RNN}}$ for all $t$. Thus, the update rule can be written as:
\begin{equation}
    \Delta \bfW^{\tn{RNN}} = -\sum_{k=1}^t \frac{d \mathcal{F}_{\tn{RNN}, t}}{d \bfg_t} \frac{\partial \bfg_t}{\partial \bfg_k} \frac{\partial \bfg_k}{\partial \bfW}
\end{equation}

Since $\bfg_k = h(\tilde{\bfg}_k) = h(\bfW\bfg_{k-1} + \bfB \bfv_k)$, and $\mathcal{F}_{\tn{RNN}, t} = \Vert \bfp_t - f(\bfF\bfg_t)\Vert_2^2$ the update rule can be written as:
\begin{equation}
    \Delta \bfW^{\tn{RNN}} = \sum_{k=1}^t \tfrac{\partial \bfg_t}{\partial \bfg_k} h'(\tilde{\bfg}_t) \bfF^{\top} f'(\bfF\bfg_t) \bm{\epsilon}_t^{\bfp} \bfg_{k-1}^{\top},
\end{equation}
concluding our proof for Equation \ref{eq:bptt-learning-rule}. The derivation for $\bfB^{\tn{RNN}}$ is similar:
\begin{equation}
    \Delta \bfB^{\tn{RNN}} = -\frac{d \mathcal{F}_{\tn{RNN}, t}}{d \bfB} = -\frac{d \mathcal{F}_{\tn{RNN}, t}}{d \bfg_t} \frac{d \bfg_t}{d \bfB},
\end{equation}
and
\begin{equation}
\begin{split}
    \frac{d \bfg_t}{d \bfB} &= \frac{\partial \bfg_t}{\partial \bfB} + \frac{\partial \bfg_t}{\partial \bfg_{t-1}}\frac{d \bfg_{t-1}}{d \bfB} \\ &= \frac{\partial \bfg_t}{\partial \bfB} + \frac{\partial \bfg_t}{\partial \bfg_{t-1}}\left(\frac{\partial \bfg_{t-1}}{\partial \bfB} + \frac{\partial \bfg_{t-1}}{\partial \bfg_{t-2}}\frac{d \bfg_{t-2}}{d \bfB} \right) \\ &= ... 
    \\ &= \sum_{k=1}^t \frac{\partial \bfg_t}{\partial \bfg_k} \frac{\partial \bfg_k}{\partial \bfB}
\end{split}
\end{equation}

Therefore, the update rule for $\bfB$ in an RNN can be written as:
\begin{equation}
    \Delta \bfB^{\tn{RNN}} = \sum_{k=1}^t \tfrac{\partial \bfg_t}{\partial \bfg_k} h'(\tilde{\bfg}_t) \bfF^{\top} f'(\bfF\bfg_t) \bm{\epsilon}_t^{\bfp} \bfv_{k}^{\top}
\end{equation}

Finally, the update rule for $\bfF^{\tn{RNN}}$ can be straightforwardly expressed as:
\begin{equation}
    \begin{split}
        \Delta \bfF^{\tn{RNN}} &= -\frac{d \mathcal{F}_{\tn{RNN}, t}}{d\bfF} \\ &= -\frac{d \mathcal{F}_{\tn{RNN}, t}}{d \hat{\bfp}_t} \frac{d \hat{\bfp}_t}{d \bfF} \\&= f'(\bfF \bfg_t)\bm{\epsilon}_t^{\bfp}\bfg_t^{\top}
    \end{split}
\end{equation}
as there is no recursive dependency.

For tPCN, at each time step $t$ the following loss is minimized with respect to $\bfW$:
\begin{equation}
    \mathcal{F}_{\textnormal{tPCN}, t} = \Vert \bfp_t - f(\bfF \bfg_t) \Vert_2^2 + \Vert \bfg_t - h(\bfW \hat{\bfg}_{t-1} + \bfB \bfv_t) \Vert_2^2
\end{equation}

Since $\hat{\bfg}_{t-1}$ is inferred through Equation \ref{eq:tpc_inf}, rather than forward-propagated by $\bfW$, the recursive dependency on $\bfW$ disappears, and thus the update rule for $\bfW$ can be locally derived as:
\begin{equation}
    \Delta \bfW^{\tn{tPCN}} = -\frac{d \mathcal{F}_{\textnormal{tPCN}, t}}{d \bfW} = h'(\tilde{\bfg}_t)\bm{\epsilon}^{\bfg}_t \hat{\bfg}_{t-1}^{\top}
\end{equation}

If we also assume that the inference dynamics in Eq.~\ref{eq:tpc_inf} have converged when the weights are updated, namely:
\begin{equation}
    \Delta \bfg_t = 0 \Rightarrow \bm{\epsilon}^{\bfg}_t = \bfF^{\top} f'(\bfF \bfg_t)  \bm{\epsilon}^{\bfp}_t,
\end{equation}
the update rule can be written as:
\begin{equation}
    \Delta \bfW^{\tn{tPCN}} = h'(\tilde{\bfg}_t) \bfF^{\top} f'(\bfF \bfg_t)  \bm{\epsilon}^{\bfp}_t \hat{\bfg}_{t-1}^{\top},
\end{equation}
which concludes our proof for Equation \ref{eq:tpc-learning-rule}. Similarly, following the same assumption of converged inference and Eq.~\ref{eq:tpc_learn}, the update rule $\Delta\bfB^{\tn{tPCN}}$ can be written as:
\begin{equation}
    \Delta\bfB^{\tn{tPCN}} = h'(\tilde{\bfg}_t) \bfF^{\top} f'(\bfF \bfg_t)  \bm{\epsilon}^{\bfp}_t \bfv_t^{\top}
\end{equation}

It can be seen that it differs from $\Delta\bfB^{\tn{RNN}}$ only in the absence of the unrolling term $\tfrac{\partial \bfg_t}{\partial \bfg_k}$. On the other hand, the update rules $\Delta\bfF^{\tn{RNN}}$ and $\Delta\bfF^{\tn{tPCN}}$ are exactly the same.

\section{Experimental setups and hyperparameters}\label{appendix:exp}
\paragraph{Place cell and trajectory parameters}

We use DoS place cell encodings throughout most of our experiments. Formally, the activity of the $i$th place cell with this encoding, given a particular location $x$ can be written as:
\begin{equation}
    K(x, C, \tau) := \exp{\left(-\frac{(x-C)^2}{\tau\xi^2}\right)}
\end{equation}
\begin{equation}
    p_i = \frac{K(x,C_i,2)}{\sum_{j=1}^{N_p}K(x,C_j,2)} - \frac{K(x,C_i,4)}{\sum_{j=1}^{N_p}K(x,C_j,4)}
\end{equation}
where $C_i$ is the center of the place cell's firing field, and $\tau$ and $\xi$ define the width of the firing field's center and surround. The table below specifies parameters defining the place cells and trajectories:

\begin{table}[!h]
\begin{center}
\begin{tabular}{| c | c | c | c |}
\hline
 $\xi$ & {\bf Path length} & {\bf Average agent speed} & {\bf Environment size}\\ \hline
 0.12m & 10 steps & \{0.02, 0.05, 0.1\}m/s & \{$1.4^2, 1.8^2, 2.0^2$\}m$^2$ \\ \hline
\end{tabular}
% \caption{Hyperparameters when training tPC models}
\end{center}
\end{table}
Specifically, at time step $t=0$, a $2$D position and a head direction scalar in $[0, 2\pi]$ are randomly initialized. At each of the subsequent time steps, a random turn angle is sampled from a normal distribution and a random speed is sampled from a Rayleigh distribution. Both values are then multiplied by $dt$ mentioned in the main text. If the simulated agent hits a border wall at this time step, its speed is slowed and its turn angle is inverted. The position of the agent is updated according to the speed and turn angle at this time step. The trajectories are simulated using parameters adapted from the code provided in \citet{sorscher2023unified}.

\paragraph{Model and training hyperparameters} 
In our experiments, we have used three models: sparsity and non-negativity constrained PCN, RNN and tPCN. The table below specifies parameters of model architectures:

\begin{table}[!h]
\begin{center}
\small
\begin{tabular}{| c | c | c | c |}
\hline
{\bf Hyperparameter} & {\bf sparse, non-neg. PCN} & {\bf tPCN} & {\bf RNN} \\ \hline
$N_p$ & 512 & 512 & 512 \\ \hline
$N_g$ & 256 & \{256, 512, 1024, 2048\} & 2048 \\ \hline
$h$ & N/A & \{$\texttt{ReLU}, \texttt{tanh}$\} & $\texttt{ReLU}$ \\ \hline
$f$ & N/A & \{$\texttt{softmax}, \texttt{tanh}$\} & $\texttt{softmax}$ \\ \hline
\end{tabular}
% \caption{Transposed Hyperparameters for tPC and RNN models}
\end{center}
\end{table}

The table below specifies hyperparameters used in training RNN and tPCN. We use $\texttt{Adam}$ optimizer for all weight updates, and plain $\texttt{SGD}$ for inference dynamics in tPCN. We found that in general, RNNs take more epochs to converge in the path integration task.
\begin{table}[!h]
\begin{center}
% \scriptsize
\begin{tabular}{| c | c | c |}
\hline
{\bf Hyperparameter} & {\bf tPCN} & {\bf RNN} \\ \hline
$N_x$ & 50000 & 50000 \\ \hline
{\bf batch size} & 500 & 500 \\ \hline
{\bf learning rate} & $10^{-4}$ & $10^{-4}$ \\ \hline
{\bf inference step size} & $10^{-2}$ & N/A \\ \hline
{\bf epochs} & 150 & 200 \\ \hline
{\bf inference iters} & 20 & N/A \\ \hline
{\bf weight decay} & $10^{-4}$ & $10^{-4}$ \\ \hline
\end{tabular}
% \caption{Transposed Hyperparameters for tPC and RNN models}
\end{center}
\end{table}

The table below specifies hyperparameters used in training the sparse, non-negative PCN. We use $\texttt{Adam}$ optimizer for all weight updates, and plain $\texttt{SGD}$ for inference dynamics.
\begin{table}[!h]
\begin{center}
\begin{tabular}{| c | c |}
\hline
{\bf Hyperparameter} & {\bf Value} \\ \hline
$N_x$ & 900 \\ \hline
{\bf batch size} & 100 \\ \hline
{\bf learning rate} & $2\times 10^{-3}$ \\ \hline
{\bf inference step size} & $10^{-2}$ \\ \hline
{\bf epochs} & 600 \\ \hline
{\bf inference iters} & 20 \\ \hline
{\bf weight decay} & $10^{-5}$ \\ \hline
$\lambda$ & 0.05 \\ \hline
\end{tabular}
% \caption{Transposed hyperparameters for tPC model training}
\end{center}
\end{table}

\paragraph{Calculation of grid scores} The following grid score calculation process is adapted from \citet{sargolini2006conjunctive} and the code of \citet{sorscher2023unified}. It is summarized below for completeness and clarity:

\begin{itemize}[leftmargin=*]
    \item Get the rate map of latent neurons (potentially hexagonal grid cells);
    \item Place one copy of the rate map on top of the other, and start moving the top copy by $\delta\in\mathbb{R}^2$. If the rate maps are hexagonal grids, for particular $\delta$'s that make the firing peaks overlap, the autocorrelation between the stationary and moved maps will be 1; otherwise, the autocorrelation will be 0. We will then have a hexagonal autocorrelation map if the rate map itself is hexagonal;
    \item We then rotate the autocorrelation map and compute the correlation between each rotated map and the original map. If the rate maps are hexagonal, the correlation as a function of rotated degrees will be sinusoidal, with 60 and 120 degrees as peaks and 30, 90 and 150 degrees as troughs.
    \item Grid score is calculated as the minimum difference between the peak and trough correlation, which in theory is a real value in range $[-2,2]$.
\end{itemize}

All experiments were performed on a single Tesla V100 GPU.