Pooling Image Datasets with Multiple Covariate Shift and Imbalance
Authors: Sotirios Panagiotis Chytas
Vishnu Suresh Lokhande
Vikas Singh
Venue: ICLR 2024 poster (2024)
Keywords: image harmonization
medical imaging
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: Use Category Theory to provide a general harmonization tool that can handle multi-equivariance and multi-invariance with respect to the images' covariates.

Abstract: Small sample sizes are common in many disciplines, 
which necessitates pooling roughly similar datasets across 
multiple sites/institutions to study weak but relevant 
associations between images and disease incidence. Such 
data often manifest shifts and imbalances in covariates 
(secondary non-imaging data). 
These issues are well-studied for classical models, but 
the ideas simply do not apply to overparameterized DNN models. 
Consequently, recent work has shown how strategies from 
fairness and invariant representation learning provides 
a meaningful starting point, but the current repertoire 
of methods remains limited to accounting for shifts/imbalances in just a couple of covariates at a time. In this paper, we show how 
viewing this problem from the perspective of Category theory 
provides a simple and effective solution that completely avoids 
elaborate multi-stage training pipelines that would otherwise be 
needed. We show the effectiveness of this approach via 
extensive experiments on real datasets. Further, we 
discuss how our style of formulation offers a unified 
perspective on at least 5+ distinct 
problem settings in vision, from self-supervised learning
to matching problems in 3D reconstruction.

Published as a conference paper at ICLR 2024
POOLING IMAGE DATASETS WITH MULTIPLE COVARI- ATE SHIFT AND IMBALANCE
Sotirios Panagiotis Chytas" UW-Madison
Vishnu Suresh Lokhande UW-Madison
Vikas Singh UW-Madison
ABSTRACT
Small sample sizes are common in many disciplines, which necessitates pooling roughly similar datasets across multiple institutions to study weak but relevant associations between images and disease outcomes. Such data often manifest shift/imbalance in covariates (i.e., secondary non-imaging data). Controlling for such nuisance variables common within standard statistical analysis, but the ideas do not directly apply to overparameterized models. Consequently, recent work has shown how strategies from invariant representation learning provides a meaningful starting point, but the current repertoire of methods is limited to counting for shifts/imbalances in just couple of covariates at a time. In this paper, we show how viewing this problem from the perspective of Category theory pro- vides a simple and effective solution that completely avoids elaborate multi-stage training pipelines that would otherwise be needed. We show the effectiveness of this approach via extensive experiments on real datasets. Further, we discuss how this style of formulation offers a unified perspective on least 5+ distinct problem settings, from self- self-super learning to matching problems in 3D reconstruction. The code is available at https//ghhuub.co////CyytaaaaaCi
INTRODUCTION
Sample sizes of medical imaging datasets at a single institution are often small due to many reasons. Ac- quiring thousands of images can be infeasible due to budget/logistics. Also, if the scope of study nar- row, only some individuals may be eligible due to in- clusion criteria (e.g., have a specific genetic risk). To improve the statistical signal in retrospective analyses of existing datasets, one option is to pool similar data across multiple sites (Thompson et al. 2014). This of- fers a chance at discovering a real statistical effect, un- detectable with small sample sizes.
Scanner
Scanner
Scanner
Sowen
Pooling data from multiple sites is common. A ma- ture body of statistical literature (e.g., covariate match- ing (Stuarti 2210) Kim & Steiner (2066), meta-analysis Thompson et al. 2014; Rücker et al., (02T)) de- scribes best practices, and mechanisms to account for the effect of a nuisance variable on the response (or target label) are well known. This allows obtaining associations between relevant predictors and the re- sponse/dependent variable of interest. For example, a standard analysis workflow may use only the images to predict a label (e.g., cognition) while controlling for nuisance variables such as race, gender or scanner type Penny al. 2007).
APOE Figure 1: Problem overview: When pool- ing image datasets from different sources, there may be differences the distribution of covari- Covariates are secondary data for each in- dividual that influences the images systemati- cally. Top row shows MR images (say, different scanners). Second row shows how the covari- ate distributions (genetic risk, age or gender) varies across scanners (but have shared sup- port). Learning representations from a pooled dataset in a manner such that covariate varia- tions are accounted for, is challenging.
AGE
SEX
*Corresponding author. Email: chytas@wi sc.edu

---

Published as a conference paper at ICLR 2024
Controlling for covariates. Deep neural network (DNN) models are now ubiquitous in medical image analysis |Ronneberger et al. 2015 Carin & Pencinal 2018). But when using such models, systematic mechanisms to deal with nuisance variables (or covariates), key concern in data pooling, are still under development. Consider pooling MR images from two different sites where participant demographies distributions (e.g., age, sex etc) across sites are similar. If the scanners the sites are different, we can include "scanner" or other continuous variables as a nuisance, and control for it within general linear model. Features in the image X deemed informative will explain variability in the response Y after the variability due to "age" and "site" has been accounted for. In other words, the image data (and not the site- site-specific artifacts) should predict the response variable. Vaamonizatiionpooliin of imaging data in a way that accounts for (or removes) the influence of non-imaging covariates is often limited to shallow pre-processing using additive/multiilicatiee batch corrections, e.g., in Combat (Johnson et al., 2007 Fortin et al. 2018).
Data pooling/harmonization
and invariance. In the context of representation learning, several re- sults have identified the link between data pooling/harmonnzation and invariant representation learn- ing or domain adaptation (Moyer et al. 2018; 2020), which have been utilized for downstream tasks (Lokhande al. 2020). Initial approaches focused image normalization techniques (e.g., his- togram matching Nyul & Udupa (1999), which allowed controlling some variations in intensities across scanners or protocols. However, such approaches cannot systematically handle covariates. Ideas based on invariant representations (Bloem-Reddy & Teh 2020 Li et al.) [2144] Arjovssy et al.) 2019] have allowed controlling for up to two covariates during representation learning. For exam- plc, one may ask that the model avoid using image features associated with two specific covariates (or in fairness terminology, sensitive attributes): age and site. Nonetheless, models that can easily remove (or control for) the influence of multiple covariates remain limited and so, either Combat- based (Johnson et al. 2207)) pre- pre-processing is used (Fortin et al. [2018) or CycleGAN-basee schemes (Zhang et al.) 2018 Nguyen et al.) 2018)) can harmonize the image data relative to a specific cat- egorical variable (say, scanner). This problem is also relevant in the longitudinal setting /auuy &
Durrleman 20222).
The problem. Different scanners can introduce systematic differences in the scans of the same person (Liu et al. 2020) and mitigation strategies continue to be topic of recent work (Zhang al.) 2023). When there is some shift/disbalance in covariates (participant data including age, sex, and so on, which influence the scans) across sites and shared support partial ("age" distribution in Fig. []. the harmonization task becomes more involved. The goal is to learn representations as the covariates were matched across the sites begin with, see Fig. D Note that in contrast covariate shift methods (Bickel et al.] 2009). the shift here is not between train and test distributions, rather between different sources of data in the training set itself.
One recent attempt in (Lokhande et al. (022)) to learn deep representations from pooled data which can also handle (Shift-imbalance) in covariates needs a two stage pipeline. The first stage obtains latent representations of the scan which is equivariant to age. The second stage enforces invariance to scanner and trained after freezing the first stage. Two covariates can be handled but an extension to many continuous (and categorical) covariates will need a complicated multi-stage model.
Main Ideas. In most learning tasks involving an equivariance or invariance criteria, the goal is to allow the model to benefit from structure and symmetry - in the data, the parameter space, or both. Frequently, one formalizes these ideas using group theory seeking invarianceecquiiariance to the action of the group (e.g., SO(n), SE(n), S.,) (Bloem-Redde & Teh 2020 Kondor & Trivedi 2018). More general criteria (beyond properties native to the chosen group) requires specialized treatment. Our starting point is based on the observation that Category theory (Fong & Spivak 2019 ManLane| 2014) provides a rich set of tools by which the "structure" (either among the co- variates the data more generally) can be expressed easily. It is known that equivariance, invariance and many group-theoretic constructs will emerge as special cases because category theory provides more abstract treatment. Once our criteria are expressed in this way, during training, the necessary constraints on the latent space fall out directly, and the formulation gracefully handles heterogeneity in many contimuous/categoricaal covariates. No ad-hoc adjustments are needed.
Contributions. (a) On the technical side, we provide general framework for imposing structure on the latent space by applying ideas from Category theory. Equivariant (and invariant) representation leamning emerge as special cases. Further, this style of formalism unifies different formulations in vision and machine learning, under the same umbrella. (b) On the practical side, we strictly
2

---

Published as a conference paper at ICLR 2024
generalize existing formulations to pool/harmonize multi-site datasets. While the existing two-stage formulation can deal with one categorical and one continuous covariate, our formulation places no restriction at all on the number of covariates. We show how the same model can also be used to heuristically approach certain hypothetical "what if"" questions and offers competitive performance on public brain imaging datasets, with strictly more functtionaliy/ffeiiility
A BRIEF REVIEW OF CATEGORY THEORY
Category theory offers a way to study abstract "structures" (Eilenberg & MacLane, 1945 MacLane 2014]. Category consists of two components (a) Objects that correspond to individual entities (e.g., scans/images), and (b) Morphisms, identified by the paths (or "arrows") between the objects.
Each Object has an identity Morphism (a self-loop often omitted when drawing the diagram).
Remark The composition of two Morphisms f S1 S2 and g S2 - S3 is well-defined Morphism and it is denoted as g S + S3 (we should read as g after f).
Example 2 Consider the Cutegory of Sets, in which the Objects are sets and the Morphisms cor- respond to functions for matchings) between sets. Since 3ʃ S1 S2, g S2 S3 then 3h = g S1 S3.
In the Category of Sets ids S S corresponds to the identity function: idg(s) s, E
Definition Functors F S + T give relationship between the Objects and the Morphisms of two different Categories: source Category S and target Category T. The conditions below hold,
gof
(i) F(ids) = idr(s) ASES
ids S
S2 1 F(S2)
7
S3
(i) F(go (F())) F())
:1 - S2.9: S2 + S3, in S
idF(S F(SI)
FUS)
Fla)
F(S:)
F(g)oF(f)
Remark 4 The reader will see why the specification above is general, but useful. The functor F takes us from the first Category to the next in way that the structure of the source Category (objects and arrows) is fully preserved in the target Category because of the two conditions above.
3 ENFORCING SYMMETRY AND STRUCTURE
Symmettyy/structure in the data parame- ters are key properties we exploit in learn- ing tasks, e.g., invariance in features (Lowe)
to_
Tolale
S
Ai
1999). compressive sensing (Candes & Wakin 2008) well as DNN training Sabour et 2017 Hinton 2021). Encod- ing symmetry/structure can involve ideas as simple rotations and cropping for data augmentation (Shorten Khoshgot- taar 2019; van Dyk & Meng [2001] Chen cct al.) 2020b) to more sophisticated con- [Cohen et al.) 2018 Fuchs et al.)
to
F
F(to
T
F(A)
F(B)
F(L))
idrup
F((o_A)
Figure 2: A Category theoretic view of CycleGAN (left) and SimCLR (right)
cepts
2020). This criteria can also be specified on the latent space, e.g., via geometricfopologicall pri- (Moor al. 2020 Ghosh et al. (020)) other conditions (Chen al. 2019 Oring al.) (021). From Fig. [] if there is a natural structure (say, in covariate space), then the structure of the repre- sentations learned on the latent space should respect that. To setup this description, first discuss how this perspective can relate distinct ideas under an abstract (but simple) formalism.
3.1 REINTERPRETING CYCLEGAN USING CATEGORY THEORY
Image translation uses a set of image pairs and trains a DNN model to map between them. Cy- cleGAN (Zhu al. 2017) is popular framework for image translation that removes the need for
3

