# Brain-like representational straightening of natural movies in robust
feedforward neural networks

Tahereh Toosi & Elias B. Issa
Department of Neuroscience, Zuckerman Mind Brain Behavior Institute
Columbia University
New York, NY, USA
<EMAIL_ADDRESS>

###### Abstract

Representational straightening refers to a decrease in curvature of visual
feature representations of a sequence of frames taken from natural movies.
Prior work established straightening in neural representations of the primate
primary visual cortex (V1) and perceptual straightening in human behavior as a
hallmark of biological vision in contrast to artificial feedforward neural
networks which did not demonstrate this phenomenon as they were not explicitly
optimized to produce temporally predictable movie representations. Here, we
show robustness to noise in the input image can produce representational
straightening in feedforward neural networks. Both adversarial training (AT)
and base classifiers for Random Smoothing (RS) induced remarkably straightened
feature codes. Demonstrating their utility within the domain of natural
movies, these codes could be inverted to generate intervening movie frames by
linear interpolation in the feature space even though they were not trained on
these trajectories. Demonstrating their biological utility, we found that AT
and RS training improved predictions of neural data in primate V1 over
baseline models providing a parsimonious, bio-plausible mechanism – noise in
the sensory input stages – for generating representations in the early visual
cortex. Finally, we compared the geometric properties of frame representations
in these networks to better understand how they produced representations that
mimicked the straightening phenomenon from biology. Overall, this work
elucidating emergent properties of robust neural networks demonstrates that it
is not necessary to utilize predictive objectives or train directly on natural
movie statistics to achieve models supporting straightened movie
representations similar to human perception that also predict V1 neural
responses.

##  1 Introduction

In understanding the principles underlying biological vision, a longstanding
debate in computational neuroscience is whether the brain is wired to predict
the incoming sensory stimulus, most notably formalized in predictive coding
(Rao & Ballard, 1999; Friston, 2009; Millidge et al., 2021), or whether neural
circuitry is wired to recognize or discriminate among patterns formed on the
sensory epithelium, popularly exemplified by discriminatively trained
feedforward neural networks (DiCarlo et al., 2012; Tacchetti et al., 2018;
Kubilius et al., 2018). Arguing for a role of prediction in vision, recent
work found perceptual straightening of natural movie sequences in human visual
perception (Hénaff et al., 2019). Such straightening is diagnostic of a system
whose representation could be linearly read out to perform prediction over
time, and the idea of representational straightening resonates with machine
learning efforts to create new types of models that achieve equivariant,
linear codes for natural movie sequences. Discriminatively trained networks,
however, lack any prediction over time in their supervision. It may not be
surprising then that large-scale ANNs trained for classification produce
representations that have almost no improvement in straightening relative to
the input pixel space, while human observers clearly demonstrated perceptual
straightening of natural movie sequences (subsequently also found in neurons
of primary visual cortex, V1 (Hénaff et al., 2019; 2021)). This deficiency in
standard feedforward ANNs might suggest a need for new models trained on
predictive loss functions rather than pure classification to emulate
biological vision.

Here, we provide evidence for an alternative viewpoint, that biologically
plausible straightening can be achieved in ANNs trained for robust
discrimination, without resorting to a prediction objective or natural movies
in training. Drawing on insights from emergent properties of adversarially-
trained neural networks in producing linearly invertible latent
representations, we highlight the link between perceptual straightening of
natural movies to invertible latent representations learned from static images
(Figure 1). We examine straightening in these robust feedforward ANNs finding
that their properties relate to those in the biological vision framework. The
contributions of this work are as follows:

  * 1.

We show that robust neural networks give rise to straightened feature
representations for natural movies in their feature space, comparable to the
straightening measured in the primate brain and human behavior, and completely
absent from standard feedforward networks.

  * 2.

We show that linearly interpolating between the start and end frames of a
movie in the output feature space of robust ANNs produces synthetic frames
similar to those of the original natural movie sequence in image space. Such
invertible linear interpolation is precisely the definition of a temporally
predictive feature representation.

  * 3.

Compared to prior models of early visual cortex, robustness to input noise
(corruption or adversarial robustness) is significantly better at explaining
neural variance measured from V1 neurons than non-robustly trained baseline
models, suggesting a new hitherto unconsidered mechanism for learning the
representations in early cortical areas that achieves natural movie
straightening.

Figure 1: Perceptual straightening of movie frames can be viewed as
invertibility of latent representations for static images. Left: straightening
of representations refers to a decrease in the curvature of the trajectory in
representation space such as a neural population in the brain or human
perceptual space, but standard ANNs do not show straightening (Hénaff et al.,
2019; 2021). Right: Invertibility of latent representation refers to
interpolation between the representation of two images (e.g. an image of a dog
and an image of a cat), where the invertible interpolations show the main
features of a dog morph into the main features of a cat. Invertible
representations emerge in robust ANNs (Engstrom et al., 2019b), obviating the
need to directly train for temporal straightening.

##  2 Related work

###  2.1 Mechanisms for producing brain-like representations

Feedforward ANNs as models of biological vision. Standard feedforward ANNs,
although lacking a number of bio-plausible features such as feedback
connections or a local learning rule (Whittington & Bogacz, 2019), still can
explain the neural variance (Schrimpf et al., 2018) recorded from rodent
(Bakhtiari et al., 2021), monkey (Yamins et al., 2014; Bashivan et al., 2019),
and human visual cortex (Khaligh-Razavi & Kriegeskorte, 2014; Cichy et al.,
2016) better than alternatives which are considered more bio-plausible by
using a prediction objective function (e.g., PredNet and CPC (Zhuang et al.,
2021; Schrimpf et al., 2020)). Thus, to learn the representations in the
brain, regardless of the bio-plausibility of mechanisms, feedforward ANNs
provide a parsimonious more tractable class of leading models for object
recognition in the visual cortex.

Models of primary visual cortex. In neuroscience, rather than rely solely on
top-down training objectives like standard ANNs do, there has been a tradition
of explaining early visual representations using more fundamental principles
such as sparse coding and predictive coding as well as invoking unsupervised
training (Olshausen & Field, 1996; Rao & Ballard, 1999). For example,
unsupervised slow feature analysis extracts the slow-varying features from
fast-varying signals in movies based on the intuition that most external
salient events (such as objects) are persistent in time, and this idea can be
used to explain the emergence of complex cells in V1 (Berkes & Wiskott, 2005).
Recent work in machine learning has attempted to blend more bottom-up
principles with top-down training by experimenting with swapping out ANN early
layers with V1-like models whose filters are inspired by neuroscience studies
(Dapello et al., 2020). This blended model turns out to have benefits for
classification robustness in the outputs. However, it remains unclear whether
there is a form of top-down training that can produce V1-like models. Such a
mechanism would provide a fundamentally different alternative to prior
proposals of creating a V1 through sparse coding or future prediction (Hénaff
et al., 2019; 2021).

###  2.2 Temporal prediction and invertibility in neural networks

Learning to predict over time. Changes in architecture, training diet
(movies), and objective (predicting future frames) have all been explored as
mechanisms to produce more explicit equivariant representations of natural
movies (Lotter et al., 2016; van den Oord et al., 2018). Directly related to
the idea of straightening, penalizing the curvature of representations of
frames was used in Learning to linearize (Goroshin et al., 2015) to learn
straightened representations from unlabeled videos. This class of models does
not need supervision which makes them more bio-plausible in nature; however,
as mentioned in the previous section, they lag behind supervised feedforward
ANNs both in terms of learning effective representations for object
recognition and in producing feature representations that predict neural data.

Learning invertible latents. In deep learning applications, invertibility is
mostly discussed in generative neural networks as a constraint to learn a
prior to address applications in signals and systems such as image de-noising,
signal compression, and image reconstruction from few and noisy measurements
or to be able to reconstruct or modify real images. Usually, invertibility is
implemented by carefully designing dedicated architectures (Jacobsen et al.,
2018b; Chen et al., 2019). However, recently it has been shown it can be
implemented in standard feedforward ANNs when they undergo training for
adversarial robustness (Engstrom et al., 2019b; c). These works showed
empirically that adversarially robust training encourages invertibility as
linear interpolation between classes (e.g., cat to dog) results in
semantically smooth image-to-image translation (Engstrom et al., 2019b) as
opposed to blurry image sequences produced by standard ANNs.

We reasoned that robust networks that encourage invertibility may also lead to
straightening as this is a property that would be related to improved
invertibility of a network, so we sought to extend prior work and study the
behavior of robustly trained networks specifically in the domain of natural
movies. We report on how these networks straighten natural movies in their
feature spaces and can invertibly reproduce movie frames in a natural
sequence.

