Genetic InfoMax: Exploring Mutual Information Maximization in High-Dimensional Imaging Genetics Studies
Authors: Yaochen Xie
Ziqian Xie
Sheikh Muhammad Saiful Islam
Degui Zhi
Shuiwang Ji
Venue: Accepted by TMLR ()

Abstract: Genome-wide association studies (GWAS) are used to identify relationships between genetic variations and specific traits. When applied to high-dimensional medical imaging data, a key step is to extract lower-dimensional, yet informative representations of the data as traits. Representation learning for imaging genetics is largely under-explored due to the unique challenges posed by GWAS in comparison to typical visual representation learning. In this study, we tackle this problem from the mutual information (MI) perspective by identifying key limitations of existing methods. We introduce a trans-modal learning framework Genetic InfoMax (GIM), including a regularized MI estimator and a novel genetics-informed transformer to address the specific challenges of GWAS. We evaluate GIM on human brain 3D MRI data and establish standardized evaluation protocols to compare it to existing approaches. Our results demonstrate the effectiveness of GIM and a significantly improved performance on GWAS.

Shuiwang Ji Published in Transactions on Machine Learning Research (05/2024)
Genetic InfoMax: Exploring Mutual Information Maximiza- tion in High-Dimensiona Imaging Genetics Studies
Yaochen Xie
ethanycxStamu.com
Texas A ABM University
Yuchao Lin Texas ABM University
kruskallón &tama edu
Ziqian Xie
mrniio.riieuuh.tmm...om
University of Teras Health Science Center
Sheikh Muhammad Saiful Islam University of Tezas Health Science Center
Degui Zhi
tmc. edu
University of Tezas Health Science Center
Sja@tam edu
Texas ABM Universily
Reviewed on OpenReview: https: // // opwnr/veew.eeeee.e net/ formm? 9UgUMFW67X
Abstract
Genomo-wide association studies (GWAS) used to identify relationships between genetic variations and specific traits. When applied to high-dimensional medical imaging data, a key step is to extract lower-dimensional, yet informative representations of the data as traits. Representation learning for imaging genetics is largely under-explored due to the unique challenges posed by GWAS in comparison to typical visual representation lcarning. In this study, we tackle this problem from the mutual information (MT) perspective by identifying key limitations of existing methods. We introduce a trans-modal learning framework Genetic InfoMax (GIM), including regularized MI estimator and novel cenetics-informed transformer to address the specific challenges of GWAS, We evaluate GIM on human brain 3D MRT data and establish standardized evaluation protocols to compare it to existing approaches. Our results demonstrate the effectiveness of GIM and a significantly improved performance on GWAS.
1 Introduction
Genomo-wide association studies (GWAS) have been an effective approach driving genetic discovery in the past 15 years (Abdellaoui et al., 2023). Given a phenotype of interest and cohort of individuals with both the measurements of the phenotype and the genotypes over markers across the genome, linear linear mixed models are built to test for the association of each marker to the phenotype and thus pinpoint the gene loci relevant to the phenotype. While the typical GWAS studics are focused on well-ostablished phenotypes, typically the risks of diseases, or well-established macro-level measurements such as heights, and BMI, or molecular- measurements such as protein and metabolomic biomarkers, there is an interest in expanding the scope GWAS to focus on phenotypes that are derived from high-dimensional complex data modality such as medical imaging data, better understand the biology of the underlying images. Such derived phenotypes sometimes called endophenotypes as they serve as intermediate between genetics and some

---

Published in Transactions on Machine Learning Research (05/2024)
discases. For example, brain images are often considered endophenotypes of brain-related discases such A Alzheimer's disease. However, there is a lack of sophisticated approaches for deriving phenotypes for GWAS. Again, taking brain imaging as an example, existing approaches mostly used traditional non-learning software to derive brain region- based volumetric or surface features. Such methods, while valuable, are reaching a limit in fully capturing phenotype complexity and can narrow the scope of genetic discoveries. In this context, the advent of deep learning techniques offers a promising avenuc for expanding the toolkit available for GWAS.
Recently deep learning approaches (Pttel et al., 2022 Xie et al. 2022 Kirchler et al. 2022 Taleb ct al. 2022
derive phenotypes from medical images by learning latent representation that captures the inherent content of the input image. However, approaches learning from imaging data alone fail to utilize the accompanying genetic data. Those approaches tend to capture patterns that are not related to genes and common patterns shared by multiple individuals. For example, Patel et al. (2022) found that representations learned by an image autoencoder are unable to fully reconstruct fine details that are individually specific. To overcome these limitations, solution is to incorporate trans-modal learning strategics that utilize the pairwise relationship between imaging and genetic data, such as trans-modal cont rastive learning (Liang et al. et al. 2021 Taleb et al. (022).. Unfortumately, the use of genetic data, including the encoding of data and capturing image-genctic relationships, still poses significant challenges. Results by Talcb et al. (2022) suggest that, despite the promising results for downstream classification of disease risk, multi-modal contrastive approaches still underperform compared to typical image-only approaches (Xie et al. 2022) on GWAS tasks. This underperformance becomes even more pronounced for
2022 Zolfaghari
3D data.
In this work, we formulate the problem as learning the representation of imaging data that shares the maximum mutual information with genetic data. By using mutual information as perspective, we are able to examine the key reasons for the failure of typical trans-modal contrastive learning for GWAS on high-dimensional imaging data. To push the limits of existing learning approaches, we propose Genetic InfoMax (GIM), a trans-modal learning framework that includes a regularized mutual information estimator and a novel transformer-basee genetic encoder. The framework addresses the issucs of dimensional collapse and non-generalizable associations in representation learning for GWAS and fully utilizes the genetic data with physical and genetic position information. Our experiments demonstrate that GIM significantly improves performance on all four evaluation metrics.
2 Problem Formulation
We study the problem of genome-wide association studies (GWAS) on high-dimeusional data. The GWAS aims to identify associations between specific genetic variants, known as single mucleotide polymorphisms (SNPs), in the genome and certain traits of interest such as the risk of disease and other biological characteristics of organisms. In particular, cach individual genetic data is denoted by G G= (g), 9... and traits of interest denoted by y. The GWAS process involves statistical tests on the sample pairs {(G,, y=)) from a large number of M individuals to identify the specific subset G9-+y G of SNPs that are associated with the target traits y. Here L is the number of SNPs and cach 9. 1, represents the number of carried variants for each individual. The values of nearby SNPs are often correlated due to their common inheritance from a shared ancestor. To account for this, it is necessary to select an independent subset of genetic information Gind C G9-y, which can be achieved by clustering and selecting the SNPs with the lowest p-value from each cluster after the statistical test, which is important for accurate analysis and interpretation of the genetic data, Practically, when conducting GWAS on high-dimensional data Y such as medical imaging, an additional step required before performing statistical tests. This step involves reducing the number of traits from the high-dimensiooal data Y to a smaller number y through experts' diagnosis or computational approaches.
Traits Computing as Representation Learning. To enable GWAS on digh-dimensional data, we are interested in computationally obtaining informative lower-dimensional traits, termed GWAS representation learning, from the high-dimensional data. Specifically, for any high-dimensional data Y to be studied, the problem is formulated as learning lower-dimensione representations of Y with a corresponding encoder fe in a self-super manner, such that a larger number of independent SNPs [Gind] can be identified from the pairs {GG,, fo (Y)))) The goal of identifying more independent SNPs requires that more information related to genetic variants is captured by the representation of Y. We focus on learning d-dimensional
2

---

