Stiefel Manifold Dynamical Systems for Tracking Representational Drift

This paper introduces the Stiefel Manifold Dynamical System (SMDS), a novel model that constrains emission matrices to evolve smoothly on the Stiefel manifold to effectively capture representational drift in neural data, outperforming traditional Linear Dynamical Systems in accuracy and dimensionality efficiency while providing new insights into the temporal dynamics of neural representations.

The Gist

Imagine you are trying to learn a new dance routine. You practice every day, and your body (the brain) is getting better at the moves. But here's the twist: every time you step onto the dance floor, the stage itself seems to rotate slightly, and the spotlights move to different angles.

If you were a robot programmed to learn this dance assuming the stage never moves, you would get confused. You'd think, "Wait, my left foot went left yesterday, but today it feels like it went right!" You'd try to learn a new, complicated set of rules for every single day, eventually getting overwhelmed and failing to see the core pattern of the dance.

This is exactly the problem neuroscientists face when studying the brain.

The Problem: The "Drifting" Brain

For a long time, scientists used a tool called a Linear Dynamical System (LDS). Think of LDS as a rigid, unchanging map. It assumes that if a specific group of neurons fires, it always means the same thing (like "move left").

But the brain is messy. Neurons don't stay static. Over time, even if an animal is doing the exact same task (like reaching for a treat), the way the neurons talk to each other changes. This is called Representational Drift.

• The Old Way (LDS): Imagine trying to map a city that is constantly rebuilding its streets. If you use a static map, you'll get lost. To make the map work, you'd have to draw a new map for every single day, requiring thousands of pages just to describe a simple trip.
• The Reality: The "traffic rules" (the underlying logic of the brain) stay the same, but the "street signs" (which neurons are lighting up) slowly rotate and shift.

The Solution: The "Stiefel Manifold" Dance Floor

The authors of this paper introduced a new model called SMDS (Stiefel Manifold Dynamical System).

Here is the analogy:
Imagine the brain's activity is a shadow cast by a 3D object (the true thought or movement) onto a 2D wall (the neurons we can measure).

• The Object: The actual thought or movement (e.g., "reach for the apple"). This stays stable.
• The Shadow: The pattern of neurons firing.
• The Drift: The light source is slowly moving around the object. The shadow changes shape and position, even though the object hasn't moved.

The old model (LDS) tried to describe the shadow as if the light source never moved. It failed.