Figure 2: ANNs show straightening of representations when robustness to noise
constraints (noise augmentation or adversarial attack) is added to their
training. Measurements for straightening of movie sequences (from (Hénaff et
al., 2019), in each layer of ResNet50 architecture under different training
regimes: supervised training (standard), no training (random parameters),
self-supervised training (Zbontar et al., 2021), supervised training with no
augmentations, supervised training with extensive augmentations, supervised
training with noise augmentation (base classifiers for RS) (Cohen et al.,
2019), and supervised training with adversarial training (Engstrom et al.,
2019a)

##  3 Methods

###  3.1 Baseline models

We consider the class of feedforward convolutional neural networks, typically
restricting to the ResNet-50 (He et al., 2015) architecture trained on
ImageNet for the main analyses. Baseline networks (not trained for robustness)
include a supervised ResNet-50/ResNet-101/ResNet-152, and self-supervised
(Barlowtwins (Zbontar et al., 2021)). We trained ResNet-50 for imagenet
classification without augmentations and with extensive augmentations (Chen et
al., 2020), labeled as SupNoAugm and SupMocoAugm, respectively. We also
consider Voneresnet (biological V1 front-end (Dapello et al., 2020)) and
ResNet-50 trained as a base network for action recognition (Chen et al., 2021)
but include these as separate examples in the Appendix since they use a
modified architecture.

Table 1: Clean accuracy and robust (attack: $L_{2},\epsilon=0.1$) accuracy for
the models used. Except for the custom models, all the other models were
obtained from the repository of the references. Note that RS here refers to
the base classifier in random smoothing without probabilistic inference.
Models | Clean accuracy | Robust accuracy | Model reference
---|---|---|---
RN50 AT $L_{2}:\epsilon=3$ | 58.50 | 57.81 | (Engstrom et al., 2019a)
RN50 AT $L_{\infty}:\epsilon=4$ | 62.80 | 61.40 | (Engstrom et al., 2019a)
RN50 AT $L_{\infty}:\epsilon=8$ | 48.29 | 47.01 | (Engstrom et al., 2019a)
RN50 RS $L_{2}:\epsilon=0.25$ | 39.40 | 36.01 | (Cohen et al., 2019)
RN50 RS $L_{2}:\epsilon=0.5$ | 23.75 | 22.21 | (Cohen et al., 2019)
RN50 RS $L_{2}:\epsilon=1$ | 10.62 | 10.17 | (Cohen et al., 2019)
RN50 Standard | 75.43 | 52.32 | (He et al., 2015)
RN50 No augmentation | 64.35 | 28.13 | custom
RN50 Extensive augmentation | 75.27 | 53.08 | custom
RN50 Self-supervised | 70.18 | 41.73 | (Zbontar et al., 2021)

###  3.2 Models trained for robustness

We consider two forms of models trained for minimizing a classification loss
$\mathcal{L}_{ce}$ in the face of input perturbations
$\delta\in\mathbb{R}^{h\times w\times c}$ subject to constraints on the
overall magnitude of perturbations in the input space, where $x$, $y$,
$\theta$ are the network input, output, and classifier parameters,
respectively:

$\begin{array}[]{c}\mathcal{L}_{ce}(\theta,x+\delta,y)\\\ \end{array}$ (1)

In adversarially trained networks, projected gradient descent from the output
space finds maximal directions of perturbation in the input space limited to
length $\epsilon$, and training entails minimizing the effect of these
perturbation directions on the network’s output (Madry et al., 2018). In
random smoothing (Lecuyer et al., 2018; Cohen et al., 2019), a supervised
network is trained but in the face of Gaussian noise added to the input space
as the base classifier before performing a probabilistic inference. In this
work, we only use the representations as learned in base classifiers without
the probabilistic inference. The perturbations in the base classifiers
$\delta$ thus can follow:

$\begin{array}[]{c}\delta_{rand}\sim\mathcal{N}(0,\sigma^{2}I)\text{,}\ \ \ \
\ \ \
\delta_{adv}:=\mathop{\arg\max}\limits_{|\delta|_{p}\leq\epsilon}\mathcal{L}_{ce}(\theta,x+\delta,y)\\\
\end{array}$ (2)

These defenses to input noise have different motivations. Adversarial
robustness provides defense against white box attacks whereas random smoothing
protects against general image corruptions. However, prior work has suggested
a connection between corruption robustness and adversarial robustness (Ford et
al., 2019). Theoretically, random smoothing leads to certified robustness
(Cohen et al., 2019) and trains a condition of invertible networks (Jacobsen
et al., 2018a), while adversarial robustness has been shown empirically to
lead to invertible latent representations in networks (Engstrom et al.,
2019b).

Figure 3: Invertibility as measured by the SSIM (Wang et al., 2004) of the
actual in-between frames (labeled as Natural Movie) and the pixel-space
projected linear interpolations between the first and the last frame labeled
Pixels, standard ResNet50, RS (ResNet50, $L_{2}:\sigma^{2}=0.5$) and AT
(ResNet50, $L_{2}:\sigma^{2}=3$). Interpolating the representations of the
first and last frames in an invertible representation space generates a
sequence of frames that are more similar to the ground-truth in-between
frames, but for a non-invertible representation, the generated frames are
blurry and more similar to the interpolation in pixel space (a.k.a. artificial
sequence). Interpolating the first and last frames in pixel space, the second
row gives exactly what was called an artificial sequence in studies of
straightening (Hénaff et al., 2019; 2021), as opposed to the natural sequence
which was the actual in-between frames.

###  3.3 Representational Metrics

Representational straightening estimates the local curvature $c$ in a given
representation $r$ of a sequence of images (natural or artificial) of length
$N$, $C_{seq}:\\{x_{t_{1}},x_{t_{2}},...,x_{t_{N}}\\}$ as the angle between
vectors connecting nearby frames, and these local estimates are averaged over
the entire movie sequence for the overall straightening in that
representational trajectory (same as (Hénaff et al., 2019)):

$\displaystyle
c_{t}=\arccos{\left({\frac{{r_{t}-r_{t-1}}}{{\left\|{r_{t}-r_{t-1}}\right\|}}}\cdot{\frac{{r_{t+1}-r_{t}}}{{\left\|{r_{t+1}-r_{t}}\right\|}}}\right)}\text{,}\
\ \ \ \ \ C_{seq}=\frac{1}{N}\sum_{t=1}^{N-1}c_{t}$ (3)

Lower curvature (angle between neighboring vectors) indicates a straighter
trajectory, and in the results, we generally reference curvature values to the
curvature in the input space (i.e., straightening relative to pixel space).
This metric has been utilized in neuroscience showing that humans tend to
represent nearby movie frames in a straightened manner relative to pixels
(Hénaff et al., 2019). This curvature metric is also closely related to
objectives used in efforts to train models with equivariance by linearizing
natural transformations in the world as an alternative to standard networks
trained for invariant object classification (Goroshin et al., 2015; Sabour et
al., 2017).

Expansion. We define the radius of a sequence of images from a movie clip as
the radial size of the minimum covering hyper-sphere circumscribing all points
representing the frames in $r$ (Gärtner, 1999). We use this measure to
supplement the geometrical characterization of a movie sequence in pixel space
and in a model’s representational spaces. Like representational straightening
values, expansion values for models in the main text are referenced to the
radius measured in pixel space or to the radius measured for the same layer in
a baseline network by simply dividing by those references. We used mini-ball,
a publicly available python package based on (Gärtner, 1999) to measure the
radius of the covering hyper-sphere.

Figure 4: Left: RS and AT are more predictive of V1 neural responses than
other non-robust models of the same architecture (ResNet50). Right: Each dot
represents a layer in ResNet50 trained under different loss function (color
codes same as left). Higher representational straightening (negative curvature
change) associates with higher V1 predictivity. Intriguingly, the highest V1
predictivity corresponds to layers that exhibit comparable straightening to
that measured from V1 neurons ($-10^{\circ}$ on average) (Hénaff et al.,
2021). Explained variance is noise-corrected and computed as in (Schrimpf et
al., 2018)

.

##  4 Results

###  4.1 Robust ANNs exhibit representational straightening

With insights from connections to invertibility (see Figure 1), we
hypothesized representational straightening of movie trajectories could be
present in robustly trained neural networks. We took the same movie stimuli
publicly available (Hénaff et al., 2019)(A.4.1, Figure 12) and the same
metrics, and we tested the same architecture, ResNet50 (He et al., 2015))
trained under different loss functions Table 1 to perform controlled head-to-
head comparisons. Figure 2 shows representational straightening of natural
movies measured in layers of ResNet50 trained under AT (Engstrom et al.,
2019a) and RS (Cohen et al., 2019) at different adversarial attack or noise
levels, respectively. Robust neural networks in contrast to other ANNs
decreased the curvature of natural movies. Straightening for artificial
sequences as measured in (Hénaff et al., 2019) (A.1, Figure 7) and other
models (A.2, Figures 9 and 8) are provided in Appendix. Importantly, although
most models, whether a standard ResNet-50 or one with a V1-like front-end, may
display an initial dip in curvature for natural movies in the very earliest
layers, this is not sustained in feature representations of later layers
except for robustly trained networks (A.2, Figure 9 vs. A.1, Figure 7) and
those trained on action recognition from temporally instructed training, which
we include here as a proxy for a movie-like training though its feedforward
architecture deviates from a ResNet50 by additional temporal processing
components (A.2, Figure 8).

Perceptual Straightening measured as invertibility of latent representations.
Next, we sought to empirically test how well robust networks can invert
natural movies given that they contain linearized feature representation of
movie frames in their high level feature spaces and given the general
conceptual benefit of linearity for invertibility Figure 1. We measured the
invertibility of each model on the same movie sequences used for measuring
straightening as follows. We linearly interpolated between latent
representations of the first and last frame of each movie and used the same
procedure as that used previously in (Engstrom et al., 2019b; a) to obtain the
pixel-space correspondence of those interpolated representations. Whereas
those generated pseudo-frames can be assessed by for their pixel-by-pixel
distance to the actual movie frame, we chose a metric, Structural Similarity
Index Measure (SSIM (Wang et al., 2004)), that utilizes intermediate-level
statistics motivated from biological vision and putatively more related to
some aspects of human perception than simple pixel space correspondence.
Figure 3 shows an example of such inverted frames for standard ResNet50, RS
($L_{2}:\sigma^{2}=0.5$) and AT ($L_{2}:\sigma^{2}=3$), and a summary of
average measured invertibility using the SSIM metric on pseudo-frames from
each model. As expected, in line with the findings of previous work (Engstrom
et al., 2019b), AT models scored relatively higher on the invertibility of
frames than a baseline discriminative model. However, what had not been
previously shown is that RS models, using merely the benefits of their
robustness to noisy augmentation (base classifier on top of learned
representation; no probabilistic inference), also exhibit higher invertibility
scores compared to standard trained models. Invertibility scores were
consistently improved in RS and AT models across a variety of movies tested
including those with relatively stationary textures and not just dynamic
objects (see A.4.4, Figure 13 for further examples and A.4.3, Table 3 for
scores across all 11 movies). Thus, RS models along with AT models exhibit
invertibility of representations for movie frames which further demonstrates
their ability to support perceptual straightening of natural movies in their
highest layers that may be functionally similar to perceptual straightening
previously measured from human subjects (Hénaff et al., 2019).

###  4.2 Random smoothing and adversarial training in explaining neural
representations in the primate visual system

Robustness to noise as a bio-plausible mechanism underlying straightening in
primary visual cortex. As shown above, straightening which is a constraint for
brain-like representations in the visual cortex manifests in robust neural
networks. Both classes of RS and AT training for robustness to $L_{2}$ norm
generate straightened representations of movie sequences. However, to
distinguish among models of object recognition, we can measure how well they
explain variance in patterns of neural activity elicited in different visual
cortical areas. Here, for all neural comparisons in our analyses, we measured
the Brain-Score (Schrimpf et al., 2018) using the publicly available online
resource to assess the similarity to the biological vision of each model,
which is a battery of tests comparing models against previously collected data
from the primate visual system (see Brain-Score.org). We found that RS and AT
models provided a better model of V1 (in terms of explained variance) compared
to non-robust models Figure 4. On other benchmarks, as we go up the ventral
stream hierarchy from V1 to IT again, keeping the layer assignment fixed
across models for proper comparison, we observed a decrease in explainability
of robust models (A.3, Figure 11), in part presumably because robust models
have lower object classification performance which is known to drive fits in
higher brain areas like V4 and IT supporting object recognition (Yamins et
al., 2014). Previous work (Dapello et al., 2020; Kong et al., 2022) linked
adversarial robustness in models to their higher Brain-Score for V1, but we
found that it may not be specifically driven by adversarial robustness per se,
rather ($L_{2}$) noise robustness is also sufficient (as in base classifiers
of RS tested here). More broadly, looking at neural fits across all models and
their layers, we find that straightening in a particular model-layer
correlates with improved explanatory power of variance in cortical area V1
(Figure 4, middle panel, each dot is a layer from a model), being even more
strongly predictive than robustness of the overall model (A3, Figure 10). The
level of straightening reached by best fitting layers of RS and AT models was
comparable to the 10 degree straightening estimated in macaque V1 neural
populations (black dashed reference line in Figure 4). This complements the
fact that robust models peak near the 30 degree straightening measured in
perception (Figure 2), suggesting that robust models can achieve a brain-like
level of straightening to V1 and perception.

Figure 5: Can straightening for a movie sequence be explained by the size of
the hyper-sphere bounding the frames (i.e. radius in pixel space)? While RS
exhibits a small but positive correlation, the rest of the models, including
AT, show negative or no correlations. A positive correlation means the smaller
the size of the bounding hyper-sphere in pixel space, the more straightened
the representation over the layers of the model.