Published in Transactions on Machine Learning Research (05/2024)
representations y fo(Y) E R with the following optimization objective
0 = arg max (())(),,))
(1)
where I(-,)) denotes the mutual information (MI) between two random variables. As computing the true value of mutual information is intractable, it becomes critical to develop an appropriate mutual information estimation under the GWAS problem setting.
Notations of Data. We instantiate our problem specifically with the 3D human brain magnetic resonance imaging (MRT) data and SNPs from the human genome. We denote the 3D brain MRI by Y Hx where H, W, D denote the three spatial dimensions, and 1 denotes the single channel of the MRI. The human genetic data G consists of N positions on the human genome with frequent variants (SNPs). Each SNP ge is represented by four-tuple (de, Ci p. pr pi )) Tn the four-tuple, di E 1, denotes the genotype of the SNP, the number of copies of the mutant allele, c E {I, 22} denotes the index of chromosome the SNP belongs to, pi E N denotes the physical position in terms of base pair (bp) of the SNP in the chromosome, and pt R+ denotes the genetic position of the SNP in terms of centimorgan (cM). Note that the genetic data is not sequence but array since two neighboring SNPs gi and gi+1 are not necessarily consecutive on the original genome and the physical (and genetic) distance between them |p, pby is meaningful. We further denote arrays consisting all genotypes, chromosomes, and positions sorted on chromosome id and physical position by d, c, pahy, and pseu, respectively.
3 What Makes Appropriate MI Estimators for GWAS?
With the goal of learning representations that capture as much information about the genetic variations as possible, our objective is maximize the MI between the representation and genetic data with an appropriate MI estimator. One commonly used approach for estimating the mutual information between multi-dimensional variables is the Jensen-Shannol Estimator (JSE) NNwozin et al. 2016). The JSE involves discriminator to distinguish heether samples of the two variables belong to the same individual or are independently sampled. Specifically, under our problem setting, the USE-based training loss is computed as
log (Po(fo(Y), G) G))
CJSE (BsE(B;,,g)
B
(Y,G)EB
(2)
E log (I-Do(fo(Y), GT){
Σ
B 1)
B ((
where B is a mini-batch paired MRI and genetic data and Do RR x x - (0, 1) is a learnable discriminator determine whether (Y) and G are from the same individual. Together with the Noise Contrastive Estimation (InfoNCE) (Bachman et al. 2019), these learning processes are also known to be in a contrastive manner across two modalities; namely, the MRI and genetic data.
To achieve desirable performance in maximizing MI and discovering genetic associations, the discriminator should meet certain requirements. First, the learnable discriminator Dj should be able to take as inputs the genetic data G and encode all the useful informa- tion from G. A well-designed genetic encoder is thus critical component of D.. It should be able to efficiently and effectively use not only the genotypes, but also their corresponding chromo- some, physical position, and genetic po- sition information. In Section 4
Dimensional
G
Non-gereralzablle
(a). Ideal Ja(Y)
(b). Reality (contrastive)
(c). With regularization
Figure 1: An illustration showing how the representation fo(Y) captures the mutual information between Y and G in different cases. The circles are the entropy of and Y, respectively, and their intersoction the mutual information I(G, Arcas indicated by squiggles in red represent the information contained in (Y).
3

---

Published in Transactions on Machine Learning Research (05/2024)
propose novel transformer-basee based genetic encoder, dubbed genetic transformer, that fully utilizes all this information based on genetic intuitions.
Second, the discriminator should make predictions based on all generalizably associated patterns, rather than memorizing noise or focusing on a small portion of associated patterns that can be easily learned. However, due to the nature of rastive learning and several differences between typical visual representation learning and GWAS representation learning, we will argue below that it is challenging to meet this requirement. The typical contrastive loss can lead to degenerated results for GWAS shown in Figure 1
3.1 Uniqueness of GWAS Representation Learning and Limitations of Contrastive Losses
To understand the limitations of typical contrastive losses in the GWAS setting, we first identify key differences between the visual representation learning problem for natural images and the GWAS representation learning on high-dimensional data. We explain how each difference can contribute to limitations or failures of typical contrastive losses in the GWAS setting, and provide empirical evidence to support our arguments.
Difference Goals of learning representations. Typical visual representations of natural images aim to capture the key semantics or class information about major objects in images. With this goal, it is acceptable for the representation to capture only semantic or class-related information even required that representations are invariant to elements such context (Zhang et al. 2021 Chang et al. 2021 and transformations (Xiao et al., 2020; Foster et al. 2020). In this case, a good representation for downstream lasks does not necessarily maximize the MI during contrastive learning. In contrast, a good representation for the GWAS purpose should capture every detail or pattern in the high-dimensional MRI data that is associated with the genes, since there no such key semantics or class information. In this case, the downstream GWAS performance is closely associated with I(fo(Y),
Limitation 1: Dimensional collapse. A recent study on visual representation learn- ing (Jing et al., 2021) iden- tifies and empirically shows
First 10 Principal Components
Representations (Pre (Pre-projection)
Embeddings (Post projection)
of Explained Variance
MI-NCE MI-JSE
Log of Singular Value
Log of Singular Value
-2
-5
that typical contrastive ap- proaches suffer from the di- mensional collapse issue, where
4
10
-5
-15
-6
-20
-10
-25
the learned representations oc-
of
MI-SEE
-30
MI-SE
-15
cupy dier-ddimnsional sub- space than their designated di- mensions. The dimensional collapse results high redun- dancy, limits the information captured by representations, and therefore leads to reduced performance in downstream tasks. Indeed, our analyses show that the dimensional collapse issue also arises in the cross-modal contrastive setting. We compare the singular values of representations learned by predictive methods and contrastive methods in Figure 2 Results indicate that the contrastive estimators NCE and JSE suffer from dimensional collapse with a dramatic drop in explained ratios and singular values. Even worse, the GWAS performance suffers more from the dimensional collapse issue due to its nature described in Difference 1, both the information of fY) and the mutual information (Y), G) is limited.
Pred
35
Pred
6
50
100
50
100
Principal Components
Singular Value Rank Index
Singular Value Rank Index
Figure 2: components (left), the singular value spectrum of learned representations (middle) and embeddings after projection (right). Comparisons are among contrastive MI estimators with InfoNCE (MI-NCE), with JSE (MI-JSE), Autoencoder (AE), and genetic data prediction (Pred).
The logarithm of the explained variance for the first 10 principal
Difference 2: Augmentation approaches and data dimensions. Contrastive learning relies on a large number of samples to more accurately estimate and maximize the mutual information between different views or modalitics. Previous studies (Chen et al. 2020 Tian et al., 2020) have shown that augnentations are crucial for contrastive learning, as they prevent representations from focusing on patterns that are irrelevant to downstream tasks and multiply the number of training samples. For higher-dimensional 3D MRI data,
4

---

Published in Transactions on Machine Learning Research (05/2024)
more samples and diverse augmentations necessary Vapnik 1999). imaging data is limited, and most augmentations used for typical visual representation learning are not applicable to medical imaging. For example, since MR are single-channel, color space augmentations are not possible. Augmentations based on rotation and flipping are not suitable for brain MRI data due to their asymmetrrc nature. Random linear transform or non-linear morph may change the shape of the elements of the image and thus are discouraged. In the case of 3D MRI, the applicable angmentations are very limited.
However, the availability of medical
Limitation 2: Non-generalizable associations. According to Huang et al. (2021), augmentations play critical role in the generalization capability of contrastive learning approaches. However, due to the limited number of applicable augmentation techniques and the dimensionality of 3D data, the discriminator tends to capture non-generalizable or false associations from the training samples such as memorizing the shape, the layout of the brain in the MRI, or specific noise in the data to identify individuals. In Figure B we evaluate the generalizability of models trained with contrastive loss by comparing the MI estimation on training and validation pairs. The remarkable discrepancy in losses between the training and validation sets suggests that the discriminator used by contrastive loss is unable to generalize to new samples, indicating that contrastive loss is a poor estimation of MI in our case.
From the perspective of mutual informa- tion, due to the dimensional collapse, a limited amount of training samples, and insufficient means of augmentations, em- pirical results show that the JSE is not an
JSE
InfoNCE
12.5
Contrastive Loss
10.0
Train Val
5-
7.5
Train
optimal estimator of mutual information. The true mutual information is hence not maximized by the learned representa- tions, leading to degraded performance in GWAS. Figure | (a-b) illustrates the ro- lationship among the brain MRI, genetic data, and the learned representation. In the ideal case shown in (a), the represen- tation should perfectly cover the mutual information between Y and G, so that the learning objective achieves its maxi- mal with
5.0
Val
2.5
3-
0.0
10
20
30
10
20
30
Training Epochs
Training Epochs
Figure 3: Generalization capability of contrastive MI estimators JSE (left) and InfoNCE (right). The learned discriminators fail to generalize to unseen pairs, leading to large discrepancy between training and validation losses.
Vfo
(3)
In practice with brain MRI data, the contrastive loss results in representations that only capture small portion of IYY, G) due to the two limitations described above, as shown in (b).
3.2 MI Estimator with Regularizations
Given the issues and limitations outlinod above, our goal to improve the representation by incorporating more generalizably associated patterns in addition to those identified by the contrastive MI estimator. However, unlike in the case of natural images where many elements are known to be non-generalizable, our limited knowledge of undiscovered genetic associations makes it difficult to determine which patterns are generalizable and which are not. As result, it is challenging to develop targeted augmentations that make the representation invariant to unwanted patterns.
To achieve our goal without requiring further knowledge, we propose to uniformly increase the total information contained in the representation by including an entropy term in the learning objective. The objective is formulated
max [(fo(Y).)) +XH[)))))))())
(4)
where i is the contrastive MI cstimation JSE, H denotes the entropy of a random variable, and a weight scalar. The entropy term encourages the representation to capture more information about Y and reduces its redundancy. A certain portion of the information can contribute to the generalizable associations,
5

---