---

Published as a conference paper at ICLR 2024
paired image pairs in the training data, and learns the map from domain (or class) to the other. Denoting image domains as type A and rype B, the goal in CycleGAN is to learn how to map a rype A image to type B image and back; which we can denote as B and to A: type A (and type B) can mean images of horses (and zebras). We will avoid defining a category for the distributions over the images for simplicity, and so training dataset made up of paired images will suffice to convey the key idea. So, for every image of type A, there is similar image of type B. We want to learn mapping (Functor) to latent space that preserves the action of change from type A B and back. This is shown in Fig. 2 (left), from which we can directly read off the following constraints,
1. F a fully-faithfie Functor meaning: 3F T S such that F-1 F(s) id, 2. Composition: to B o to A(B.) = idg, (B.) = B, to A to ((A,) = idA, (A,) A,
The first pair of constraints is implicit in autoencoder-like models. The second pair of constraints define the cycle consistency conditions in CycleGANN We note that different formalism for Cy- cleGAN was described in (Gavranović 22000) by considering CycleGAN as a schema where cycle- consistencies enforce composition invariants in that Category.
3.2 REINTERPRETING SIMCLR USING CATEGORY THEORY
SimCLR (Chen et al. 2020b) is a self-supervised learning method where one attempts to map the images to a latent space which is invariant to transformations such as rotations, croppings, and so on. This can be instantiated using an invariant Functor (see Fig. 2 Unlike CycleGAN, such an invariant Functor F, is no longer fully-faithful (i.e., FF-! such that F o F- = idz).
3.3 REINTERPRETING OTHER FORMULATIONS USING CATEGORY THEORY
(a) Latent space interpolation. A number of approaches seek to perform interpolation in latent space (Chen al. 2019 Oring et al. 202I)) informed by "covariates" associated with the original images (e.g., age, sex, hair length). The recipe described in the two examples above lends itself di- rectly to describing such formulations, and will be subsumed by our formulation in the next section. (b) Rotation and Permutation synchronization. Rotation synchronization is a central problem in Cryo-EM image reconstruction (Marshall et al.) 2022 Bendory et al. 2022 Tagare al.. 2008). In computer vision, permutation synchronization Pachauri et al. 2013) seeks to find replicable matches of keypoints across many images of the same 3D scene, e.g., for 3D reconstruction from datasets such as Tourism (Li al. 2022, [Brdal et al. 2021 Birdal & Şimşekli, 2019) as well in robotics (Leonardos et al. 2017). Since the synchronization task always involves consistency con- straints over the symmetric Group, and Groups
Category, the reinterpretation is immediate. (c) Equivariance and Invariance. Many problems in learning benefit from equivariance and invari- ance. Recall that equivariance is formally defined in terms of action of Group G.
Definition 5 (Equivariance) A mapping f - T defined over measarable Borel spaces S and T is said to be G-equivariant under the action of group G iff
f[g 8) = g f((), UEC G
(1)
A more general definition of equivariance follows directly from the Functor's definition.
Definition 6 A Functor F S + T is defined as a mapping from S to T such that
F(ids) = id VS S
(2)
F(g(S:)) = (F(())) Vs.g S₁ + S2 S
(3)
If S, are Borel spaces, and g, F()) belong to Group G, then we obtain the group-theoretic equivariance definition. Category theory gives more general result with no restrictions on the form of the relationships (or Morphisms) g. Similarly, invariance (e.g., in Group theory) is defined as
Definition 7 A mapping f S + T defined over measurable Borel spaces S and T is said to be G-invariant under the action of group G iff
s( = f(s), gEG
(4)
In Category theory, this results in a special type of Functor:
Definition 8 (Invariance) Functor + T invariant to the Morphism g : S + S2 if F(g) = (s)-
4

---

Published as a conference paper at ICLR 2024
3.4 A SIMPLE SANITY CHECK OF BENEFITS ON MNIST DATASET
Consider the case where the objects of the source Category S consist of MNIST images (LeCun) jtt al., 1998). Since these images represent integers on the number line, we can ask whether we can learn latent space which allows algebraic manipulations, with one (or more) basic operations de- fined on it. If our operation of interest was "addition", a primitive we will need for counting would be the "+1" operation. To keep things simple, we will focus this operation where the source Category S consists of Morphisms that represent the "+1" difference between two subsequent dig- its. It is easy check that we have indeed defined a Category since each Object has an identity Morphism and the Morphisms compose. The target Category T (i.e., our latent space) can consist of vectors v Rn objects. Orthogonal linear mappings W |R=* as Morphisms will suffice. In our pooling/harmonization problem, the morphisms will reflect traversing the axis of each covariate.
For this MNIST example, our goal is to impose the "+1" operation in the latent space, see Fig. 3 We de- fine our loss to express precisely the equivalence for
0
1
2
3
4
5
6
?
8
9
F--1
F
F-1
subsequent digits as well as a subset of their +1 com-
W
W-!
positions that can be read off of Fig. [3] The model is
W
W W.
nof presented data for "-1" operations but discovers due to the special structure of W (orthogonal, so W- exists). Then, after training. when presented a "seed" image of 6, the model can (forward/backward) traverse the latent space: query "6 - 4" starts from the latent code for 6, hops backwards 4 times and generates an image of 2. All images except "6" shown in Fig. were generated in this way. It is worth making a brief comment on the generality here. If we had & different morphisms, say, "+1" as well as rotation, scaling and shear, we could apply them in sequence and generate the corresponding images, even if such training data were not shown to the model (more experiments can be found in (A). This feature will allow dealing with multiple covariates easily our experiments. A similar solution has been provided in Keeurt et al. (2023), although their experiments were restricted to simple datasets and dealt only with 2D object dispositions. Here, we already showed how to learn something more complicated (addition) and, in the next section, we will show how to adapt this idea to complicated covariates.
Figure 3: MNIST Example: Modelling the re- lationships between the digits as linear mappings in the target Category.
4 A CATEGORY THEORY INSPIRED FORMULATION FOR DATA POOLING
Overview. The preceding section provides us all the necessary modules. Our task involves using brain MR images to predict diseased or healthy controls status. Our covariates will include sec- ondary (non-imaging) data pertaining to the participants, including scanner type, site, age, sex, and genotype status. Some of these are ordinal/continuous such as genotype (APOE; three risk types) and age, whereas others are categorical such as scanner type, site and sex. For the categorical co- variates, our formulation will seek to enforce invariance on the learned representations, i.e., asking that the latent representations be devoid of information pertinent to nuisance variables like scanner and site. It turns out that sex is associated with Alzheimer's 's disease (AD) because two-thirds of those diagnosed are women (Mielke 2018). But if our goal is understand which image features are relevant for the disease (and not simply to maximize accuracy), it makes sense to control for sex as a nuisance variable and add it separately at the last layer, if needed. For ordinal/continuous co- variates, we see both shifts and imbalances across sites or scanner (different sites/scanners may not cover exactly the same age range of participants or include the same number of individuals for each genotype value). Equivariance will allow adjusting for these shifts, i.e., coordinate system where the latent representations are "aligned" (modulo their covariate induced morphism). The diagram in Fig. 4 will setup overall formulation, and we use the example below to describe its operation.
Example 9 Assume we want to control only for age when learning the latent representations of the images for objects/participanss S, and S2. S1, S2 are identical for all covariates but have different ages. Their scans are also different. If S2 is years older than St, in S, we have S2 z fS1) (i.e., equal up to isomorphism. In practice, isomorphism can be as strict as lz-norm, cosine similarity as in -ased models (Radford et al.) [02T), or even distribution-
measures such as MMD). For someone, 2x years older. similar to Fig. 3 we will compose fa twice. In T. latent representations Tı, T2 correspond to F(Si) and F(S2). We seek to learn a functor F by learning
5

---

Published as a conference paper at ICLR 2024
Is = F(fz) such that T2 _ 9x((T)). S1 and S2 differed in two covariates, by amounts z and resp., the morphism from S1 to S2 would involve composing fx and hy (if h denoted the morphisms for the second covariate). The morphisms in the Category T would compose similarly.
Setting up loss function. We will use simple mor- phisms which correspond to linear transformations W E R (this is design choice). The pair of Func- tors (F, F correspond to an autoencoder, while the pair (I, C*) corresponds to appropriate classifier (or regressor). Simply walking the paths from Fig. f pro- vides all necessary constraints which we write as dis- tinct terms in our loss function,
S
T
F))
51
4
F(11)
C1
F(D=W
W
S
F(S2)
R
(x) idz gives the reconstruction loss:
(5) Figure 4: diagrammatic representation of equivariance with respect to a single covariate whose change corresponds to "f" in the orig- inal data space. Our goal is to preserve this (6) structure in the latent space (Category T) in which model FJ)) linear transforma- tion E R=x= many practical cases the downstream goal is to classify the latent repre- sentations, which we formulate with the Func- tor C T Free(N) (or Free((R) for regression task, etc.)
C.
Σ 1
o ())
AES
(ii) the prediction loss (given labels y):
Lp Σ CE (y (C (C)FF(()
(1,8)S
emss-entropy
(iii) Finally, preserving the morphisms, in its simplest form, is a structure- preserving loss,
structure-peeserving
L. =
Σ 81
W F(a)) - F(82)
(7)
The training for multiple covariates (Alg. is direct generalization of the above formulation.
Algorithm Structure preserving training of Funetors
Input: Parameters: (0, W), Multipliers:
Data: input/output: yS, E (B xn, 187)) and covariates: C C Nmxe
for ep in epochs do
L, =. - (Fo, (s) C. E CE(y,, (Co, (o,) C. c C'
for (81, pairs((S c) do
C. += l|Wa . Fo, (s1) Fo, (s2)]l2
C= + +e6 AaC. step(C, (01, 02 0s, Wi,..., We))
// update
Function pairs (S, c):
pairs 11
10
for (i, m x m do
11
if St is paired with Sj then pairs. (S., c₁)
13
return pairs
14
Enriching the latent space with structure provides some information about "hypotheticals" Given that we are unlikely to be provided different versions of sample (one for each composition of the known Morphisms), a structure preeserrne- latent space gives us a way to:
ii generate new data samples, by evaluating the expression (F-l o F()) - (S). answer questions, by evaluating the expression (C F()) (S).
Role of Category Theory. There is clear advantage of expressing the problem in Category Theory. enables a more concise and straightforward representation of the problem (sometimes, apparent in hindsight), as well as additional capabilities compared to earlier works (Lokhande et al. 2022).
5 EXPERIMENTAL EVALUATIONS
Rationale. Our motivation stems from pooling brain imaging datasets, particularly in scenarios where covariate distributions vary across sites. In the ADNI brain imaging study (Mueller et al.

---

Published as a conference paper at ICLR 2024
while acquisition protocols are consistent among 50+ sites, demographic distributions differ. The diversity of scanners across sites, along with the impact of scanner upgrades on analysis tasks (Chen et al. 2020a Assffrr et al. 2021 Lee et al. 2019), further complicates the scenario. Scanner variations, even within the same study (e.g.. ADNN-3), introduce controlled nuisances for standard analyses. Beyond ADNI, we also perform analysis on the ADCP dataset Lokkand([ al. (2022). Our objective is to incorporate this functionality into representation learning for iigrrsssiooollsssiiiattinn
2005).
This problem setting underscores limitations in existing methods. Recent VAE or GAN-based ap- proaches (Ren et al. 202T] Hu et al. 2023 Moyer et al. 2020 Bashyam et al. 2022) target scanner invariance, often accommodating binary or multiple scanner variables. However, handling variations in other covariates (c.g., age, sex, APOE) remains challenging (Husain et al.) (2021) (see SB.I).
Baselines. The following baselines address distributional differences in the data while maintaining associations with covariates: 1. Naive: Pooling data without pre/post processing. 2. MMD (Li al. Minimizes Maximum Mean Discrepancy (MMD) for invariance but lacks equivariance. 3. CAT (Xie et al. 2017): Achieves invariance with a discriminator but lacks equivariant mappings.
and minimizes MMD for each subgroup. 5. RM (Motiian et al., 2017): Matches similar samples from different scanners for invariance. 6. GE (Group Equivariance) (Lokhande et al. 2022): Uses group theory to minimize MMD while sceking equivariance with respect to one covariate. GE achieves the best overall accuracy but has limitations: it uses expensive matrix exponentials and relies on a two-stage training, handling only two covariates.
SS (Zhou et al. 2017): Divides the population into subgroups (Subsampling-2
Evaluations. As the classification task, we will predict Alzheimer's ' disease (AD)/Control normals (CN) labels using the invariant and/or equivariant representations. We will use the following metrics to systematically assess each algorithm. or not participant has AD. Accuracy is reported in order ensure that the model, despite its extra constraints, maintains all the required information that an input image carries. (b) M.MD: MMD is a measure of distributional differences, defined as the Euclidean distance of the kernel embeddings means (typically an RBF kernel). (c) ADV: an alternative to measure invariance is by training a model to predict the nuisance covariate (e.g.. "Scanner") using the latent representations. If obtain latent representations devoid of this information, then the accuracy should be almost random. The above two metrics identify invariance with respect to scanners. To check equivariance to covariates, we use two measures: Minimum distance (D) and Cosine similarity (CS),
(a) Accuracy (ACC): Test set accuracy of whether
(a) Assume that a value of C1 (for a covariate c) corresponds to an individual $1, and Morphism W E Rax= in the latent space gives a change in the covariate from e₁ to cz. Then, the Minimum Distance (similar to covariate matching) identifies the closest individual s such that & 02:
|| F(s1) F(s)
D(s1; ca) = min
(8)
w:de
(b) A low minimum distance by itself is insufficient. So, we also calculate the cosine similarity (CS) between all the other covariates of s and $2, where $2 = arg min || FF(1)) F(s)]l2
Low Minimum Distance (D) and high Cosine Similarity (CS) is desirable.
Setting. We model the Functor F using modified ResNet (He al. 2016). and the Functor C using a Fully-Connected Neural network. For our experiments, using linear mappings W E Rax= for the Morphisms in the target Category S (i.e., latent space) was sufficient but this can be easily upgraded. All the experiments are result of 5-fold cross validation procedure. In the following sub-sections, we summarize our experimental findings. We gradually increase the complexity of the experiments (number of constraints that we simultaneously optimize).
Can Category theory constraints help impose invariance to scanner? The following objective derived directly from the formulation is used,
L=
Σ AI
X- -((s)))))F((zz)
cross-cntropy
V "11nng from differrnnt Scanners
Our results (Table 2 show that we obtain invariant representations without degrading the model's accuracy, consistent with the literature on invariant representation learning (Moyer et al. 2018). MMD decreases by more than 80% compared to the best baseline, while accuracy is 1.3% higher than the naive model, suggesting stronger generalization to new samples (more experiments on tabular data can be found in (B.2 and ablation studies on À and latent space size in B.3).
7

---

Published as a conference paper at ICLR 2024
ADNI MMD(x10")
ADCP
ACC
ADV 4
ACC 1
MMM((x10))
ADV 4
Random
64
49
74
42
80(26)
27(1.6)
59
83(4.4)
90(es.7)
(90.))4
MMD Li
2014
80(2.6)
27(1.8)
59:3
84(6.5)
86(11.0)
49(11.9)
CAr et a
74(3.6)
27(1.5)
61(2.1)
(2(6.1)
85(12.3)
66(06.9)
Zhou et
2017
81.7.7)
26(1.6)
57(2.1)
(2(3.5)
88(14.6)
51(06.7)
RAK Mobaan
2017
78(3.8)
22(0.61
52
84(b.3)
77(15.8)
40(04.7)
GR Lochantle
2022
77(1.8)
16(7.2)
50(1.2)
81(18)
70,22.3)
(0(07)))
Ours (INV)
(1.3.3)
(2
62(2.0)
86
34
(0(44)2)
Ours (1 cov)
82(2.6)
11:2.9)
53(1.2)
86(5.9)
39.03.6)
18(06.7)
Ours (2 cov)
82(2.2)
10(0.6)
51(5.0)
85(7.7)
40(04.3)
44(06.1)
Ours (5 cuv)
80(2.3)
11(5.3)
52(1-79
Table 2: Quantitative Results on ADNI (Mueller et al., 2005) and ADCP from LLkhhande et al., 2022). The mean accuracy (ACC) and invariance evaluated by MMD and ADV are shown. The standard de- viations are in parenthesis. The baseline corresponds to only invariance to Scanner, cov corresponds equivariance with respect to a covariates (along with invariance to Scanner). .MMD measure is significantly reduced without any drop in accuracy. The ADV measure is close to the random baseline (desirable).
Can model be invariant to scanner and remain equivariant to covariates? Besides invariance, we seek equivariance to specific covariates in some cases. Here, we examine the simplest case of equivariance to a single covariate. We test three different covariates: 1. Age 2. Sex 3. APOE Al. Since age is continuous, we discretize it into bins of 10 years (so we have enough samples in each
bin). APOE Al (and A2) are genotypes in ADNI that are associated with AD (Husain et 2021 Sienski et al. 2021). Using our formulation, we derive the following loss function, Al
L
Cp cross entropy #1:92E5
Σ
A||
Cop F(+))= F(22)
(10)
Cey Fag represent the nuisance covariate
Note that the invariance term is implicit in this case because for two individuals 81, 82 with G2s we recover the invariance term from (13).
Ca2
While this specific problem setting is not new (see MMotiaan et al. 2017)) here we show that the simplicity of our method comes with certain advantages. First, get improved results and further, our model has fewer parameters and smaller runtime complexity (Tables) 2
Remark 10
In Fig. we examine the performance of the naive model where we infer W post-training, GE (Lokhande al.] 2022) which defines different linear transformation W for each age difference, and our model (Ours). In all three models, we model the age increase (e.g., 60 to 65) by applying a linear transformation in the latent space. Then, find the closest point for that age (e.g., 65) in the latent space w.r.t. & norm (D) and cosine sim- ilarity (CS). The transformed vector should be in the neighborhood of latent vectors with that age (c.g., (5), and D should be small and CS should be high. The results show that such a morphism is not accurate enough in the naive or even the GE algorithm. The naive model tends to "spread" the latent vectors without preserving structure, so it cannot capture equivariance in the latent space (clear when we examine D), while GE achieves equivariance but our results suggest improvements.
Encoder
Classifier
Parameters
Morphisms
Runtime
GE Ours
O(E) O(E)
O(C)
Ole P) Ole R. O(R)
O(C)
O(e)
Table Parameters/ Runtime comparison with respect to the number of covariates. e the num- ber of covariates and P the number of unique pairs for each covariate (exponential growth). E E/C is the number of parameters in the Encoder and Classifier respectively (independent of c). R stands for sin- gle training runtime. Our method's complexity is better both in parameters and runtime.
Assessing hypothetical samples. Enriching the latent space with structure provides information about questions such What would be the change in AD if APOE AI had different value? Recall that APOE A1 is ordinal covariate with values (f2, 4)) and associated with AD (Husain al., 2 Sienski et al.) 2021). For each sub-group (individuals with specific value for APOE Al), increase/decrease its value using morphisms and then classify the new vector. An increase in APOE Al results in higher probability of AD and vice-versa (Fig. 6), as expected (Husain al. 2021).
RELATED WORK
Data pooling, Fairness, disentanglement, and Invariant Representation Learning. Gen- eral approaches for analyzing data from different sites involve meta analysis (Thomp-
8

---

Published as a conference paper at ICLR 2024
al. 2014 Rücker al.
2021 When data transfer is feasible, pooling can
be approached using the Johnson-h Neyman technique (e.g.. when ANCOVA is inapplicable). Some tools from statistical genomics have been deployed for brain image data pooling (Johnson et al. 2007). also see (Garcia-Dias al. 2020). For deep models, the link between invariant represen- lation learning/fairness and data pooling was seen in (Moyer et al.) 2018 Jaiswal al. (2019 Banijamali et (0117). and has been the defacto approach [Moyer et al., 20200] Lokhande et al. 2022] Peng al., 2017: Liu al.) |002] Pham et al.) (2023) to disentangle the
Naive
123
1.0
0.9
0.01
0.8
CS
influence of nuisance variables (but provably doing so is difficult
Age
Locatello et al. 2.22)).
APOE
APOE A2
Category theory in machine learning. The applications of Category theory in applied disciplines are somewhat limited. But more recently, these ideas have been successfully applied to machine learning specifically, (Fong et 2019 Cruttwell [rutt//elttttelll et al. 2022 Fong & John- sson| 2019; Barbiero et al., (003)) model the learning process gradient-based learning algorithms using Lenses and Gavranović ett al. (2024) uses Monads. Separately, (Wilson & 'Zanasi 2020) used Category theory to offer framework for training Boolean cir- cuits. The results in (Gavranović (2020) modeled CycleGAN using Functors and showed how the lormulation can be used for insert- ing/deleting objects from an image. A summary of developments is in (Shiebler al. 202T).
Sex
Education
0.00
0.01
0.90
0.95
1.00
other fields (e.g. (Fong& & Spivak 2019)).
CS
Figure 5: (Leff)) Minimum Dis- tance (D) and Cosine Similar- ity (CS) in for age equivariance, compared with Naive and GE. (Right) Minimum Distance (D) and Cosine Similarity (CS) for 5-covariate cquivariance. Results are consistently good.
CONCLUSIONS
Imposing structure on latent spaces learned by DNN models is being actively studied, for problems ranging from disentanglement to in- terpretability. Using data pooling problem in brain imaging as a motivation, we discuss how category theory provides a precise framework for such tasks. Not only does such an approach significantly simplify recent works in the litera- ture, but also offers perspective unifying an array of standalone approaches/models. Due to its rather abstract nature, the application of category theory vision/machine learning is rather limited. Nonetheless, we believe that the models/experimmnts in this paper provide evidence that these ideas can inform other challenging problems in our field, including explainability, modu- larity and interpretability.
0.70
0.70
x(DD)
0,90
0.50
0.30
0.25
0:25
0.00
0,00
0.00
RIAD)
0.90
1150
0.30
0.00
1100
APOE
APOE
APOE
Figure 6: Increasing APOE Al in the latent space us- ing morphisms leads higher probability of AD and vice versa. Such manipulations in latent space are fea- sible through learned morphisms.
Limitations. We point out a few key caveats. Since the goal was to demonstrate the benefits of the formalism rather than maximize accuracy attained for each downstream task, the choice of most modules was kept quite simple (linear transformations, ResNets etc). Our construction offers great deal of flexibility upgrade these modules, as needed. However, due to the nature of our algorithm (i.e.. imposing constraints on the latent space), we cannot directly benefit from contemporary archi- tectures such U-Net (Ronneberger et al. 2015) Zuo al. 2021 Zhang al. 2018) since their skip-connections provide shortcuts that will avoid the category-theory informed constraints. This means that, in contrast to some other works that harmonize two sets of images (mostly, with one cat- egorical nuisance variable), the proposed method is best suited for obtaining well-behaved latent space with multiple nuisance variables/covariate shifts, for broad range of downstream tasks.
Acknowledgments. The authors thanks Veena Nair and Vivek Prabhakaran from UW Health for help with the ADCP dataset. Research was supported by NIH grants RFIAG059312 and RF RFLAG059869, as well as NSF award CCF 1918211 and the Bodosaki scholarship.