Does the geometry of movie frame representations in pixel space dictate
straightening in downstream representations? The connection between two
properties of the same representation manifold, robustness to independently
sampled noise, and straightened trajectories of smooth input temporal
sequences, is not immediately clear. Because robustness is achieved by adding
noise bounded by a norm ($L_{2}$, $L_{2}$, or $L_{\infty}$) in pixel space, a
natural question is whether the radius of the bounding hyper-sphere of the
frames of the tested movies in pixel space (see _Expansion_ in Methods) was
correlated with the measured straightening in feature space in each layer of
the robustly trained models (Figure 5; also see A.5, Figure 14). We found,
however, that there seemed to be different mechanisms at play for RS versus AT
in terms of achieving straightening. RS models showed (small but) positive
correlations, which means the smaller the ball containing all the frames of
the movie in input space, the larger the straightening effect for the
representations of frames of that movie in the model. While in AT models we
see the opposite (negative) or no correlation. These divergent patterns
underscore differences between these models and suggest that geometric size in
pixel space is not strongly constraining the degree to which a movie can be
straightened.

Geometry of movie frame representations in feature space is relevant for
capturing neural representations in V1 Between different RS models tested on
different input noise levels, RS $L_{2}:\sigma^{2}=0.5$ stands out as it gives
a better model of V1 than those using smaller or larger magnitude input noise
(Figure 4). For this model, we found that in addition to its intermediate
level of straightening, the expansion score of movie frames, which is the
radial size in its representation normalized to size in the same layer of a
baseline ResNet50, was highest compared to the other RS models (Figure 6,
middle panel; measures are referenced to layers in a standard ResNet50 to
highlight relative effect of robustness training rather than effects driven by
hierarchical layer). This demonstrates a potential trade-off between improving
straightening in a representation while avoiding too much added contraction of
movies by robust training relative to standard training. This balance seems to
be best achieved for $\sigma^{2}=0.5$, where we also see the significantly
higher predictivity of V1 cortical data (Figure 6, right panel). The best AT
model also shows little contraction of movies coupled with high straightening
(A.5, 15).

Figure 6: Geometric characteristics, straightening, and curvature, of RS
models related to V1 explainability. $\Delta$ means quantity is referenced to
the same measure in a standard ResNet50.

##  5 Discussion

We have demonstrated novel properties of robust neural networks in how they
represent natural movies. Conceptually, this work establishes a seemingly
surprising connection between disparate ideas, robust discriminative networks
trained on static images, on one hand, to work learning to linearize by
training on natural movies, on the other. These modeling paths could both
result in linearized or straightened, natural movie representations (Figure
1). From a machine learning perspective, the invertibility and concomitant
representational straightening of robust networks suggest that they learn
explainable representations of natural movie statistics. Biologically, the
emergence of straightening in these networks as well as their ability to
better explain V1 data than baselines relatively lacking in straightening
Figure 4 provides new insights into potential neural mechanisms for previously
difficult-to-explain brain phenomena.

Biological constraints could lend parsimony to selecting among models, each
with a different engineering goal. On the face, RS by virtue of utilizing
Gaussian noise instead of engineered noise gains traction over adversarial
training as a more simple, and powerful way of achieving robustness in ANNs,
which is in line with a long history of probabilistic inference in the visual
cortex of humans (Pouget et al., 2013). Indeed, looking across the range of
robust models tested, the best-fitting model of V1 was not necessarily the
most robust but tended toward more straightened representations that also
showed the least contracted representations – consistent with a known
dimensionality expansion from the sensory periphery to V1 in the brain (Field,
1994). Future work exploring a wider variety of robustness training in
conjunction with more bioplausible architectures, objectives, and training
diets may yet elucidate the balance of factors contributing to biological
vision.

At the same time, our work does not directly address how straightened
representations in the visual system may or may not be utilized to influence
downstream visual perception and behavior, and this connection is an important
topic for future work. On the one hand, for supporting dynamical scene
perception, behaviors that predict (extrapolate) or postdict (interpolate)
scene properties over time (e.g., object position) may be supported by
straightened natural movie representations. Indeed, both explanations,
prediction and postdiction, have been invoked to account for psychophysical
phenomena like the flash-lag illusion which presents an interesting test case
of how the brain processes complex stimuli over time (Eagleman & Sejnowski,
2000). However, even for relatively stationary scenes such as those containing
textures, we observed benefits for straightening and invertibility in robustly
trained networks (see A.4, Tables 2 and 3). Further work is needed to explore
how spatially local versus global features in the presence of simple versus
complex motion are affected in their relative straightening by model training.

#### Acknowledgments

This work was supported by a Klingenstein-Simons fellowship, Sloan Foundation
fellowship, and Grossman-Kavli Scholar Award as well as a NVIDIA GPU grant and
was performed using the Columbia Zuckerman Axon GPU cluster. We thank all
three reviewers for their constructive feedback that led to an improved final
version of the paper.

## References

  * Bakhtiari et al. (2021) Shahab Bakhtiari, Patrick J Mineault, Tim Lillicrap, Christopher C Pack, and Blake Aaron Richards.  The functional specialization of visual cortex emerges from training parallel pathways with self-supervised predictive learning.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), _Advances in Neural Information Processing Systems_ , 2021.  URL https://openreview.net/forum?id=t1czgrQOrwW.
  * Bashivan et al. (2019) Pouya Bashivan, Kohitij Kar, and James J DiCarlo.  Neural population control via deep image synthesis.  _Science_ , 364(6439):eaav9436, May 2019.
  * Berkes & Wiskott (2005) P Berkes and L Wiskott.  Slow feature analysis yields a rich repertoire of complex cell properties.  _J. Vis._ , 5(6):9–9, July 2005.
  * Chen et al. (2021) Chun-Fu Chen, Rameswar Panda, Kandan Ramakrishnan, Rogerio Feris, John Cohn, Aude Oliva, and Quanfu Fan.  Deep analysis of cnn-based spatio-temporal representations for action recognition.  In _The IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , June 2021.
  * Chen et al. (2019) Ricky T. Q. Chen, Jens Behrmann, David K Duvenaud, and Joern-Henrik Jacobsen.  Residual flows for invertible generative modeling.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), _Advances in Neural Information Processing Systems_ , volume 32. Curran Associates, Inc., 2019.  URL https://proceedings.neurips.cc/paper/2019/file/5d0d5594d24f0f955548f0fc0ff83d10-Paper.pdf.
  * Chen et al. (2020) Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He.  Improved baselines with momentum contrastive learning.  March 2020.
  * Cichy et al. (2016) Radoslaw Martin Cichy, Aditya Khosla, Dimitrios Pantazis, Antonio Torralba, and Aude Oliva.  Comparison of deep neural networks to spatio-temporal cortical dynamics of human visual object recognition reveals hierarchical correspondence.  _Sci. Rep._ , 6(1), September 2016.
  * Cohen et al. (2019) Jeremy M Cohen, Elan Rosenfeld, and J. Zico Kolter.  Certified adversarial robustness via randomized smoothing, 2019.  URL https://arxiv.org/abs/1902.02918.
  * Dapello et al. (2020) Joel Dapello, Tiago Marques, Martin Schrimpf, Franziska Geiger, David D. Cox, and James J. DiCarlo.  Simulating a primary visual cortex at the front of cnns improves robustness to image perturbations.  _bioRxiv_ , 2020.  doi: 10.1101/2020.06.16.154542.  URL https://www.biorxiv.org/content/early/2020/10/22/2020.06.16.154542.
  * DiCarlo et al. (2012) James J DiCarlo, Davide Zoccolan, and Nicole C Rust.  How does the brain solve visual object recognition?  _Neuron_ , 73(3):415–434, February 2012.
  * Eagleman & Sejnowski (2000) David M. Eagleman and Terrence J. Sejnowski.  Motion Integration and Postdiction in Visual Awareness.  _Science_ , 287(5460):2036–2038, March 2000.  doi: 10.1126/science.287.5460.2036.  URL https://www.science.org/doi/full/10.1126/science.287.5460.2036.  Publisher: American Association for the Advancement of Science.
  * Engstrom et al. (2019a) Logan Engstrom, Andrew Ilyas, Hadi Salman, Shibani Santurkar, and Dimitris Tsipras.  Robustness (python library), 2019a.  URL https://github.com/MadryLab/robustness.
  * Engstrom et al. (2019b) Logan Engstrom, Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Brandon Tran, and Aleksander Madry.  Adversarial robustness as a prior for learned representations.  June 2019b.
  * Engstrom et al. (2019c) Logan Engstrom, Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Brandon Tran, and Aleksander Madry.  Learning perceptually-aligned representations via adversarial robustness.  In _ArXiv preprint arXiv:1906.00945_ , 2019c.
  * Field (1994) David J. Field.  What Is the Goal of Sensory Coding?  _Neural Computation_ , 6(4):559–601, July 1994\.  ISSN 0899-7667.  doi: 10.1162/neco.1994.6.4.559.  URL https://doi.org/10.1162/neco.1994.6.4.559.
  * Ford et al. (2019) Nic Ford, Justin Gilmer, Nicolas Carlini, and Dogus Cubuk.  Adversarial examples are a natural consequence of test error in noise.  January 2019.
  * Friston (2009) Karl Friston.  The free-energy principle: a rough guide to the brain?  _Trends Cogn. Sci._ , 13(7):293–301, July 2009\.
  * Gärtner (1999) Bernd Gärtner.  Fast and robust smallest enclosing balls.  In Jaroslav Nešetřil (ed.), _Algorithms - ESA’ 99_ , pp. 325–338, Berlin, Heidelberg, 1999. Springer Berlin Heidelberg.  ISBN 978-3-540-48481-3.
  * Goroshin et al. (2015) Ross Goroshin, Michael F Mathieu, and Yann LeCun.  Learning to linearize under uncertainty.  In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett (eds.), _Advances in Neural Information Processing Systems_ , volume 28. Curran Associates, Inc., 2015.  URL https://proceedings.neurips.cc/paper/2015/file/eefc9e10ebdc4a2333b42b2dbb8f27b6-Paper.pdf.
  * He et al. (2015) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  _CoRR_ , abs/1512.03385, 2015.  URL http://arxiv.org/abs/1512.03385.
  * Hénaff et al. (2019) Olivier J Hénaff, Robbe L T Goris, and Eero P Simoncelli.  Perceptual straightening of natural videos.  _Nat. Neurosci._ , 22(6):984–991, June 2019.
  * Hénaff et al. (2021) Olivier J Hénaff, Yoon Bai, Julie A Charlton, Ian Nauhaus, Eero P Simoncelli, and Robbe L T Goris.  Primary visual cortex straightens natural video trajectories.  _Nat. Commun._ , 12(1):5982, October 2021.
  * Jacobsen et al. (2018a) Jörn-Henrik Jacobsen, Jens Behrmann, Richard Zemel, and Matthias Bethge.  Excessive invariance causes adversarial vulnerability.  November 2018a.
  * Jacobsen et al. (2018b) Jörn-Henrik Jacobsen, Arnold Smeulders, and Edouard Oyallon.  i-RevNet: Deep invertible networks.  February 2018b.
  * Khaligh-Razavi & Kriegeskorte (2014) Seyed-Mahdi Khaligh-Razavi and Nikolaus Kriegeskorte.  Deep supervised, but not unsupervised, models may explain IT cortical representation.  _PLoS Comput. Biol._ , 10(11):e1003915, November 2014.
  * Kong et al. (2022) Nathan C L Kong, Eshed Margalit, Justin L Gardner, and Anthony M Norcia.  Increasing neural network robustness improves match to macaque V1 eigenspectrum, spatial frequency preference and predictivity.  _PLoS Comput. Biol._ , 18(1):e1009739, January 2022.
  * Kubilius et al. (2018) Jonas Kubilius, Martin Schrimpf, Aran Nayebi, Daniel Bear, Daniel L K Yamins, and James J DiCarlo.  CORnet: Modeling the neural mechanisms of core object recognition.  September 2018.
  * Lecuyer et al. (2018) Mathias Lecuyer, Vaggelis Atlidakis, Roxana Geambasu, Daniel Hsu, and Suman Jana.  Certified robustness to adversarial examples with differential privacy.  February 2018.
  * Lotter et al. (2016) William Lotter, Gabriel Kreiman, and David Cox.  Deep predictive coding networks for video prediction and unsupervised learning.  May 2016.
  * Madry et al. (2018) Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu.  Towards deep learning models resistant to adversarial attacks.  In _International Conference on Learning Representations_ , 2018.  URL https://openreview.net/forum?id=rJzIBfZAb.
  * Millidge et al. (2021) Beren Millidge, Anil Seth, and Christopher L Buckley.  Predictive coding: A theoretical and experimental review.  July 2021.
  * Olshausen & Field (1996) B A Olshausen and D J Field.  Emergence of simple-cell receptive field properties by learning a sparse code for natural images.  _Nature_ , 381(6583):607–609, June 1996.
  * Pouget et al. (2013) Alexandre Pouget, Jeffrey M Beck, Wei Ji Ma, and Peter E Latham.  Probabilistic brains: knowns and unknowns.  _Nat. Neurosci._ , 16(9):1170–1178, September 2013.
  * Rao & Ballard (1999) R P Rao and D H Ballard.  Predictive coding in the visual cortex: a functional interpretation of some extra-classical receptive-field effects.  _Nat. Neurosci._ , 2(1):79–87, January 1999.
  * Sabour et al. (2017) Sara Sabour, Nicholas Frosst, and Geoffrey E Hinton.  Dynamic routing between capsules.  October 2017.
  * Schrimpf et al. (2018) Martin Schrimpf, Jonas Kubilius, Ha Hong, Najib J. Majaj, Rishi Rajalingham, Elias B. Issa, Kohitij Kar, Pouya Bashivan, Jonathan Prescott-Roy, Franziska Geiger, Kailyn Schmidt, Daniel L. K. Yamins, and James J. DiCarlo.  Brain-score: Which artificial neural network for object recognition is most brain-like?  _bioRxiv preprint_ , 2018.  URL https://www.biorxiv.org/content/10.1101/407007v2.
  * Schrimpf et al. (2020) Martin Schrimpf, Jonas Kubilius, Michael J Lee, N Apurva Ratan Murty, Robert Ajemian, and James J DiCarlo.  Integrative benchmarking to advance neurally mechanistic models of human intelligence.  _Neuron_ , 2020.  URL https://www.cell.com/neuron/fulltext/S0896-6273(20)30605-X.
  * Tacchetti et al. (2018) Andrea Tacchetti, Leyla Isik, and Tomaso A Poggio.  Invariant recognition shapes neural representations of visual input.  _Annu. Rev. Vis. Sci._ , 4:403–422, September 2018.
  * van den Oord et al. (2018) Aaron van den Oord, Yazhe Li, and Oriol Vinyals.  Representation learning with contrastive predictive coding.  July 2018.
  * Wang et al. (2004) Zhou Wang, A.C. Bovik, H.R. Sheikh, and E.P. Simoncelli.  Image quality assessment: from error visibility to structural similarity.  _IEEE Transactions on Image Processing_ , 13(4):600–612, 2004.  doi: 10.1109/TIP.2003.819861.
  * Whittington & Bogacz (2019) James C.R. Whittington and Rafal Bogacz.  Theories of error back-propagation in the brain.  _Trends in Cognitive Sciences_ , 23(3):235–250, 2019.  ISSN 1364-6613.  doi: https://doi.org/10.1016/j.tics.2018.12.005.  URL https://www.sciencedirect.com/science/article/pii/S1364661319300129.
  * Yamins et al. (2014) Daniel L. K. Yamins, Ha Hong, Charles F. Cadieu, Ethan A. Solomon, Darren Seibert, and James J. DiCarlo.  Performance-optimized hierarchical models predict neural responses in higher visual cortex.  _Proceedings of the National Academy of Sciences_ , 111(23):8619–8624, 2014.  doi: 10.1073/pnas.1403112111.  URL https://www.pnas.org/doi/abs/10.1073/pnas.1403112111.
  * Zbontar et al. (2021) Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny.  Barlow twins: Self-supervised learning via redundancy reduction.  March 2021.
  * Zhuang et al. (2021) Chengxu Zhuang, Siming Yan, Aran Nayebi, Martin Schrimpf, Michael C. Frank, James J. DiCarlo, and Daniel L. K. Yamins.  Unsupervised neural network models of the ventral visual stream.  _Proceedings of the National Academy of Sciences_ , 118(3):e2014196118, 2021.  doi: 10.1073/pnas.2014196118.  URL https://www.pnas.org/doi/abs/10.1073/pnas.2014196118.

