# Comparing high-dimensional neural recordings by aligning their low-
dimensional latent representations

Max Dabagia111School of Computer Science at the Georgia Institute of
Technology., Konrad P. Kording222Depts of Bioengineering and Neuroscience at
the University of Pennsylvania., Eva L. Dyer 333Department of Biomedical
Engineering at the Georgia Institute of Technology<EMAIL_ADDRESS>
###### Abstract

Many questions in neuroscience involve understanding of the responses of large
populations of neurons. However, when dealing with large-scale neural
activity, interpretation becomes difficult, and comparisons between two
animals, or across different time points becomes challenging. One major
challenge that we face in modern neuroscience is that of correspondence e.g.
we do not record the exact same neurons at the exact same times. Without some
way to link two or more datasets, comparing different collections of neural
activity patterns becomes impossible. Here, we describe approaches for
leveraging shared latent structure across neural recordings to tackle this
correspondence challenge. We review algorithms that map two datasets into a
shared space where they can be directly compared, and argue that alignment is
key for comparing high-dimensional neural activities across times, subsets of
neurons, and individuals.

##  1 Introduction

The study of brain function and cognition relies on the measurement and
interpretation of changes in the brain’s activity. Whether between states of
disease and health (Frere & Slutsky,, 2018), across novice- and expert-level
performance on a complex task (Sadtler et al.,, 2014), or in wakefulness
versus sleep (Hengen et al.,, 2016), we need ways to compare neural activity
patterns and draw conclusions about how two neural recordings differ.

Historically, changes in neural activity patterns have been examined at the
level of single cells, by characterizing how the firing rate (or tuning
curves) of each cell shifts across conditions (Hubel & Wiesel,, 1959; Saxena &
Cunningham,, 2019) . However, we know that the brain consists of highly
interconnected networks of neurons (Bullmore & Sporns,, 2009), and thus the
impact of different disease states or conditions is likely to have complex
effects across an entire circuit. Without modeling these higher-order
dependencies between neurons, it may be impossible to model the collective
impact of a given perturbation or transformation of neural state (Golub et
al.,, 2018; Kaufman et al.,, 2014; Russo et al.,, 2020).

At first glance, it may be tempting to attempt to compare two recordings of
neural activity by pairing off similar neurons across recordings and then
comparing the changes in the activity across these pairs. However, in many (if
not most) cases we may be technically unable to do so (as we can not generally
choose which cells we record), or it may be fundamentally impossible to do so
(the brain changes over time, and different individuals might not share
neurons that are identical). The alternative is to summarize each recording in
a way which is _invariant to the specific subset of neurons that are polled._
Many existing approaches build invariance through the computation of metrics
that capture essential properties of the pairwise interactions between neurons
(like clustering or density of connections) (Soudry et al.,, 2015;
Nonnenmacher et al.,, 2017; Brinkman et al.,, 2018), or extract a global
property of the network such as criticality (Chialvo,, 2010; Ma et al.,,
2019). While these approaches can be informative when we know the right
metrics to compare across conditions, a data-driven strategy to identify
couplings across multiple units would provide an even richer multi-dimensional
picture of changes in neural circuits.

One potential approach to tackle this correspondence challenge is to first
find a representation of neural activity which is more abstract than the
firing of individual units, and then compare different neural datasets through
their representations (Figure 1). We can construct suitable representations by
approximating the activity of each neuron as driven by a combination of a
smaller number of _latent factors_ – signals which describe coordinated
activity across the population (Cunningham & Yu,, 2014). Once each dataset has
been fitted with a latent model, we hope that when the circuits are performing
similar tasks, their latent models will agree with each other, allowing them
to be compared. A substantial body of evidence shows that there are many
settings in which latent factors are stable (or can be tracked) across subsets
of neurons, across time, and even individuals (Churchland et al.,, 2012; Mante
et al.,, 2013; Kaufman et al.,, 2014; Dyer et al.,, 2017; Gallego et al.,,
2018; Pandarinath et al.,, 2018a; Williamson et al.,, 2019; Gallego et al.,,
2020; Degenhart et al.,, 2020), suggesting that these models may be a powerful
tool for comparison.

The roadmap for the rest of this article is as follows: In Section 2, we
discuss settings where latent space alignment is needed to effectively compare
multi-dimensional neural activities across times, subsets of neurons, and
individuals. After introducing the idea behind latent space alignment, in
Section 3, we describe and compare different approaches for aligning neural
datasets. In Section 4, we discuss the challenges that must be overcome to
compare recordings at larger scales and across more severe perturbations.
Finally, in Section 5, we end with prospective analyses and discoveries that
these approaches will make possible.

##  2 Why we need alignment

Figure 1: Comparing neural activity patterns generated across two conditions
through the alignment of their latent spaces. In this example, we show the
generation of two latent spaces (bottom) where datasets A and B share common
structure but also have unique differences that can be detected.

### Forming a latent state space model

To build a population-level view of neural activity, we start by representing
the firing rates of all measured neurons as a point in a multi-dimensional
state space, where each axis represents the firing rate of one of the recorded
neurons (Figure 2a) (Paninski et al.,, 2010). Viewed this way, the dimension
of this space is simply the number of neurons sampled in the recording. If
every neuron operated completely independently this would be the end of the
story, and moreover neuroscience would be in trouble, because understanding a
neural circuit would require the characterization of every single neuron
involved in it. Fortunately, there is a substantial body of evidence that this
is not the case (Mazor & Laurent,, 2005; Mante et al.,, 2013; Kaufman et al.,,
2014). Intuitively, neurons interact with each other through both excitation
and inhibition, which naturally leads to correlation in their activity (Zohary
et al.,, 1994; Abbott & Dayan,, 1999). When many neurons are correlated with
each other, it is reasonable to expect that their activities might be
explainable by a smaller number of latent variables.

The simplest methods for building such a model identify each latent variable
with a particular pattern of neural activity, such that that any observation
can be approximated as a weighted sum of these factors (Tsodyks et al.,, 1999;
Sadtler et al.,, 2014; Luczak et al.,, 2015). Principal Components Analysis
(PCA) is probably the most widely-used technique (Kaufman et al.,, 2014;
Harvey et al.,, 2012; Cowley et al.,, 2016), which chooses latent variables
which retain the most variance of the observations. There are many extensions
of this approach that either relax PCA’s assumptions (e.g. that the latent
variables must be orthogonal) (Child,, 1990), or incorporate additional
structure into the solution (e.g., adding a non-negativity constraint) (Lee &
Seung,, 1999; Hyvärinen,, 2013) to better model neural data. More
sophisticated manifold learning strategies can be used to identify low-
dimensional structures that may be nonlinearly embedded in neural activity,
and “flatten” them to recover the latent factors (Stopfer et al.,, 2003;
Ganmor et al.,, 2015).

The methods mentioned above assume that time points are independent and thus
don’t necessarily make use of temporal dynamics. Rather than modeling each
point independently, we can instead model neural activity as being governed by
a dynamical system to find a low-dimensional representation which preserves
coherent dynamics (Gao et al.,, 2017; Kobak et al.,, 2019; Churchland et al.,,
2012; Brunton et al.,, 2016). A state-of-the-art approach for learning
nonlinear latent dynamics, called LFADS (Pandarinath et al.,, 2018a), provides
a deep learning solution to this problem that builds a realistic generative
model of neural spiking using a sequential autoencoder. This approach can then
be used to identify stable low-dimensional structure in neural activity across
multiple individuals performing the same motor task. Incorporating temporal
information can yield more stable latent factors, and crucially more
interpretable ones.

### Why can’t latent spaces be compared directly?

The critical assumption for using a latent factor model to compare two neural
recordings is that their latent spaces provide an adequate representation of
what is driving the circuit’s activity. For the purposes of comparing two sets
of neurons, we assume that if they encode similar information there should
exists views of their latent spaces in which their activity looks similar. The
objective is to find these views and put both datasets into a common reference
frame.

Unfortunately, even when they exist, finding views of two neural recordings
that situate similar brain states in nearby parts of the latent space is
typically not possible without further intervention or assumptions. This is
due to the fact that each dimensionality reduction technique relies on
assumptions about how information is encoded to differentiate signal from
noise and infer latent factors. Without a complete understanding of the
brain’s encoding strategies, these assumptions are at best approximations. For
example, PCA associates latent factors with the patterns of activity which
account for as much of the population’s variance as possible. When the
patterns of interest are roughly equally important, relatively insignificant
changes in the activity of a few neurons can reorder these patterns, leading
to latent models which encode the same information but must be transformed to
match each other (Figure 2). In more sophisticated methods, the ordering of
factors is typically not unique, and comes down to particularities of the
algorithm. Moreover, common factors may be obscured by additional neural
activity, which is unrelated to the information shared across the recordings.

Adding to these challenges, in long-term recordings, instabilities arise from
the difficulty of tracking particular neurons (Grill et al.,, 2009; McCreery
et al.,, 2010; Fu et al.,, 2016) or matching them across time after sorting,
and this can also introduce changes in which factors are dominant. In
comparisons across individuals, or in acute recordings, direct comparisons
across neurons is impossible. All of these challenges make direct comparisons
across latent factors difficult and highlights the need for alignment: we must
ensure that similar brain states in the two datasets are placed in the same
part of the space.

Figure 2: Perturbations to neural activity can radically alter the latent
space.  In (a), we show an example of unperturbed neural activities of 5
neurons, with the firing rates of two neurons displayed on the left and after
the population is transformed into 2D with principal components analysis
(PCA). In (b-c), we show examples of two different perturbations that are
often observed in neural recordings. In (b), the baseline firing rate of the
fifth neuron is increased by 1 spike per time interval (highlighted in pink).
In (c), the activity of the second neuron is replaced with random noise,
simulating its death or loss of contact with the electrode. In all cases the
latent space displays the same structure, although the perturbed datasets have
experienced an information-preserving transformation which must be reversed
via alignment. To the right in (b) and (c), we show each perturbed latent
space after alignment with Hierarchical Wasserstein Alignment (HiWA) (Lee et
al.,, 2019) to the match the distribution obtained in the first baseline
dataset.

##  3 Finding correspondence across distinct neural recordings

In this section, we discuss different approaches for alignment. These
approaches rely on different assumptions about the shared structure of the two
datasets, or partial information about correspondences, either across neurons
or across time.

### Approach 1. Using the global distribution of latent states for alignment

If we assume that the latent spaces are comparable but they may have a
different ordering of their latent factors, then this introduces the need for
building in rotational invariance, or finding the best rotation of the latent
factors that aligns the source and target data maximally (Figure 2b-c, right).
One way to pose this problem is to find a mapping between the neural
representations which _aligns_ them – that is, makes them maximally similar
under some measure of similarity. The space of mappings to consider should
account for the invariances we require, and may be tailored to the individual
problem.

Distribution Alignment Decoding (DAD) (Dyer et al.,, 2017), uses density
estimation to infer the distribution of neural activity in the latent space,
and then, searches over the space of rotations to identify the transformation
that best matches the two distributions based upon their Kullback-Liebler (KL)
divergence. This idea has been applied to achieve nearly unsupervised neural
decoding of movement, by aligning the low-dimensional projections of neural
activity to a distribution of movements without known correspondences between
the two. Hierarchical Wasserstein Alignment (HiWA) (Lee et al.,, 2019)
improves upon this strategy by leveraging the tendency of neural circuits to
constrain their low-dimensional activity to clusters or multiple low-
dimensional subspaces. Combined with with a more geometric formulation of
distance drawn from optimal transport (Villani,, 2008), HiWA is able to more
quickly and robustly recover the correct rotation to align latent spaces
across neural recordings.