---

Published as a conference paper at ICLR 2024
REFERENCES
Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk minimiza-
tion. Technical report, arXiv:1907.02893. 2019. URL http::/leen.botter wttp//ppere.e...o.eeeeeoeee
ar jrjosky- o
2019
Miriam T Ashford, Garrett Miller, Rema Raman, Michael C Donohue, Ozioma C Okonkwo, Monica Rivera Mindt, Rachel L Nosheny, Ronald C Petersen, Paul S Aisen, and Mike W Weiner. The screening and enroll- ment of underrepresented ethnoracial and educational populations in the alzheimer's 's disease neuroimaging http://ooorgg11110000/0000.00000 https://alz jourall-oorraaaaa wnw.nnniine.... librarrbbrrr.... wi eyy wwt.co////ooooo www.//oo//aaaaooooooooooo abs 10..:..00222111
initiative. Alzheimer's & Dementia, 17($10):057424, 2021. doi:
URL
Ershad Banijamali, Amir-Hossein Karimi, Alexander Wong, and Ali Ghodsi. JADE: joint autoenenders for dis-entanglement. CoRR, ahs/1711,09163. 2017.
wwwpr//////a///// /11/ 1.9 63
Pietro Barbiero, Stefano Fioravanti, Francesco Giannini, Alberto Tonda, Pietro Lio, and Elena Di Lavore. Categorical foundations of explainable ai: A unifying formalism of structures and semantics. arXiv preprint arXiv:2304. 14094, 2023.
Vishnu M. Bashyam, Jimit Doshi, Guray Erus, Dhivya Srinivasan, Ahmed Abdulkadir, Ashish Singh, Mo- hamad Habes, Yong Fan, Colin L. Masters, Paul Maruff, Chuanjun Zhuo Henry Völzke, Sterling C. Johnson, Jurgen Fripp, Nikolaos Koutsouleris, Thcodore D. Satterthwaite, Daniel H. Wolf, Raquel E. Gur, Ruben C. Gur, John C. Morris, Marilyn S. Albert, Hans J. Grabe, Susan M. Resnick, Nick R. Bryan, Katharina Wit- tfeld, Robin Bülow, David A. Wolk, Haochang Shou, Ilya M. Nasrallah, Christos Davatzikos, and The iSTAGING and PHENOM consortia. Deep generative medical image harmonization for improving cross- site generalization in deep learning predictors. Journal of Magnetic Resonance Imaging, 55(3)9088-16, 2022. doi: https://doi.or// 1002/mmri.27908 URL [ttp::ttp///ttc http::s://oonnnellbaaraa wasietyywwieyy. wiwey com/ do17
aps 10. 1002/ jmri 1. 27908
Barry Beeker and Ronny Kohavi.
Adult.
UCI
Machine Learning Repository, 1996. DOI:
Tamir Bendory, Yuehaw Khoo, Joe Kileel, Oscar Mickelin, and Amit Singer. Autocorrelation analysis for cryo-em with sparsity constraints: Improved sample complexity and projection-based algorithms, 2022.
URL
10531
Steffen Bickel, Michael Brückner, and Tobias Scheffer. Diseriminative learning under covariate shift. Journal of Machine Learning Research, 10(75):2137-25 2155, 2009. URL http::/ wwtp://w...ppppp.s..eees
Dickel09a. html
Tolga Birdal and Umut Şimşekli. Probabilistic permutation synchronization using the riemannian structure of the birkhoff polytope. In Conference Computer Vision and Pattern Recognition (CVPR), 2019.
Tolga Birdal, Vladislav Golyanik, Christian Theobalt, and Leonidas Guibas. Quantum permutation synchro- nization. In Computer Vision and Pattern Recognition (CVPR), 2021.
Benjamin Bloem-Reddy and Yee Whye Teh. Probabilistic symmetry and invariant neural networks. ArXiv, 0bs/1901.06088 2020.
Emmanuel J. Candes and Michael B. Wakin. An introduction to compressive sampling. Magazine, 25(2):21-30, 2008. doi: 10.1 10.10000PP20007994495
samplnnn IEEE Signal Processing
Lawrence Carin and Michael J. Pencina. On Deep Learning for Medical Image Analysis. JAMA, 320(11): 193, 09 2018. ISSN 0098-7484, doi: 10. (0.1001/jamaa2018. 133316. [tt.tttss:////// org/ CC. 1001/jama 2018. 13316
Andrew A Chen, Joanne C Beer, Nicholas J Tustison, Philip A Cook, Russell T Shinohara, Haochang Shou, Alzheimer's Discase Neuroimaging Initiative, et al. Removal of scanner effects in covariance improves multivariate pattern analysis in neuroimaging data. bioRxiv, pp. 858415, 2020a.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. arXiv preprint orXiv:2002.05700 2020b.
Ying-Cong Chen, Xiaogang Xu, Zhuotao Tian, and Jiaya Jia. Homomorphic latent space interpolation for un- paired image-to-imase translation. In 2019 IEEE/CVF Conference Computer Vision and Pattern Recog- nition (CVPR), pp. 2403 2403-2411, 2411, 2019. doi: 10.1 199/ 00051..
Taco S. Cohen, Mario Geiger, Jonas Köhler, and Max Welling. Spherical CNNs. In International Conference on Learning Representations, 2018. URL https: htttss://ppeeeeeeeeeeee http:////ppeeeeeeeeeee wwwp://wwwwwweeeenn.ne...m http://oorumm..oootaaaaeaaa d-Hkbd5x%Rb)
10

---

Published as a conference paper at ICLR 2024
Geoff S. H. Cruttwell, Bruno Gavranovi' Neil Ghani, Paul Wilson, and Fabio Zanasi. Categorical foundations of gradient- based learning. In ESOP, 2022.
Samuel Filenherg and Saunders Macl aan. General theory of natural equivalences. Transactions of the Amer-
ican Mathematical Society, 58(2):33-2294, 1945. ISSN 00029947. URL http:/ http://ww..stt..ooo..o..oooooo.t jstor. org/
stable/ 1990244
Brendan Fong and Michael Johnson. Lenses and learners. In PLW, 2019.
Brendan Fong and David I. Spivak. An Invitation to Applied Category Theory: Seven Sketches in Composition- ality. Cambridge University Press, 2019. doi: 10.1017/978 10007/978100668000 108668804.
Brendan Fong, David I. Spivak, and Rémy Tuyéras. Backprop as functor: A compositional perspective on supervised learning. 2019 34th Annual ACM/IEEE Symposium on Logic in Computer Science (LJCS), 1-13, 2019.
Shan-Philippe Fortin, Nicholas Cullen, Yvette Sheline, Warren D. Taylor, Irem Aselcioglu, Philip A. Cook, Phil Adams, Crystal Cooper. Maurizio Fava, Patrick J. McGrath, Melvin Mclnnis. Mary L. Phillips, Madhukar H. Trivedi, Myrna M. Weissman, and Russell T. Shinohara. Harmonization of cortical thick- ness measurements acToss scanners and sitcs. Neurolmage, 167:104-120, 2018. ISSN 1053-8119. doi: https://doi.rg//// memm://m........mmmmmmmm........ 2017.11..2
wttp://wwwwwccceeeeereree..... com/
Fabian Fuchs, Daniel Worrall, Volker Fischer, Max Welling. seansorrrnssssmmer::: 3d roto-translation equivariant attention networks. In H. Larochelle. M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1970-1981. Cur- Associates, Inc., 2020. URL 2223la7cebaa89dl 30 212.::30000332222255333333333
Rafael Garcia-Dian Cristina Scarpazza, Lea Baecker, Sandra Vieira. Walter Hugo Lopez Pinaya, Aiden P. Corvin, Alberto Redolfi, Barnaby Nelson, Benedicto Crespo-Facorro, Colm Medonald, Diana Tordesillas- Gutiérrez, Dara M. Cannon, David Mothersill, Dennis Hernaus, Derek Morris, Esther Setién-Suero, Gary Donohoe, Giovanni Battista Frisoni, and Andrea Mechelli. Neuroharmony; A new tool for harmonizing volumetric mri data from unseen scanners. Neuroimage, 220, 2020.
Bruno Gavranović. Learning functors using gradient descent. Electronic Proceedings in Theoretical Computer Science, 323:230 245, sep 2020. doi: 10.4204/eptes.323.15 URL /attps://do// http:///...oo.//...// 42048 2Feptes 323.15
Bruno Gavranović, Paul Lessard, Andrew Dudzik, Tamara von Glehn, João GM Araújo, and Petar Veličković. Categorical deep learning: An algebraic theory of architectures. arXiv preprint arXiv:2402. 153332 2024.
Ghosh, M. M. Sajjadi, Vergari, M. Black, and B. Schölkopf. From variational to deterministic autocncoders. In 8th International Conference on Learning Representations (ICLR), April 2020. URL https: htt.ieewweeeeeeeeeeeeeeeeeeeeeee wttpe//oormmmmmllgtttpppp
*cqual contribution.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning image recognition. In 2016 IEEE Conference on Computer Vision and Puttern Recognition (CVPR), pp. 770-778, 2016. doi: N0.1109/CVP..201..22
Geoffrey E. Hinton. How to represent phot-whole hierarchies in neural network. CoRR, abx/2102.12627,
2021. htttsstts: /tte:///xxiii wwtpw/wwwgg///o/////oo/o 12627
Hans Hofmann. Statlog (German Credit Data).
UCI Machine Learning Repository, 1994. DOI:
https://doi.
Fengling Hu, Alfredo Lucas, Andrew A. Chen, Kyle Coleman, Hamnah Horng, Raymond W.S. Ng. Nicholas Tustison. Kathryn A. Davis, Haochang Shou, Mingyao Li, Russell T. Shinohara, and The Alzheimer's Dis- ease Neuroimaging Initiative. Deepcombat: A statistically motivated, hyperparamette--robus.. deep learning approach to harmonization of neuroimaging data. bioRxiv, 2023. doi: 04.1100///223...44333396 URL
https: www. wttp://wowox....o//oooteeooooooe http:///www...aooooo////////////
Mohammed Amir Husain, Benoit Laurent, and Mélanie Plourde. Apoe and alzheimer's disease: From lipid transport to physiopathology and therapeutics. Frontiers in Neuroxcience, 15, 2021. ISSN 1662-453X. 23006000022
doi: 10.338%/fnins.202 202 630502. URL htttps:///wwwt www.:ffoootttiiio wttp://ww.w.ttttttaeeeeoeee 3389/
Enins 2021. 630502]
Ayush Jaiswal, Yuehua Wu, Wael AhdAlmageed, and P. Natarajan. Unified adversarial invariance. ArXiv, abs/1905.03629, 2019.
11

---

Published as a conference paper at ICLR 2024
W Evan Johnson, Cheng Li, and Ariel Rabinovic. Adjusting batch effects in microarray expression data using empirical bayes methods. Biostatistics, 8(1):118-18 127, 2007.
Hamza Keurti, Hsiao-Ru Pan, Michel Besserve, Benjamin F Grewe, and Bernhard Schölkopf. Homomorphism AutoEncoder learning group structured representations from observed transitions. Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceed- ings of the 40th International Conference on Machine Learning. volume 202 of Proceedings of Machine Learning Research, pp. 16215,16215. PMLR, 23-29 Jul 2023. URL | htttpp:://proceedingg. press/v202/ hww220//eeuti2aaaa.
Yongnam Kim and Peter Steiner. Quasi-experimental designs for causal inference. Educational Psychologist, 51(3-4):395- 405, 2016. doi: 10. 11w.0//4661222221122212277 URL htttps://do..ogg/// 00461520. 2016, 1207177
PMID: 30100637.
Risi Kondor and Shubhendu Trivedi. On the generalization cquivariance and convolution in neural networks to the action of compact groups. ArXiv, abs/ 180.003609, 2018.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- bbased learning applied to document recognition. Proc. IEEE, 86:2278-2324, 1998.
Hyunwoo Lee, Kunio Nakamura, Sridar Narayaman, Robert A Brown, Douglas L. Arnold, Alzheimer's Dis- ease Neuroimaging Initiative, et al. Estimating and accounting for the effect of mri scanner changes on longitudinal whole-brain volume change measurements, Neuroimage, 184:555 565, 2019.
Spyridon Leonardos, Xiaowei Zhou, and Kostas Danilldiis. Distributed consistent data association via permu- tation synchronization. 2017 IEEE International Conference Robotics and Automation (ICRA), pp. 2645-2652, 2017.
doi: 10. R0LL0//CAA22221.7999990000
Shaohan Li. Yunpeng Shi, and Gilad Lerman. Fast, accurate and memory- partial permutation syn- chronization. In 2022 IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp- 15714-15722, 2022.
doi: 10.1 200CVVRR5688.2002..0121
Yujia Li, Kevin Swersky, and Richard Zemel. Learning unbiased features. arXiv preprint arXiv:1412.5244, 2014.
Sirui Liu, Bo Hou, Yiwei Zhang, Tianye Lin, Xiaoyuan Fan, Hui You, and Feng Feng. Inter-scanner repro- ducibility of brain volumetry: influence of automated brain segmentation software. BMC Neuroscience, 21: 35. 09 2020. doi: 1111///1188--11111111111 020-00285-11
Xiao Liu, Pedro Sanchez, Spyridon Thermos, Alison Q O'Neil, and Sotirios A. Tsaftaris. Learning dis- entangled representations in the imaging domain. Medical Image Analysis, 80:102516, 2022. ISSN 1361- 8415. doi: https://ooi.org//00 1016////0//.....0...ofo... media. 2022. 102516. WRLhhtts::///ww http://wwwww.ciiceeeeeeee... com/
Sicience/artticle/p1////
841 02200122001155 633
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Rätsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. A sober look at the unsupervised learning of disentangled representations and their evalu- ation. J. Mach. Learn. Res., 21(1), jun 2022. ISSN 1532-4435.
Vishnu Suresh Lokhande. Aditya Kumar Akash, Sathya N Ravi, and Vikas Singh. Fairalm: Augmented la- grangian method for training fair models with little regret. In European Conference Computer Vision, pp. 365-381, Springer. 2020.
Vishnu Suresh Lokhande, Rudrasis Chakraborty, Sathya N. Ravi, and Vikas Singh. Equivariance allows handling multiple nuisance variables when analyzing pooled neuroimaging datasets. 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10422-1043) 2022. doi: 10.1109/
CVPR52688.2022.0111
David G Lowe. Object recognition from local scale-invariant features. In ICCV, volume 2, pp. 1150-1157, 1999.
Saunders MacLane. 9781461
Categories for the working mathematician. Springer, New York, NY, 2014. ISBN
Nicholas F. Marshall, Oscar Mickelin, Yunpeng Shi, and Amit Singer. Fast principal component analysis for cryo-em images, 2022. URL https::/aaxxi..ottt wttp.////...b.2....1...
Michelle M. Mielke. Sex and gender differences in alzheimer disease dementia. Psychiotric Times, 35(1 I):
14-15, November 2018. ISSN 0893-29055
12