##  Appendix A Appendix

###  A.1 straightening for both natural and artificial sequences

Figure 7: ANNs show straightening of representations when robustness to noise
constraints (noise augmentation or robustness to adversarial attack) is added
to their training. Counterclockwise from top left, measurements for
straightening of movie sequences (from (Hénaff et al., 2019), natural
sequence: green, artificial sequence: black) in each layer of ResNet50
architecture under different training regimes: supervised training (standard),
supervised training with adversarial training ($L_{2},\sigma^{2}=3$) (Engstrom
et al., 2019a) and supervised training with noise augmentation
($L_{2},\sigma^{2}=0.5$) (Cohen et al., 2019). The top right shows
straightening for artificial (open circles) and natural (closed circles)
sequences using ResNet architecture with no training (random parameters),
self-supervised training (Chen et al., 2020), or additional layers.

###  A.2 straightening in other architectures

Figure 8: Straightening for ResNet50 and ResNet101 trained as the base
architecture for action recognition (Chen et al., 2021) They were trained on
video clips for action recognition. Although these models were not trained for
straightening or predicting the next frame, they exhibit small but measurable
straightening for natural movies. However, the curvatures for artificial
sequences were not increased as much as the curvature increase for artificial
sequences in robust neural networks (Figure 7). Figure 9: Lack of
straightening for natural movies in ResNet50 trained with a biologically-
inspired model of V1 in the front-end (Dapello et al., 2020). Vone-ResNet50
exhibits robustness to adversarial attack, but the fact that it does not
exhibit straightening (except for the front-end) provides further evidence
that adversarial robustness does not always accompany straightening.

###  A.3 More on neural data predicitivity

Figure 10: Clean accuracy (left) and robust accuracy (right) vs. V1
predictivity (same color convention as used in main text). Figure 11: Brain-
score (data and metric publicly available (Schrimpf et al., 2018)) for the
models used in this study. The bar plot Figure 4 is a summary of this plot.
Since the focus of this work was on V1 which is the only visual area for which
neural straightening has been measured (Hénaff et al., 2021), we measured the
brain-score for more layers for V1.

###  A.4 Movie characteristics and additional interpolation examples

####  A.4.1 Movies

Figure 12: Movies used to evaluate straightening and invertibility (11 frames
each).

####  A.4.2 Table for average straightening

Table 2: Average curvature change (re: pixels) for each movie. RN stands for
ResNet. The architecture of all robust models used was ResNet50.

| water | carn. | walk. | dogv. | egomo. | chiron. | bees | leaves | smile | chiron. | prair.
---|---|---|---|---|---|---|---|---|---|---|---
RN18 | 45.29 | 36.25 | -27.58 | 12.39 | -34.82 | 13.30 | -5.94 | -58.03 | 9.28 | 0.69 | 48.45
RN34 | 44.79 | 36.63 | -28.43 | 12.39 | -35.67 | 12.39 | -5.81 | -58.65 | 12.08 | 0.43 | 48.14
RN50 | 48.85 | 37.33 | -28.13 | 14.00 | -34.68 | 13.45 | -5.66 | -58.53 | 10.79 | 0.76 | 49.72
RN50 Self-sup | 52.05 | 40.32 | -26.73 | 18.98 | -33.19 | 14.37 | -5.68 | -58.22 | 19.42 | 0.52 | 49.90
RN50 MocoAugm | 49.03 | 36.15 | -26.04 | 16.75 | -32.88 | 13.82 | -5.22 | -57.69 | 14.00 | 0.57 | 49.30
RN50 NoAugm | 46.38 | 33.64 | -24.79 | 15.89 | -33.58 | 14.71 | -5.10 | -57.32 | 8.87 | 0.77 | 48.31
RN101 | 48.81 | 37.73 | -29.28 | 12.51 | -35.03 | 12.68 | -6.17 | -59.23 | 11.87 | -1.05 | 49.63
RN152 | 48.91 | 40.38 | -29.09 | 13.27 | -35.53 | 12.23 | -6.60 | -59.75 | 13.80 | -1.00 | 50.68
AT $\bm{L_{2}:\epsilon=3}$ | 17.43 | 2.78 | -40.56 | -3.61 | -55.42 | -2.54 | -13.44 | -73.01 | -25.09 | -13.06 | 30.34
AT $\bm{L_{\infty}:\epsilon=4}$ | 25.93 | 16.51 | -36.55 | 3.48 | -50.46 | 5.45 | -12.90 | -68.74 | -11.49 | -7.91 | 38.46
AT $\bm{L_{\infty}:\epsilon=8}$ | 25.64 | 13.15 | -40.79 | 1.10 | -54.99 | 4.11 | -15.59 | -71.48 | -17.99 | -10.23 | 39.09
RS $\bm{L_{2}:\epsilon=0.25}$ | 24.97 | 12.23 | -31.81 | 5.38 | -45.86 | 5.10 | -9.37 | -63.84 | -9.08 | -5.22 | 34.23
RS $\bm{L_{2}:\epsilon=0.5}$ | 16.25 | 1.34 | -34.65 | 1.76 | -51.28 | 0.80 | -12.57 | -70.21 | -14.35 | -10.09 | 24.26
RS $\bm{L_{2}:\epsilon=1}$ | 10.65 | -0.61 | -42.40 | -7.27 | -58.04 | -4.68 | -16.57 | -79.82 | -23.55 | -14.62 | 18.03

####  A.4.3 Invertibility measure for each movie

Table 3: Average SSIMs for each movie. RN stands for ResNet. Architecture of
all robust models used was ResNet50.