A deep learning-driven solution to the alignment problem (Farshchian et al.,,
2018) draws inspiration from generative adversarial networks (GANs)
(Goodfellow et al.,, 2020), which have received some attention in the
computational neuroscience community both for generative tasks like producing
realistic neural activity (Arakaki et al.,, 2017; Molano-Mazon et al.,, 2018),
or to aid in downstream tasks like decoding visual stimuli (St-Yves &
Naselaris,, 2018). Here, the idea is to start with a latent model for one
dataset from which one can reconstruct neural activity effectively, and then
transform new datasets to match by minimizing their reconstruction loss using
the same latent model. This approach is promising for brain-machine interfaces
(BMIs), where the priority is correctly decoding target variables and less
about interpreting the differences in the learned embeddings. On the other
hand, when using latent representations to study a neural circuit, it might be
desirable for distinctive new features of the shifted distribution to be
preserved, but a sufficiently powerful alignment network which has “overfit”
to the target dataset will attenuate them in favor of matching the original
dataset as closely as possible.

### Approach 2. Using partial neuron correspondences for alignment

This same idea of rotation-invariant manifold alignment has recently been
demonstrated in an online BMI setting (Degenhart et al.,, 2020). This method
assumes a sufficient number of neurons are captured across subsequent
recordings, so that only a few coordinate axes of the neural manifold need to
be adjusted. The stable units are estimated to be those which have
approximately the same relationship with the latent factors, while the
perturbed units are remapped using a constrained transformation. Although less
general than distribution alignment because it assumes the underlying
transformation to be a permutation of a few of the coordinate axes and that
the identities of some neurons are known, this assumption also decreases the
computational requirements, making this solution more practical in a realistic
online BMI setting. When used to incrementally update the decoder for intra-
cortical multi-electrode arrays (Chestek et al.,, 2011; Vaidya et al.,, 2014)
(and their successors), this strategy is likely to be a crucial building block
for stable long-term BMIs, especially when used in conjunction with covariance
inference techniques used in calcium imaging (Turaga et al.,, 2013;
Nonnenmacher et al.,, 2017) to stitch together recordings across large
populations to fill in manifolds.

### Approach 3. Using temporal correspondence for alignment

The methods described in the last section assume that we have common global
distributional or subspace structure between datasets to align them. However,
there are many cases where this assumption may be insufficient or incorrect.
Instead, when we have some amount of temporal correspondence between datasets,
either through simultaneous recordings (Sponberg et al.,, 2015), or by using
the temporal structure which arises from multiple trials of the same task
(Gallego et al.,, 2020), then we can use information to embed multiple
datasets jointly (Figure 3). A straightforward approach for doing so is
canonical correlation analysis (CCA) (Hotelling,, 1936; Fujiwara et al.,,
2013; Dmochowski et al.,, 2018; Gallego et al.,, 2020), and more sophisticated
nonlinear extensions (Lai & Fyfe,, 2000; Huang et al.,, 2006; Andrew et al.,,
2013). This technique is a generalization of PCA, where rather than choosing a
single factor (a pattern of neural activity) at a time which maximizes
variance, it chooses a factor for _each dataset_ so that the covariance
between these components is maximized. CCA reduces to PCA when the datasets
being compared are the same. The problem is that CCA requires exact
correspondence between moments in time across datasets, which, unless they
were made simultaneously, is difficult to satisfy. Nonetheless, due to its
simplicity and efficiency, CCA has been widely used in neuroscience, even when
this key assumption is only loosely satisfied (such as when each recording is
a trial of a precisely timed task) (Gallego et al.,, 2018, 2020).

### Approach 4. Using common dynamics for alignment

Temporal information can be used more generally by finding a latent space in
which the dynamics of each dataset are similar (Pandarinath et al.,, 2018b).
For instance, in LFADS (Pandarinath et al.,, 2018a), the latent
representations of dynamics from many recording sessions can be matched to
each other by learning a shared latent space for all sessions jointly in which
temporal dynamics are consistent. This method develops the mappings between
each recording session (dataset) and the latent space by first finding the
trial-averaged responses for each unit and then learning a weight matrix
(input map) to match these trial-averaged responses. This means that given
training data consisting of the direction (or target) of the reach, all of the
trajectories in the same direction can be utilized to improve the quality of
the aligned latent representations, which alleviates the need to maintain
contact with the same neurons across recordings. In principle, however, this
approach could also be fully unsupervised, or use partial information about
clustering of behaviors (like HiWA) to find alignment of dynamics. By
leveraging shared task structure across datasets, methods like LFADS show that
it is possible to stitch together dynamics from different populations of
neurons and decode intended movement variables.

Figure 3: Alignment Paradigms. In both (a-b), we compare distribution
alignment with strategies that leverage temporal structure and dynamics to
align data into a common latent space. In (a), we show an example where in the
baseline condition (blue), neural activity traverses across three subspaces in
the latent space (T1, T2, T3); in the second condition (pink), the individual
makes a very different traversal through the latent space, jumping between
different subspaces and not adhering to the same dynamics as in the first
case. However, the overall global structure of the distribution is similar in
both cases and thus manifold alignment techniques can be applied successfully.
In (b), we show an example where the temporal structure is preserved between
both datasets by having a consistent sequence of transitions between activity
modes. In this case, we do not need to assume the same overall global
structure, but that dynamics are consistent.

### Comparison of various approaches

Given how challenging it can be to read out the same information when it has
been encoded in multiple unknown complex dynamical systems, each of the
methods discussed previously relies on a different set of assumptions to make
the problem tractable. Distribution alignment exploits similarities between
the data embedded in their latent spaces to bring them into agreement; this
typically requires that there are sufficient similarities in the behavior over
time and that a sufficient number of samples are available to decipher this
structure (Figure 3a). This places a hard requirement on the latent structure;
if the similarities are absent, too few, or too imprecise, the alignment
problem becomes ill-posed. However, when this latent structure exists, these
methods can align recordings separated by large amounts of time or even from
different individuals. In contrast to this more global approach, approaches
that leverage dynamics in the progression of neural states (Figure 3b), may
take a more local view of modeling neural activity patterns which places
constraints on how two datasets must evolve over time. When additional
information is available which allows datasets to be matched to each other,
such as temporal or partial neuron correspondence, alignment techniques can
leverage this information to succeed in cases too ambiguous for manifold
alignment.

##  4 Challenges

An essential tool for comparing neural datasets is extracting latent
representations that preserve key structures within them. However, using
dimensionality reduction approaches to study neural population dynamics can be
daunting because they are typically unsupervised and the latent factors
learned by these methods can be hard to interpret. Even with supervised labels
or training data (like movement kinematics, sleep or behavioral state), it can
be hard to know when a latent space model is sufficient to describe the
structure of the data.

Furthermore, it is unclear when latent variables should be considered
intrinsic to the functionality of a neural population, or when they are just a
concise description of activity. This is an instance of the classic
statistical conundrum of differentiating causation versus correlation. Even
between neurons with correlated activity, it is practically impossible to
determine whether the correlation arises because one neuron drives the
activity of another, or if there is another _hidden_ neuron driving both of
them (Mehler & Kording,, 2018). For example, it is not understood whether the
ability to decode movement patterns from activity in motor cortex is because
neurons have been found which are designed to encode movement, or because
their activity is highly correlated with movement as a side-effect of their
true functional roles. Most dimensionality reduction techniques are agnostic
to this distinction, yet hopefully we can distill neural activity down to a
space that reveals its true underlying causes (Lillicrap & Kording,, 2019).
Moving forward, careful design of experiments that can modify neural response
properties, either through either direct manipulation (Gradinaru et al.,,
2010; Roth,, 2016) or new learning paradigms (Shenoy & Kao,, 2021), are needed
to better interpret neural representations, and understand how circuit and
systems-level mechanisms modulate or change the distribution of the neural
state space.

Most of the alignment methods covered in the previous section assume that
there is observable low-dimensional structure that persists across multiple
recordings. However, there are many settings where changes in underlying brain
states may occur due to sources of modulation like attention and engagement
(Mitchell et al.,, 2009; Cohen & Maunsell,, 2009; Zhang et al.,, 2011; McAdams
& Maunsell,, 1999), or shift as units die or electrodes move (Polikov et al.,,
2005; Grill et al.,, 2009; McCreery et al.,, 2010). In more challenging
settings where common structure may not be easily ascertained, some supervised
data or labels could be provided to the methods to help ground alignment when
common structure between latent spaces is not sufficiently similar to find a
map between them. More transient modulations are harder to overcome, but
methods developed for dealing with them for supervised decoding might provide
the starting point for a way forward (Fan et al.,, 2014; Sederberg et al.,,
2019; Whiteway et al.,, 2019). Additionally, methods for partial domain
adaptation (Cao et al.,, 2018, 2019; Lin et al.,, 2020) may be utilized to
facilitate matching in these conditions. Methods that can account for
realistic shifts in neural data will allow for comparisons, while also
modeling state-dependent changes that may modulate the neural distribution
(Feulner & Clopath,, 2021).

As the algorithms to perform these types of analyses grow more sophisticated
and specialized, it will be important to apply them to neural datasets that
are generated across sufficiently diverse conditions. After all, carefully
obtained datasets of this nature are the only way of knowing if the approaches
we use actually work in practice. In machine learning, common datasets used
across the community have often spurred progress, from simple benchmarks like
MNIST (LeCun et al.,, 1998) to more substantial datasets like the ImageNet
corpus (Russakovsky et al.,, 2015). In neuroscience, open-access brain imaging
datasets and tools like those provided in the Allen Institute’s
optophysiology-based Brain Observatory (de Vries et al.,, 2020), and their
more recent high-density survey with Neuropixel probes (Siegle et al.,, 2019),
could become benchmarks by which to compare methods, as well as fertile ground
for more in-depth studies. Moving forward, widespread data sharing like this
will allow for more reproducible results and pave the way for more powerful
and generalizable algorithms.

##  5 Outlook

In current brain machine interfaces (BMIs), decoders rely on stable mappings
from neural activity to the variables of interest (e.g., position, speed, or
direction of intended movement), which often degrade over time (Pandarinath et
al.,, 2018b). Thus with the ability to directly compare latent representations
and align them, it will be possible to correct for changes in neural
distributions to develop more robust neural decoding interfaces. These
controllers might be able to adapt to shifts in neural activity due to
cognitive state without collection of new training data, which is crucial for
a controller which will be chronically implanted. We emphasize that latent
space alignment should not be considered only a convenience for analysis (and
label transfer), but also a tool for improving the weaknesses of current BMIs
and neural prosthetics.

While many of the approaches we discussed use linearity assumptions to
simplify the challenge of representation alignment, the computations and
associated latent spaces present within both small circuits and across larger
neural populations are likely to lie in spaces that linear models cannot
appropriately capture. In these cases, more complex models like unions of
subspaces and manifolds (Elhamifar & Vidal,, 2011, 2013; Dyer et al.,, 2013)
and statistical tests based upon these models may be used to compare more
complex low-dimensional structures. By coupling multi-manifold models with
matrix completion methods (Bishop & Byron,, 2014; Liu et al.,, 2017), it could
also be possible to stitch together manifolds from many brain areas and start
to fill in the gaps across recordings that are collected from different sites
of contact in the brain across non-overlapping populations of neurons. Thus,
through combining the techniques described in this article with more complex
latent space models, we may be able to compare whole brain activity patterns,
even when individual recordings provide only a small view of the bigger
picture.

While understanding of the impact of neurological and neurodegenerative
diseases has advanced tremendously in recent years (Frere & Slutsky,, 2018;
Wingo et al.,, 2021), a description of the disruptions they cause to the
functionality of neural circuits has not yet been possible. It is highly
likely that when such a description is made, it will be through the study of
shifts in the activity of many neurons. The alignment strategies outlined
here, and more sophisticated methods built on them, will provide new avenues
to compare the latent spaces that the brain occupies in health and provide
signatures of decline as it progresses into disease.