---

Published as a conference paper at ICLR 2024
Michael Moor, Max Hor,, Bastian Rieck, and Karsten Borgwardt. Topological autoencoders. In Hal Daumé III and Aarti Singh (eds.), Proceedings of 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 7045-70 7054. PMLR, 13-18 Jul 2020. URL
//proccedings.. mlr. prwps /vl1iv1ii /mmorooo2.a. .0mmm
Sacid Motiann, Marco Piccirilli, Donald A. Adjeroh, and Gianfranco Doretto. Unified deep supervised domain adaptation and generalization. 2017 IEFE International Conference Computer Vision (ICCV), pp. 5716- 5726, 2017.
Daniel Moyer, Shuyang Gao, Rob Brekelmans, Aram Galstyan, and Greg Ver Steeg. Invariant representations without adversarial training. Advances Neural Information Processing Systems, 31, 2018
Daniel Mayer, Greg Ver Steeg, Chantal M. W. Tax, and Paul M. Thompson. Scanner invariant representations for diffusion mri harmonization. Magnetic Resonance in Medicine, 84(4):2174-2182, 2020. doi: https:// doi.orge Noi.org/10.J0 101.0..00000.mmmm00 28243. URL StEtps://onlinee www.iinnnibbbrry.e www.ee......l.e...o.o...oo.oo wwwp////.occ.eeeeoooo/ooe/oo..et abs//ab//ii 10.2202/
28243
Susanne Mueller, Michael Weiner, Leon Thal, Ronald Petersen, Clifford Jack, William Jagust, John Tro- janowski, Arthur Toga, and Laurel Beckett. Ways toward an early diagnosis in alzheimer's disease: The alzheimer's disease neuroimaging initiative (adni). Altheimer's dementia the journal of the Alzheimer's 's. Association, 1:55 66, 08 2005. doi: 10. jal.220.0
Harrison Nguyen, Richard W. Morris, Anthony W. Harris, Mayuresh S. Korgoankar, and Fabio Ramos. Cor- recting differences in multi-site neuroimaging data using generative adversarial networks, 2018.
László G. Nyúl and Jayaram K. Udupa. On standardizing the image intensity scale. Magnetic Resonance Medicine, 42(6): 1072-10078,,, 1999. doi: https://oii.rg/// otg/O/...ooo/O../00 (802///CCD1152225544119999999999
1072::
AID WDD-MRM1133.0.00:2221 URL htt //ol Abbarbrrr
Tey com wtw.dd///bbbbsase 1002/ MRM 13333 0.3.3....O 00.8
3B2-M -
Alon Oring, Zohar Yakhini, and Yacov Hel-Or. Autoencoder image interpolation by shaping the latent space. In Marina Meila and Tong Zhang (eds.). Proceedings the 38th International Conference Machine Learning, volume 139 Proceedings of Machine Learning Research, pp. 8281-8290. PMLR, 18-24 Jul 2021. URL httltttps://prcceeeeee
mlr. wttppppwwsss.....p.poo
html
Deepti Pachauri, Risi Kondor, and Vikas Singh. Solving the multi-way matching problem by permutation synchronization. Advances in neural information processing systems, 26, 2013.
Xi Peng. Xiang Yu, Kihyuk Sohn, Dimitris N. Metaxas, and Manmohan Chandraker. Reconstruction cronstructioonbbsee disentanglement for pose-invariant face recognition. 2017 IEEE International Conference Computer Vision (ICCV), pp. 1632 1641, 2017.
William D. Penny, Karl J. Friston, John Ashburner, Stefan Kiebel, and Thomas Nichols, Statistical parametric mapping: The analysis of functional brain images. 2007.
Thai-Hoang Pham, Xueru Zhang, and Ping Zhang. Faireess and accuracy under domain generalization. arXiv preprint arXiv:2301 2.ti:2301..332... 13323, 2023.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Tlya Sutskever. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, 2021. URL/https: api,
Mengwei Ren, Neel Dey, James Fishbaugh, and Guido Gerig. Segmentation--enommmliice deep feature modu- lation for unpaired image harmonization. IEEE Transactions on Medical Imaging, 101651551-15500, 2021. doi: 10.1 109/TM1.222 3059972..
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Nassir Navab, Joachim Hornegger. William M. Wells. and Alejandro F. Frangi (eds.), Medical Image Computing and Computer-Assister Intervention MICCAI 2015, pp. 234-241, Cham, 2015. Springer International Publishing. ISBN 978-3-319-2457442
Gerta Rücker, Susanne Schmitz, and Guido Schwarzer. Component network meta analysis compared to a matching method in a disconnected network: A case study. Biometrical Journal, 63(2):477 461, 2021. doi: https://doi.or//100 100./////.......mmm.o 201 900339. URL https: // on nel ibra ry.. wi ley. com /oo / /abs/
10. 10002/0m 20000000339 900339
13