| water | carn. | walk. | dogv. | egomo. | chiron. | bees | leaves | smile | chirono. | prair.
---|---|---|---|---|---|---|---|---|---|---|---
RN18 | 0.47 | 0.34 | 0.21 | 0.22 | 0.13 | 0.28 | 0.31 | 0.20 | 0.29 | 0.31 | 0.34
RN34 | 0.46 | 0.34 | 0.21 | 0.22 | 0.13 | 0.27 | 0.30 | 0.20 | 0.29 | 0.30 | 0.34
RN50 | 0.47 | 0.34 | 0.22 | 0.22 | 0.14 | 0.28 | 0.30 | 0.20 | 0.29 | 0.31 | 0.34
RN50 Self-sup | 0.49 | 0.34 | 0.22 | 0.23 | 0.14 | 0.29 | 0.33 | 0.20 | 0.33 | 0.33 | 0.34
RN50 MocoAugm | 0.44 | 0.32 | 0.20 | 0.21 | 0.12 | 0.25 | 0.27 | 0.19 | 0.27 | 0.28 | 0.31
RN50 NoAugm | 0.43 | 0.32 | 0.20 | 0.21 | 0.12 | 0.26 | 0.32 | 0.19 | 0.28 | 0.28 | 0.32
RN101 | 0.48 | 0.35 | 0.22 | 0.23 | 0.14 | 0.28 | 0.31 | 0.21 | 0.30 | 0.30 | 0.35
RN152 | 0.49 | 0.36 | 0.22 | 0.23 | 0.14 | 0.29 | 0.31 | 0.21 | 0.30 | 0.31 | 0.36
AT $\bm{L_{2}:\epsilon=3}$ | 0.76 | 0.59 | 0.52 | 0.67 | 0.22 | 0.62 | 0.58 | 0.35 | 0.81 | 0.63 | 0.59
AT $\bm{L_{\infty}:\epsilon=4}$ | 0.72 | 0.56 | 0.43 | 0.60 | 0.25 | 0.57 | 0.53 | 0.32 | 0.69 | 0.55 | 0.56
AT $\bm{L_{\infty}:\epsilon=8}$ | 0.78 | 0.59 | 0.45 | 0.66 | 0.29 | 0.65 | 0.56 | 0.35 | 0.78 | 0.61 | 0.62
RS $\bm{L_{2}:\epsilon=0.25}$ | 0.53 | 0.48 | 0.35 | 0.47 | 0.18 | 0.44 | 0.44 | 0.29 | 0.57 | 0.46 | 0.42
RS $\bm{L_{2}:\epsilon=0.5}$ | 0.60 | 0.49 | 0.37 | 0.51 | 0.22 | 0.44 | 0.45 | 0.30 | 0.54 | 0.49 | 0.46
RS $\bm{L_{2}:\epsilon=1}$ | 0.61 | 0.50 | 0.35 | 0.53 | 0.24 | 0.44 | 0.48 | 0.30 | 0.52 | 0.49 | 0.48

####  A.4.4 Additional movie interpolation examples

Figure 13: Three more example for interpolations for movies: chironomous,
dogville, and egomotion, respectively. The gray dot in the middle of all
frames is known as fixation spot where subjects (humans or monkeys) are
instructed to keep their gaze toward during the experiment.

###  A.5 Expansion metric vs. straightening across models and layers

Figure 14: For each layer in each model, the expansions (re: pixels) and
curvature change (re: pixels) were plotted for first row: all movies, second
row: average over movies. Figure 15: Geometric straightening and expansion of
AT models and their relation to V1 explainability. $\Delta$ means quantity is
referenced to the same measure in a standard ResNet50.

###  A.6 Reproducibility information

Almost all data (models, movies, and metrics) used in this work are publicly
available and we provided references to them in the text (for instance see
Table 1). We will release the code to reproduce the main results in this work
at
https://github.com/toosi/BrainLike_Straightening
and we provide pointers to the publicly available resources used in this work
as listed below.
Movies and images. We used the same movies used in the original studies on
human perception and monkey primary visual cortex (Hénaff et al., 2019; 2021)
which are available from the first author’s Github as referenced in their
papers. Images used to measure the clean accuracy and robust accuracy were
taken from ImageNet validation set.
Models. All the models used in this study were from the ResNet family and
checkpoints for the main robust models are publicly available as references in
the main text (Table 1). The checkpoints for the only two custom-trained
models (supervised with no augmentations and supervised with Moco
augmentation) will be made publicly available along with the code.
Neural predictivity metric. We used brain-score, which is a publicly available
benchmark to evaluate how well a model predicts variance in neural data
(Schrimpf et al., 2018).


# A Federated Cox Model with Non-Proportional Hazards

Dekai Zhang
Department of Computing
Imperial College London
London, SW7 2BX
<EMAIL_ADDRESS>
&Francesca Toni
Department of Computing
Imperial College London
London, SW7 2BX
<EMAIL_ADDRESS>
&Matthew Williams
Department of Surgery & Cancer
Imperial College London
London, SW7 2BX
<EMAIL_ADDRESS>

###### Abstract

Recent research has shown the potential for neural networks to improve upon
classical survival models such as the Cox model, which is widely used in
clinical practice. Neural networks, however, typically rely on data that are
centrally available, whereas healthcare data are frequently held in secure
silos. We present a federated Cox model that accommodates this data setting
and also relaxes the proportional hazards assumption, allowing time-varying
covariate effects. In this latter respect, our model does not require explicit
specification of the time-varying effects, reducing upfront organisational
costs compared to previous works. We experiment with publicly available
clinical datasets and demonstrate that the federated model is able to perform
as well as a standard model.

##  1 Introduction

Estimating how long patients might live for is a key task in clinical
medicine, and is a common question from patients. Survival analysis is the
statistical branch used to perform these estimates, which can range in its
application from predicting death following diagnosis to loan defaults or
machine part failures. Amongst survival models, the Cox model (Cox, 1972) is
one of the most widely used.

Machine learning techniques have received attention for their potential to
improve upon the performance of the Cox model. Many recent efforts (Katzman et
al., 2018; Luck et al., 2017; Kvamme et al., 2019) have exploited neural
networks (NNs) to model more complex relationships as well as enable typically
unsupported input data types such as images (Zhu et al., 2017a, b; Li et al.,
2019; Bello et al., 2019). Notwithstanding, the Cox model has remained the
standard in survival analysis (Wang et al., 2019). Indeed, the adoption of
machine learning has progressed haltingly in many areas of healthcare (Kelly
et al., 2019).

One challenge lies in the distributed nature of healthcare data (Wiens et al.,
2019). In much of machine learning, data are centralised, whereas privacy
concerns often result in secure data “silos” in healthcare. Federated learning
(FL) accommodates this decentralised data environment and has shown promise in
clinical contexts (Rieke et al., 2020).

Despite a fast emerging literature in FL, there has been scant work on
federated survival analysis. Andreux et al. (2020) propose a federated Cox
model that is closest to this paper. The standard Cox model is, however,
limited in that it can only correctly model proportional hazards. We take an
alternative approach allowing us to embed time-varying covariate effects
(_non_ -proportional hazards) directly in the architecture, potentially
reducing organisational setup costs for federations. Such effects are relevant
to adapt models for patients such as those with breast cancer where the
proportional hazards assumption has been shown to be violated (Bellera et al.,
2010; Coradini et al., 2000; Gore et al., 1984).

In the following, we briefly discuss relevant background and highlight related
work (Section 2), before defining our model (Section 3), instantiating it with
different hazards assumptions and presenting our experiments with real-world
clinical datasets (Section 4). Section 5 concludes with potential directions
for future work.

##  2 Background and Related Work

The promise of greater control over data ownership and enhanced privacy that
FL affords has generated interest in the healthcare community. Few works,
however, have investigated the intersection between survival analysis and FL.
We present background on each of these areas separately before discussing
their intersection relevant for this work.

####  2.0.1 Background on Federated Learning (FL).

FL is a framework for decentralised data that cannot be shared due to their
sensitive content or prohibitive communication costs (McMahan et al., 2017).
In the context of healthcare, patient data may be kept in this way by the
clinical unit (e.g., the hospital) at which the patient was treated. In the
following, we will simply refer to these data-keeping units as federation
members or centres. Typically (and in this paper), the federated objective is
to minimise $\mathcal{L}_{F}(X,\phi)$ with respect to $\phi$ with:

$\mathcal{L}_{F}(X,\phi)\\!=\\!\sum_{k\in K}w_{k}\mathcal{L}_{k}(X_{k},\phi)$
(1)

where $\mathcal{L}_{F}$ represents the global loss: an average of the local
losses $\mathcal{L}_{k}$ computed by the federation members in $K$ on their
own data $X_{k}$ weighted by $w_{k}$, where $\phi$ represents the model
parameters. Typically, each member customises $\phi$ for a number of local
optimisation rounds before aggregating the customised $\phi$ for a new global
consensus model.

####  2.0.2 Background on Survival Analysis.

Survival analysis estimates the time to an event for a population $N$ with
data $\mathcal{D}=\\{(x_{i},t_{i},s_{i})\\}_{i\in N}$ where each person $i$
has covariates $x_{i}=(x_{i1},...,x_{ip})^{\top}$, a time of observation
$t_{i}$ and an indicator $s_{i}\in\\{0,1\\}$ which equals 1 if $i$ has
experienced the event or 0 if not, i.e., if $i$ is _censored_.

The _Cox model_ (Cox, 1972) is one of the most widely used survival models. It
defines a hazard function $h$, which expresses the rate of failure at time $t$
subject to survival until then as follows:

$\begin{split}h(t|x_{i})&=P(T=t|T>t-1)\\\ &=h_{0}(t)exp[g(x_{i})]\text{ with
}g(x_{i})=\beta^{\top}x_{i}\end{split}$ (2)

where $h_{0}(t)$ is some baseline hazard and where
$\beta=(\beta_{1},...,\beta_{p})^{\top}$ is a coefficient vector. Later works
replace the linear predictor $\beta^{\top}x_{i}$ with NNs $g_{\phi}(x_{i})$,
demonstrating competitive performance (Faraggi and Simon, 1995; Katzman et
al., 2018). The coefficients are estimated by minimising the negative partial
log-likelihood given by:

$-\sum_{i\in N}s_{i}[g(x_{i})-\log(\sum_{j\in R_{i}}exp[g(x_{j})])]$ (3)

where $R_{i}=\\{j\in N:t_{j}\geq t_{i}\\}$ denotes the individuals who are
still _at risk_ when $i$ experiences the event.

In a federated setting, this loss generally cannot be decomposed into local
losses due to the logarithmic term, as the risk set $R_{i}$ can contain
individuals from centres other than the one of $i$. This therefore does not
match the formulation of Eq. 1. The hazard function also assumes proportional
hazards (PH) – differences in covariates result in constant proportional
differences in hazards. Over long time horizons, this can be restrictive
(Kvamme et al., 2019; Bellera et al., 2010; Antolini et al., 2005).

####  2.0.3 State-of-the-Art in Federated Survival Analysis.

Our work is situated in the intersection of federated learning and survival
analysis and proposes a novel framework. A handful of works have already
proposed such frameworks of which we provide a brief overview here. The works
of Lu et al. (2015) and Dai et al. (2020) embody one approach which relies on
substantial sharing of summary statistics over the local datasets in every
training iteration. This differs in spirit from FL where more abstract
parameters are shared and, often, infrequently so. Moreover, their models are
based on linear predictors and do not address integration with NNs.

Recent work by Andreux et al. (2020) is closest to our approach. Their model
exploits a discretisation of the Cox model (also by Cox (1972)) with an NN-
based predictor:

$\dfrac{h(t|x)}{1-h(t|x)}=\dfrac{h_{0}(t)}{1-h_{0}(t)}exp[g_{\phi}(x_{i})]$
(4)

which can be rewritten in a sigmoid form:

$h(t|x_{i})=\dfrac{1}{1+exp[-(\alpha_{t}+g_{\phi}(x_{i}))]}$ (5)

where $\alpha_{t}=\log(\frac{h_{0}(t)}{1-h_{0}(t)})$.

They follow Craig et al. (2021) in estimating this function like a logistic
regression with negative log-likelihood:

$-\sum_{i\in
N}\sum_{k=1}^{t_{i}}[y_{ik}\log[h(k|x_{i})]+(1-y_{ik})\log[(1-h(k|x_{i}))]]$
(6)

where $y_{ij}=\mathds{1}\\{t_{j}=t_{i},s_{i}=1\\}$. Importantly, this loss
does not depend on risk sets and is therefore separable – each centre’s loss
only depends on local data – recovering the federated objective (Eq. 1).

Andreux et al. (2020) demonstrate that a federation of this model can draw
even in performance with a model trained on pooled data. This is, however,
only shown with aggregation after every local optimisation round – a setup
that may need to differ in practice (Rieke et al., 2020) – and assuming PH, as
their predictor $g_{\phi}(x_{i})$ is time-invariant.

We note that non-PH can be admitted to their model by including time
interactions (giving $g_{\phi}(x_{i},f(t))$) – an approach sometimes taken in
standard Cox models – as demonstrated on pooled data by Craig et al. (2021).
This could, however, introduce a dependency on the specification of $f(t)$ and
its interactions. Crucially, this may add to the organisational setup costs of
a federation: even though interactions could be learned, $f(t)$ needs to be
fully specified and agreed upon in advance. In contrast, we follow Gensheimer
and Narasimhan (2019) by making the choice between PH and non-PH a binary
decision over the architecture of the output layer.

##  3 Model

We build upon the discretised Cox model and detail how the PH assumption is
relaxed and formulate the federated objective. We describe a discretisation
procedure and an optional interpolation scheme for smooth predictions. Lastly,
we outline two complementary performance metrics.

##### Non-Proportional Hazards.

We use a discretised Cox model (Eq. 5) but parameterised with a time-varying,
NN-based predictor $g_{\phi,t}(x)$. Following Gensheimer and Narasimhan
(2019), we allow for non-PH by fully connecting the output layer to the
previous layer. The output layer thus encodes time-varying covariate effects
in time-specific weights. A sigmoid is used to retrieve the hazard rates.

For PH, the output component is split into a first layer with a single neuron
and no bias. The output of this neuron is passed into a second layer with as
many neurons as time steps. This captures time-varying baseline hazards in the
second layer and time-invariant covariate effects in the first. The difference
in components is illustrated in Figure 1.

(a) Non-PH

(b) PH

Figure 1: Output components for 3 time steps.

##### Federated Objective.

To conform to a federated formulation (Eq. 1), we split the objective (Eq. 6):

$\begin{split}\mathcal{L}_{F}&=\sum_{k\in
K}w_{k}\mathcal{L}_{k}(X_{k},\phi)\\\ &=-\sum_{k\in
K}\frac{|N_{k}|}{|N|}\sum_{i\in
N_{k}}\sum_{j=1}^{t_{i}}\Big{[}y_{ij}\log[h(j|x_{i})]+(1-y_{ij})\log[(1-h(j|x_{i}))]\Big{]}\end{split}$
(7)

where $y_{ij}=\mathds{1}\\{t_{j}=t_{i},s_{i}=1\\}$. Each centre calculates
$\mathcal{L}_{k}(X_{k},\phi)$ on its own subset of the population $N_{k}$. We
adapt the FedAvg algorithm of McMahan et al. (2017) to minimise this loss
(Algorithm 1).

Algorithm 1 Procedure to optimise the federated objective.

1: Initialise global model with $\phi_{0}$

2: for each round $t=1,...,T$ do

3: for each centre $k=1,...,K$ in parallel do

4: Send $\phi_{t-1}$ to centre k

5: for $\text{each local round }b=1,...,B$ do

6: Local update
$\phi_{t}^{k}\leftarrow\phi_{t}^{k}-\lambda\nabla\mathcal{L}_{k}(X_{k},\phi_{t}^{k})$

7: end for

8: Receive $\phi_{t}^{k}$ from centre k

9: end for

10: Aggregate $\phi_{t}\leftarrow\sum_{k\in K}\frac{|N_{k}|}{|N|}\phi_{t}^{k}$

11: end for

12: return $\phi_{T}$

##### Discretisation.

The model operates on discretised time, so that $t$ indexes into a set of
intervals $[\tau_{t-1},\tau_{t})$. Following Kvamme and Borgan (2019) we
discretise on Kaplan-Meier quantiles. Defining the survival curve
$S(\tau)=S(\tau-1)(1-h(\tau))$, the quantiles
$\\{\tau_{1},\tau_{2},...,\tau_{m}\\}$ can be obtained as:

$S(T=\tau_{j})-S(T=\tau_{j+1})=\dfrac{1-S(T=\tau_{max})}{m}$ (8)

for $j=\\{0,1,...,m-1\\}$. This discretisation procedure yields a set of steps
$\\{\tau_{1},\tau_{2},...,\tau_{m}\\}$ where each step results in the same
decrease in survival (an illustration is provided in Appendix A).

##### Interpolation.

To smooth step-wise predictions, we use constant density interpolation (Kvamme
and Borgan, 2019). Letting $\widetilde{S}(\tau)$ denote the interpolation of
the survival curve $S(\tau)$, we then have:

$\widetilde{S}(\tau)\\!=\\!S(\tau_{j-1})+\left[S(\tau_{j})-S(\tau_{j-1})\right]\dfrac{\tau-\tau_{j-1}}{\tau_{j}-\tau_{j-1}}$
(9)

for a given time $\tau\in(\tau_{j-1},\tau_{j}]$. Intuitively, the step
survival curve is linearly interpolated between any adjacent steps, resulting
in constant densities in the corresponding interval (an illustration is
provided in Appendix A).

##### Performance Metrics – Concordance.

We use the time-dependent concordance index (Antolini et al., 2005), or
c-index, which is a discriminative measure for how well the model ranks the
relative survival between patient pairs, expressed as:

$P(S(t_{i}|x_{i})<S(t_{i}|x_{j})\ \&\ t_{i}<t_{j}\ \&\ s_{i}=1)$ (10)

which is estimated as follows:

$\widehat{c}=\dfrac{\sum_{i\in N}\sum_{j\in N,j\neq i}conc_{ij}}{\sum_{i\in
N}\sum_{j\in N,j\neq i}comp_{ij}}$ (11)

$comp_{ij}=\mathds{1}\\{t_{i}<t_{j}\ \&\ s_{i}=1\\}+\mathds{1}\\{t_{i}=t_{j}\
\&\ s_{i}=1\ \&\ s_{j}=0\\}$ (12)

$conc_{ij}=\mathds{1}\\{S(t_{i}|x_{i})<S(t_{i}|x_{j})\\}\ comp_{ij}$ (13)

##### Performance Metrics – Calibration.

While the c-index measures the discriminative performance of the model, it
does not measure how well _calibrated_ these estimates are (an illustration is
provided in Appendix A).

As a measure of calibration, we follow Graf et al. (1999) who propose a Brier
score for use with censored data defined as follows:

$BS(t)=\dfrac{1}{|N|}\sum_{i\in
N}w_{i}(t)\Big{(}y_{i}(t)-h(t|x_{i})\Big{)}^{2}$ (14)

$w_{i}(t)=\begin{cases}s_{i}/G(t),&\text{if $t_{i}\leq t$}\\\ 1/G(t),&\text{if
$t_{i}>t$}\end{cases}$ (15)

where $y_{i}(t)=\mathds{1}\\{t_{i}=t\\}$ and $G(t)$ is the Kaplan-Meier
estimate of the censoring distribution (i.e., estimated on
$\\{(x_{i},t_{i},1-s_{i})\\}_{i\in N}$). To measure calibration across the
entire time horizon, we numerically integrate the Brier score using 100 time
points (Kvamme et al., 2019).

##  4 Experiments

We introduce the datasets we experiment on, describe the setup of a simulated
federation and instantiate our model with three different linearity and
hazards assumptions, and present our results.111For source code, see
https://github.com/dkaizhang/federated-survival.

###  4.1 Datasets

Dataset |  Size |  Features |  Prop. censored |  Last event
---|---|---|---|---
METABRIC |  1,904 |  9 |  42% |  355 days
SUPPORT |  8,873 |  14 |  32% |  1,944 days
GBSG |  2,232 |  7 |  43% |  83 days
Table 1: Overview of datasets.

We experiment on three clinical datasets (Table 1; for Kaplan-Meier curves see
Appendix A) made available by Katzman et al. (2018), namely the Molecular
Taxonomy of Breast Cancer International Consortium (METABRIC), the Study to
Understand Prognoses Preferences Outcomes and Risks of Treatment (SUPPORT),
and the Rotterdam tumour bank and German Breast Cancer Study Group (GBSG).
METABRIC and GBSG both relate to breast cancer patients, a group for whom non-
PH have been noted (Bellera et al., 2010; Coradini et al., 2000), while
SUPPORT presents serious hospitalisations for a second application area.

###  4.2 Setup

We simulate two federated data cases: In the first, data are randomly
distributed (“IID”), simulating the case of each centre seeing a similar
sample of the patient population. In the second, data are stratified on the
time to event (“Non-IID”), simulating the case that each centre sees a non-
overlapping quantile of the population – from centre 1 seeing only the
shortest survivals leading to centre 4 with the longest survivals (Figure 2).
For comparability, we maintain the total number of local training rounds at
100. A pooled data baseline is provided (“Pooled” – no distinction between
local and global rounds). In all cases, 80% of the overall data are split, if
federated, and used for training, while 20% are held out for evaluation.

(a) IID

(b) Non-IID

Figure 2: Event time distribution for IID and non-IID data using
stratification by event time on METABRIC.

We instantiate the model with different choices for $g(x)$ – with a linear
predictor or with an NN, with and without PH (Table 2). For baselines, we
considered the works of Andreux et al. (2020) and Craig et al. (2021). The
former assumes PH, however, while the latter is a pooled data model. Both
require upfront agreement on a specification of $f(t)$ to include non-PH,
adding to the setup costs of a federation. We further note that no
implementations of these models are available. We therefore provide the NN PH
model to approximate the model of Andreux et al. (2020) – a federated NN-based
Cox model with PH – and the Linear PH model as a standard baseline.

Model | Predictor
---|---
Linear PH | $\beta^{\top}x$
NN PH | $g_{\phi}(x)$
NN nonPH | $g_{\phi,t}(x)$
Table 2: Model choices.

Architectures (Figure 1) are implemented in PyTorch 1.8.0 (Paszke et al.,
2019) with two hidden layers of 32 neurons for NN models and none for the
linear model. Optimisation uses Adam with grid-sought learning rates
($10^{-1}$ to $10^{-5}$ on 20% of the training data) and a batch size of 256.
Base case discretisation uses 10 time steps.

###  4.3 Results