The new model (SMDS) realizes the light source is moving. It assumes the object (the brain's internal logic) is stable, but the projection (how that logic shows up in neurons) is slowly rotating.

How SMDS Works (The Magic Trick)

The paper uses some fancy math (Stiefel Manifolds), but you can think of it like this:

1. The "Orthogonal" Rule: The model forces the "light source" to rotate smoothly, like a dancer spinning on a stage. It can't just jump randomly; it has to glide. This ensures the model doesn't get confused by noise.
2. Shared Rules: The model learns the dance steps (the dynamics) once and uses them for the whole session. It doesn't relearn the steps every time the stage rotates.
3. Tracking the Drift: Because it knows the rules are stable, it can measure exactly how much the stage rotated between Trial 1 and Trial 100.

What They Found

The team tested this on both computer simulations and real brain recordings from monkeys and rats.

• Better Accuracy: SMDS predicted the brain's activity much better than the old models. It needed fewer "dimensions" (fewer pages in the map) to explain the same amount of data.
• The "Important" vs. "Unimportant" Drift: This was the coolest discovery.
  • When they looked at the parts of the brain responsible for the most important things (like the actual movement of the hand), those parts drifted very little. They stayed stable.
  • The parts of the brain that were less critical for the task drifted a lot.

The Metaphor: Imagine a band playing a song. The melody (the important part) stays perfectly in tune. But the drummer might change their rhythm slightly, and the guitarist might tweak their tone. The song is still the same song, but the "texture" of the performance changes. SMDS figured out that the brain keeps the "melody" (task-relevant info) rock-solid while letting the "texture" (less important neural noise) drift around.

Why This Matters

This isn't just about math; it's about understanding how we learn and remember.

• For Brain-Computer Interfaces (BCIs): If you want to build a robotic arm controlled by your thoughts, you need a decoder that doesn't break every time your brain drifts. SMDS could help build systems that adapt to these changes automatically.
• For Understanding Memory: It suggests the brain has a clever way of keeping important information stable while letting the "background noise" shuffle around, perhaps to make room for new learning without overwriting old memories.

In short: The brain is like a jazz band. The song (the task) stays the same, but the musicians (neurons) improvise and shift their positions. The old models tried to force the musicians to stand still. The new model (SMDS) understands the music, tracks the improvisation, and realizes that the song is what really matters.

Technical Summary

1. Problem Statement

Representational Drift: In systems neuroscience, it is well-established that neural population activity often resides on low-dimensional manifolds. However, a major challenge in modeling these dynamics is representational drift: the phenomenon where the mapping between latent neural states and observed neural activity changes gradually over time (across trials or sessions), even when the underlying latent dynamics and behavior remain stable.

Limitations of Existing Models:

• Linear Dynamical Systems (LDS): The standard tool for modeling neural data assumes a static emission matrix (the mapping from latent states to observations). This assumption fails to capture non-stationarity, leading to poor model fit and an inability to recover true latent dynamics when drift is present.
• PCA and Subspace Tracking: While Principal Component Analysis (PCA) can detect overall subspace rotation, it fails to track drift per dimension because principal components can flip signs or swap ordering between trials, making dimension-specific drift unidentifiable.
• Existing Non-Stationary Models: Previous approaches often rely on discrete switching or complex condition-dependent parameters that do not explicitly model the smooth evolution of the emission subspace on a geometric manifold.

2. Methodology: Stiefel Manifold Dynamical System (SMDS)

The authors propose SMDS, a novel probabilistic state-space model that generalizes LDS by allowing the emission matrix to evolve smoothly over trials while maintaining the stability of the underlying latent dynamics.

Core Mathematical Framework

• Stiefel Manifold Constraint: The emission matrix $C^{(k)}$ for trial $k$ is constrained to lie on the Stiefel manifold $\text{Stiefel}(N, D)$, the space of all $N \times D$ orthonormal matrices ($C^\top C = I$). This ensures the mapping remains a valid rotation/orthonormal projection, addressing identifiability issues inherent in standard LDS.
• Latent Displacement Variables: Instead of learning a new matrix for every trial, SMDS models the evolution of the emission matrix via a latent "displacement" variable $z^{(k)}$.
  • $z^{(k)}$ parameterizes a skew-symmetric matrix $B^{(k)}$.
  • The emission matrix is updated via the Cayley transform: $C^{(k)} = U_{\text{base}} f_{\text{Cay}}(B^{(k)}) O_{\text{readout}}$.
  • This formulation separates rotations within the subspace (which do not change the subspace) from rotations of the subspace (drift on the Grassmann manifold).
• Dynamics:
  • Latent States ($x_t$): Evolve according to a standard linear Gaussian dynamic ($x_{t+1} = Ax_t + b + q_t$) shared across all trials.
  • Displacements ($z^{(k)}$): Evolve across trials as a random walk ($z^{(k+1)} = z^{(k)} + e^{(k)}$), where the variance of the noise term controls the drift rate. Crucially, drift rates are learned separately for each dimension of the displacement space, allowing different axes to drift at different speeds.

Inference Algorithm

Since the joint posterior over latent states and displacements is intractable, the authors developed a Variational Expectation-Maximization (VEM) algorithm:

1. E-step (Coordinate Ascent):
   • Latent States: Given current displacement estimates, update the within-trial latent states using Kalman Smoothing (exact inference for linear Gaussian systems).
   • Displacements: Given current latent state estimates, update the posterior over displacements using an Extended Kalman Smoother (EKS). The emission function is non-linear (due to the Cayley transform), requiring linearization (Jacobian) for the EKS.
2. M-step: Update model parameters ($A, Q, R$, etc.) in closed form using the sufficient statistics from the E-step.
3. Regularization: To ensure the first-order Taylor approximation in the EKS remains valid, the drift rate parameter $\tau^2_z$ is clipped after each M-step.

3. Key Contributions

1. Novel Model Class: Introduction of SMDS, the first model to explicitly constrain emission matrices to the Stiefel manifold and model their smooth evolution as a latent dynamical process.
2. Identifiability and Quantification: By enforcing orthonormality, SMDS uniquely identifies the latent subspace and allows for the quantification of drift via Grassmann distance (the geodesic distance between subspaces).
3. Dimension-Specific Drift Analysis: Unlike PCA, SMDS can track drift for individual latent dimensions, revealing that not all dimensions drift at the same rate.
4. Efficient Inference: Development of a scalable VEM algorithm combining Kalman smoothing and Extended Kalman smoothing for non-linear manifold dynamics.

4. Results

The authors validated SMDS on both synthetic and real neural datasets, comparing it against standard LDS and Conditionally Linear Dynamical Systems (CLDS).

A. Simulated Data

• Recovery of Ground Truth: In non-stationary settings, SMDS successfully recovered the true latent dynamics and the true drift pattern. Standard LDS failed to recover the true dynamics, requiring higher latent dimensions to compensate for drift and yielding lower held-out log-likelihood.
• Dimensionality: SMDS saturated at the true latent dimensionality ($D$), whereas LDS log-likelihood continued to increase with $D$, indicating overfitting to the drift.

B. Macaque Neural Data (Center-Out Reaching)

• Dataset: 96-channel M1 recordings over 750 trials (~24 mins).
• Performance: SMDS achieved higher held-out log-likelihood than LDS and CLDS with fewer latent dimensions.
• Drift Findings:
  • Gradual Drift: SMDS detected a gradual drift of $13^\circ$–$50^\circ$ over the session.
  • Stability of Task-Relevant Dimensions: Dimensions explaining high neural variance and those critical for predicting behavioral variables (cursor velocity) exhibited significantly less drift.
  • Instability of Noise Dimensions: Dimensions with low explanatory power drifted more rapidly.

C. Rodent Neural Data (Directional Licking)

• Dataset: Two-photon calcium imaging (132 ROIs) over 155 trials (~31 mins).
• Performance: Consistent with macaque data, SMDS outperformed LDS and CLDS in log-likelihood and required fewer dimensions.
• Drift Findings:
  • Drift ranged from $37^\circ$–$67^\circ$.
  • Similar to the macaque data, dimensions encoding task-relevant information (lick direction) were more stable than those encoding less relevant variance.

5. Significance and Implications

• Understanding Neural Coding: The results suggest that the brain maintains stable representations of task-relevant information despite underlying neuronal turnover or noise. The "drift" is not uniform; it preferentially affects dimensions that are less critical for the task.
• Tool for Neuroscience: SMDS provides a principled framework for quantifying representational drift on short timescales (minutes), a regime previously difficult to analyze due to a lack of interpretable computational tools.
• Beyond Stationarity: This work challenges the assumption of static neural mappings in dynamical systems, offering a more realistic model for how the brain processes information over time. It opens avenues for investigating the causes of drift (e.g., learning, plasticity, or experimental artifacts) by correlating drift rates with specific neural or behavioral features.

In summary, SMDS represents a significant advancement in neural data analysis, bridging differential geometry and dynamical systems to reveal that while neural representations drift, the core computational substrates for behavior remain remarkably stable.