---

Published as a conference paper at ICLR 2024
Sara Sabour, Nicholas Frosst, and Geoffrey E. Hinton. Dynamic routing between capsules. In Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS' 17, pp. 3859-3869, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN
Benoit Sauty and Stanlcy Durrleman. Progression models for imaging data with longitudinal variational auto encoders. In Linwei Wang, Qi Dou, P. Thomas Fletcher, Stefanic Speidel, and Shuo Li (cds.), Medical Image Computing and Computer Assisted Intervention MICCAI 2022, pp. 3-13, Cham, 2022. Springer Nature Switzerland. ISBN 978-3-031-164314
Dan Shichler, Bruno Gavranovi'c, and Paul Wilson. abs/2106.07032, 2021.
Category theory in machine learning. ArXiv,
Connor Shorten and Taghi Khoshgoftaar. A survey image data augmentation for deep learning. Journal of Big Data, 6, 07 2019. doi: N..L186940557-019-019777
Grzegorz Sienski, Priyanka Narayan, Julia Maeve Bonner, Nora Kory, Sebastian Boland, Aleksandra A. Ar- czewska, William T. Ralvenius, Leyla Akay. Elana Lockshin, Liang He, Blerta Milo, Agnese Graziosi, Valeriya Baru, Caroline A. Lewis, Manolis Kellis, David M. Sabatini, Li-Huei Tsai, and Susan Lindquist. Apoe4 disrupts intracellular lipid homeostasis in human ipsc-derived glia. Science Translational Medicine, 13(583), 2021, doi: 10. 126//citransmmmdd..ae aa74564 URL. htll/httss: //www//www....////////....w// http///wwwcc..co.....occoo.cooot www.:asciencc... ence. org. /abs/ 10. 1 126/scccc.. /et/siitaannn. w 4524564
Elizabeth A. Stuart. Matching Methods for Causal Inference: A Review and Look Forward. Statistical Sci- ence, 25(1):1 21, 2010. doi: 10.1214/99-STS3 URL http:tttt///o...og////aa/a 1214/09- STS313
Hemant D. Tagare, Frederick Sigworth, and Andrew Barthel. Fast, adaptive
align- ment for cryo-em. In Proceedings of the 11th International Conference on Medical Image Computing and Computer-Assisted Intervention, Part II, MICCAI '08, pp. 855-862, Berlin, Heidelherg. 2008 Springer- Verlag. ISBN 9783540859895 doi: 10.1007//78-3-555 0-5-7788-----888899----- 1-193, URL (ntpps: //t::///oo.o.//ooo org/IC.
Paul Thompson, Jason Stein, Sarah Medland, Derrek Hibar, Alejandro Arias Arias- Vásquez, Miguel E Renteria, Roberto Toro, Neda Jahanshad, Gunter Schumann, and Barbara Franke. The enigma consortium: Large- scale collaborative analyses of neuroimaging and genetic data. Brain imaging and behavior, 8:153-188 01 2014.
David A van Dyk and Xiao-Li Li Meng. The art data augmentation, Journal of Computationall and Graphical Statistics, 10(1):1-50, 2001. doi: 008/100616001554188855 URL 10618 8600160015300 52418584
and
Paul Wilson and Fabio Zanasi. Reverse derivative ascent: categorical approach to learning boolean circuits. In ACT, 2020.
Qizhe Xie, Zihang Dai, Yulun Du Eduard H. Hovy, and Graham Neubig. Controllable invariance through adversarial feature learning. NIPS, 2017.
Ronggian Zhang, Lindsay D. Oliver, Aristotle N. Voineskos, and Jun Young Park. RELIEF: A structured multivariate approach for removal of latent inter-scanner effects. Imaging Neuroscience, 1:1-16, 08 2023. 00011
ISSN 2837-6056. doi: 10. 1..1162/ma...a00001. URL ptt: / wttp::/io.rr.1...........
Zizhao Zhang, Lin Yang, and Yefeng Zheng. Translating and segmenting multimodal medical volumes with cycle- and chape-consistency generative adversarial network. In 2018 IEEE/CVF Conference Computer Vision and Pattern Recognition, pp- 9242-9251, 2018. doi: 10.1 109/CVPR.22 PP..P..011..00000 00963.
Hao Henry Zhou, Yilin Zhang, Vamsi K. Ithapu, Sterling C. Johnson, Grace Wahba, and Vikas Singh. When can multi-site datasets be pooled for regression? Hypothesis tests, fg-consistency and neuroscience applications. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 4170-4179. PMIRR, 06-11 Aug 2017. URL https:/://////.. pttps////ococeeeggg..... m.r. press /V70 zhoul .h.mll
Jun-Yan Zhu, Taesung Park, Phillip Isola, and Alexei A. Efros. Unpaired image-to-image translation using cycle-consiseee adversarial networks. In 2017 IEEE International Conference on Computer Vision (ICCV), pp. 2242 2251, 2017. doi: 10.1109/CCV.2217.2244
Lianrui Zuo, Blake E. Dewey, Yihao Liu, Yufan He, Scott D. Newsome, Ellen M. Mowry, Susan M. Resnick, Jerry L. Prince, and Aaron Carass. Unsupervised mr harmonization by learning disentangled representations using information bottleneck theory. Neurolmage, 183:1, 2021. ISSN 1053-8119, doi: http:://bb..rrg//l...........mmp Sctiance//rriicle//11//113311111
18569. https::s:: :/ www. sci enceddeect.coomm
14