Red and blue boxes represent the Published in Transactions on Machine Learning Research (05/2024)
Down-sampling
Window Attention
Shifted Window Attention
Value
Value
Embedddings
within
Embeddiinge
within window
Key
Attantion
Genotype and chromosome d; C/
map
map
X
Query
Query
Relative physical posttions
Relative physical positions
Genetic distances
Genetic distances
SNP embeddings
Updated SNP embeddings
(glte piplle
Figure A twan-transformer block in the proposed genetic encoder. windows and shifted windows.
illustrated in Figure (c). The entropy term serves regularization to the estimated MI to improve its generalizability. From the optimization aspect of view, the objective is considered as adding a Lagrange multiplier to maximize the entropy H(Jo(Y)), subject to the constraint that the estimated mutual information (Y)o(Y) G) achieves its maximum. When multiple patterns can be used to identify individuals, the entropy term encourages the model to capture as many of them as possible, instead capturing the easiest ones.
There are various methods estimate and optimize the entropy, such as minimizing the off off-diagonal values in the covariance matrix of the representation (Zbontar et al. 2021). However, these estimations require large mini-batch size, which is not suitable in our case due to memory constraints caused by the 3D data and MRI encoder. As an alternative, we use the reconstruction of MRI data as a proxy to maximize the entropy. The loss then computed as
A
C(B:0,6,0) LISE B:
|Y - hofo)(f((Y)) (Y,G)EB
Σ
(5)
B
where hy is a deterministic decoding head used to reconstruct the MRI from the representation fo(Y). Compared to other proxies discussed by Zbbontar et al. (2021). the reconstruction term is less sensitive small mini-batch sizes. To justify the reconstruction term, we have
H(fo(Y)) I(fo(Y) + HUMYY)),
(6)
where the conditional entropy ((fo(Y) is reduced to when the encoder is deterministic. Intuitively, given the input Y, the learned embedding ((Y) will not carry additional information about Y. We then consider the maximization of the mutual information I(fo(Y Y). According to |Wang and Isola 2000 the
objective consists of two components: 1) the alignment of fY( to Y, which is achieved by maxinizing a positive pair belong to the same individual, and 2) the uniformity of
the likelihood that fY) and (Y)),, which is achieved by maximizing the likelihood that (Y and any other Y' are independent. To reduce the batch-size sensitivity and optimize memory cost for extremely high-dimensional MRI data, we avoid introducing additional encoders and discriminators used by conventional MI estimators. Instead, we use the reconstruction term to maximize the log- likelhoodo that fY( and Y in positive pair belong to the same individual, corresponding to the first component. The uniformity of fo is addressed in T(fo(Y) G) by maximizing the likelihood that fY(Y) and any other G' are independent, inluitively ensuring that the embeddings any two different individuals are distinguishable from each other. From the perspective of Wang and Isola (2020). the two terms in LJSE( B; 0, Φ) and the reconstruction term aim at optimizing three propertics of fY()); namely, the alignment to G, the uniformity, and the alignment Y.
4 Genetics-Informed Transformer
A typical current genotyping microarray of the human genome includes more than 650k SNPs, with the physical spacing between any two consecutive SNPs being inconsistent. Genetic encoders developed in existing

---

Published in Transactions on Machine Learning Research (05/2024)
studies (van Hilten et al., 2021 Taleb et al. 2022) based on convolutional neural networks (CNNs) and multi-layer perceptrons (MIPss) are incapable of handling the unstructured genetic data with extremely large sizes. To address this, we develop an effective genetics-informer transformer to encode genetic data in accordance with the following objectives:
1. Signíficant optimized computational cost, in recognition of certain biological assumptions.
2. Information aggregation among SNPs from arbitrary positions in the genome, considering multiple genetics dependency measurements.
3. Flexibilitty to accept any segments or subsets of the genetic data as input, thereby facilitating cropping Nownsampling-baset augmentations on the genetic data.
An overview of the proposed transformer block is shown in Fig. 1 where allention operators with shifting windows (Liu et al. 2021) are used to enable efficient computation, and the aggregation is specialized with both physical and genetic distances of SNPs. In the transformer, two blocks are connected by down-sampling with altention-based pooling operators, and the initial SNP embeddings are computed based on both the genotypes and the chromosome each SNP belongs. Finally, attention- based readout is used to compute the global representation.
Window attention in swin-transformer The 1D trin-trannformer performs self-attention operations within each window split. from the entire genetic array to enable efficient computing. Tt. contains two components; those are, window attention and shifted window attention, as shown in Fig. 4 Given input SNP embeddings H E R where L denotes the number of SNPs and q denotes the embedding dimension, window attention first splits H into set of windows (H, E )=i,,, size w. A sclf-attention block is then applied to cach H; to update the SNP embeddings by aggregating information within that window. The updated windows are then merged back following the order when splitting H, forming H' as the final output of the window attention component. In the following shifted window attention component, H' is first shifted by a length of [w/2]. Then similar splitting and self-attention are performed as in window attention to update SNP embeddings from each individual window. Finally, the updated windows are merged back, and the merged sequence is also shifted back by a length of [ur/2]. There exists biological assumption that strong and informative dependencies between SNPs exist only when they are within certain distance. Hence, compared with performing global attention all 650k marker positions, performing allention within windows in our proposed met aggregates similar information, but largely reduces the computing cost.
where each window has
Aggregation based on multiple dependencies. Multiple attention heads are computed to capture various types of dependencies among markers. Specifically, the computing of attention scores captures dependencies from three perspectives; those are, the SNP embeddings reflecting potential co-mutation, the encodings of the physical positions of SNPs on chromosome indicating local dependencies among SNPs, and the encodings of genetic positions of SNPs measuring genetic linkages. Formally, the attention score af between the i-th and j-th SNPs for the k-th attention head is computed as
1 (g., 9;) So, (hr, )) Iāz (v."h .rghrr + SKBP (v." 1 "
where h; = Hi E R2 denotes the SNP embedding. We compute the three individual components of attention scores that can capture different genetic dependencies as
(he, (h w (h.) wh:r) Sto (opthy rpphor (ep Wwhk (o
(8)
JkBF (r - - (1( 23 D)' Pog -p 120 Wh+t 1 -1p-0 )o rbf
where e.(-) denotes the position encoding (Devlin et al. are trainable projections. 1 L is an indicator vector where all elements are is if the condition holds, and are Os otherwise. The function r denotes distance expansion with radial basis functions (RBF) (Schitt
k 2019 Wh., Whr WkJ Wh. and rbf rbf
et al. 2017). Denoting s = 1p; gen gen the term r ((p, gen P, p) in the above equation is computed as
r(s) = exp (s - /)/ (ss-))/o²)
(fo. At) (ee R
(9)
7

---

Published in Transactions on Machine Learning Research (05/2024)
where (to.... .te) is set of non-negative real numbers ranging from 0 to preset threshold. The asymmetric projections in all Sh: Spe and /RBF functions indicate that Oi,j does not necessarily equal to Oj,,, leading to more expressive models to capture dependencies. In addition, the computing of each attention head is based on a combination of three types of dependencics, which enables information aggregation among different SNPs based on genelic dependencies. By doing this, the complicated genetic dependencies of the input, genome data can be captured. Additional details about the transformer are provided in Appendix B
5 Related Work
Dimension reduction for GWAS. When doing cenomo-wide association studies, people oftentimes find themselves dealing with high-dimensional quantitative traits. In order to reduce computational cost and rodundancy, and in the hope of finding meaningful underlying patterns, many works perform dimension reduction of the high dimensional traits before doing GWAS, including principal component analysis (Zhao
ct al.
2021 Yano et al. 2019 Ma et al. 2021). independent component analysis (Elliott et al.) 2018 Pearlson 2015) and non-negative matrix factorization (Wen et al.) 2022). These approaches are effective in
et al. capturing lincar dependencies but are less capable of identifying complicated traits from imaging data.
Deep learning-hased approaches. Recently, several works used unsupervised learning to characterize high-dimensional medical data. iGWAS (Xie et al., 2022) applied contrastive learning between multiple images of the same person to reveal potential genctic signals, ContIG (Taleb et al. 2022) applied contrastive learning ween medical imaging data and genetic data to learn the feature representation. DeepEndo autoencoder (Patel et al. 2022) used a convolutional autoencoder to reduce the dimensionality of the imaging data and found genetic associations of these extracted phenotypes. TransferGWAS (Kirchler et al.) 2022) used both supervised task and reconstruction task to learn the feature representation. Specifically, ContIG (Tleeb et al.] 2022) is the first to use contrastive learning between images and genetics on the GWAS problem. However, there are distinguishable differences between the work and ours. First, ContIG aims to learn general representation for multiple downstream tasks, such as classifications of the risk of several diseases. With this goal, ContIG treats the problem a typical visual representation learning task. On the contrary, our study focuses on the representation lcarning specifically for GWAS. Second, our approaches are built upon the grounding of mutual information maximization, whereas ContIG is grounded by contrastive learning for unlabelled data. Third, our work focuses on a more challenging setting with 3D MRI data, where typical contrastive approaches may fail.
Mutual Information Maximization. Previous research has employed mutual information maximization pretext task for representation learning on various data types, including images (Hjelm et al. 2019). videos (ijelm and Bachman 2020), and graphs (Veličković et al.) 2019 Stärk et al. (022).. However, these studies primarily focus on classification or regression as downstream tasks. Our work presents unique challenges and goals of mutual information maximization under the GWAS setting and we are the first to examine GWAS from a mutual information perspective.
6 Experiments
We use the brain imaging dataset from UK Biobank (Sullow et al., 2015) in this study, it is currently the largest public brain imaging dataset. Specifically, we use TT-weighted MRT imaging data accessed on October 15, 2021. We register and pre-process the MRI data into the shape of 182x218x 182. For representation learning, we split the MRI-genetic data pairs into 4,597 training and 1,533 validation pairs based on ethnicity. A detailed description of the processing and split is provided in Appendix A
GWAS Evaluation Metrics We involve three metrics to evaluate the representations Iearned by different models; namely, the number of loci discovered by GWAS, the estimated mutual information, and the heritability of the representations. All three metrics are computed on a testing dataset that is unseen during the representation lcarning process and measures the quality of representation for GWAS purposes. To enable efficient evaluation, we obtain the first 10 principal components of representations and compute all metrics the 10-dimensional vectors. We provide detailed descriptions of each evaluation metric in Appendix A
8

---

Published in Transactions on Machine Learning Research (05/2024)
Table 1: Comparisons of quantitative evaluation results on the test set. # Loci are counted under Bonferroni corrected p-value threshold (cutoff) of 50-9. "Unique" refers to the number of loci discovered by method that is NOT discovered by any methods in other groups.
# #Iddependent Loci
Methods
Allf
Uniquet
Estimated MIf
Heritability h²
Random Init
14
1.2178 + 0.0128
0.0756 = 0.0656
Autoencoder (Patel al.
2022
1
1.3124 + 0.0062
0.3121
0.0769
Predictive
Autoencoder-attennioo
23
4
1.3138 + 0.0061
0.2984 + 0.0773
Gen Prediction
10
0
1.2482 + 0.0116
0.0918 = 0.1110
Barlow Twins (Zbontar al.
2021
11
1
1.3008 + 0.0072
0.0814 0.0636
Contrastive
SimCLR (Chen et al. 2020
15
1
1.2355 + 0.0131
0.1448 = 0.1128
SimCLR-JSE
7
1.3062 + 0.0068
0.1604 0.1151
InfoNCE (ContIG, Taleb et al. (2022))
11
0
1.2268 + 0.0142
0.1334 = 0.0588
Trans-Modal
Decorrelated InfoNCE
3
1.2274 + 0.0128
0.0527 0.0349
Contrastive
GIM (Ours)
40
15
1.3782 + 0.0010
0.3807 + 0.0517
6.1 Quantitative Results and Analysis
We compare our approach with multiple baseline approaches four groups; namely, MRI encoders that are randomly initialized, trained by predictive approaches, contrastive approaches with MRI data only, and trans-modal contrastive approaches that involve genetic data. The baselines include existing straightforward
training schemes, namely, Autoencoder (Kirchler et al. 2022 Patel et al. 2022 Gen Prediction that uses genetic data prediction as a pretext task, Barlow-Twins (Zbontar et al., 2021 SimCLR (Chen et al. 2020), and ContIG (Taleb et al. 2022). We additionally include their variants Autoencoder-atteetiioo that uses the same MRI encoder as ours, SimCLR-JSE. where the contrastive objective in SimCLR is replaced by JSE, and Decorrelated InfoNCE, where a decorrelation term added to the contrastive loss. We describe the implementation details of GIM and baselines in Appendix B
Comparisons among representations learned by different methods in terms of the three metrics shown in Table 1 The results indicate that the proposed learning framework with a regularized MI estimator and genetic transformer significantly improves the quality of learned representation in terms of the number of discovered loci and the heritability. The improved MI of our methods on test pairs also suggests a st.onggr generalization capability. Additionally, we have the following observations.
The level of mutual information on test pairs agrees with # loci and heritability. The results suggest that higher mutual information the test set implies higher heritability and more loci discovered. It justifies our formulation of learning representation for GWAS as the problem of maximizing mutual information.
Typical trans-modal contrastive approaches fail for MRI data. Trans-modal contrastive learning with typical contrastive loss performs fairly well on 2D retina imaging (Tleb et al. 2022) but suffers more from the performance reduction on the liggher-dimenssonall 3D data. the 3D MRT case, we found that the simplest Autoencoder approach performs even better than contrastive and typical trans-modal contrastive approaches.
GIM has consistent performance. To further examine the robustness of GIM, we obtain additional results by training 3 additional GIM models independently at different randomizations. They resulted in 43, 42, and 46 discovered loci, respectively, which indicates consistent and significant improvement compared to baseline approaches.
6.2 Ablations and Additional Results
We perform additional quantitative studies to demonstrate the effectiveness of GTM. Visualizations for
representation uniformity and Manhattan plot for GWAS are provided in Appendices C D and
Effectiveness of individual proposed components. To show the effectiveness and necessity of both the proposed learning objective and the genetic transformer, track the change in the number of all loci, unique

---

Published in Transactions on Machine Learning Research (05/2024)
Table Change in the number of loci and heritability when incrementally adding components to the models.
Methods
# Loci
Change
Unique Loci
Change
1²2
Significant Gene
Base-Contrastiver
11
0
0.1334
0.0588
+ Regularization
29
+18
6
+6
0.3390
0.0627
CENPW
+ Genetic Transformer
32
+3
10
+4
0.3773
= 0.0384
WNT16
+ Random Cropping
36
+4
13
+3
0.3723
0.0558
ITPR3
+ Co-training
40
441
15
4-2
0.3807 1 0.0517
MSRB3
loci, the heritability score, and newly discovered genes with the highest significance when incrementally adding each component. Table 2 shows the results of adding the regularization to the objective, replacing the MLP encoder with the genetics-informed transformer, and performing random cropping the genetic data. The results suggest that adding each component generally increases the useful information carried by the representations, leading to more loci discovered. We discuss the new gene discovered by each component and the change of h? in Appendix D
GWAS representation for Teiweihtted MRI. We additionally apply GIM to a second modality, namely the T2-weighted MRI. Sim- ilarly, compute GWAS on the first 10 principle components of the learned representation on the test set. In contrast to the results for T1, we observe that learning informative representations is less challenging for contrastive methods in the T2 case. In addition, con- trastive methods equipped with NCE generally perform better than their JSE counterparts. This result consistent with those presented in Taleb et al. (2022). Nevertheless, results in Table 3 show a consis- tent out-performance of GIM over baselines, indicating generalizable effectiveness.
Table 3: # Loci for T2- 12-weighted
MRI at cutoffs 5e-9 and 1e-8.
Methods
5e-9
1e-8
Autoencoder
29
48
Barlow-Twin
14
SimCLE
21
28
SiRCLR-JSE
15
20
ContIG
22
29
GIM (Ours)
38
62
Downstrcam classification performance. Although this is not our primary goal, we perform additional evaluation with three downstream binary classification lasks, namely stroke, cancer, and biological sex, from UK Biobank. We follow Taleb et al. (0222 to perform linear evaluation protocol and compute their ROC-AUC. Specifically, we obtain MRI embedding from individual MRI encoders. We then freeze all embeddings and perform linear classification upon the embeddings for the three tasks. Since ROC-AUC is computed on the entire validation set, we repeat, the classification 5 random training-validatin splits and show their average ROC-AUC score. We include ablations for GIM with without the presence of regularized loss (Reg.), Genetics-Infone The results shown in Table 4 indicate that GIM achieves better performance compared to its contrastive counterparts. For cancer and biological sex classification, we identify continuous performance gain when adding individual components, whereas, for stroke, the regularized loss and more intensive cropping bring the most improvement. For Autoencoder, we observe significantly better performance on biological sex classification. The underperformance of cross-modal approaches is due to the exclusion of the sex chromosome in the given SNP data, and thus the MI maximization discourages the model from capturing biological sex-related information from MRI data.
Transformer (GT), and random cropping (Cropping) with different probabilities.
Table 4: ROC-AUC of downstream classification.
Methods
Reg.
GT
Cropping
Stroke
Cancer
Sex
Autoencoder
0.6748 0.0026
0.5828 + 0.0005
0.9039 + 0.0003
SimCLR
0.6329 0.0015
0.5797 = 0.0005
0.7739 = 0.0217
ContIG
0.6034 0.0001
0.5385 + 0.0038
0.5927 + 0.0002
GIM
0.6563 0.0069
0.5738 + 0.0006
0.7515 + 0.0057
GIM
0.6510 = 0.0002
0.5789 - 0.0007
0.8601 = 0.0010
GIM
p = 0.2
0.6387 + 0.0180
0.5836 + 0.0013
0.8650 + 0.0016
GIM
p 0.5
0.6772 + 0.0250
0.5898 + 0.0009
0.8460 - 0.0161
10

---

Published in Transactions on Machine Learning Research (05/2024)
7 Limitation and Future Directions
Current GWAS relies statistical testing based on linear models to identify significant loci and compute the heritability scores. However, these methods can be limited in their ability to capture non-lincar and complex associations, leading to missed loci and inconsistencies between mutual information and other metries. To address these limitations, new approaches to discovering non-linear associations are needed.
Regarding mutual information optimization, our current mutual information optimization solution aims to improve representations by capturing more generalizable associations, but it is still limited in reducing information on non-generalizable ones due to lack of knowledge about non-generalizable patterns. With specific domain knowledge, it is possible to further improve these representations.
The interpretability is a further step towards understanding the association between gene and brain phenotypes and is another important future direction of this work. In fact, the learning of informative representations with "good" enough encoders is a crucial prerequisite of interpretability. Once such an encoder model becomes available, we are able to adapt existing visual interpretation voxel-level explanation approaches to identify significant regions in the brain that are associated with certain genes. The interpretability encoder side is cqually important. While capturing some genetic information, it lacks the resolution to trace back the ronttioiionns from individual SNP markers. Future work is needed to develop more elficient and high-resolution genetic encoder.
the genetic
8 Conclusions
In this work, we have investigated the differences and limitations of GWAS representation learning compared to typical visual representation learning and have presented Genetic InfoMax, GWAS representation learning framework. We have establishod standardized evaluation protocols to benchmark existing and our approaches. Our experiments demonstrate a significant boost in GWAS performance by GTM.
Acknowledgement
This work was supported by the National Institutes of
Health grant U01 AG0701122 12.
Broader Impact Statement
The use of individualized genetic data and imaging data in medical research raises ethical concerns related privacy, informed consent, and potential discrinination. The data used in this study were from the UK Biobank under the approved Project 24247. The Committee for the Protection of Human Subjects of The University of Texas Health Science Center at Houston gave ethical approval for this work under HSC-SBMI- 20- 1323. We followed the best practice and confirmed that all necessary patient/participan obtained and the appropriate institutional forms have been archived, and that any patient/pani http//paaiiipttssammpnn identifiers included were not known to anyone (e.g., hospital staff, patients, or participants themselves) outside the research group so cannot be used to identify individuals. Participants enrolled in the UK Biobank have provided broad consent, expressly permitting the use of their data for biomedical rescarch purposes that encompass the objectives of this study, ensuring compliance with ethical standards and the intended broad impact of our research.
consent has been
The work focuses on the development of new methods for uncovering genetic factors contributing to brain morphology as measured by brain MRI images. Our model training sets include people from diverse ancestry backgrounds. The GWAS cohorts include primarily reportep British white individuals. Therefore, the gencralizability of the genetic findings to other populations and people of different ancestrics is not testod. Future studies of more diverse samples are warranted. Since the brain morphological features derived are not of immediate medical use, the risk of perpetuating existing healthcare disparities is indirect.
11

---

Published in Transactions on Machine Learning Research (05/2024)
References
Abdel Abdellaoui, Loic Yengo, Karin JH Verweij, and Peter M Visscher. 15 years of gwas discovery: Realizing the promise. The American Journal of Human Genelics, 2023.
Stephanie L Adams, Laurent Benayoun, Kathy Tilton, Olivia R Chavez, Jayandra J Himali, Jan Krzysztof Blusztajn, Sudha Seshadri, and Ivana Delalle. Methionine sulfoxide reductase-b3 (msrb3) protein associates with synaptic vesicles and its expression changes in the hippocampi of alzheimer's disease patients. Journal of Alzheimer's Disease, 60(1):43-566 2017.
Nil Aygün, Angela L Elwell, Dan Liang, Michael Lafferly, Kerry E Cheek, Kenan P Courtney, Jessica Mory, Ellie Hadden-Ford, Oleh Krupa, Luis de la Torre Ubieta, et al. train-trait-associatee variants impact cell-type-specific gene regulation during neurogenesis. The American Journal of Human Genetics, 108(9): 1647-1668, 2021.
Philip Bachman, R Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views. Advances in neural information processing systems, 32, 2019.
Brendan Bulik-Sullivan, Hilary K Finucane, Verneri Anttila, Alexander Gusev, Felix R. Day, Po-Ru Loh, Laramie Duncan, John RB Perry, Nick Patterson, Elise B Robinson, al. An atlas of genetic correlations across human discases and traits. Neture genetics, 47(11):1236-122) 2015.
Judit Cabana--ominguee, Bårbara. Torrico, Andreas Reif, Noèlia BernändezCasttillo and Bru Cormand. Comprehensive exploration of the genetic contribution of the dopaminergic and serotonergic pathways to psychiatric disorders. Translational Psychiatry, 12(1):11, 2022.
Chun-Hao Chang, George Alexandru Adam, and Anna Goldenberg. Towards robust classification model by counterfactual and invariant data generation. In Proceedings of the IEEE/CVP Conference on Computer Vision and Pottern Recognition, pages 15212 15221, 2021.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Proceedings of the International Conference on Machine Learning, 2020.
Jacob Devlin, Ming-Wae Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North Americon Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 4171-4186, 2019.
Lloyd T Elliott, Kevin Sharp, Fidel Alfaro-Almage Sinan Shi, Karla L Miller, Gwenaëlle Douaud, Jonathan Marchini, and Stephen M Smith. Genome-wide association studios of brain imaging phenotypes in uk biobank. Nature, 562(7726):210-216 2018.
Adam Foster, Rattana Pukdee, and Tom Rainforth. Improving transformation invariance in contrastive representation learning. In International Conference on Learning Representations, 2020.
Francesca Gori, Ulf Lerner, Claes Ohlsson, and Roland Baron. A new wnt on the bone: Wnt16, cortical bone thickness, porosity and fractures. BoneKEy reports, 4:669, 2015.
R Devon Hjelm and Philip Bachman. Representation learning with video deep infomax. ar preprint arXiv:2007.1392 132778 2020.
R Devon Hjelm, Alex Fedorov, Samuel Tavoie-Marchildon) Karan Grewal, Phil Bachman, Adam Trischler, and Yoshma Bengio. Learning deep representations by mutual information estimation and maximization. International Conference on Learning Representations, 2019.
Weiran Huang, Mingyang Yi, and Xuyang Zhao. Towards the generalization of contrastive self-suud learning. Xio preprint arXiv:2111 211.00044, 2021.
12

---

Published in Transactions on Machine Learning Research (05/2024)
Mark Jenkinson, Peter Bannister, Michael Brady, and Stephen Smith. Improved optimization for the robust and accurate linear registration and motion correction of brain images. Neuroimage, 17(2):825-85 881, 2002.
Li Jing, Pascal Vincent, Yann LeCun, and Yuandong Tian. Understanding dimensional collapse in contrastive self-supervised learning. In International Conference on Learning Represenlations, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiu:1412.6900 2014.
Matthias Kirchler, Stefan Konigorski, Matthias Norden, Christian Meltendorf, Marius Kloft, Claudia Schur- mann, and Christoph Lippert. transfergwas: Gwas of images using deep transfer learning. Bioinformatics, (5(14):3621-3628 2022.
Weixin Liang, Yuhui Zhang, Yongchan Kwon, Serena Yeung, and James Zou. Mind the gap: Understanding the modality gap in multi-modal contrastive reprosentation learning. In Advances in Neural Information Processing Systems, 2022.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hicrarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Compuler Vision (ICCV), 2021.
Po-Ru Loh, George Tucker, Brendan K Bulik-Sullvaan, Bjarni J Vilhjalmsson, Hilary K Finucane, Rany M Salem, Daniel Chasman, Paul M Ridker, Benjamin M Neale, Bonnie Berger, et al. Efficient bayesian mixed-model analysis increases association power in large cohorts. Nature genetics, 47(3):284 290, 2015.
Nya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. ar Xiw preprinl arXiv: 608. 039883 2016.
Langlang Ma, Chunyan Qing, Minyan Zhang, Chaoying Zou, Guangtang Pan, and Yaou Shen. Gwas with a pca uncovers candidate genes for accumulations microelements in maize scedlings. Physiologia Plantarum, 172(4):2170 2180, 2021.
Jacqueline MacArthur, Emily Bowler, Maria Cerezo, Laurent Gil, Peggy Hall, Emma Hastings, Heather Junkins, Aoife McMahon, Annalisa Milano, Joannella Morales, et al. The new nhgri-ebi catalog of published genome-wide association studies (gwas catalog). Nucleic acids research, 45(D1) DD966-D9DDD 2017.
Carolina Medina-Gomez, Benjamin H Mullin, Alessandra Chesi, Vid Prijatelj, John P Kemp, Chen Shochat- Carvalho, Katerina Trajanoska, Carol Wang, Raimo Joro, Tavia E Evans, et al. Genome wide association metanalysis of skull bone mineral density identifies loci relevant for osteoporosis and craniosynostosis. MedRatv, pages 2021-11, 2021.
Adam C Naj, Christiane Reitz, Farid Rajabli, Gyungah R. Jun, Penelope Benchek, Giuseppe Tosto, Jin Sha, Congcong Zhu, Nicholas A Kushch, Wan-Ping Lee, et al. Multi-ancostry genome-wide association analysis of late-onset alzheimer's discase (load) in 60,941 individuals identifies novel cross-ancestry association in Irre4c. Alzheimer's & Dementio, 18:e065822, 2022.
Sebastian Nowozin, Botond Cseke, and Ryota Tomioka. f-GAN: Training generative neural samplers using variational divergence minimization. In Advances in Neural Information Processing Systems, pages 271-279, 2016.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilankurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An imperative style, high-performanc deep learning library. In Advances in Neural Information Processing Systems, pages 8024-8035. 2019.
Khush Patel, Ziqian Xie, Hao Yuan, Sheikh Muhammad Saiful Islam, Wanheng Zhang, Assaf Gottlieb, Han Chen, Luca Giancardo, Alexander Knaack, Evan Fletcher, et al. New phenotype discovery method by unsupervised deep representation learning empowers genetic association studies of brain imaging. medRziv, pages 2022-182 12, 2022.
13

---

Published in Transactions on Machine Learning Research (05/2024)
Godfrey D Pearlson, Jingyu Liu, and Vince D Calhoun. An introductory review of parallel independent component analysis (p-ica) and guide to applying p-ica to genetic data and imaging phenotypes to identify disease-associated biological pathways and systems in common complex disorders. Frontiers in genetics, 6: 276, 2015.
Shaun Purcell, Benjamin Neale, Kathe Todd-Brown, Lori Thomas, Manuel AR Ferreira, David Bender, Julian Maller, Pamela Sklar, Paul IW De Bakker, Mark Daly, et al. Plink: a tool set for whole-genome association and population-based linkage analyses. The American journal of haman genetics, (1) ((33:559-55)) 575, 2007.
Julius Rönkkö, Svetlana Molchanova, Anya Revah-Politi, Elaine M Pereira, Mari Auranen, Jussi Toppila, Jouni Kvist, Anastasia Ludwig, Julika Noumann, Goert Bultynck, et al. Dominant mutations in itpr3 cause mharcot-marie-oooth disease. Annals of Clinical and Translational Newrology, 7(10):1962 1972 2020.
K. T. Schütt, P.-J. Kindermans, H. E. Sauceda, S. Chmiela, A. Tkatchenko, and K.-R. Müller. Schnet: A continuous-file convolutional neural network for modeling quantum interactions. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pages 992 992-1002, 2017.
Stephen M Smith, Gwenaëlle Douaud, Winfield Chen, Taylor Hanayik, Fidel Allao-Almagee Kevin Sharp, and Lloyd T Elliott. An expanded set of genome-wide association studies of brain imaging phenotypes in uk biobank. Nature neuroscience, 24(5):737 745, 2021.
Hannes Stärk, Dominique Beaini, Gabriele Corso, Prudencio Tossou, Christian Dallago, Stephan Günnemann, and Pictro Liò. 3d infomax improves gnns for molecular property prodiction. In International Conference on Machine Learning, pages 20479 20502. PMLR, 2022.
Cathie Sudlow, John Gallacher, Naomi Allen, Valerie Beral, Paul Burton, John Danesh, Paul Downey, Paul Elliott, Jane Green, Martin Landray, et al. Uk biobank: an open access resource for identifying the causes wide range of complex diseases of middle and old age. PLoS medicine, :e 2015.
Benjamin B Sun, Stephanie J Loomis, Fabrizio Pizzagalli, Natalia Shatokhina, Jodie N Painter, Christopher N Foley, Megan Jensen, Donald G McLaren, Sai Spandana Chintapalli, et al. Genetic map of regional sulcal morphology in the human brain from uk biobank data. Nature Communications, 13(1):6071, 2022.
Aiham Taleb, Matthias Kirchler, Remo Monti, and Christoph Lippert. ContIG: Self-supervised multimodal contrastive learning for medical imaging with genetics. In Proceedings of the IEEE/CVF Conference on Compuler Vision and Pallern Recognition, pages 20908 20921, 2022.
Yonglong Tian, Chen Sun, Ben Poole, Dilip Krishnan, Cordelia Schmid, and Phillip Isola. What makes for good views for contrastive learning? orXiv preprint arXiv:2005. 1055..1022 10243, 2020.
Arno van Hilten, Steven A Kushner, Manfred Kayser, M Arfan Ikram, Hicab HH Adams, Caroline CW Klaver, Wiro J Niessen, and Gennady V Roshchupkin. Gennet framework: interpretable deep learning for predicting phenotypes from genetic data. Communications Biology, 4(1):1-9, 2021.
V.N. Vapnik. An overview of statistical learning theory. IEEE Transactions on Neural Networks, 10(5): 988 999, 1999.
Petar Veličkovié, William Fedus, William L. Hamilton, Pietro Liò, Yoshua Bengio, and Devon Hjelm. Deep graph infomax. In International Conference on Learning Representations, 2019.
Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. In International Conference on Machine Learning, pages 9929 9939. PMLR, 2020.
Xiao-Lan Wang and Lianjian Li. Cell type-specific potential pathogenic genes and functional pathways in alzheimer's disease. BMC neurology, 21(1):1-18, 2021.
Zhengyang Wang, Yaochen Xie, and Shuiwang Ji. Global voxel transformer networks for augmented microscopy. Nature Machine Intelligence, 3:161-171, 2021. doi: 10.
14

---

Published in Transactions on Machine Learning Research (05/2024)
Junhao Wen, Ilya Nasrallah, Ahmed Abdulkadir, Theodore Satterthwaite, Guray Erus, Timothy Robert- Fitzgerald, Ashish Singh, Aristeidis Sotiras, Aleix Boquet-Pujadas, Zhijian Yang, et al. Mega-analysis of brain structural covariance, genetics, and clinical phenotypes. Research Square, 2022.
Tete Xiao, Xiaolong Wang, Alexei A Efros, and Trevor Darrell. What should not be contrastive in contrastive learning. In International Conference Learning Representations, 2020.
Ziqian Xie, Tao Zhang, Sangbae Kim, Jiaxiong Lu, Wanheng Zhang, Cheng- Hui Lin, Man-Ru Wu, Alexander Davis, Roomasa Channa, Luca Giancarlo, et al. igwas: imago-based genomo-wide association of self- supervised deep phenotyping of human medical images. medRaiv, 2022.
Kenji Yano, Yoichi Morinaka, Fanmiao Wang, Peng Huang, Sayaka Takehara, Takaaki Hirai, Aya Ito, Eriko Koketsu, Mayuko Kawamura, Kunihiko Kotake, et al. Gwas with principal component analysis identifies a gene comprchensively controlling rice architccture. Proceedings of the National Academy of Sciences, 116 (42):21262 21267, 2019.
Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self-supervised learning via redundancy reduction. International Conference on Machine Learning, pages 22310-12320 PMLR, 2021.
Michael Zhang, Nimit Sharad Sohoni, Hongyang R Zhang, Chelsea Finn, and Christopher Ré. Correct-n- contrast: A contrastive approach for improving robustness to spurious correlations. In NewrIPS 2021 Workshop on Distribution Shifts: Connecting Methods and Applications, 2021.
Bingxin Zhao, Tengfei Li, Yue Yang, Xifeng Wang, Tianyou Luo, Yue Shan, Ziliang Zhu, Di Xiong, Mads E Hauberg, Jaroslav Bendl, et al. Common genetic variation influencing human white matter microstructure. Science, reb(6333),, 2021.
Mohammadreza Zolfaghari, Yi Zhu, Peter Gehler, and Thomas Brox. Crossclr: Cross-modal contrastive learning for multi-modal video representations. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 1450-1459, October 2021.
15

---

Published in Transactions on Machine Learning Research (05/2024)
A Data Processing and Split
All MRIs were linearly registered (affine registration with 12 DOF) to standard MN1152 space using the UKBiobanl-provided transformation matrix with FSL FLIRT (Jenkinson et al. 2002) and all the outputs are of shape 182x218x182. A large portion of the UKBiobank population are white British. In order to maximize the power of genetic discovery and avoid the complication of population stratification, the genetic association study was only done on the white British (UKBiobank data field 21000 and 22006) cohort. So we selected 6,130 images from subjects of mixed ethnicitics (all non white British samples plus a small number of random white British samples) not overlapped with the samples for the genetic discovery to do the training and validation, among which 4,597 was randomly selected for training and 1,533 for validation. We used two quality metries "inverted contrast-to-noiss ratio" (UKBiobank data field 25735) and "Discrepancy between T2 FLAIR brain image and T1 brain image" (UKBiobank data field 25736) to ensure the quality of the training data.
Details about the evaluation metrics are provided below.
Number of Loci We perform genome-wide scans over 658,720 directly genotyped SNPs and on 28,489 white British participants unseen during training. We use BOLT-LMM (Version 2.3.4) (Loh et al. 2015 for running GWAS. Age, gender, and the first 10 ancestral principal components are used as covariates. We use Bonferroni corrected p-value threshold of 5e-9 and a minor allele frequency threshold of 1% to get the significant SNPs and filter out the rare variants. We then cluster the significant SNPs into loci using a 250 kb window, which is approximately 0.25 cM (Elliott et al. 2018). The number of loci indicates the amount of genetic contribution the learned features.
Heritability measures the proportion of variation of the feature explained by the genetic factors. It provides insight into the genetic basis of feature. A higher heritability indicates that the representation is better associated with the genetic data. The heritability is computed using LDSC v1.0.1 (Bulik-Sullivan et al. 2015
Mutual Information We estimate the mutual information between MRI representations and genetic data on the test. set to explicitly demonstrate that the proposed objective adds to the generalizability of captured associations to unseen pairs. We train individual JSE--ased mutual information estimators with the same architecture described in Appendix B for different methods. We train the MI estimator until the contrastive loss converges and take the opposite of the converged value as the MI estimation for each model.
B Implementation Details
MRI encoders The MRI encoder is constructed 3D convolutional network consisting of three residual blocks (Wang et al. 2021) connected by two downsampling operators with stride convolutions. The numbers of channel maps are 32, 64, and 128, respectively for the three blocks. The final representations are 128- dimensional computed by a dense layer upon flattened feature maps. When computing the reconstruction loss, we include additional 128-dimensional vectors computed from a multi-head attentive readout from feature maps, and the reconstruction is performed on the 256-dimensional representation after concatenation. However, dimensions from atlention are not used in GWAS compulation. This is to further prevent. the encoder from learning too detailed patterns, possibly noise, that are non-generalizable to the test set.
blocks connected by two down-sampling operators with a down- wown-sampling rate of 10. The positional encoding for SNP physical positions is 128. The embedding dimensions are 32, 64, and 128 for the three blocks, respectively. The window size to perform attention is 10 and the number heads is 4 for all self-attention operators. The downsampling operator computes the attention with learnable queries within each window, where the window size equal to the downsampling size. The global pooling operators compute attention with learnable queries among all positions at multiple scales and resolutions.
Genetic encoders Our genetic encoder consists of three 1D
1 Applied Biosystems UK BILEVE Axiom Array, UKRiobank data field 22438
16

---

Published in Transactions on Machine Learning Research (05/2024)
Multi-modal InfoNCE
SimCLR
GIM
Figure 5: Visualization of learned representations with t-SNE. Forming chusters not desired in the GWAS setting. Instead, higher uniformity better utilizes the space capacity and could lead to better GWAS discovery.
Training The models are implemented with PyTorch (Paszke et al. 2019 and are trained on single Nvidia A100 GPU. The training is performed with the Adam optimizer (Kingma and Ba 2014), cosine annealing scheduler (Loshchilov and Hutter 2016) with a starting learning rate of 0.001 and the mini-batch size of 12. We simply set À in the objective to I and do not exhaustively tune it. During training, we randomly crop the 3D MRI into smaller patches of size [160, 160, 160). We first train the models with the genetic encoder frozen for 200 epochs, then include the augmentation on genetic data and continue training for an additional 100 epochs, and finally co-train both MRI and genetic encoders and projection heads for 50 epochs with augmentation genetic data. To perform augmentation, the genetic data has probability of 0.2 to be randomly cropped and probability of 0.8 to be evenly down-sampled into length of 65,000.
Baseline approaches For Gen Prediction, we apply a linear layer to the output representation as prediction head h IR9 > |G] to predict the class of each genotype in the genetic data and optimize the following loss.
LGenProd = Cross-Entropy (ho (No(oo)))))))) (Jo()))
For bascline trans-modal contrastive methods, follow the architecture, training loss, and training settings in Taleb et al. (2022). For the MRI augmentations, perform the random flipping and rotation the x-2 plane, along with the random 3D patching. However, found the flipping and rotation do not help the GWAS performance in the 3D MRI case. For correlated InfoNCE, we compute the covariance matrix of Iearned MRI representations and minimize its difference with the identity matrix,
Ldecor || -
where & is the normalized MRI representations the mini-bated Since the mini-batcl size small due to memory constrain, the covariance estimation can be less accurate, still leading reduced performance.
C Distribution of learned representations
We visualize the distribution of representations learned by trans-modal InfoNCE, SimCLR, and GIM, respectively, with t-SNE: in Figure 5 Compared to baseline approaches, GTM learns representations that more uniformly distributed in the space. According to the discussion the difference between learning goals, the goal of our representation learning is not to form clusters for downstream classification purposes but to uniformly encode as much information about the genetic ic data as possible (Wang and Tsola 2020). Under this setting, clusters of representations are not desired and may harm the GWAS performance due to reduced capacity for other characteristics.
Gene mapping and query for previous known associations
We mapped significant SNPs to genes using Plink v1.9 (Purcell al. 2007 and presented the genes that associated with the most significant new SNP of each model in the ablation study in Table CENPW
17

---

Published in Transactions on Machine Learning Research (05/2024)
is known to associated with neurogenesis [Aygin al. 2021 and cortical morphology Sun et al, 2022 WNT16 with skull and brain shape (Gori et al.
2015 Medina-Gomez et al. 2021 ITPR3 with neuropathy et al., 2022 and MSRB3 with 2017p. For cach unique locus in the best-performing model, we did a range query Alzheimer's (Adams et al.
(Rönkkö et al. 2020) and many psychiatric disorders Slabana-Domingun
using the GWAS Catalog MacArthur et al. 2017] API for the previously identified brain-related associations We also queried each locus in the result of the Big40 study (Smith et al. 2021; Elliott al.] 2018), which uses thousands of conventional image-derived phenotypes to do GWAS and found locus not presented in the Big40 study in Chromosome 2, base pair 218466221 to 218604356 (in hg19 coordinate). This locus is mapped to DTRC9, which has been shown to be associated with Alzheimer's disease (Naj et al., 2022) Wang and Li 2021). This showcases the potential of our method in capturing features missed by the traditional expert-defined pipelines.
and the result shown in Figure (
For the heritability, we see that the score generally increases when adding each component to the model. The only exception is the random cropping without co-training of MRI and genetic encoder. This is due that without co-training, the genetic encoder has limited capability of adapting to the cropped genetie data from different locations. However, the cropping augmentation still help the learning of better MRI encoder that identifies more loci, and co-training leads to further improvement.
Number of Unique Loci mapped to Brain Related Phenotypes
brain measurement cortical thickness brain volume
cortical surface area
neuroimaging measurement educational attainment PHF -tu
self reported educational attainment
cognitive decline
Alzhelmer disease
neurafibrillary tangles
autism spectrum disorder
brain connectivity
white matter microstructure
total cortical area
mathematica ability
schizon schizophrenia
Brain Phenotypes
lateral ventricle volume change
Parkinson disease
Grey matter detsy
late-onset Alzheimers disease
white matter hyperintensity
isthmus cingulate coftex
brain stem volume
prefrontal cortex growth
microglial activation
cerebral amylold anglopathy
lifestyle l measurement
cerebral cortex growth
thalamus volume
brain amygdala volume
episodic memory
brain
cerebellar
neuroticism
oppositional defiant disorder
antisocial behavlour
borderline personality disorder
functional brain measurement
nucleus accumbens volume
depressive sym ptom
right ventricular systolic volume
Number Unique Loci mapped
Figure 6: The number of unique loci found by the best model associated with subset of brain-related traits in the GWAS Catalog.
E Manhattan Plot of GIM results
The Manhattan Plot for loci discovered by GIM is shown in Figure
18

---

Published in Transactions on Machine Learning Research (05/2024)
30
25
20
15
10
2
3
4
5
8
9
10
11
12
13
14
15
16
Figure 7:
Aggregated Manhattan plot of the 10 PCs from the best model, the grey line represents the Bonferroni corrected p-value threshold of 5e-9.
19


Reviews:
[{"review_id": "5sGpMLYY5K", "reviewer": "Reviewer_iD54", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n•\tThe paper is in general well written.\n\n•\tThe paper appear to improve upon current state-of-the-art for relating high-dimensional traits (structural MRI) to GWAS.  \n\n•\tThe dataset is used is large (UK Biobank).\n\n•\tThe approach appears sound.\n\nWeaknesses\n\n•\tComplete lack of error bars in the assessments making it very hard to judge the significance of the approach when compared to baselines (I therefore rate claim and evidence insufficient).\n\n•\tThe derivation although strongly grounded in mutual information ends up being a rather simple standard loss function trading of auto-encoder reconstruction and discriminator losses with rather limited novelty here.\n\n•\tTransformers are well known to provide improved context learning and the inclusion of such architecture is also of limited novelty.\n\n•\tHow to tune the regularization \\lambda is very unclear and it is not discussed how it influences results nor how the tuning of \\lambda is practically done in the experimentation.\n\nAdditionally, I find the relation between entropy and least squares reconstruction regularization to not be well elaborated and unclear in the presentation although this is part of the central contributions apart from the proposed use of transformers. I.e., what is H(f_theta(Y)|Y) and what does the crossing arrow and zero here mean (this term goes to zero?).  I assume I(f_theta(Y),Y) is the mutual information but how it that estimated here (is it using the  JSE og NCE estimators?) how exactly this relates to the reconstruction regularization should further be sharpened as the text below equation 6 is unclear and could be better elaborated upon.\n\nMinor comment:\nIs there a mistake in table 2.1 in which Change of +5 should be +3?\nTable 3 and 4 does the autoencoder correspond to setting \\lambda->infty?\n\nIn general this is a well written paper, however, I find the technical contribution rather limited and the results unconvincing due to the lack of error bars. I.e., how do results change according to random initialization (and potentially changes to training and test splits)? Such uncertainty assessment is necessary to adequately judge the current results.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a mutual information inspired framework for the mapping of 3D imaging data to GWAS, by use of a loss function that provides a proxy for mutual information optimization being composed of two losses, i.e. a discriminator loss as well as an autoencoder (AE) reconstruction loss with a regularization term \\lambda trading of its impact. The paper further incorporates a transformer based framework for the efficient learning of relevant contexts. \n\nSpecifically, the paper considers the Jensen-Shannon Estimator (JSE) to map high-dimensional traits to optimally correspond to genetic profiles of individuals which is related to the noise-contrastive estimation (NCE) optimization in which a discriminator should optimally differentiate corresponding pairs from non-corresponding pairs. However, such loss is found to potentially degenerate focusing on undesirable non-generalizable properties (i.e., referred to as dimensional collapse) which is proposed resolved using the inclusion of the AE loss.\n\nThe approach is compared to existing frameworks considering a large UK Biobank datasets of structural MRI and GWAS finding that the approach provides favorable loci identification and downstream task predictions of stroke, cancer and gender in general outperforming existing approaches.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The work can be substantially improved in particular by providing:\n\nImproved description of the relation to the regularization by the AE and mutual information – present description below equation 5 are found unclear and insufficient.\n\nInclusion of error bars on all results for statistical assessment of significance of the results.\n\nCareful description of how to tune \\lambda and how the value of \\lambda was assessed in the experimentation.\n\nSystematic contrasting the procedure to ContIG. ContIG also uses contrastive loss – is this identical to the considered InfoNCE – to me these appears to be the same please clarify. It would be good to further evaluate how the ContIG contrastive loss compare to the present JSE loss in the experimentation and if JSE change performance compared to the contrastive loss considered in the context of ContIG. The paper should thus compare training with this loss in the results as opposed to the JSE loss to highlight the benefit if any of JSE.\n\nInclusion of the model ablations in Table 2 also in the downstream task experimentation provided in Table 4 regarding classification of stroke, cancer and gender.", "broader_impact_concerns": "Broader impact is only briefly discussed including appendix F. It would be good here to expand upon this in terms of how the use of the developed modeling procedures can impact understanding of GWAS and its relation to traits as well as potentially lead to misinterpretations if used incorrectly etc.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "FGO1ipHPp1", "reviewer": "Reviewer_V9SE", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\nMotivation and Background: The paper thoroughly discusses the unique challenges posed by GWAS and identifies the limitations of contrastive losses in GWAS representation learning. It reviews related work and elaborates on the distinctions between this work and previous research, which is useful for advancing understanding.\nMethodological novelty: The paper proposes a novel and effective trans-modal learning framework for addressing the challenges of GWAS in high-dimensional medical imaging data.\n\nWeaknesses:\nThe clarity of presentation can be improved. Some concepts could be explained in a clearer way. \nThe use of English needs improvement and should be proofread by a native English speaker. In particular, there are a lot of language and mathematical typos and mistakes.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper discusses the unique challenges of Genome-Wide Association Studies (GWAS) in the context of high-dimensional imaging phenotype data and proposes a novel framework, Genetic InfoMax (GIM), which extracts low-dimensional representations from high-dimensional medical imaging data. These representations capture as much information about genetic variations as possible, thereby improving the performance of GWAS. GIM includes a regularized MI estimator, approximated by the Jensen-Shannon Estimator (JSE), and a novel genetics-informed transformer. Experiments on human brain 3D MRI data demonstrate the effectiveness of GIM and a significantly improved performance on GWAS.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "This paper contains some language and mathematical typos and mistakes. Below are some examples:\nSection 1：\n\"linear or linear mixed model\" should be \"linear or linear mixed models\".\n\"and thus limited\" should be \"thus limiting\".\n\nSection 2:\nThe sub-sentence \"the GWAS process involves statistical tests on the sample pairs\" should be started as a new sentence; otherwise, the whole sentence is not grammatically correct (comma splice or a run-on sentence).\n\"each gi \\in {0, 1, 2}, representing\" should be \"and each gi \\in {0, 1, 2} represents\".\n\"is formulated as to learn\" should be \"is formulated as learning\".\n\"two neighbor SNPs gi and gi+1 are not necessarily to be consecutive\" should be \"two neighboring SNPs, gi and gi+1, are not necessarily consecutive\".\n\nSection 3:\n\"two variables\" should be \"two sets of variables\" or \"two multivariates\".\n\"also presents\" should be \"also arises\" or \"is also present\".\n\"is formulated as to learn\" should be \"is formulated as learning\".\n\"both the information of fθ(Y ) and the mutual information I(fθ(Y ),G) is limited\" should be \"... are limited\".\nShould \"sufficient means of augmentations\" be \"insufficient means of augmentations\"?\n\"in the right-hand side (RHS) of Eq. (2)\" should be \"on the right-hand side (RHS) of Eq. (2)\".\n\nSection 4：\nLine 1 below Eq. (7): \"denote\" should be \"denotes\".\nLine 1 below Eq. (9): \"ranging from 0 and a preset threshold\" should be \"ranging from 0 to a preset threshold\" or \"ranging between 0 and a preset threshold\".\n\nSection 6.1：\n\"transmodal contrastive approaches that involves genetic data\" should be \"transmodal contrastive approaches that involve genetic data\".\n\"The baselines include existing or straightforward training schemes\" should be followed by a semicolon (:).\n\nTable 2, column \"Change\": \"+5\" should be \"+3\".", "broader_impact_concerns": "The use of genetic data and imaging data in medical research raises concerns related to privacy. But in the paper there is no discussion on how to address this.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ATRnL8Dt3C", "reviewer": "Reviewer_3bDZ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\nThe paper is well written.\nIt clearly states the considered problem, motivates and explains the chosen approach, and describes the experiments and results.\nMany baselines are included in the experiments and the proposed method outperforms them on the considered metrics.\nAn additional experiment evaluates the usefulness of the learned representations for relating imaging features to other phenotypes -- stroke, cancer and sex.\n\n## Weaknesses\n\nI think the motivation for tackling this problem should be stated in more detail.\nThe baselines only contain deep learning approaches.\nI believe including simple baselines that rely on \"region-based volumetric or surface features\" and linear dimensionality reduction would be helpful to situate the reported metrics. \nToo few details are included about the procedure for the \"downstream classification performance\" evaluation.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors present a method for extracting features from 3D MRI scans that have strong associations with genetic data (SNPs).\nThe method relies on maximizing a combination of (i) the Mutual Information between the image representations and the genetic data as estimated by the Jensen-Shannon Estimator (JSE) and (ii) a term measuring how well the original image can be reconstructed from its lower-dimensional representation (low mean squared error).\nThe JSE consists in evaluating how well a model can discriminate pairs of (image features, genetic features)  that belong vs do not belong to the same individual.\n\nThe authors conduct experiments on the UKBiobank.\nImage representations are learned using 4.5K individuals for training and 1.5K individuals for validation.\nThen, representations are extracted for 28K unseen individuals and used to perform GWAS, looking for associations between genetic data and loadings on the 10 first principal components of the image representations.\nThey report the number of loci detected by the GWAS, the heritability of the PC loadings (proportion of variation explained by genetic features), and the estimated MI (with the same JSE approach) between genetic and imaging features.\n\nThe authors find that for those metrics (number of loci, estimated MI, heritability), their method outperforms various baselines that also rely on deep learning.\nMoreover, they observe that several of the genes identified by the GWAS on their image features are known in the literature to have associations with phenotypes related to brain structure or function.\n\nAnother experiment shows that the extracted features perform on par with several alternative representations for classifying stroke, cancer and sex on the UKBioBank.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "## Critical\n\nThe authors mention that \"typical GWAS studies are focused on well-established phenotypes, typically the risks of diseases\".\nIt is easy to imagine why identifying associations between genetic variants and risks of diseases can be useful, eg to monitor more closely individuals at a higher risk.\nHowever, for a reader (such as myself) who is not familiar with GWAS or genomics in general, it is less obvious why it is important to identify associations between genetic variants and abstract, high-dimensional features extracted from brain images with deep learning.\nExtracting features as representations for downstream prediction tasks is useful but the authors point out that here it is not the primary goal.\nHow are the resulting features and associations to be interpreted and used, what are the potential applications or insights?\nI think providing more context and motivation in the introduction would be very useful.\n\nThe authors mention that non-learning features have limited expressiveness and that dimensionality reduction such as PCA, ICA, NMF are less capable of identifying complex traits.\nTo support these observations, could the authors add examples of these approaches to the baselines?\n\nRelating those features to other phenotypes relevant to psychology or clinical practice is interesting, and for that the downstream classification task in 6.2 is very interesting.\nHowever the paper does not provide enough information about the training and evaluation procedure to interpret the results in Table 4.\nCould the authors add more details to that section?\n\n## Minor\n\nFor the final GWAS, a PCA is applied to the learned features and the first 10 are kept.\nDoes this limit the expressiveness of those features, could the authors explain why it is not a limitation and how the number 10 was chosen.\n\nOne of the metrics, the MI estimated with JSE, is exactly what the method is designed to maximize.\nIt is therefore somewhat expected that it scores higher on that metric.\nI do not think this is a limitation especially since other metrics are included but I would suggest including a short comment about this.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Please do a thorough pass on the writing.
Decision: Accept (poster)