## Code availability

Demonstrations of the alignment methods DAD (Dyer et al.,, 2017) and HiWA (Lee
et al.,, 2019) and implementations in Python and MATLAB are available at:
nerdslab.github.io/neuralign/

## Acknowledgements

We thank Benjamin Dewey (https://benjamindewey.tumblr.com/) for creating the
awesome illustration in Figure 1, and Chethan Pandarinath and Mehdi Azabou for
helpful feedback on this manuscript. This work was supported by
NIH-1R01EB029852, as well as generous gifts from the Sloan Foundation and
McKnight Foundation.

## References

  * Abbott & Dayan, (1999) Abbott, Larry F, & Dayan, Peter. 1999.  The effect of correlated variability on the accuracy of a population code.  Neural Computation, 11(1), 91–101.
  * Andrew et al., (2013) Andrew, Galen, Arora, Raman, Bilmes, Jeff, & Livescu, Karen. 2013.  Deep canonical correlation analysis.  Pages 1247–1255 of: International Conference on Machine Learning.
  * Arakaki et al., (2017) Arakaki, Takafumi, Barello, Gregory, & Ahmadian, Yashar. 2017.  Capturing the diversity of biological tuning curves using generative adversarial networks.  arXiv preprint arXiv:1707.04582.
  * Bishop & Byron, (2014) Bishop, William E, & Byron, M Yu. 2014.  Deterministic symmetric positive semidefinite matrix completion.  Pages 2762–2770 of: Advances in Neural Information Processing Systems.
  * Brinkman et al., (2018) Brinkman, Braden AW, Rieke, Fred, Shea-Brown, Eric, & Buice, Michael A. 2018.  Predicting how and when hidden neurons skew measured synaptic interactions.  PLoS Computational Biology, 14(10), e1006490.
  * Brunton et al., (2016) Brunton, Bingni W, Johnson, Lise A, Ojemann, Jeffrey G, & Kutz, J Nathan. 2016\.  Extracting spatial–temporal coherent patterns in large-scale neural recordings using dynamic mode decomposition.  Journal of Neuroscience Methods, 258, 1 – 15.
  * Bullmore & Sporns, (2009) Bullmore, Ed, & Sporns, Olaf. 2009.  Complex brain networks: graph theoretical analysis of structural and functional systems.  Nature reviews neuroscience, 10(3), 186–198.
  * Cao et al., (2018) Cao, Zhangjie, Ma, Lijia, Long, Mingsheng, & Wang, Jianmin. 2018.  Partial adversarial domain adaptation.  Pages 135–150 of: Proceedings of the European Conference on Computer Vision (ECCV).
  * Cao et al., (2019) Cao, Zhangjie, You, Kaichao, Long, Mingsheng, Wang, Jianmin, & Yang, Qiang. 2019\.  Learning to transfer examples for partial domain adaptation.  Pages 2985–2994 of: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition.
  * Chestek et al., (2011) Chestek, Cynthia A, Gilja, Vikash, Nuyujukian, Paul, Foster, Justin D, Fan, Joline M, Kaufman, Matthew T, Churchland, Mark M, Rivera-Alvidrez, Zuley, Cunningham, John P, Ryu, Stephen I, et al. 2011\.  Long-term stability of neural prosthetic control signals from silicon cortical arrays in rhesus macaque motor cortex.  Journal of neural engineering, 8(4), 045005.
  * Chialvo, (2010) Chialvo, Dante R. 2010.  Emergent complex neural dynamics.  Nature Physics, 6(10), 744–750.
  * Child, (1990) Child, Dennis. 1990.  The essentials of factor analysis.  Cassell Educational.
  * Churchland et al., (2012) Churchland, Mark M, Cunningham, John P, Kaufman, Matthew T, Foster, Justin D, Nuyujukian, Paul, Ryu, Stephen I, & Shenoy, Krishna V. 2012.  Neural population dynamics during reaching.  Nature, 487(7405), 51.
  * Cohen & Maunsell, (2009) Cohen, Marlene R, & Maunsell, John HR. 2009.  Attention improves performance primarily by reducing interneuronal correlations.  Nature Neuroscience, 12(12), 1594.
  * Cowley et al., (2016) Cowley, Benjamin R, Smith, Matthew A, Kohn, Adam, & Byron, M Yu. 2016.  Stimulus-driven population activity patterns in macaque primary visual cortex.  PLoS Computational Biology, 12(12), e1005185.
  * Cunningham & Yu, (2014) Cunningham, John P, & Yu, Byron M. 2014.  Dimensionality reduction for large-scale neural recordings.  Nature Neuroscience, 17(11), 1500.
  * de Vries et al., (2020) de Vries, Saskia EJ, Lecoq, Jerome A, Buice, Michael A, Groblewski, Peter A, Ocker, Gabriel K, Oliver, Michael, Feng, David, Cain, Nicholas, Ledochowitsch, Peter, Millman, Daniel, et al. 2020\.  A large-scale standardized physiological survey reveals functional organization of the mouse visual cortex.  Nature Neuroscience, 23(1), 138–151.
  * Degenhart et al., (2020) Degenhart, Alan D, Bishop, William E, Oby, Emily R, Tyler-Kabara, Elizabeth C, Chase, Steven M, Batista, Aaron P, & Byron, M Yu. 2020.  Stabilization of a brain–computer interface via the alignment of low-dimensional spaces of neural activity.  Nature Biomedical Engineering, 1–14.
  * Dmochowski et al., (2018) Dmochowski, Jacek P, Ki, Jason J, DeGuzman, Paul, Sajda, Paul, & Parra, Lucas C. 2018.  Extracting multidimensional stimulus-response correlations using hybrid encoding-decoding of neural activity.  NeuroImage, 180, 134–146.
  * Dyer et al., (2013) Dyer, Eva L, Sankaranarayanan, Aswin C, & Baraniuk, Richard G. 2013.  Greedy feature selection for subspace clustering.  The Journal of Machine Learning Research, 14(1), 2487–2517.
  * Dyer et al., (2017) Dyer, Eva L, Azar, Mohammad Gheshlaghi, Perich, Matthew G, Fernandes, Hugo L, Naufel, Stephanie, Miller, Lee E, & Körding, Konrad P. 2017.  A cryptography-based approach for movement decoding.  Nature Biomedical Engineering, 1(12), 967.
  * Elhamifar & Vidal, (2011) Elhamifar, Ehsan, & Vidal, René. 2011.  Sparse manifold clustering and embedding.  Pages 55–63 of: Advances in neural information processing systems.
  * Elhamifar & Vidal, (2013) Elhamifar, Ehsan, & Vidal, Rene. 2013.  Sparse subspace clustering: Algorithm, theory, and applications.  IEEE transactions on pattern analysis and machine intelligence, 35(11), 2765–2781.
  * Fan et al., (2014) Fan, Joline M, Nuyujukian, Paul, Kao, Jonathan C, Chestek, Cynthia A, Ryu, Stephen I, & Shenoy, Krishna V. 2014.  Intention estimation in brain–machine interfaces.  Journal of Neural Engineering, 11(1), 016004.
  * Farshchian et al., (2018) Farshchian, Ali, Gallego, Juan A, Cohen, Joseph P, Bengio, Yoshua, Miller, Lee E, & Solla, Sara A. 2018.  Adversarial Domain Adaptation for Stable Brain-Machine Interfaces.  arXiv preprint arXiv:1810.00045.
  * Feulner & Clopath, (2021) Feulner, Barbara, & Clopath, Claudia. 2021.  Neural manifold under plasticity in a goal driven learning behaviour.  PLOS Computational Biology, 17(2), e1008621.
  * Frere & Slutsky, (2018) Frere, Samuel, & Slutsky, Inna. 2018.  Alzheimer’s disease: from firing instability to homeostasis network collapse.  Neuron, 97(1), 32–58.
  * Fu et al., (2016) Fu, Tian-Ming, Hong, Guosong, Zhou, Tao, Schuhmann, Thomas G, Viveros, Robert D, & Lieber, Charles M. 2016.  Stable long-term chronic brain mapping at the single-neuron level.  Nature Methods, 13(10), 875.
  * Fujiwara et al., (2013) Fujiwara, Yusuke, Miyawaki, Yoichi, & Kamitani, Yukiyasu. 2013.  Modular encoding and decoding models derived from Bayesian canonical correlation analysis.  Neural Computation, 25(4), 979–1005.
  * Gallego et al., (2018) Gallego, Juan A, Perich, Matthew G, Naufel, Stephanie N, Ethier, Christian, Solla, Sara A, & Miller, Lee E. 2018.  Cortical population activity within a preserved neural manifold underlies multiple motor behaviors.  Nature Communications, 9(1), 4233.
  * Gallego et al., (2020) Gallego, Juan A, Perich, Matthew G, Chowdhury, Raeed H, Solla, Sara A, & Miller, Lee E. 2020.  Long-term stability of cortical population dynamics underlying consistent behavior.  Nature Neuroscience, 23(2), 260–270.
  * Ganmor et al., (2015) Ganmor, Elad, Segev, Ronen, & Schneidman, Elad. 2015.  A thesaurus for a neural population code.  eLife, 4, e06134.
  * Gao et al., (2017) Gao, Peiran, Trautmann, Eric, Yu, Byron M, Santhanam, Gopal, Ryu, Stephen, Shenoy, Krishna, & Ganguli, Surya. 2017.  A theory of multineuronal dimensionality, dynamics and measurement.  BioRxiv, 214262.
  * Golub et al., (2018) Golub, Matthew D, Sadtler, Patrick T, Oby, Emily R, Quick, Kristin M, Ryu, Stephen I, Tyler-Kabara, Elizabeth C, Batista, Aaron P, Chase, Steven M, & Yu, Byron M. 2018.  Learning by neural reassociation.  Nature Neuroscience, 21(4), 607.
  * Goodfellow et al., (2020) Goodfellow, Ian, Pouget-Abadie, Jean, Mirza, Mehdi, Xu, Bing, Warde-Farley, David, Ozair, Sherjil, Courville, Aaron, & Bengio, Yoshua. 2020.  Generative adversarial networks.  Communications of the ACM, 63(11), 139–144.
  * Gradinaru et al., (2010) Gradinaru, Viviana, Zhang, Feng, Ramakrishnan, Charu, Mattis, Joanna, Prakash, Rohit, Diester, Ilka, Goshen, Inbal, Thompson, Kimberly R, & Deisseroth, Karl. 2010.  Molecular and cellular approaches for diversifying and extending optogenetics.  Cell, 141(1), 154–165.
  * Grill et al., (2009) Grill, Warren M, Norman, Sharon E, & Bellamkonda, Ravi V. 2009.  Implanted neural interfaces: biochallenges and engineered solutions.  Annual Review of Biomedical Engineering, 11, 1–24.
  * Harvey et al., (2012) Harvey, Christopher D, Coen, Philip, & Tank, David W. 2012.  Choice-specific sequences in parietal cortex during a virtual-navigation decision task.  Nature, 484(7392), 62.
  * Hengen et al., (2016) Hengen, Keith B, Pacheco, Alejandro Torrado, McGregor, James N, Van Hooser, Stephen D, & Turrigiano, Gina G. 2016.  Neuronal firing rate homeostasis is inhibited by sleep and promoted by wake.  Cell, 165(1), 180–191.
  * Hotelling, (1936) Hotelling, Harold. 1936.  Relations between two sets of variates.  Biometrika, 28(3-4), 321–377.
  * Huang et al., (2006) Huang, Su-Yun, Lee, Mei-Hsien, & Hsiao, Chuhsing Kate. 2006.  Kernel canonical correlation analysis and its applications to nonlinear measures of association and test of independence.  Institute of Statistical Science: Academia Sinica, Taiwan.
  * Hubel & Wiesel, (1959) Hubel, David H, & Wiesel, Torsten N. 1959.  Receptive fields of single neurones in the cat’s striate cortex.  The Journal of physiology, 148(3), 574–591.
  * Hyvärinen, (2013) Hyvärinen, Aapo. 2013.  Independent component analysis: recent advances.  Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences, 371(1984), 20110534.
  * Kaufman et al., (2014) Kaufman, Matthew T, Churchland, Mark M, Ryu, Stephen I, & Shenoy, Krishna V. 2014\.  Cortical activity in the null space: permitting preparation without movement.  Nature Neuroscience, 17(3), 440.
  * Kobak et al., (2019) Kobak, Dmitry, Pardo-Vazquez, Jose L, Valente, Mafalda, Machens, Christian K, & Renart, Alfonso. 2019.  State-dependent geometry of population activity in rat auditory cortex.  eLife, 8, e44526.
  * Lai & Fyfe, (2000) Lai, Pei Ling, & Fyfe, Colin. 2000.  Kernel and nonlinear canonical correlation analysis.  International Journal of Neural Systems, 10(05), 365–377.
  * LeCun et al., (1998) LeCun, Yann, Bottou, Léon, Bengio, Yoshua, & Haffner, Patrick. 1998.  Gradient-based learning applied to document recognition.  Proceedings of the IEEE, 86(11), 2278–2324.
  * Lee & Seung, (1999) Lee, Daniel D, & Seung, H Sebastian. 1999.  Learning the parts of objects by non-negative matrix factorization.  Nature, 401(6755), 788.
  * Lee et al., (2019) Lee, John, Dabagia, Max, Dyer, Eva, & Rozell, Christopher. 2019.  Hierarchical Optimal Transport for Multimodal Distribution Alignment.  Pages 13453–13463 of: Advances in Neural Information Processing Systems.
  * Lillicrap & Kording, (2019) Lillicrap, Timothy P, & Kording, Konrad P. 2019.  What does it mean to understand a neural network?  arXiv preprint arXiv:1907.06374.
  * Lin et al., (2020) Lin, Chi-Heng, Azabou, Mehdi, & Dyer, Eva L. 2020.  Making transport more robust and interpretable by moving data through a small number of anchor points.  arXiv preprint arXiv:2012.11589.
  * Liu et al., (2017) Liu, Guangcan, Liu, Qingshan, & Yuan, Xiaotong. 2017.  A new theory for matrix completion.  Pages 785–794 of: Advances in Neural Information Processing Systems.
  * Luczak et al., (2015) Luczak, Artur, McNaughton, Bruce L, & Harris, Kenneth D. 2015.  Packet-based communication in the cortex.  Nature Reviews Neuroscience, 16(12), 745–755.
  * Ma et al., (2019) Ma, Zhengyu, Turrigiano, Gina G, Wessel, Ralf, & Hengen, Keith B. 2019.  Cortical circuit dynamics are homeostatically tuned to criticality in vivo.  Neuron, 104(4), 655–664.
  * Mante et al., (2013) Mante, Valerio, Sussillo, David, Shenoy, Krishna V, & Newsome, William T. 2013\.  Context-dependent computation by recurrent dynamics in prefrontal cortex.  Nature, 503(7474), 78.
  * Mazor & Laurent, (2005) Mazor, Ofer, & Laurent, Gilles. 2005.  Transient dynamics versus fixed points in odor representations by locust antennal lobe projection neurons.  Neuron, 48(4), 661–673.
  * McAdams & Maunsell, (1999) McAdams, Carrie J, & Maunsell, John HR. 1999.  Effects of attention on orientation-tuning functions of single neurons in macaque cortical area V4.  Journal of Neuroscience, 19(1), 431–441.
  * McCreery et al., (2010) McCreery, Douglas, Pikov, Victor, & Troyk, Philip R. 2010.  Neuronal loss due to prolonged controlled-current stimulation with chronically implanted microelectrodes in the cat cerebral cortex.  Journal of Neural Engineering, 7(3), 036005.
  * Mehler & Kording, (2018) Mehler, David Marc Anton, & Kording, Konrad Paul. 2018.  The lure of causal statements: Rampant mis-inference of causality in estimated connectivity.  arXiv e-prints, arXiv–1812.
  * Mitchell et al., (2009) Mitchell, Jude F, Sundberg, Kristy A, & Reynolds, John H. 2009.  Spatial attention decorrelates intrinsic activity fluctuations in macaque area V4.  Neuron, 63(6), 879–888.
  * Molano-Mazon et al., (2018) Molano-Mazon, Manuel, Onken, Arno, Piasini, Eugenio, & Panzeri, Stefano. 2018.  Synthesizing realistic neural population activity patterns using generative adversarial networks.  arXiv preprint arXiv:1803.00338.
  * Nonnenmacher et al., (2017) Nonnenmacher, Marcel, Turaga, Srinivas C, & Macke, Jakob H. 2017.  Extracting low-dimensional dynamics from multiple large-scale neural population recordings by learning to predict correlations.  Pages 5702–5712 of: Advances in Neural Information Processing Systems.
  * Pandarinath et al., (2018a) Pandarinath, Chethan, O’Shea, Daniel J, Collins, Jasmine, Jozefowicz, Rafal, Stavisky, Sergey D, Kao, Jonathan C, Trautmann, Eric M, Kaufman, Matthew T, Ryu, Stephen I, Hochberg, Leigh R, et al. 2018a.  Inferring single-trial neural population dynamics using sequential auto-encoders.  Nature Methods, 1.
  * Pandarinath et al., (2018b) Pandarinath, Chethan, Ames, K Cora, Russo, Abigail A, Farshchian, Ali, Miller, Lee E, Dyer, Eva L, & Kao, Jonathan C. 2018b.  Latent factors and dynamics in motor cortex and their application to brain–machine interfaces.  Journal of Neuroscience, 38(44), 9390–9401.
  * Paninski et al., (2010) Paninski, Liam, Ahmadian, Yashar, Ferreira, Daniel Gil, Koyama, Shinsuke, Rad, Kamiar Rahnama, Vidne, Michael, Vogelstein, Joshua, & Wu, Wei. 2010.  A new look at state-space models for neural data.  Journal of Computational Neuroscience, 29(1-2), 107–126.
  * Polikov et al., (2005) Polikov, Vadim S, Tresco, Patrick A, & Reichert, William M. 2005.  Response of brain tissue to chronically implanted neural electrodes.  Journal of Neuroscience Methods, 148(1), 1–18.
  * Roth, (2016) Roth, Bryan L. 2016.  DREADDs for neuroscientists.  Neuron, 89(4), 683–694.
  * Russakovsky et al., (2015) Russakovsky, Olga, Deng, Jia, Su, Hao, Krause, Jonathan, Satheesh, Sanjeev, Ma, Sean, Huang, Zhiheng, Karpathy, Andrej, Khosla, Aditya, Bernstein, Michael, et al. 2015\.  Imagenet large scale visual recognition challenge.  International Journal of Computer Vision, 115(3), 211–252.
  * Russo et al., (2020) Russo, Abigail A, Khajeh, Ramin, Bittner, Sean R, Perkins, Sean M, Cunningham, John P, Abbott, Laurence F, & Churchland, Mark M. 2020.  Neural trajectories in the supplementary motor area and motor cortex exhibit distinct geometries, compatible with different classes of computation.  Neuron, 107(4), 745–758.
  * Sadtler et al., (2014) Sadtler, Patrick T, Quick, Kristin M, Golub, Matthew D, Chase, Steven M, Ryu, Stephen I, Tyler-Kabara, Elizabeth C, Yu, Byron M, & Batista, Aaron P. 2014.  Neural constraints on learning.  Nature, 512(7515), 423.
  * Saxena & Cunningham, (2019) Saxena, Shreya, & Cunningham, John P. 2019.  Towards the neural population doctrine.  Current Opinion in Neurobiology, 55, 103–111.
  * Sederberg et al., (2019) Sederberg, Audrey J, Pala, Aurelie, Zheng, He JV, He, Biyu J, & Stanley, Garrett B. 2019.  State-aware detection of sensory stimuli in the cortex of the awake mouse.  PLoS Computational Biology, 15(5), e1006716.
  * Shenoy & Kao, (2021) Shenoy, Krishna V, & Kao, Jonathan C. 2021.  Measurement, manipulation and modeling of brain-wide neural population dynamics.  Nature Communications, 12(1), 1–5.
  * Siegle et al., (2019) Siegle, Joshua H, Jia, Xiaoxuan, Durand, Séverine, Gale, Sam, Bennett, Corbett, Graddis, Nile, Heller, Greggory, Ramirez, Tamina K, Choi, Hannah, Luviano, Jennifer A, et al. 2019\.  A survey of spiking activity reveals a functional hierarchy of mouse corticothalamic visual areas.  bioRxiv, 805010.
  * Soudry et al., (2015) Soudry, Daniel, Keshri, Suraj, Stinson, Patrick, Oh, Min-hwan, Iyengar, Garud, & Paninski, Liam. 2015.  Efficient “shotgun” inference of neural connectivity from highly sub-sampled activity data.  PLoS Computational Biology, 11(10), e1004464.
  * Sponberg et al., (2015) Sponberg, Simon, Daniel, Thomas L, & Fairhall, Adrienne L. 2015.  Dual dimensionality reduction reveals independent encoding of motor features in a muscle synergy for insect flight control.  PLoS Comput Biol, 11(4), e1004168.
  * St-Yves & Naselaris, (2018) St-Yves, Ghislain, & Naselaris, Thomas. 2018.  Generative Adversarial Networks Conditioned on Brain Activity Reconstruct Seen Images.  Pages 1054–1061 of: 2018 IEEE International Conference on Systems, Man, and Cybernetics (SMC).  IEEE.
  * Stopfer et al., (2003) Stopfer, Mark, Jayaraman, Vivek, & Laurent, Gilles. 2003.  Intensity versus identity coding in an olfactory system.  Neuron, 39(6), 991–1004.
  * Tsodyks et al., (1999) Tsodyks, Misha, Kenet, Tal, Grinvald, Amiram, & Arieli, A. 1999.  Linking spontaneous activity of single cortical neurons and the underlying functional architecture.  Science, 286(5446), 1943–1946.
  * Turaga et al., (2013) Turaga, Srini, Buesing, Lars, Packer, Adam M, Dalgleish, Henry, Pettit, Noah, Hausser, Michael, & Macke, Jakob H. 2013.  Inferring neural population dynamics from multiple partial recordings of the same neural circuit.  Pages 539–547 of: Advances in Neural Information Processing Systems.
  * Vaidya et al., (2014) Vaidya, Mukta, Dickey, Adam, Best, Matthew D, Coles, Josh, Balasubramanian, Karthikeyan, Suminski, Aaron J, & Hatsopoulos, Nicholas G. 2014.  Ultra-long term stability of single units using chronically implanted multielectrode arrays.  Pages 4872–4875 of: 2014 36th Annual International Conference of the IEEE Engineering in Medicine and Biology Society.  IEEE.
  * Villani, (2008) Villani, Cédric. 2008.  Optimal transport: old and new.  Vol. 338.  Springer Science & Business Media.
  * Whiteway et al., (2019) Whiteway, Matthew R, Socha, Karolina, Bonin, Vincent, & Butts, Daniel A. 2019.  Characterizing the nonlinear structure of shared variability in cortical neuron populations using latent variable models.  arXiv preprint arXiv:1904.10441.
  * Williamson et al., (2019) Williamson, Ryan C, Doiron, Brent, Smith, Matthew A, & Yu, Byron M. 2019.  Bridging large-scale neuronal recordings and large-scale network models using dimensionality reduction.  Current Opinion in Neurobiology, 55, 40–47.
  * Wingo et al., (2021) Wingo, Aliza P, Liu, Yue, Gerasimov, Ekaterina S, Gockley, Jake, Logsdon, Benjamin A, Duong, Duc M, Dammer, Eric B, Robins, Chloe, Beach, Thomas G, Reiman, Eric M, et al. 2021\.  Integrating human brain proteomes with genome-wide association data implicates new proteins in Alzheimer’s disease pathogenesis.  Nature Genetics, 1–4.
  * Zhang et al., (2011) Zhang, Ying, Meyers, Ethan M, Bichot, Narcisse P, Serre, Thomas, Poggio, Tomaso A, & Desimone, Robert. 2011.  Object decoding with attention in inferior temporal cortex.  Proceedings of the National Academy of Sciences, 108(21), 8850–8855.
  * Zohary et al., (1994) Zohary, Ehud, Shadlen, Michael N, & Newsome, William T. 1994.  Correlated neuronal discharge rate and its implications for psychophysical performance.  Nature, 370(6485), 140.


# MASTER: Market-Guided Stock Transformer for Stock Price Forecasting

Tong Li1, Zhaoyang Liu2, Yanyan Shen1, Xue Wang2, Haokun Chen2, Sen Huang2
This work was done during her internship at Alibaba Group.Corresponding
author.

###### Abstract

Stock price forecasting has remained an extremely challenging problem for many
decades due to the high volatility of the stock market. Recent efforts have
been devoted to modeling complex stock correlations toward joint stock price
forecasting. Existing works share a common neural architecture that learns
temporal patterns from individual stock series and then mixes up temporal
representations to establish stock correlations. However, they only consider
time-aligned stock correlations stemming from all the input stock features,
which suffer from two limitations. First, stock correlations often occur
momentarily and in a cross-time manner. Second, the feature effectiveness is
dynamic with market variation, which affects both the stock sequential
patterns and their correlations. To address the limitations, this paper
introduces MASTER, a MArkert-Guided Stock TransformER, which models the
momentary and cross-time stock correlation and leverages market information
for automatic feature selection. MASTER elegantly tackles the complex stock
correlation by alternatively engaging in intra-stock and inter-stock
information aggregation. Experiments show the superiority of MASTER compared
with previous works and visualize the captured realistic stock correlation to
provide valuable insights.

## Introduction

Stock price forecasting, which utilizes historical data collected from the
stock market to predict future trends, is a vital technique for profitable
stock investment. Unlike stationary time series that often exhibit regular
patterns such as periodicity and steady trends, the dynamics in the stock
price series are intricate because stock prices fluctuate subject to multiple
factors, including macroeconomic factors, capital flows, investor sentiments,
and events. The mixing of factors interweaves the stock market as a correlated
network, making it difficult to precisely predict the individual behavior of
stocks without taking other stocks into account.

Most previous works (Feng et al. 2019; Xu et al. 2021; Wang et al. 2021, 2022;
Wang, Qu, and Chen 2022) in the field of stock correlation have relied on
predefined concepts, relationships, or rules and established a _static_
correlation graph, e.g., stocks in the same industry are connected to each
other. While these methods provide insights into the relations between stocks,
they do not account for the real-time correlation of stocks. For example,
different stocks within the same industry can experience opposite price
movements on a particular day. Additionally, the pre-defined relationships may
not be generalizable to new stocks in an evolving market where events such as
company listing, delisting, or changes in the main business happen normally.
Another line of research (Yoo et al. 2021) follows the Transformer
architecture (Vaswani et al. 2017), and use the self-attention mechanism to
compute _dynamic_ stock correlations. This data-driven manner is more flexible
and applicable to the time-varying stock sets in the market. Despite different
schemes for establishing stock correlations, the existing methods generally
follow a common two-step computation flow. As depicted in Figure 1, the first
step is using a sequential encoder to summarize the historical sequence of
stock features, and obtain stock representation, and the second step is to
refine each stock representation by aggregating information from correlated
stocks using graph encoders or attention mechanism. However, such a flow
suffers from two limitations.

Figure 1: The framework of existing works. The dashed lines represent the
underlying momentary and cross-time stock correlations, which reside between
some ($stock_{1},time_{1}$), ($stock_{2},time_{2}$) pairs.

First, existing works distill an overall stock representation and blur the
time-specific details of stock sequence, leading to weakness in modeling the
de-facto stock correlations, which often occurs _momentarily_ and in a _cross-
time_ manner (Bennett, Cucuringu, and Reinert 2022). To be specific, the stock
correlation is highly dynamic and may reside in misaligned time steps rather
than holding true through the whole lookback period. This is because the
dominating factors of stock prices constantly change, and different stocks may
react to the same factors with different delays. For instance, upstream
companies’ stock prices may react faster to a shortage of raw materials than
those of downstream companies, and individual stocks exhibit a lot of catch-up
and fall-behind behaviors.

Since the stock correlation may underlie between every stock pair and time
pair, a straightforward way to simulate the momentary and cross-time
correlation is to gather the $\tau\times|\mathcal{S}|$ feature vectors for
pair-wise attention computation, where $\tau$ is the lookback window length
and $\mathcal{S}$ is the stock set. However, in addition to the increased
computational complexity, this approach faces practical difficulties because
the stock forecasting task is in intense data hunger. Intuitively, there are
only around 250 trading days per year, producing limited observations on
stocks. When the model adopts such a large attention field with insufficient
training samples, it often struggles to optimize and may even fall into
suboptimal solutions. Although clustering approaches like local sensitive
hashing (Kitaev, Kaiser, and Levskaya 2020) have been proposed to reduce the
size of the attention field, they are sensitive to initialization, which is a
fatal issue in a data-hungry domain like stock forecasting. To address these
challenges, we propose a novel stock transformer architecture specifically
designed for stock price forecasting. Rather than directly modeling the
$\tau\times|\mathcal{S}|$ attention field or using clustering-based
approximation methods, our model aggregates information from different time
steps and different stocks alternately to model realistic stock correlation
and facilitate model learning.

Another limitation of existing works is that they ignore the impact of varying
market status. In long-term practice with the market variation, one essential
observation by investors is that the features come into effect and expire
dynamically. The effectiveness of features has an influence on both the intra-
stock sequential pattern and the stock correlation. For instance, in a bull
market, the correlations among stocks are more significant due to the
investors’ optimism. Traditional investors repeatedly conduct statistical
examination on to select effective feature, which is exhaustive and face a gap
when integrated with learning-based methods. To save the human efforts, we are
motivated to equip our stock transformer with a novel gating mechanism, which
incorporates the market information to perform automatically feature
selection. We name the proposed method MASTER, standing for MArket-Guided
Stock TransformER. To summarize, our main contributions are as follows.

$\bullet$ We propose a novel stock transformer for stock price forecasting to
effectively capture the stock correlation. To the best of our knowledge, we
are the first to mine the momentary and cross-time stock correlation with
learning-based methods.

$\bullet$ We introduce a novel gating mechanism that integrates market
information to automatically select relevant features and adapt to varying
market scenarios.

$\bullet$ We conducted experiments to validate the designs of our proposed
method and demonstrated its superiority compared to baselines. The
visualization results provided valuable insights into the real-time dynamics
of stock correlations.

Figure 2: Overview of the MASTER framework.

## Methodology

### Problem Formulation

The indicators of each stock $u\in\mathcal{S}$ are collected at every time
step $t\in[1,\tau]$ to form the feature vector $x_{u,t}\in\mathbf{R}^{F}$.
Following existing works on stock market analysis (Feng et al. 2018; Sawhney
et al. 2020; Huynh et al. 2023), we focus on the prediction of the change in
stock price rather than the absolute value. The return ratio, which is the
relative close price change in $d$ days, is
$\tilde{r}_{u}=(c_{u,\tau+d}-c_{u,\tau+1})/c_{u,\tau+1}$, where $c_{u,t}$ is
the closing price of stock $u$ at time step $t$, and $d$ represents the
predetermined prediction interval. The return ratio normalizes the market
price variety between different stocks in comparison to the absolute price
change. Since stock investment is to rank and select the most profitable
stocks, we perform daily Z-score normalization of return ratio to encode the
label with the rankings, $r_{u}={\rm Norm}_{\mathcal{S}}(\tilde{r}_{u})$, as
in previous work (Yang et al. 2020).

######  Definition 1 (Stock Price Forecasting)

Given stock features $\\{x_{u,t}\\}_{u\in\mathcal{S},t\in[1,\tau]}$, the stock
price forecasting is to jointly predict the future normalized return ratio
$\\{r_{u}\\}_{u\in\mathcal{S}}$.

### Overview

Figure 2 depicts the architecture of our proposed method MASTER, which
consists of five steps. (1) Market-Guided Gating. We construct a vector
representing the current market status $m_{\tau}$ and leverage it to rescale
feature vectors by a gating mechanism, achieving market-guided feature
selection. (2) Intra-Stock Aggregation. Within the sequence of each stock, at
each time step, we aggregate information from other time steps to generate a
local embedding that preserves the temporal local details of the stock while
collecting all important signals along the time axis. The local embedding
$h_{u,t}$ will serve as relays and transport the collected signals to other
stocks in subsequent modules. (3) Inter-Stock Aggregation. At each time step,
we compute stock correlation with attention mechanism, and each stock further
aggregates the local embedding of other stocks. The aggregated information
$z_{u,t}$, which we refer to as temporal embedding, contains not only the
information of the momentarily correlated stocks at $t$, but also preserves
the personal information of $u$. (4) Temporal Aggregation. For each stock, the
last temporal embedding queries from all historical temporal embedding and
produce a comprehensive stock embedding $e_{u}$. (5) Prediction. The
comprehensive stock embedding is sent to prediction layers for label
prediction. We elaborate on the details of MASTER step by step in the
following sub-sections.

### Market-Guided Gating

#### Market Status Representation

First, we propose to combine information from two aspects into a vector
$m_{\tau}$ to give an abundant description of the current market status. (1)
Market index price. The market index price is a weighted average of the prices
of a group of stocks $\mathcal{S^{\prime}}$ by their share of market
capitalization. $\mathcal{S^{\prime}}$ is typically composed of top companies
with the most market capitalization, representing a particular market or
sector, and may differ from user-interested stocks in investing $\mathcal{S}$.
We include both the current market index price at $\tau$ and the historical
market index prices, which is described by the average and standard deviation
in the past $d^{\prime}$ days to reveal the price fluctuations. Here,
$d^{\prime}$ specifies the referable interval length to introduce historical
market information in applications. (2) Market index trading volume. The
trading volumes of $\mathcal{S^{\prime}}$ reveals the investors involvement,
reflecting the activity of the market. We include the average and standard
deviation of market index trading volume in the past $d^{\prime}$ days, to
reveal the actual size of the market. $\mathcal{S^{\prime}}$ and $d^{\prime}$
are identical to the aforementioned definitions. Now we present the market-
guided stock price forecasting task.

######  Definition 2 (Market-Guided Stock Price Forecasting)

Given $\\{x_{u,t}\\}_{u\in\mathcal{S},t\in[1,\tau]}$ and the constructed
market status vector $m_{\tau}$, market-guided stock price forecasting is to
jointly predict the future normalized return ratio
$\\{r_{u}\\}_{u\in\mathcal{S}}$.

#### Gating Mechanism

The gating mechanism generates one scaling coefficient for each feature
dimension to enlarge or shrink the magnitude of the feature, thereby
emphasizing or diminishing the amount of information from the feature flowing
to the subsequent modules. The gating mechanism is learned by the model
training, and the coefficient is optimized by how much the feature contributes
to improve forecasting performance, thus reflect the feature effectiveness.

Given the market status representation $m_{\tau},|m_{\tau}|=F^{\prime}$, we
first use a single linear layer to transform $m_{\tau}$ into the feature
dimension $F=|x_{u,t}|$. Then, we perform Softmax along the feature dimension
to obtain a distribution.

$\alpha(m_{\tau})=F\cdot{\rm softmax}_{\beta}(W_{\alpha}m_{\tau}+b_{\alpha}),$

where $W_{\alpha}$, $b_{\alpha}$ are learnable matrix and bias, $\beta$ is the
temperature hyperparameter controlling the sharpness of the output
distribution. Softmax compels a competition among features to distinguish the
effective ones and ineffective ones. Here, a smaller temperature $\beta$
encourages the distribution to focus on certain dimension and the gating
effect is stronger while a larger $\beta$ makes the distribution incline to
even and the gating effect is weaker. Note that we enlarge the value at each
dimension by $F$ times as the scaling coefficient. This operation compare the
generated distribution with a uniform distribution where each dimension is
$1/F$, to determine whether to enlarge or shrink the value. The intuition to
generate coefficients from $m_{\tau}$ is that the effectiveness of features
are influenced by market status. For example, if the model learns moving
average (MA) factor is useful during volatile market periods, it will
emphasize MA when the market becomes volatile again. Under the same
$m_{\tau}$, $\alpha$ are shared for $\\{x_{u,t}\\}$, $u\in\mathcal{S}$,
$t\in[1,\tau]$, in that we incorporate the most recent market status to
perform unified feature selection. The rescaled feature vectors are
$\tilde{x}_{u,t}=\alpha(m_{\tau})\circ x_{u,t}$, where $\circ$ is the Hadamard
product.

### Intra-Stock Aggregation

In MASTER, we use intra-stock aggregation followed by inter-stock aggregation
to break down the large and complex attention field. Although the entire
market is complicated with diverse behaviours of individual stocks, the
patterns of a specific stock tend to be relatively continuous. Therefore, we
perform intra-stock aggregation first due to its smaller attention field and
simpler distribution. In our proposed intra-stock aggregation, the feature at
each time step aggregate information from other time steps and form a local
embedding. Compared with existing works which initially mix the feature
sequence into one representation (Yoo et al. 2021), we maintain a sequence of
local embedding which are advised with the important signals in sequence
through intra-stock aggregation while reserve the local details.

We first send the rescaled feature vectors to a feature encoder and transform
them into the embedding space, $y_{u,t}=f(\tilde{x}_{u,t})$, $|y_{u,t}|=D$. We
simply use a single linear layer as $f(\cdot)$. Then, we apply a bi-
directional sequential encoder to obtain the local output at each time step
$t$. Inspired by the success of transformer-based models in modeling
sequential patterns, we instantiate the sequential encoder with a single-layer
transformer encoder (Vaswani et al. 2017). Each feature vector at a particular
time step is treated as a token, and we add a fixed $D$-dimensional sinusoidal
positional encoding $p_{t}$ to mark the chronically order in the look back
window.

$Y_{u}=||_{t\in[1,\tau]}\textsf{LN}(f(\tilde{x}_{u,t})+p_{t}),$

where $||$ denotes the concatenation of vectors and LN the layer
normalization. Then, the feature embedding at each time step queries from all
time steps in the stock sequence. We introduce multi-head attention
mechanisms, denoted as MHA$(\cdot)$, with $N_{1}$ heads to perform different
aggregations in parallel. We also utilize feed forward layers, FFN$(\cdot)$,
to fuse the information obtained from the multi-head attention.

$Q^{1}_{u}=W^{1}_{Q}Y_{u},\quad K^{1}_{u}=W^{1}_{K}Y_{u},\quad
V^{1}_{u}=W^{1}_{V}Y_{u},$
$H^{1}_{u}=||_{t\in[1,\tau]}h_{u,t}=\textsf{FFN}^{1}(\textsf{MHA}^{1}(Q^{1}_{u},K^{1}_{u},V^{1}_{u})+Y_{u}),$

where FFN is a two-layer MLP with ReLU activation and residual connection. As
a result, the local embedding $h_{u,t}$ both reserve the local details and
encode indicative signals from other time steps.

### Inter-Stock Aggregation

Then, we consider aggregating information from correlated stocks. Compared
with existing works that distill an overall stock correlation, we establish a
series of momentary stock correlation corresponding to every time step.
Instead of using pre-defined relationships that face a mismatch with the
proximity of real-time stock movements, we propose to mine the asymmetric and
dynamic inter-stock correlation via attention-mechanism. The quality of the
correlation will be measured by its contribution to improving the forecasting
performance, and automatically optimized by the model training process.

Specifically, at each time step, we gather the local embedding of all stocks
$H^{2}_{t}=||_{u\in\mathcal{S}}h_{u,t}$ and perform multi-head attention
mechanism with $N_{2}$ heads.

$Q^{2}_{t}=W^{2}_{Q}H^{2}_{t},\quad K^{2}_{t}=W^{2}_{K}H^{2}_{t},\quad
V^{2}_{t}=W^{2}_{V}H^{2}_{t},$
$Z_{t}=||_{u\in\mathcal{S}}z_{u,t}=\textsf{FFN}^{2}(\textsf{MHA}^{2}(Q^{2}_{t},K^{2}_{t},V^{2}_{t})+H^{2}_{t}).$

With the residual connection of FFN, the temporal embedding $z_{u,t}$ is
encoded with both the information from momentarily correlated stocks and the
personal information of stock $u$ itself. Our stock transformer is able to
model the cross-time correlation of stocks, as shown in Figure 2 (Right). The
local details of $y_{v,j}$ can first be conveyed to $h_{v,i}$ by the intra-
stock aggregation of stock $v$, and then transmitted to $z_{u,i}$ by inter-
stock aggregation at time step $i$, hence modeling the correlation from any
$(v,j)$ to $(u,i)$. We further visualize and explain the captured cross-time
correlation in the experiments section.

### Temporal Aggregation

In contrast with existing works which obtain one embedding for each stock
after modeling stock correlation (Feng et al. 2019), our approach involves
producing a series of temporal embedding $z_{u,t},t\in[1,\tau].$ Each
$z_{u,t}$ is encoded with information from stocks that are momentarily
correlated with $(u,t)$. To summarize the obtained temporal embeddings and
obtain a comprehensive stock embedding $e_{u}$, we employ a temporal attention
layer along the time axis. We use the latest temporal embedding $z_{u,\tau}$
as the query vector, and compute the attention score $\lambda_{u,t}$ in a
hidden space with transformation matrix $W_{\lambda}$,

$\lambda_{u,t}=\frac{\exp(z^{T}_{u,t}W_{\lambda}z_{u,\tau})}{\sum_{i\in[1,\tau]}\exp(z^{T}_{u,i}W_{\lambda}z_{u,\tau})},\quad
e_{u}=\sum_{t\in[1,\tau]}\lambda_{u,t}z_{u,t}.$

### Prediction and Training

Finally, the stock embedding $e_{u}$ is fed into a predictor $g(\cdot)$ for
label regression. We use a single linear layer as the predictor, and the
forecasting quality is measured by the MSE loss. In each batch, MASTER is
jointly optimized for all $u\in\mathcal{S}$ on a particular prediction date.
And a training epoch is composed of multiple batches correspond to different
prediction dates in the training set.

$\hat{r}_{u}=g(e_{u}),\quad
L=\sum_{u\in\mathcal{S}}\textsf{MSE}(r_{u},\hat{r}_{u}).$

### Discussions

#### Relationships with Existing Works

Modeling stock correlations has long been an indispensable research direction
for stock price prediction. Today, many researchers and quantitative analysts,
still opt for linear models, support vector machines, and tree-based methods
for stock price forecasting (Nugroho, Adji, and Fauziati 2014; Chen and
Guestrin 2016; Kamble 2017; Xie et al. 2013; Li et al. 2015; Piccolo 1990).
The aggregation of correlation information within and between stocks is often
achieved through feature engineering, which relies heavily on manual expertise
and constantly faces the risk of factor decay. Inspired by the success of
neural sequential data analysis, researchers are driven to take into account
the stock feature sequences and learn the temporal correlation automatically.
They design various sequential models, such as RNN-based (Feng et al. 2019;
Sawhney et al. 2021; Yoo et al. 2021; Huynh et al. 2023), CNN-based (Wang et
al. 2021), and attention-based models(Liu et al. 2019; Ding et al. 2020), to
mine the internal temporal dynamics of a stock. Recent research focus on the
modeling of stock correlation, which add a correlation module in posterior to
the sequential model as illustrated in Figure 1. They propose to use graph-
based (Feng et al. 2019; Xu et al. 2021; Wang et al. 2021, 2022), hypergraph-
based (Sawhney et al. 2021; Huynh et al. 2023) and attention-based (Yoo et al.
2021; Xiang et al. 2022) modules to build the overall stock correlation and
perform joint prediction. Our MASTER is dedicated to momentary and cross-time
stock correlation mining. To do so, we develop a novel model architecture as
in Figure 2 that is genuinely different from all existing methods.
Furthermore, MASTER is specialized for stock price forecasting, which is
distinct in data form and task properties from existing transformer-based
models in spatial-temporal data (Bulat et al. 2021; Cong et al. 2021; Xu et
al. 2020; Li et al. 2023) or multivariate time series domains (Zhang and Yan
2022; Nie et al. 2022).

#### Complexity Analysis

We now analyze the computation complexity of our proposed method. Let
$M=|\mathcal{S}|$, the market-guided gating rescale $M\times\tau$ feature
vectors of dimension $F$. In intra-stock aggregation, the calculation amount
of pair-wise attention is $\tau^{2}$ for each stock at each attention head. In
inter-stock aggregation, the calculation amount is $M^{2}$ at each time step
and each attention head. In temporal aggregation, we compute $\tau$ attention
scores for each stock. The overall computation complexity is
O$(FM\tau+N_{1}M\tau^{2}D^{2}+N_{2}M^{2}\tau D^{2}+M\tau D^{2})$, where
$M\gg\tau$. Therefore, MASTER is of ${\rm O}(N_{2}M^{2}\tau D^{2})$ time
complexity. Compared with directly operating on the $M\times\tau$ attention
field with $N$ attention heads, which is in O$(NM^{2}\tau^{2}D^{2})$, we
reduce the computation cost by about $\tau$ times and achieve modeling cross-
time correlations between stocks more efficiently. The overall parameters to
be trained in MASTER are transformation matrices
$W^{1}_{Q},W^{1}_{K},W^{1}_{V},W^{2}_{Q},W^{2}_{K},W^{2}_{V},W_{\lambda}$,
which is in shape $D\times D$, and parameters in MLP layers
$\alpha,f,\textsf{FFN}^{1},\textsf{FFN}^{2}$ and $g$.

## Experiments

In this section, we conduct experiments to answer the following four research
questions:

  * •

RQ1 How is the overall performance of MASTER compared with state-of-the-art
methods?

  * •

RQ2 Is the proposed stock transformer architecture effective for stock price
forecasting?

  * •

RQ3 How do hyper-parameter configurations affect the performance of MASTER?

  * •

RQ4 What insights on the stock correlation can we get through visualizing the
attention map?

#### Datasets

We evaluate our framework on the Chinese stock market with CSI300 and CSI800
stock sets. CSI300 and CSI800 are two stock sets containing 300 and 800 stocks
with the highest capital value on the Shanghai Stock Exchange and the Shenzhen
Stock Exchange. The dataset contains daily information ranging from 2008 to
2022 of CSI300 and CSI800. We use the data from Q1 2008 to Q1 2020 as the
training set, data in Q2 2020 as the validation set, and the last ten
quarters, i.e., Q3 2020 to Q4 2022, are reserved as the test set. We apply the
public Alpha158 indicators (Yang et al. 2020) to extract stock features from
the collected data. The lookback window length $\tau$ and prediction interval
$d$ are set as $8$ and $5$ respectively. For market representation, we
constructed $63$ features with CSI300, CSI500 and CSI800 market indices, and
refereable interval $d^{\prime}=5,10,20,30,60$.

#### Baselines

We compare the performance of MASTER with several stock price forecasting
baselines from different categories. $\bullet$ XGBoost (Chen and Guestrin
2016): A decision-tree based method. According to the leaderboard of Qlib
platform (Yang et al. 2020), it is one of the strongest baselines. $\bullet$
LSTM (Graves and Graves 2012), GRU (Cho et al. 2014), TCN (Bai, Kolter, and
Koltun 2018), and Transformer (Vaswani et al. 2017): Sequential baselines that
leverage vanilla LSTM/GRU/Temporal convolutional network/Transformer along the
time axis for stock price forecasting. $\bullet$ GAT (Veličković et al. 2017):
A graph-based baseline, which first use sequential encoder to gain stock
presentation and then aggregate information by graph attention networks111More
discussion is provided in the supplementary materials.. $\bullet$ DTML (Yoo et
al. 2021): A state-of-the-art stock correlation mining method, which follows
the framework in Figure 1. DTML adopts the attention-mechanism to mine the
dynamic correlation among stocks and also incorporates the market information
into the modeling.

#### Evaluation

We adopt both ranking metrics and portfolio-based metrics to give a thorough
evaluation of the model performance. Four ranking metrics, Information
Coefficient (IC), Rank Information Coefficient (RankIC), Information Ratio
based IC (ICIR) and Information Ratio based RankIC (RankICIR) are considered.
IC and RankIC are the Pearson coefficient and Spearman coefficient averaged at
a daily frequency. ICIR and RankICIR are normalized metrics of IC and RankIC
by dividing the standard deviation. Those metrics are commonly used in
literature (e.g., Xu et al. 2021 and Yang et al. 2020) to describe the
performance of the forecasting results from the value and rank perspectives.
Furthermore, we employ two portfolio-based metrics to compare the investment
profit and risk of each method. We simulate daily trading using a simple
strategy that selects the top 30 stocks with the highest return ratio and
reports the Excess Annualized Return (AR) and Information Ratio (IR) metrics.
AR measures the annual expected excess return generated by the investment,
while IR measures the risk-adjusted performance of an investment.

Dataset | Model | IC | ICIR | RankIC | RankICIR | AR | IR
---|---|---|---|---|---|---|---
CSI300 | XGBoost | $0.051\pm 0.001$ | $0.37\pm 0.01$ | $0.050\pm 0.001$ | $0.36\pm 0.01$ | $0.23\pm 0.03$ | $1.9\pm 0.3$
LSTM | $0.049\pm 0.001$ | $0.41\pm 0.01$ | $0.051\pm 0.002$ | $0.41\pm 0.03$ | $0.20\pm 0.04$ | $2.0\pm 0.4$
GRU | $0.052\pm 0.004$ | $0.35\pm 0.04$ | $0.052\pm 0.005$ | $0.34\pm 0.04$ | $0.19\pm 0.04$ | $1.5\pm 0.3$
TCN | $0.050\pm 0.002$ | $0.33\pm 0.04$ | $0.049\pm 0.002$ | $0.31\pm 0.04$ | $0.18\pm 0.05$ | $1.4\pm 0.5$
Transformer | $0.047\pm 0.007$ | $0.39\pm 0.04$ | $0.051\pm 0.002$ | $0.42\pm 0.04$ | $0.22\pm 0.06$ | $2.0\pm 0.4$
GAT | $0.054\pm 0.002$ | $0.36\pm 0.02$ | $0.041\pm 0.002$ | $0.25\pm 0.02$ | $0.19\pm 0.03$ | $1.3\pm 0.3$
DTML | $0.049\pm 0.006$ | $0.33\pm 0.04$ | $0.052\pm 0.005$ | $0.33\pm 0.04$ | $0.21\pm 0.03$ | $1.7\pm 0.3$
MASTER | $0.064^{*}\pm 0.006$ | $0.42\pm 0.04$ | $0.076^{*}\pm 0.005$ | $0.49\pm 0.04$ | $0.27\pm 0.05$ | $2.4\pm 0.4$
CSI800 | XGBoost | $0.040\pm 0.000$ | $0.37\pm 0.01$ | $0.047\pm 0.000$ | $0.42\pm 0.01$ | $0.08\pm 0.02$ | $0.6\pm 0.2$
LSTM | $0.028\pm 0.002$ | $0.32\pm 0.02$ | $0.039\pm 0.002$ | $0.41\pm 0.03$ | $0.09\pm 0.02$ | $0.9\pm 0.2$
GRU | $0.039\pm 0.002$ | $0.36\pm 0.05$ | $0.044\pm 0.003$ | $0.39\pm 0.07$ | $0.07\pm 0.04$ | $0.6\pm 0.3$
TCN | $0.038\pm 0.002$ | $0.33\pm 0.04$ | $0.045\pm 0.002$ | $0.38\pm 0.05$ | $0.05\pm 0.04$ | $0.4\pm 0.3$
Transformer | $0.040\pm 0.003$ | $0.43\pm 0.03$ | $0.048\pm 0.003$ | $0.51\pm 0.05$ | $0.13\pm 0.04$ | $1.1\pm 0.3$
GAT | $0.043\pm 0.002$ | $0.39\pm 0.02$ | $0.042\pm 0.002$ | $0.35\pm 0.02$ | $0.10\pm 0.04$ | $0.7\pm 0.3$
DTML | $0.039\pm 0.004$ | $0.29\pm 0.03$ | $0.053\pm 0.008$ | $0.37\pm 0.06$ | $0.16\pm 0.03$ | $1.3\pm 0.2$
MASTER | $0.052^{*}\pm 0.006$ | $0.40\pm 0.06$ | $0.066\pm 0.007$ | $0.48\pm 0.06$ | $0.28^{*}\pm 0.02$ | $2.3^{*}\pm 0.3$
Table 1: Overall performance comparison. The best results are in bold and the second-best results are underlined. And * denotes statistically significant improvement (measured by t-test with p-value $<$ 0.01) over all baselines. Model | IC | ICIR | RankIC | RankICIR | AR | IR
---|---|---|---|---|---|---
(MA)STER | $0.064\pm 0.003$ | $0.43\pm 0.02$ | $0.074\pm 0.004$ | $0.48\pm 0.04$ | $0.25\pm 0.03$ | $2.1\pm 0.3$
(MA)STER-Bi | $0.058\pm 0.005$ | $0.38\pm 0.04$ | $0.066\pm 0.008$ | $0.41\pm 0.05$ | $0.19\pm 0.03$ | $1.6\pm 0.2$
Naive | $0.041\pm 0.008$ | $0.30\pm 0.05$ | $0.046\pm 0.007$ | $0.32\pm 0.04$ | $0.18\pm 0.05$ | $1.6\pm 0.6$
Clustering | $0.044\pm 0.003$ | $0.36\pm 0.02$ | $0.049\pm 0.005$ | $0.39\pm 0.04$ | $0.18\pm 0.04$ | $1.7\pm 0.3$
Table 2: Experiments on CSI300 to validate the effectiveness of proposed stock
transformer architecture. The best results are in bold and the second-best
results are underlined.

#### Implementation

We implemented MASTER222Code and supplementary materials are at
https://github.com/SJTU-Quant/MASTER with PyTorch and build our methods based
on the open-source quantitative investment platform Qlib (Yang et al. 2020).
For DTML, we implement it based on the original paper since there is no
official implementation publicly. For other baselines, we use their Qlib
implementations. For hyperparameters of each baseline method, the layer number
and model size are tuned from $\\{1,2,3\\}$ and $\\{128,256,512\\}$
respectively. The learning rate $lr$ is tuned among
$\\{10^{-i}\\}_{i\in\\{3,4,5,6\\}}$, and we selected the best hyperparameters
based on the IC performance in the validation stage. For hyperparameters of
MASTER, we tune the model size $D$ and learning rate $lr$ among the same range
as the baselines, and the final selection is $D$=$256$, $lr$=$10^{-5}$ for all
datasets; we set $N_{1}$=$4$, $N_{2}$=$2$ for all datasets and $\beta$=$5$ and
$\beta$=$2$ for CSI300 and CSI800 respectively. More implementation details of
baseline methods are summarized in the supplementary materials. Each model is
trained for at most $40$ epochs with early stopping. All the experiments are
conducted on a server equipped with Intel(R) Xeon(R) Platinum 8163 CPU, 128GB
Memory, and a Tesla V100-SXM2 GPU (16GB Memory). Each experiment was repeated
5 times with random initialization and the average performance was reported.

### Overall Performance (RQ1)

The overall performance are reported in Table 1 MASTER achieves the best
results on 6/8 of the ranking metrics, and consistently outperforms all
benchmarks in the portfolio-based metrics. In particular, MASTER achieve
$13\%$ improvements in ranking metrics and $47\%$ improvements in portfolio-
based metrics compared to the second-best results on the average sense. Note
that ranking matrics are computed with the whole set and portfolio-based
metrics mostly consider the 30 top-performed stocks. The achievements in both
types of metrics imply that MASTER is of good predicting ability on the whole
stock set without sacrificing the accuracy of the important stocks. The
significant improvements cast light on the importance of stock correlation
modeling, so each stock can also benefit from the historical signals of other
momentarily correlated stocks. We also observe all methods gain better
performance on CSI300 over CSI800. We believe it is because CSI300 consists of
companies with larger capitalization whose stock prices are more predictable.
When compared to the existing stock correlation method (i.e., DTML), MASTER
outperforms in all 6 metrics, which tells our proposed Market-Guided Gating
and aggregation techniques are more efficient in mining cross-stock
information than existing literature.

Figure 3: The average and standard deviation of metrics with different
$(N_{1},N_{2})$ combinations on CSI300. Figure 4: MASTER performance with
varying $\beta$. The horizontal dash lines are performance without market-
guided gating. Figure 5: The correlation towards three target stocks on Aug
19th, 2022. The y-axis is time steps in the lookback window and the x-axis is
source stocks. _Avg._ denotes the evenly distributed value.

### Stock Transformer Architecture (RQ2)

We validate the effectiveness of our specialized stock transformer
architecture by experiments on four settings. (1) (MA)STER, which is our stock
transformer without the gating. (2) (MA)STER-Bi, in which we substitute the
single-layer transformer encoder with a bi-directional LSTM to evince that the
effectiveness of our proposed architecture is not coupled with strong
sequential encoders. (3) Naive, which directly performs information
aggregation among $\tau\times|\mathcal{S}|$ tokens. (4) Clustering, in which
we adapt the Local Sensitive Hashing (Kitaev, Kaiser, and Levskaya 2020) to
allocate all tokens into $10$ buckets by similarity and perform aggregation
within each group, which is a classic task-agnostic technique to reduce the
scale of the attention field. For a fair comparison, in (3) and (4), we first
use the same transformer encoder to extract token embedding and then use the
same multi-head attention mechanism as in our stock transformer, so the only
difference is the attention field. Due to resource limits, we only conduct
experiment on CSI300 dataset. The results in Table 2 illustrate the efficacy
of our tailored stock transformer architecture, which performs intra-stock
aggregation and inter-stock aggregation alternatively.

### Ablation Study (RQ3)

First, we conduct ablation study on $(N_{1},N_{2})$ combination. The results
of CSI300 are shown in Figure 3 and the results on CSI800 are similar. The
difference among head combinations is not significant compared with the
inherent variance under each setting. In the studied range, most settings
consistently performed better than the baselines.

Second, we study the influence of temperature $\beta$ in the gating mechanism.
As explained before, a smaller $\beta$ forces a stronger feature selection
while a larger $\beta$ turns off the gating effect. Figure 4 shows the
performance with varying $\beta$. The CSI300 is a relatively easier dataset
where most features are quite effective, so the temperature is expected to be
larger to relax the feature selection, while more powerful feature selection
intervention is needed for the sophisticated CSI800 dataset whose $\beta$ of
the best performance is smaller.

### Visualization of Attention Maps (RQ4)

We show how MASTER captures the momentary and cross-time stock correlation
that previous methods are not expressive enough to model. Figure 5 shows the
inter-stock attention map at different time steps in the lookback window. We
choose three representative stocks as the target and sample $100$ random
stocks as sources for visualization. The highlighted part is scattered instead
of exhibiting neat strips, implying that the correlation is momentary rather
than long-standing. Also, the inter-stock correlation is sparse, with only a
few stocks having strong correlations toward the target stocks. Figure 6
displays the correlation between stock pairs to show how the correlation
resides in time. From source stock $v$ to target stock $u$, we compute
$I_{u\leftarrow v}[i,j]=\mathrm{S}^{1}_{v}[i,j]\mathrm{S}^{2}_{i}[u,v]$ as the
$\tau\times\tau$ correlation map, while $\mathrm{S}^{1}$ and $\mathrm{S}^{2}$
are the intra-stock and inter-stock attention map. First, the highlighted
blocks are not centered on the diagonal, because the stock correlation is
usually cross-time rather than temporally aligned. Second, the left two
figures are totally different, illustrating that correlation is highly
asymmetric between $u\leftarrow v$ and $v\leftarrow u$. Third, the importance
of mined correlation changes slowly when the lookback window slides to
forecast on different dates. For example, blocked regions in the right two
figures correspond to the same absolute time scope of different prediction
dates, whose patterns are to a certain degree similar.

Figure 6: Cross-time correlation of stock pairs on Aug 19th and 25th, 2022.
The x-axis is the source time steps and the y-axis is the target time steps.

## Conclusion

We introduce a novel method MASTER for stock price forecasting, which models
the realistic stock correlation and guide feature selection with market
information. MASTER consists of five steps, market-guided gating, intra-stock
aggregation, inter-stock aggregation, temporal aggregation, and prediction.
Experiments on the Chinese market with $2$ stock universe shows that MASTER
achieves averagely $13\%$ improvements on ranking metrics and $47\%$ on
portfolio-based metrics compared with all baselines. Visualization of
attention maps reveals the de-facto momentary and cross-time stock
correlation. In conclusion, we provide a more granular perspective for
studying stock correlation, while also indicating an effective application of
market information. Future work can explore to mine stock correlations of
higher quality and study other uses of market information.

## Acknowledgements

The authors would like to thank the anonymous reviewers for their insightful
reviews. This work is supported by the National Key Research and Development
Program of China (2022YFE0200500), Shanghai Municipal Science and Technology
Major Project (2021SHZDZX0102), and SJTU Global Strategic Partnership Fund
(2021SJTU-HKUST).

## References

  * Bai, Kolter, and Koltun (2018) Bai, S.; Kolter, J. Z.; and Koltun, V. 2018.  An empirical evaluation of generic convolutional and recurrent networks for sequence modeling.  _arXiv preprint arXiv:1803.01271_.
  * Bennett, Cucuringu, and Reinert (2022) Bennett, S.; Cucuringu, M.; and Reinert, G. 2022.  Lead–lag detection and network clustering for multivariate time series with an application to the US equity market.  _Machine Learning_ , 111(12): 4497–4538.
  * Bulat et al. (2021) Bulat, A.; Perez Rua, J. M.; Sudhakaran, S.; Martinez, B.; and Tzimiropoulos, G. 2021.  Space-time mixing attention for video transformer.  _Advances in neural information processing systems_ , 34: 19594–19607.
  * Chen and Guestrin (2016) Chen, T.; and Guestrin, C. 2016.  Xgboost: A scalable tree boosting system.  In _Proceedings of the 22nd acm sigkdd international conference on knowledge discovery and data mining_ , 785–794.
  * Cho et al. (2014) Cho, K.; Van Merriënboer, B.; Gulcehre, C.; Bahdanau, D.; Bougares, F.; Schwenk, H.; and Bengio, Y. 2014.  Learning phrase representations using RNN encoder-decoder for statistical machine translation.  _arXiv preprint arXiv:1406.1078_.
  * Cong et al. (2021) Cong, Y.; Liao, W.; Ackermann, H.; Rosenhahn, B.; and Yang, M. Y. 2021.  Spatial-temporal transformer for dynamic scene graph generation.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , 16372–16382.
  * Ding et al. (2020) Ding, Q.; Wu, S.; Sun, H.; Guo, J.; and Guo, J. 2020.  Hierarchical Multi-Scale Gaussian Transformer for Stock Movement Prediction.  In _IJCAI_ , 4640–4646.
  * Feng et al. (2018) Feng, F.; Chen, H.; He, X.; Ding, J.; Sun, M.; and Chua, T.-S. 2018.  Enhancing stock movement prediction with adversarial training.  _arXiv preprint arXiv:1810.09936_.
  * Feng et al. (2019) Feng, F.; He, X.; Wang, X.; Luo, C.; Liu, Y.; and Chua, T.-S. 2019.  Temporal relational ranking for stock prediction.  _ACM Transactions on Information Systems (TOIS)_ , 37(2): 1–30.
  * Graves and Graves (2012) Graves, A.; and Graves, A. 2012.  Long short-term memory.  _Supervised sequence labelling with recurrent neural networks_ , 37–45.
  * Huynh et al. (2023) Huynh, T. T.; Nguyen, M. H.; Nguyen, T. T.; Nguyen, P. L.; Weidlich, M.; Nguyen, Q. V. H.; and Aberer, K. 2023.  Efficient integration of multi-order dynamics and internal dynamics in stock movement prediction.  In _Proceedings of the Sixteenth ACM International Conference on Web Search and Data Mining_ , 850–858.
  * Kamble (2017) Kamble, R. A. 2017.  Short and long term stock trend prediction using decision tree.  In _2017 International Conference on Intelligent Computing and Control Systems (ICICCS)_ , 1371–1375. IEEE.
  * Kitaev, Kaiser, and Levskaya (2020) Kitaev, N.; Kaiser, Ł.; and Levskaya, A. 2020.  Reformer: The efficient transformer.  _arXiv preprint arXiv:2001.04451_.
  * Li et al. (2023) Li, L.; Duan, L.; Wang, J.; He, C.; Chen, Z.; Xie, G.; Deng, S.; and Luo, Z. 2023.  Memory-Enhanced Transformer for Representation Learning on Temporal Heterogeneous Graphs.  _Data Science and Engineering_ , 8(2): 98–111.
  * Li et al. (2015) Li, Q.; Jiang, L.; Li, P.; and Chen, H. 2015.  Tensor-based learning for predicting stock movements.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 29.
  * Liu et al. (2019) Liu, J.; Lin, H.; Liu, X.; Xu, B.; Ren, Y.; Diao, Y.; and Yang, L. 2019.  Transformer-based capsule network for stock movement prediction.  In _Proceedings of the first workshop on financial technology and natural language processing_ , 66–73.
  * Nie et al. (2022) Nie, Y.; Nguyen, N. H.; Sinthong, P.; and Kalagnanam, J. 2022.  A Time Series is Worth 64 Words: Long-term Forecasting with Transformers.  In _The Eleventh International Conference on Learning Representations_.
  * Nugroho, Adji, and Fauziati (2014) Nugroho, F. S. D.; Adji, T. B.; and Fauziati, S. 2014.  Decision support system for stock trading using multiple indicators decision tree.  In _2014 The 1st International Conference on Information Technology, Computer, and Electrical Engineering_ , 291–296. IEEE.
  * Piccolo (1990) Piccolo, D. 1990.  A distance measure for classifying ARIMA models.  _Journal of time series analysis_ , 11(2): 153–164.
  * Sawhney et al. (2021) Sawhney, R.; Agarwal, S.; Wadhwa, A.; Derr, T.; and Shah, R. R. 2021.  Stock selection via spatiotemporal hypergraph attention network: A learning to rank approach.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 35, 497–504.
  * Sawhney et al. (2020) Sawhney, R.; Agarwal, S.; Wadhwa, A.; and Shah, R. R. 2020.  Spatiotemporal hypergraph convolution network for stock movement forecasting.  In _2020 IEEE International Conference on Data Mining (ICDM)_ , 482–491. IEEE.
  * Vaswani et al. (2017) Vaswani, A.; Shazeer, N.; Parmar, N.; Uszkoreit, J.; Jones, L.; Gomez, A. N.; Kaiser, Ł.; and Polosukhin, I. 2017.  Attention is all you need.  _Advances in neural information processing systems_ , 30.
  * Veličković et al. (2017) Veličković, P.; Cucurull, G.; Casanova, A.; Romero, A.; Lio, P.; and Bengio, Y. 2017.  Graph attention networks.  _arXiv preprint arXiv:1710.10903_.
  * Wang et al. (2021) Wang, H.; Li, S.; Wang, T.; and Zheng, J. 2021.  Hierarchical Adaptive Temporal-Relational Modeling for Stock Trend Prediction.  In _IJCAI_ , 3691–3698.
  * Wang et al. (2022) Wang, H.; Wang, T.; Li, S.; Zheng, J.; Guan, S.; and Chen, W. 2022.  Adaptive long-short pattern transformer for stock investment selection.  In _Proceedings of the Thirty-First International Joint Conference on Artificial Intelligence_ , 3970–3977.
  * Wang, Qu, and Chen (2022) Wang, Y.; Qu, Y.; and Chen, Z. 2022.  Review of graph construction and graph learning in stock price prediction.  _Procedia Computer Science_ , 214: 771–778.
  * Xiang et al. (2022) Xiang, S.; Cheng, D.; Shang, C.; Zhang, Y.; and Liang, Y. 2022.  Temporal and Heterogeneous Graph Neural Network for Financial Time Series Prediction.  In _Proceedings of the 31st ACM International Conference on Information & Knowledge Management_, 3584–3593.
  * Xie et al. (2013) Xie, B.; Passonneau, R.; Wu, L.; and Creamer, G. G. 2013.  Semantic frames to predict stock price movement.  In _Proceedings of the 51st annual meeting of the association for computational linguistics_ , 873–883.
  * Xu et al. (2020) Xu, M.; Dai, W.; Liu, C.; Gao, X.; Lin, W.; Qi, G.-J.; and Xiong, H. 2020.  Spatial-temporal transformer networks for traffic flow forecasting.  _arXiv preprint arXiv:2001.02908_.
  * Xu et al. (2021) Xu, W.; Liu, W.; Wang, L.; Xia, Y.; Bian, J.; Yin, J.; and Liu, T.-Y. 2021.  Hist: A graph-based framework for stock trend forecasting via mining concept-oriented shared information.  _arXiv preprint arXiv:2110.13716_.
  * Yang et al. (2020) Yang, X.; Liu, W.; Zhou, D.; Bian, J.; and Liu, T.-Y. 2020.  Qlib: An ai-oriented quantitative investment platform.  _arXiv preprint arXiv:2009.11189_.
  * Yoo et al. (2021) Yoo, J.; Soun, Y.; Park, Y.-c.; and Kang, U. 2021.  Accurate multivariate stock movement prediction via data-axis transformer with multi-level contexts.  In _Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining_, 2037–2045.
  * Zhang and Yan (2022) Zhang, Y.; and Yan, J. 2022.  Crossformer: Transformer utilizing cross-dimension dependency for multivariate time series forecasting.  In _The Eleventh International Conference on Learning Representations_.