---

Published as a conference paper at ICLR 2024
A A WARM-UP ANALYSIS
A.1 MULTI-EQUIVARIANCE
In this section, we evaluate the performance of our approach when we require the latent space to be equivariant with respect different transformations. We consider dataset of images and the goal is to create representations that are equivariant with respect to image rotations as well as image scaling.
We use the MNIST dataset LeCun et al. (1998) as a toy example to illustrate the key idea. We assume that the latent space is the Category with the following characteristics:
(a) Objects: vectors E R"
Morphisms: orthogonal linear transformations W R WTW (the identity morhpism is the identity matrix T)
We learn an autoencoder (i.e., a pair of fully faithful Functors F, F-1) and two matrices W,, W, that represent the rotation and scale morhpism respectively in IR.. In this experiment, we set n 128. During training, we provide our model with an image from MNIST, along with a rotated version (of the same image) and separately, scaled one. An important note is that we do not provide any images that are both rotated and scaled during training. We define
1. Rotate to be counter-clockwise rotation by 5 degrees. This means that rotation by e.g. 15 degrees corresponds to "Rotate Rotate Rotate" in the source Category of images and to W in the latent space of IR"
2. Scale the operation of "zoom-out" in an image. To achieve this behaviour, we first pad the input image (I extra pixel each side) and then resize it back to the original shape.
According to this definition, the inverse operation (Scale ') is defined zooming into the image.
Our training objective consists of the following loss terms: - 1)2 WWTW- 10) -)²
L L, + XC, + ( -
(11)
The first two terms are suggested directly by our framework while the last two constraints ensure that the latent space morphisms have the desired form and are case-specific.
Although during training we provide images that are only counter-clockwis rotated (up to 50 de- grees) scaled (up 10 paddings), our model is able to map the images to the latent space and preserve the two morphisms, which leads to the following abilities:
(a) We can generate realistic images that are counter-clockwis. rotated by more than 50 degrees by continuously applying the linear transformation W, in the latent space and mapping the new vector back the original Category of images using the decoder (F-1).
(b) We can generate realistic images that are scaled by more than 10 pads by continuously applying the linear transformation W, the latent space and mapping the new vector back to the original Category of images using the decoder (F-1).
(c) We can generate realistic images that are clockwise rotated by applying the inverse linear trans- formation W-1 to the latent representation of an image and then map the new vector back to the original Category of images using the decoder (F-1).
(d) We can generate realistic images that are zoomed-in in by applying the inverse linear transforma- tion W-1 to the latent representation of an image and then mapping the new vector back to the original Category of images using the decoder 11-
(e) Finally, we can generate realistic images that are both rotated and scaled by applying both linear transformations W, W, (= W, W,) in the latent representation of an image and then mapping new vector back to the original Category of images using the decoder (F-1).
15

---

Published as a conference paper at ICLR 2024
0
12
Rotation
4
12
16
Rotation
9
9
9
0
2
7
7
r
9
9
0
C
2
7
7
9
9
a
C
8
2
7
7
T
y
9
4
A
C
12
2
7
+
?
16
Scale
Scale
Figure 7: Composition of rotation and scaling. The transtormations were applied to the original image's latent vector (v) in the latent space according the composition rule WFW? where i, denote the degree of rotation and scaling respectively. In red we depict the original image, in blue depict the only transformations that were "visible" during training.
A.1.1 QUALITATIVE RESULTS
In Figure show the generated images when apply both transformations to latent vector. The 2 main takeaways from that figure are:
(a) Our model is able to generalize beyond the presented transformations depicted by the row/column in which scale/rotation was 0 respectively. We observe that even after applying W W20 (twice much as the data presented during training) we get realistic results that fol- low the rotation and scaling rules.
(b) Our model is able compose the transformation in the latent space and generate images that follow both transformations, although no composed images were presented during training.
In Fig. 8 the effect of the inverse linear transformation shown. Although there were such im- ages during training, our model is able to understand the two transformations due to the categorical structure imposed on the latent space. As a result, we can generate clockwise rotated and zoomed- in images with no quality drop (zoom-in fails after a significant amount of constant applications due the space constraints while rotation remains realistic even for extreme number of linear transformations).
A.1.2 QUANTITATIVE RESULTS
In Fig. 9 we quantify what we observed already. Rotation provides consistently good results inde- pendent of the direction, while inverse scaling (zoom-in) fails after some steps, due to the fact that the depicted digit is too big fit in the image. On the contrary, we can continuously apply the scale operation (zoom-out) up to the point which the digit no larger than a couple pixels, with no decrease in the quality. Also, the results do not vary much increase the latent dimension from 32 to 256.
A.2 EQUIVARIANCE TO ALGEBRAIC MORPHISMS
Here we show that the same formulation, without any change, able to create a mapping to the latent space which is equiavariant to much more abstract and complex morphism (algebraic addition). Specifically, we consider the MNIST dataset and our goal is to map the images to a latent space that preserves the relationships between the digits. We use the same Category as above for our latent space and we wish to enforce the following condition: if for two images, I, In with labels 4, l the corresponding latent representations 01, ; U2, then U2 111--1111
16

---

Published as a conference paper at ICLR 2024
-16
8
0
8
16
Rotation
3
2
7
7
4
4
V
3
7
7
7
J
3
3
3
3
3
9
8
8
8
8
8
6
16
8
0
8
16
Scale
R
H
R
8
8
8
3
8
T
6
6
6
6
6
6
7
7
7
7
7
7
7
9
9
9
9
Figure 8: Rotation and Scaling in both directions. The transformations were applied to the original image's latent vector (v) in the latent space according the morphisms Wr W.1, W., W- In red we depict the original image, in blue we depict the only transfor- mations that were "visible" during training. Inverse scaling, as expected, fails after an adequate number of steps, since the digit can no longer fit in the image. On the contrary, rotation is more robust and performs well even for large rotations in both directions.
Rotate
0.20
Scale
0.15
MSE
0.10
0.05
0.00
20
-15
-10
-5
Matrix power
Figure 9: MSE error of the predicted image against the ground truth rotate it using Wr, W W-1 and scale using W., W.-1 The line intensity represents the size of the latent space
(32, 64, 128, 256) and it shows that the results are accurate enough even for small latent spaces.
17

---

Published as a conference paper at ICLR 2024
0
4
8
3
4
5
7
8
9
0
4
2
3
4
5
6
7
8
9
0
7
3
3
5
3
6
7
0
6
4
3
4
5
6
7
8
9
0
2
3
4
5
6
7
8
9
Figure 10: Images generated when we apply W and W-1 to the latent vector of the image in red.
The images the left of the red image were generated by using the linear transformation W-1 on e, where v is the latent vector of the original image (red), and the images to the right were generated by using the linear transformation W on v. While there might exist some failed cases (e.g. row 3), in most cases we are able to generate realistic images for each digit.
We use the same setting as in the previous experiment (but here we learn a single matrix W). During training, provide our model with pairs of data that corresponds to images (71, T2) with 21 1 = 12. This means that, during training, we do not optimize for W-1 simultaneously (just like in the previous experiment). In this experiment we set 32.
We use the same AutoEncoder architecture and the objective in this case is:
L = L. + XL. + +W((W> -))
(12)
which means that we use the exact same constraints that our formulation has suggested along with the constraint that W orthogonal.
A.2.1 QUALITATIVE RESULTS
Figure shows some of the generated images we obtain when we apply the linear transformations W. WTT
in the latent vector of image. We can observe that, altough we do not explicitly train model for the inverse transformation (W-1), it is able generate realistic images due to the orthogonality constraint we impose in our model.
Importance of orthogonality Besides the structure preserving constraints we use another con- straint in our models; the orthogonality constraint. In Fig. 11 we depict the results when such constraint is not enforced in the model. While still we can generate realistic images when we ap- ply the transformation W in the latent space, the results are less ideal when we apply the inverse transformation W-1, Since we provide no constraints on W during training, it ends up being a non-singular matrix. This means that can only approximate the inverse matrix W- and, as the results indicate, this estimation not close to the ideal. In contrast, the orthogonality constraint forces the matrix have well-defined inverse (W-1= WT))
Smooth image interpolation Since W represents the abstract morphism "add one" we can sume that the matrix W represents the abstract morphism "add a" Va IR. This provides us with simple way interpolate between the digits and create smooth transitions.
In the supplementary material we have included gif files that show a complete transition from 0 to with step of 0.2. The results show that, although goal was preserve specific relationship of the original data into the latent space, a by- by-pruducct of our method is the ability image interpolation and smooth transitions.
18