In this section, we first compare the performance of the three models trained
in a centralised fashion on pooled data against their federated performance on
decentralised data. We next provide additional experiments exploring the
impact of the discretisation grid chosen for the base case. We report averaged
5-fold cross-validation performance throughout.

####  4.3.1 Federated Performance.

On pooled data, the NN nonPH model outperforms or ties in concordance (Table
4) and ties with the best in calibration (Table 4), indicating a gain from the
relaxation of the PH assumption. For METABRIC and GBSG this aligns with
Bellera et al. (2010); Coradini et al. (2000); Gore et al. (1984) who find
non-PH amongst this patient group.

Comparing this to the federated setting with IID data, all three models
maintain their performance (within one standard deviation) in concordance and
calibration when aggregation is frequent. First hints of performance
degradation amongst the NN-based models occur as aggregation becomes very
infrequent (rightmost columns) while the Linear PH model appears largely
unaffected. This observation is noteworthy, as infrequent aggregation will be
a likely feature in practice given communication costs. This indicates a
potential trade-off between model complexity and achievable aggregation
frequency to support its training.

In practice, data are likely to be non-IID across centres. The results show
that the performances of all three models suffer when this is the case.
Generally, the NN-based models experience the most severe losses in
performance and are largely outperformed by the Linear PH model. When
aggregation is infrequent, the NN-based models on SUPPORT and GBSG effectively
approach a no-skill predictor in concordance (average c-index of 0.5).
Further, performance losses under infrequent aggregation of the NN-based
models are, as would be expected, worse than under IID data. On SUPPORT, the
NN-based models exhibit much greater performance differences than on the other
datasets. In this respect, we note that SUPPORT has much longer survival times
than METABRIC or GBSG (Table 1), so that stratification by event time likely
results in a more significantly different partition of the data for the former
than for the latter two.

|  | METABRIC | SUPPORT | GBSG
---|---|---|---|---
|  | Global / local rounds | Global / local rounds | Global / local rounds
Data | Model | 100 / 1 | 20 / 5 | 1 / 100 | 100 / 1 | 20 / 5 | 1 / 100 | 100 / 1 | 20 / 5 | 1 / 100
Pooled | Linear PH | 63.5 $\\!\pm\\!$ 1.4 |  |  | 57.2 $\\!\pm\\!$ 1.0 |  |  | 66.5 $\\!\pm\\!$ 2.1 |  |
NN PH | 64.0 $\\!\pm\\!$ 0.6 |  |  | 60.8 $\\!\pm\\!$ 0.6 |  |  | 66.2 $\\!\pm\\!$ 2.6 |  |
NN nonPH | 66.7 $\\!\pm\\!$ 2.1 |  |  | 61.5 $\\!\pm\\!$ 1.2 |  |  | 66.6 $\\!\pm\\!$ 1.9 |  |
IID | Linear PH | 63.9 $\\!\pm\\!$ 0.8 | 64.0 $\\!\pm\\!$ 2.2 | 63.7 $\\!\pm\\!$ 1.7 | 57.2 $\\!\pm\\!$ 0.8 | 57.2 $\\!\pm\\!$ 0.8 | 57.2 $\\!\pm\\!$ 0.4 | 66.5 $\\!\pm\\!$ 1.5 | 66.3 $\\!\pm\\!$ 0.6 | 66.3 $\\!\pm\\!$ 1.4
NN PH | 63.8 $\\!\pm\\!$ 1.6 | 62.5 $\\!\pm\\!$ 1.9 | 63.3 $\\!\pm\\!$ 1.2 | 60.6 $\\!\pm\\!$ 1.0 | 60.9 $\\!\pm\\!$ 0.9 | 58.3 $\\!\pm\\!$ 1.4 | 67.4 $\\!\pm\\!$ 1.7 | 67.1 $\\!\pm\\!$ 1.2 | 63.5 $\\!\pm\\!$ 2.6
NN nonPH | 65.4 $\\!\pm\\!$ 1.9 | 65.7 $\\!\pm\\!$ 1.4 | 61.5 $\\!\pm\\!$ 2.1 | 62.1 $\\!\pm\\!$ 0.7 | 62.4 $\\!\pm\\!$ 0.4 | 56.7 $\\!\pm\\!$ 2.3 | 66.5 $\\!\pm\\!$ 1.0 | 66.4 $\\!\pm\\!$ 0.9 | 62.8 $\\!\pm\\!$ 1.9
Non-IID | Linear PH | 59.2 $\\!\pm\\!$ 3.0 | 59.8 $\\!\pm\\!$ 2.1 | 61.0 $\\!\pm\\!$ 1.3 | 55.4 $\\!\pm\\!$ 1.5 | 56.1 $\\!\pm\\!$ 0.9 | 56.2 $\\!\pm\\!$ 0.5 | 61.1 $\\!\pm\\!$ 0.9 | 61.9 $\\!\pm\\!$ 3.0 | 56.5 $\\!\pm\\!$ 6.9
NN PH | 60.9 $\\!\pm\\!$ 1.1 | 59.4 $\\!\pm\\!$ 2.5 | 57.3 $\\!\pm\\!$ 4.2 | 57.0 $\\!\pm\\!$ 1.4 | 56.7 $\\!\pm\\!$ 1.1 | 52.9 $\\!\pm\\!$ 0.8 | 61.3 $\\!\pm\\!$ 4.3 | 61.8 $\\!\pm\\!$ 2.9 | 53.0 $\\!\pm\\!$ 6.2
NN nonPH | 57.9 $\\!\pm\\!$ 2.9 | 59.6 $\\!\pm\\!$ 3.9 | 54.6 $\\!\pm\\!$ 5.3 | 50.8 $\\!\pm\\!$ 0.6 | 50.9 $\\!\pm\\!$ 0.9 | 50.8 $\\!\pm\\!$ 1.2 | 57.6 $\\!\pm\\!$ 1.9 | 55.8 $\\!\pm\\!$ 2.0 | 52.0 $\\!\pm\\!$ 2.5

Table 3: C-index (rebased to 100) – mean and standard deviation. Higher values
are better.

|  | METABRIC | SUPPORT | GBSG
---|---|---|---|---
|  | Global / local rounds | Global / local rounds | Global / local rounds
Data | Model | 100 / 1 | 20 / 5 | 1 / 100 | 100 / 1 | 20 / 5 | 1 / 100 | 100 / 1 | 20 / 5 | 1 / 100
Pooled | Linear PH | 16.4 $\\!\pm\\!$ 0.6 |  |  | 20.9 $\\!\pm\\!$ 0.5 |  |  | 18.0 $\\!\pm\\!$ 0.5 |  |
NN PH | 16.8 $\\!\pm\\!$ 0.7 |  |  | 19.6 $\\!\pm\\!$ 0.4 |  |  | 18.2 $\\!\pm\\!$ 0.8 |  |
NN nonPH | 16.4 $\\!\pm\\!$ 0.8 |  |  | 19.6 $\\!\pm\\!$ 0.4 |  |  | 18.0 $\\!\pm\\!$ 0.4 |  |
IID | Linear PH | 16.3 $\\!\pm\\!$ 1.1 | 16.7 $\\!\pm\\!$ 0.9 | 16.5 $\\!\pm\\!$ 0.7 | 20.9 $\\!\pm\\!$ 0.6 | 20.9 $\\!\pm\\!$ 0.3 | 20.9 $\\!\pm\\!$ 0.6 | 18.1 $\\!\pm\\!$ 0.2 | 18.1 $\\!\pm\\!$ 0.2 | 18.1 $\\!\pm\\!$ 0.3
NN PH | 17.2 $\\!\pm\\!$ 1.1 | 18.1 $\\!\pm\\!$ 0.5 | 18.1 $\\!\pm\\!$ 1.1 | 19.7 $\\!\pm\\!$ 0.7 | 19.8 $\\!\pm\\!$ 0.2 | 22.6 $\\!\pm\\!$ 1.3 | 17.7 $\\!\pm\\!$ 0.3 | 17.8 $\\!\pm\\!$ 0.4 | 19.8 $\\!\pm\\!$ 1.9
NN nonPH | 16.3 $\\!\pm\\!$ 1.3 | 16.5 $\\!\pm\\!$ 0.6 | 19.1 $\\!\pm\\!$ 0.8 | 19.5 $\\!\pm\\!$ 0.7 | 19.4 $\\!\pm\\!$ 0.4 | 21.1 $\\!\pm\\!$ 0.6 | 18.3 $\\!\pm\\!$ 0.5 | 18.2 $\\!\pm\\!$ 0.6 | 21.7 $\\!\pm\\!$ 1.5
Non-IID | Linear PH | 18.3 $\\!\pm\\!$ 0.8 | 18.0 $\\!\pm\\!$ 0.6 | 19.3 $\\!\pm\\!$ 0.7 | 22.9 $\\!\pm\\!$ 0.7 | 22.1 $\\!\pm\\!$ 0.3 | 21.4 $\\!\pm\\!$ 0.3 | 20.4 $\\!\pm\\!$ 0.5 | 20.1 $\\!\pm\\!$ 0.5 | 22.1 $\\!\pm\\!$ 0.3
NN PH | 18.3 $\\!\pm\\!$ 0.5 | 18.1 $\\!\pm\\!$ 0.2 | 19.7 $\\!\pm\\!$ 1.0 | 22.6 $\\!\pm\\!$ 0.5 | 21.9 $\\!\pm\\!$ 0.2 | 24.0 $\\!\pm\\!$ 1.8 | 20.5 $\\!\pm\\!$ 0.6 | 20.6 $\\!\pm\\!$ 0.4 | 22.3 $\\!\pm\\!$ 0.5
NN nonPH | 20.9 $\\!\pm\\!$ 0.3 | 20.1 $\\!\pm\\!$ 0.5 | 21.1 $\\!\pm\\!$ 1.0 | 26.2 $\\!\pm\\!$ 0.4 | 25.3 $\\!\pm\\!$ 0.5 | 25.3 $\\!\pm\\!$ 0.7 | 23.6 $\\!\pm\\!$ 0.9 | 22.4 $\\!\pm\\!$ 0.8 | 22.7 $\\!\pm\\!$ 1.0

Table 4: Integrated Brier scores (rebased to 100) – mean and standard
deviation. Lower values are better.

####  4.3.2 Impact of Discretisation Fineness.

Figure 3: Model performance with increasing discretisation fineness. Federated
models were trained with 100 global and 1 local rounds. Performance decreases
on smaller METABRIC and GBSG datasets with mixed results on the larger SUPPORT
dataset.

We re-train models on finer time grids using 100 global and 1 local rounds. A
finer grid on METABRIC (Figure 3 upper panel) and GBSG (Figure 3 lower panel)
did not improve performance and, in fact, appears to degrade performance.
Notably, the Linear PH model becomes a no-skill predictor in terms of
concordance on the non-IID GBSG case. The results are less conclusive on
SUPPORT (Figure 3 middle panel), as a finer time grid appears to result in a
minor to no increase in concordance at the expense of a loss in calibration.

A finer time grid can be expected to result in a trade-off between closer
approximation of true (smooth) survival and a reduction in data available in
any given time step. An increase from 10 to 20 time steps, for instance,
halves the number of available data points to estimate a given step. The
latter effect appears to dominate on the smaller METABRIC and GBSG datasets,
and less so for the approximately 4-times larger SUPPORT dataset.