---

Published as a conference paper at ICLR 2024
/
2
3
9
5
7
8
9
4
5
6
7
8
9
2
3
9
5
6
7
8
9
8
9
Figure 11: Without the orthogonality constraint, the inverse matrix may not exist and its approxi- mation leads to generated images that do not resemble real digits (the application the matrix though still leads realistic images).
W
W
Mean Squared Error
0001 0,002 0.003 0,004
128
20
7
7
7
Labels difference grrverated
loriginal)
Figure 12: Mean Squared Error of the generated images for different sizes of the latent space (n)). The application of the inverse has an only marginally higher MSE due to the orthogonality of the matrix W.
A.2.2 QUANTITATIVE RESULTS
Besides the qualitative results that indicate that our model successfully learned the relationship be- tween the depicted digits, we can also quantify how accurate our model is. In Fig. 12 we show how the Mean Squared Error (MSE) changes when we continuously apply the linear transformations W. W-1 to the latent vector and generate an image. In order calculate it, we found the minimum MSE between the generated image and any image that depicts the specific digit. We see clearly that our model is able to perfectly generate an image of the new digit, no matter it is smaller (W-1 ap- plied) or greater (W applied), even after applying multiple times the linear transformations. Another important aspect is the robustness of our model to different values of (latent space dimension).
B EXPERIMENTAL EVALUATIONS IN DEPTH
B. DATASET DETAILS
The ADNI dataset can be obtained from sttpps:/addi..oni. www.looii......u.... was launched in 2004.
The ADNI project
This specific dataset consists of 449 MRI scans, along with dataset of covariates for each individual participant. Each MRI was obtained in one of three different scanners ((GE Medical Systems, Philips Medical Systems, SIEMENS)) and some of the covariates include {Age, Sex, Years of Education, APOE A1, APOE A2, Marital status, Race) along with lot of cognitive battery scores such (MMSE, RAVLT, Ecog).
19

---

in the case of invariance, and Published as a conference paper at ICLR 2024
ACC
MMD(x100)
ADV
Naive
74(0.9)
7.7(0.8)
62(3.1)
MMDLi et al.
(2014)
73(1.5)
1.5(0.3)
66(0.04)
CAfXie et al.
2017
76(1.3)
1.2(2.4)
65(0.01)
SS Zhou
(017)
76(0.9)
1.5(0.0)
70(6.9)
RMMotnan et al. (2017)
74(2.1)
7.5(0.9)
66(4.2)
GoLohhande et al. (2022)
75(3.3)
2.7(0.6)
54(1.1)
Ours (inv)
74(0.6)
0.7(0.1)
51(2.4)
Ours (1 cov)
74(1.4)
1.3(0.9)
55(1.4)
Table 3: Quantitative Results German Hofmann 1994).. The accuracy (ACC) and invariance as evaluated by M.D and ADV are shown. The standard deviations are in parenthesis. The baseline inv corresponds only invariance to foreigner. cov corresponds equivariance with respect to the age covariates (plus invariance to foreigner). MMD and ADV are significantly reduced without any significant drop accuracy.
neuroimaging data, "APOE" typically refers to the apolipoprotein gene. The APOE gene involved in encoding protein that plays crucial role in the metabolism of lipids (fats) in the body, including cholesterol. This gene has different variants, or alleles, known as APOE c2, APOE e3, and APOE e4. In the context of neuroimaging and neuroscience, the APOE gene has been of particular interest because one of its alleles, APOE e4, is considered significant genetic risk factor for Alzheimer's disease Husain al. (2021) Sienski al. (202T).
The second dataset (ADCP) was shared with us by the authors of Lokhande et al. (2022).
The data for ADCP collected through an NiH-sponsored Alzheimer's Disease Connectome Project (ADCP) UOI AG051216. The study inclusion criteria for AD (Alzheimer's disease) / MCI (Mild Cognitive Impairment) patients consisted of age between 55-90 years, willing and able to undergo all procedures, retaining decisional capacity at the initial visit, and meet criteria for probable AD or MCI. The MRI images were acquired at three distinct sites.
Besides the medical image datasets, we evaluate our performance in two tabular datasets; the Ger- man(Hofmann 1994) and the Becker & Kohavi (996).. The German consists of 58 features, with some of them being foreigner and age, and our goal is to predict the consumers' default loans. The Adult consists of 100 with sensitive attributes such as sex and age. Our goal here is predict whether someone's income is higher than $50K.
B.2 RESULTS ON TABULAR DATA
Tables contains all 3 metrics for German and Adult datasets, respectively. We observe that we have accuracy similar to the other methods, while we have significantly lower value of MMD and ADV, meaning that we obtain less biased embedding space and, as result, predictions.
In of our experiments we used latent space size 32 and 1 1-hidden-layee (of size 32) feedfor- ward network to map the input features to R2
B.3 EFFECT OF LAGRANGE MULTIPLIER (A) AND LATENT SPACE DIMENSION (n)
Here, we evaluate the effect of the latent space dimension well as the effect of the Lagrange multiplier a reminder, the objective was defined as:
==
Cp
+
Σ
A
F(s1) F(s1)- F(sz)
(13)
$1.45S
cross-entropy
V $1.32 from different Scanners
20

---

Published as a conference paper at ICLR 2024
ACC
MMD(x109)
ADV 4
Naive
84(0.1)
9.8(0.3)
83
MMD et al.
2014)
(4
3.1(0.3)
83(0.1)
CAIXie et al.
2017
84(0.04)
2.2(2.4)
81
SSZhou
2017)
84.0.1)
1.5(0.2)
83(0.2)
RMMonnan et al. (2017)
84(0.3)
4.8(0.7)
82(0.4)
LoLkhaddde al. (2022)
83
7.1(0.6)
75(1.4)
Ours (inv)
83
0.6(0.4)
74(1.)
Ours (1 cov)
84(0.1)
1.1(0.9)
76(1.3)
Table 4: Quantitative Results Adult Becker & Kohavi (199)). The mean accuracy (ACC) and invariance as evaluated by M.M and ADV are shown. The standard deviations are in parenthesis. The baseline inv corresponds only invariance gender, cov corresponds equivariance with respect to the age covariates (plus invariance to gender). Similarly to German, MMD and ADV are significantly reduced without any significant drop in accuracy.
C= Cp cp + Σ
W F()) F(s)) F(s2)
(14)
cross-entropy 81 $1SS
Cay Sayy represent the nuisance covariate
the case of equivariance.
Tables 6 10 present the results for 1-covariate equivariance, for both datasets. We can observe that (i.e. ~ 0.01) lead a low MMD value (much lower than the existing methods in most of the cases) with no Accuracy drop. In fact, in most of the cases, the effect of invariance leads to better generalization capabilities (i.e. higher Accuracy).
moderate values
Tables show the same results but in the case of 2-covariate cquivariance. Although, in theory, such problem is harder to solve, this does not translate to lower values MMD and Accuracy. In fact, MMD is slightly lower in this case, since we enforce invariance to each subgroup separately (i.e. a less "aggressive" form of invariance).
Finally, in Table and Fig. 1 present the results for 5-covariate equivariance, for the ADNI dataset. The results show, undoubtedly, that we have presented a general method which has no problem in handling multiple nuisance covariates, in contrast existing approaches that either can not handle such scenarios, or exhibit a high performance drop.
Only invariance
ACC
MMD
A
16
32
64
16
32
64
0.001
81(2.3)
81(1.))
84(3.9)
18(4.4)
21
24(0.8)
0.01
(2.3)
(9
81.33)
06(2.0)
17(5.0)
16(4.8)
0.1
81.0)
78(3.8)
83
02(2.4)
07(3.3)
13(8.4)
Effect of the Lagrange multiplier and size of the latent space Accuracy and Maximum Mean Discrepancy (MMD), where enforce only invariance (with respect to the site) to the latent space. Higher values lead to lower MMD (more invariant representations). In all the combinations of (A, n)) Accuracy remains high (in most of the cases higher than the baselines) while MMD is significantly lower than the baselines in most cases.
Table 5: ADNI (Mueller al., 2005
21

---

Published as a conference paper at ICLR 2024
Age
Sex
APOE A1
X
16
32
64
16
32
64
16
32
64
0.001
(9(3.3)
80().7)
(9(4.0)
80(3.3)
81(3.0)
80(3.4)
79(3.9)
80(1.7)
79(3.2)
MMD ACC
0.01
(2().5)
80(3.2)
(3.4)
80(4.5)
80(2.2)
(1.3))
82(1.7)
80(3.4)
81(5.0)
0.1
79(1.)
79:3.2)
(4.9.9)
78(3.9)
77(L7)
77(.7)
78(4.2)
77(1.1)
74(1.))
0.001
(7.4.4)
(7.6))
(7.0.)
27(1.2)
26(1.))
(7(6.)
28(0.7)
27(0.7)
27(0.1)
0.01
(1(2.9)
15(6.5)
19(3.6)
17(6.3)
13(2.9)
20(6.1)
11(4.7
15(2.8)
22(4.0)
0.1
(4(7.7)
09(8.1)
16(2.5)
(3(4.4)
07(6.0)
15(5.1)
04(0.6)
11(4.4)
14(1.4)
Effect of the Lagrange multiplier À and size of the latent space n to Accuracy and to Maximum Mean Discrepancy (MMD) in the case single covariate equivariance. No Accuracy drop is observed, while MMD is significantly lower than the baselines, in most of the experiments.
Table 6: ADNI (Mueller et al.
2005):
Age-APOE AI
Age-Sex
APOE A A1-Sex
A
16
32
64
16
32
64
16
32
64
0.001
80(1.)
79(3.2)
80(3.4)
82(4.2)
82(2.3)
82(2.2)
83(1.)
(1.2)2
80(2.8)
MMD ACC
0.01
80(1.3)
81(1.7)
82(1.7)
80(1.1)
81(3.6)
82(3.0)
82(2.2)
79(3.0)
80(3.6)
0.1
78(2.2)
77(1.9)
75(2.8)
77(2.8)
77(0.0)
77(2.3)
77(1.2)
79(2.3)
76(0.6)
0.001
24(2.8)
26(0.44
17(0..)
22(1.1)
27(2.9)
27(0.4)
25(3.7)
26(8.0)
26(4.0)
0.01
07(1.2)
14(7.0)
18(5.4)
(9(0.4)
17(3.9)
16(1.7)
10(6.6)
16(5.7)
15(6.0)
0.1
04(1.7)
(7.9..)
10(7.0)
2(0.7)
04(2.2)
08(3.4)
03(2.6)
09(9.0)
09(1.0)
Table 7: ADNI (Mueller et al.
2005):
Effect of the Lagrange multiplier and size of the latent
space to Accuracy and to Maximum Mean Discrepancy (MMD) in the case of two-covariate equivariance.
Age-Sex-APOE AI-APOE A2-Education
ACE
MMD
A
16
32
64
16
32
64
0.001
82(4.2)
81(1.7)
82(1.1)
11(5.3)
21
25(3.4)
0.01
80(2.3)
78(2.3)
77(1.7)
05(5.3)
09
14(3.7)
0.1
78(1.9)
76(2.2)
78(0.6)
01(0.2)
03
07(5.6)
Table 8: ADNI (Mueller et al. Accuracy and Maximum Mean Discrepancy (MMD), in the case of equivariance with respect to five variables.
2005):
Effect of the Lagrange multiplier À and size of the latent space
Only invariance
ACC
MMD
16
64
16
64
0.001
82(2.2)
82(2.2)
83
48(5.9)
53(4.8)
60(7.0)
0.01
83(2.2)
86(8.0)
885(3.9)
35(6.1)
34(1.2)
54(4.1)
0.1
(3.8)
(2(5.9)
85(7.7)
15(2.5)
25(2.2)
28(2.0)
Table 9: ADCP: Effeet of the Lagrange multiplier À and size of the latent space to Accuracy and Maximum Mean Discrepancy (MMD), where we enforce only invariance (with respect to the site) to the latent space.
22

---

Published as a conference paper at ICLR 2024
Age
Sex
A
16
32
64
16
32
64
0.001
85(3.8)
83(2.2)
86(5.9)
83(4.4)
83(4.4)
83(5.8)
MMD ACC
0.01
(9(2.2)
(8.00))
86(5.9)
(2(4.))
81
85(3.8)
0.1
77(3.9)
(9(1.4)
83(44)
81(0.2)
78(8.0)
83(14)
0.001
37(4.0)
(5.6)
57(5.3)
(7.88)
52(8.4)
55(3.0)
0.01
30(12)
(5
(9(3.6)
(0(0.0)
31(122
34(5.8)
0.1
27(3.2)
27(5.8)
23(5.4)
23(3.3)
21(7.0)
46(1.)
Table 10: ADCP: Effect of the Lagrange multiplier À and size of the latent space Accuracy and to Maximum Mean Discrepancy (MMD) in the case of single covariate equivariance.
Age-Sex
ACC
MMD
A
16
32
64
16
32
64
0.001
83(5.8)
82(4.4)
(2.6.6)
(7(2.1)
42
49(4.
0.01
(5
82(9.5)
85(3.8)
29(3.3)
32(1.2)
35(6.1)
0.1
79(2.2)
(9.8.7)
83(4.4)
24(4.))
27(4.8)
22(2.8)
Table 11: ADCP: Effect the Lagrange multiplier and size of the latent space n to Accuracy and Maximum Mean Discrepancy (MMD), in the case of two-covariates equivariance.
À 0.1
À 0.01
A 0.001
Age
APOE A1
APOE A2
Sex
Education
0.90
0.95
1.00
0.90
0.95
1.00
0.90
0.95
1.00
CS
CS
CS
Age
APOE A1
APOE A2
Sex
Education
0.00
0.01
0,000
0.025
0.00
0,05
D
D
D
Figure 13: Effect of the Lagrange multipler À on the equivariance metries CS and D. As expected, and D negatively correlated, while CS remains high for different values X.
23


Reviews:
[{"review_id": "rLc7sFJ1cc", "reviewer": "Reviewer_KLWx", "summary": "Authors present a Category Theory-based formalism for covariate shifts\nand a method derived from this formalism for pooling data with\ncovariate shifts for simple tasks. Images are mapped to a latent space\n- using AE but I guess this can be other things - and covariate shifts\nare modeled as transformations in the latent space. Invariance to\nsome covariate shifts, e.g., using different acquisition devices, are\nmodeled by enforcing similar latent space\nrepresentations. Equivariance is obtained by enforcing a chosen\ntransformation. Experimens show that the resulting model can retain\naccuracy while removing differences between samples with covariate\nshift in the latent space in cases where invariance is\ndesired. When equivariance is desired, authors show that the latent\nspace transformations are also good.", "questions": "1. I encourage authors to reconsider their claims, especially\n   concerning counterfactuals and models requirements of overlapping\n   covariate values between different sets.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The general formalism Category Theory provides is elegant and\n   contains a class of algorithms as special cases. This has the\n   potential of inventing new directions for data pooling and\n   addressing covariate shifts.\n2. Authors provide a great introduction to the theory and maps the\n   relevant aspects to practical problems pertinent to covariate\n   shifts observed in application in neuroimaging.\n3. Experimental results are motivating. Authors compare with some\n   recent work to convince us that the model proposed is indeed\n   useful.\n4. The article is not necessarily providing a novel technique but\n   rather discusses a theory that views already applied methods as a\n   special case and thus allows further generalization. For instance,\n   https://link.springer.com/chapter/10.1007/978-3-031-16431-6_1\n   applied similar ideas for longitudinal imaging data.", "weaknesses": "1. The current method requires overlapping covariate values between\n   data sets coming from different centers. Equation (9) seems crucial\n   for the methods to work. However, until (9) the reader is left with\n   the opinion that Category Theory would magically work for\n   non-overlapping covariate values between different data sets. This\n   should be clearly discussed early on in the article to avoid any\n   confusions and - to be honest - false hope. Given two data sets\n   whose covariate values do not overlap, the current method may not\n   work at all. \n2. The emphasis on couterfactuals is rather questionable. There is no\n   decoupling of the effects of different covariates on data. In other\n   words, a transformation in the latent space may correspond to\n   changes in multiple covariates. A well curated data set may avoid\n   this problem, however, generally this will not hold. Therefore, it\n   is difficult to discuss about counterfactuals and even\n   interventions. The latent space transformations will not be able to\n   account a missing causal model linking covariates, observed\n   images and labels. I recommend authors to reconsider claims\n   regarding counterfactuals.\n3. Figure 5 is not explained at all. This seems like an important\n   figure and requires further explanation in the text.\n4. I find the claim that this approach simplifies recent\n   works in the literature not substantiated. Furthermore, without\n   explicit disentanglement nor the capability to do interventions, I\n   am not sure to what authors may be referring.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ai3EtKhnO6", "reviewer": "Reviewer_SjKQ", "summary": "The paper explores an interesting direction for giving a defined structure to the latent space by wrapping an understanding around it derived from category theory. The motivation is to find commonality among datasets in terms of the latent space and aid in data pooling / data harmonisation, something that is needed when data collection is expensive and data from multiple sources need to be combined.  The hope is that the theory is able to define structure and invariances/equivariances between them,, and invariance to shifts in feature distributions for clinical and demographic features when dealing with image data would be easily drawn. For quantifying  the shift, minimum distance and cosine sim are employed. The paper also describes a reinterpretaion of two differentiable models from adversarial and self-supervised learning families, and walks through a demonstration of latent space operation sequences on MNIST before developing notions of losses and evaluations on data. It reads almost like a position paper with few evaluations.", "questions": "While converting to measurable spaces, does the scale of the Borel interval have an effect? That's something that may merit digging into. \n\nWhile evaluating, accuracy seems just an indicator and not really a measure of distributional discrepancy like mmd and adv. Do I understand it right?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "Even though elements of the ideas exist, as elaborated in literature, this is perhaps a first formulation of latent space equivalences in terms of category theory. The formulation expressed, developed and proved mathematically appears sound on a few reads of the text.", "weaknesses": "Both classes of functor mappings are linear or affine, given properties of identity, distributivity and commutativity. The development of these ideas for supra-linear mappings may present an uncertainty at the fundamental that is unknown at the moment.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6zJokNqlcx", "reviewer": "Reviewer_Fvtj", "summary": "In this study, the authors introduced the category theory to reinterpret the problem of medical data pooling.  The study firstly elaborated on the MNIST toy dataset for explaining the morphism composition imposed on the latent space, then conducted experiments on ADNI and ADCP datasets by testing on the scanner parameter. Compared to the closely-related GE method, the proposed approach achieved both low time complexity and better performance reflected on ACC and MMD measurements.", "questions": "Please see the weakness section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The introduction of category theory to the medical data pooling is novel and interesting.\n\n2. Despite the complicated topic, the paper is well-written and well-structured, therefore quite easy to understand.\n\n3. Both the runtime and quantitative results are improved over the GE method, which are quite impressive.", "weaknesses": "1. Though it is novel and interesting to introduce category theory to this problem, I am not sure if it is really necessary to do so. As far as I understand and please correct me if I am wrong, there are no additional theoretical novelties other than quoting existing definitions to the paper.\n\n2. In terms of the scanner parameter discussed in the experiments, it is also very common to apply different kinds of normalization tricks to resolve the issue. Unfortunately, I didn't see relevant discussion or experiments presented in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Btb0PVlZ3q", "reviewer": "Reviewer_zBw9", "summary": "This paper aims to provide a general harmonization tool that can handle multi-equivariance and multi-invariance with respect to the images' covariates.", "questions": "None", "limitations": "", "rating": 5, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The problem to be solved in the paper is more interesting.", "weaknesses": "I don't see any obvious Weaknesses", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)