##  5 Conclusion

We present a federated Cox model that relaxes the proportional hazards (PH)
assumption and demonstrate its ability to maintain concordance and calibration
relative to a pooled baseline under various linearity and PH assumptions.
Compared to prior work, this federation scheme encodes the decision between PH
and non-PH in a binary choice over the output layer, rather than requiring
upfront agreement on a specification of $f(t)$. We note that our model is not
restricted to a particular data type or network architecture excepting the
output component. Future work could adapt the model for image-based federated
survival predictions.

The decrease in performance on non-IID data (even if pathologically derived in
this paper) represents a challenge to the application of federated learning in
practice. Extensions could include exploring methods accounting for
statistical heterogeneity (Yang et al., 2020; Li et al., 2018) or other
federation topologies which maintain locally specialised models trained in a
peer-to-peer fashion (Rieke et al., 2020). While the heterogeneity in this
paper was derived from label stratification, other types of heterogeneity,
such as covariate shifts, could be explored: for image-based survival
predictions, differences in acquisition protocols could provide one such
avenue.

#### Acknowledgements.

This work was supported by the UKRI CDT in AI for Healthcare
http://ai4health.io (Grant No. EP/S023283/1)

##  Appendix A Additional Figures

(a) Kaplan-Meier-based discretisation.

(b) Constant density interpolation.

Figure 4: Discretisation and interpolation. Figure 5: Two sets of survival
estimates with correct ranking (green above blue) but poor calibration given
under- / overestimation of true survival curves. Figure 6: Kaplan-Meier
estimates with 95% confidence interval.

## References

  * Cox [1972] David R. Cox.  Regression Models and Life-Tables.  _Journal of the Royal Statistical Society. Series B (Methodological)_ , 34(2):187–202, 1972.  ISSN 0035-9246.  doi: 10.1111/j.2517-6161.1972.tb00899.x.  URL https://www.jstor.org/stable/2985181.
  * Katzman et al. [2018] Jared L. Katzman, Uri Shaham, Alexander Cloninger, Jonathan Bates, Tingting Jiang, and Yuval Kluger.  DeepSurv: Personalized treatment recommender system using a Cox proportional hazards deep neural network.  _BMC Medical Research Methodology_ , 18(1):1–11, 2018.  ISSN 14712288.  doi: 10.1186/s12874-018-0482-1.
  * Luck et al. [2017] Margaux Luck, Tristan Sylvain, Héloïse Cardinal, Andrea Lodi, and Yoshua Bengio.  Deep Learning for Patient-Specific Kidney Graft Survival Analysis.  (Nips 2017), 2017.  URL http://arxiv.org/abs/1705.10245.
  * Kvamme et al. [2019] Havard Kvamme, Ornulf Borgan, and Ida Scheel.  Time-to-event prediction with neural networks and cox regression.  _Journal of Machine Learning Research_ , 20:1–30, 2019\.  ISSN 15337928.
  * Zhu et al. [2017a] Xinliang Zhu, Jiawen Yao, and Junzhou Huang.  Deep convolutional neural network for survival analysis with pathological images.  _Proceedings - 2016 IEEE International Conference on Bioinformatics and Biomedicine, BIBM 2016_ , (2):544–547, 2017a.  doi: 10.1109/BIBM.2016.7822579.
  * Zhu et al. [2017b] Xinliang Zhu, Jiawen Yao, Feiyun Zhu, and Junzhou Huang.  WSISA: Making survival prediction from whole slide histopathological images.  _Proceedings - 30th IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2017_ , 2017-Janua:6855–6863, 2017b.  doi: 10.1109/CVPR.2017.725.
  * Li et al. [2019] Hongming Li, Pamela Boimel, James Janopaul-Naylor, Haoyu Zhong, Ying Xiao, Edgar Ben-Josef, and Yong Fan.  Deep convolutional neural networks for imaging data based survival analysis of rectal cancer.  _arXiv_ , pages 1–4, 2019.  ISSN 23318422.
  * Bello et al. [2019] Ghalib A. Bello, Timothy J.W. Dawes, Jinming Duan, Carlo Biffi, Antonio de Marvao, Luke S.G.E. Howard, J. Simon R. Gibbs, Martin R. Wilkins, Stuart A. Cook, Daniel Rueckert, and Declan P. O’Regan.  Deep-learning cardiac motion analysis for human survival prediction.  _Nature Machine Intelligence_ , 1(2):95–104, 2019\.  ISSN 25225839.  doi: 10.1038/s42256-019-0019-2.
  * Wang et al. [2019] Ping Wang, Yan Li, and Chandan K. Reddy.  Machine learning for survival analysis: A survey.  _ACM Computing Surveys_ , 51(6), 2019.  ISSN 15577341.  doi: 10.1145/3214306.
  * Kelly et al. [2019] Christopher J. Kelly, Alan Karthikesalingam, Mustafa Suleyman, Greg Corrado, and Dominic King.  Key challenges for delivering clinical impact with artificial intelligence.  _BMC Medicine_ , 17(1):1–9, 2019.  ISSN 17417015.  doi: 10.1186/s12916-019-1426-2.
  * Wiens et al. [2019] Jenna Wiens, Suchi Saria, Mark Sendak, Marzyeh Ghassemi, Vincent X. Liu, Finale Doshi-Velez, Kenneth Jung, Katherine Heller, David Kale, Mohammed Saeed, Pilar N. Ossorio, Sonoo Thadaney-Israni, and Anna Goldenberg.  Do no harm: a roadmap for responsible machine learning for health care.  _Nature Medicine_ , 25(September), 2019.  ISSN 1546170X.  doi: 10.1038/s41591-019-0548-6.  URL http://dx.doi.org/10.1038/s41591-019-0548-6.
  * Rieke et al. [2020] Nicola Rieke, Jonny Hancox, Wenqi Li, Fausto Milletarì, Holger R. Roth, Shadi Albarqouni, Spyridon Bakas, Mathieu N. Galtier, Bennett A. Landman, Klaus Maier-Hein, Sébastien Ourselin, Micah Sheller, Ronald M. Summers, Andrew Trask, Daguang Xu, Maximilian Baust, and M. Jorge Cardoso.  The future of digital health with federated learning.  _npj Digital Medicine_ , 3(1):1–7, 2020.  ISSN 23986352.  doi: 10.1038/s41746-020-00323-1.  URL http://dx.doi.org/10.1038/s41746-020-00323-1.
  * Andreux et al. [2020] Mathieu Andreux, Andre Manoel, Romuald Menuet, Charlie Saillard, and Chloé Simpson.  Federated Survival Analysis with Discrete-Time Cox Models.  _arXiv_ , pages 1–21, 2020.  ISSN 23318422.
  * Bellera et al. [2010] Carine A. Bellera, Gaëtan MacGrogan, Marc Debled, Christine Tunon De Lara, Véronique Brouste, and Simone Mathoulin-Pélissier.  Variables with time-varying effects and the Cox model: Some statistical concepts illustrated with a prognostic factor study in breast cancer.  _BMC Medical Research Methodology_ , 10, 2010.  ISSN 14712288.  doi: 10.1186/1471-2288-10-20.
  * Coradini et al. [2000] Danila Coradini, Maria Grazia Daidone, Patrizia Boracchi, Elia Biganzoli, Saro Oriana, Gianluigi Bresciani, Cinzia Pellizzaro, Gorana Tomasic, Giovanni Di Fronzo, and Ettore Marubini.  Time-dependent relevance of steroid receptors in breast cancer.  _Journal of Clinical Oncology_ , 18(14):2702–2709, 2000.  ISSN 0732183X.  doi: 10.1200/JCO.2000.18.14.2702.
  * Gore et al. [1984] Sheila M. Gore, Stuart J. Pocock, and Gillian R. Kerr.  Regression Models and Non-Proportional Hazards in the Analysis of Breast Cancer Survival Author.  _Journal of the Royal Statistical Society. Series C (Applied Statistics)_ , 33(2):176–195, 1984.
  * McMahan et al. [2017] H. Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Arcas.  Communication-efficient learning of deep networks from decentralized data.  _Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, AISTATS 2017_ , 54, 2017.
  * Faraggi and Simon [1995] David Faraggi and Richard Simon.  A neural network model for survival data.  _Statistics in Medicine_ , 14(1):73–82, 1995\.  ISSN 10970258.  doi: 10.1002/sim.4780140108.
  * Antolini et al. [2005] Laura Antolini, Patrizia Boracchi, and Elia Biganzoli.  A time-dependent discrimination index for survival data.  _Statistics in Medicine_ , 24(24):3927–3944, 2005\.  ISSN 02776715.  doi: 10.1002/sim.2427.
  * Lu et al. [2015] Chia Lun Lu, Shuang Wang, Zhanglong Ji, Yuan Wu, Li Xiong, Xiaoqian Jiang, and Lucila Ohno-Machado.  WebDISCO: A web service for distributed cox model learning without patient-level data sharing.  _Journal of the American Medical Informatics Association_ , 22(6):1212–1219, 2015.  ISSN 1527974X.  doi: 10.1093/jamia/ocv083.
  * Dai et al. [2020] Wenrui Dai, Xiaoqian Jiang, Luca Bonomi, Yong Li, Hongkai Xiong, and Lucila Ohno-Machado.  VERTICOX: Vertically Distributed Cox Proportional Hazards Model Using the Alternating Direction Method of Multipliers.  _IEEE Transactions on Knowledge and Data Engineering_ , 4347(c):1–1, 2020.  ISSN 1041-4347.  doi: 10.1109/tkde.2020.2989301.
  * Craig et al. [2021] Erin Craig, Chenyang Zhong, and Robert Tibshirani.  Survival stacking: casting survival analysis as a classification problem.  _arXiv_ , pages 1–17, 2021.  URL http://arxiv.org/abs/2107.13480.
  * Gensheimer and Narasimhan [2019] Michael F. Gensheimer and Balasubramanian Narasimhan.  A scalable discrete-time survival model for neural networks.  _PeerJ_ , pages 1–19, 2019.  ISSN 21678359.  doi: 10.7717/peerj.6257.
  * Kvamme and Borgan [2019] Havard Kvamme and Ornulf Borgan.  Continuous and discrete-time survival prediction with neural networks.  _arXiv_ , 2019.  ISSN 23318422.
  * Graf et al. [1999] Erika Graf, Claudia Schmoor, Willi Sauerbrei, and Martin Schumacher.  Assessment and comparison of prognostic classification schemes for survival data.  _Statistics in Medicine_ , 18(17-18):2529–2545, 1999.  ISSN 02776715.  doi: 10.1002/(sici)1097-0258(19990915/30)18:17/18¡2529::aid-sim274¿3.0.co;2-5.
  * Paszke et al. [2019] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  PyTorch: An imperative style, high-performance deep learning library.  _Advances in Neural Information Processing Systems_ , 32(NeurIPS), 2019.  ISSN 10495258.
  * Yang et al. [2020] Lixuan Yang, Cedric Beliard, and Dario Rossi.  Heterogeneous data-aware federated learning.  _arXiv_ , (1), 2020.  ISSN 23318422.
  * Li et al. [2018] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith.  Federated Optimization in Heterogeneous Networks.  2018\.  URL http://arxiv.org/abs/1812.06127.