SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Duy M. H. Nguyen\*1,2,3, Hoang Nguyen\*3, Nghiem T. Diep\*3, Tan N. Pham\*4, Tri Cao\*3, Binh T. Nguyen\*4, Paul Swoboda\*5, Nhat Ho\*6, Shadi Albarqouni\*7,8, Pengtao Xie\*9,10, Daniel Sonntag\*3,11, Mathias Niepert\*+1,2
<sup>1</sup>University of Stuttgart, <sup>2</sup>IMPRS for Intelligent Systems
$^{3}$ German Research Center for Artificial Intelligence,  $^{4}$ University of Science - VNUHCM,
<sup>5</sup>Max Planck Institute for Informatics, <sup>6</sup>University of Texas at Austin <sup>7</sup>Helmholtz Munich,
<sup>8</sup>University Hospital Bonn, <sup>9</sup>UC San Diego, <sup>10</sup>MBZUAI, <sup>11</sup>Oldenburg University.
Obtaining large pre-trained models that can be fine-tuned to new tasks with limited annotated samples has remained an open challenge for medical imaging data. While pre-trained deep networks on ImageNet and vision-language foundation models trained on web-scale data are prevailing approaches, their effectiveness on medical tasks is limited due to the significant domain shift between natural and medical images. To bridge this gap, we introduce LVM-Med, the first family of deep networks trained on large-scale medical datasets. We have collected approximately 1.3 million in medical images from 55 publicly available datasets, covering a large number of organs and modalities such as CT, MRI, X-ray, and Ultrasound. We benchmark several state-of-the-art self-supervised algorithms on this dataset and propose a novel self-supervised contrastive learning algorithm using a graph matching formulation. The proposed approach makes three contributions: (i) it integrates prior pair-wise image similarity metrics based on local and global information; (ii) it captures the structural constraints of feature embeddings through a loss function constructed via a combinatorial graph-matching objective; and (iii) it can be trained efficiently end-to-end using modern gradient-estimation techniques for black-box solvers. We thoroughly evaluate the proposed LVM-Med on 15 downstream medical tasks ranging from segmentation and classification to object detection, and both for the in and out-of-distribution settings. LVM-Med empirically outperforms a number of state-of-the-art supervised, self-supervised, and foundation models. For challenging tasks such as Brain Tumor Classification or Diabetic Retinopathy Grading, LVM-Med improves previous vision-language models trained on 1 billion masks by  $6 - 7\%$  while using only a ResNet-50. We release pre-trained models at this link https://github.com/duyhominhnguyen/LVM-Med.
Constructing large-scale annotated medical image datasets for training deep networks is challenging due to data acquisition complexities, high annotation costs, and privacy concerns [1, 2]. Vision-language pretraining has emerged as a promising approach for developing foundational models that support various AI tasks. Methods such as CLIP [3], Align [4], and Flava [5] propose a unified model trained on large-scale image-text data, showing exceptional capabilities and performance across
*Corresponding authors, †Co-Senior authors.
various tasks. However, their effectiveness in the medical domain still remains unclear. A recent work SAM [6] trains large vision models on over one billion annotated masks from 11M natural images, enabling interactive segmentation. Nevertheless, SAM's zero-shot learning performance is moderate on other datasets [7, 8], highlighting the need for fine-tuning to achieve satisfactory results [9].
Figure 1: (left) Overview of the body organs and modalities in our collected dataset; (right) LVM-Med unifies and extends contrastive and instance-based self-supervised learning approaches by specifying graph's properties.
To facilitate the development of foundation models in the medical domain, we make two major contributions. First, we have curated a vast collection of 55 publicly available datasets, resulting in approximately 1.3 million medical images covering various body organs and modalities such as CT, MRI, X-ray, ultrasound, and dermoscopy, to name a few. Second, we propose LVM-Med, a novel class of contrastive learning methods, utilizes pre-trained ResNet-50 and a ViT network SAM[10]. We evaluate various instances of LVM-Med relative to popular supervised architectures and vision-language models across 15 medical tasks. To our best knowledge, this is the first time such a large-scale medical dataset has been constructed and used to investigate the capabilities of SSL algorithms.
LVM-Med incorporates a second-order graph-matching formulation, which subsumes and extends a large class of contrastive SSL methods. Given a batch of images, two random transformations are applied to each image, and the resulting transformed images are then fed to an image encoder. The embedding vectors obtained from images in a batch are used to construct two graphs where vertices represent pairs of transformed images generated from the same original one. Through solving a graph-matching problem [11, 12], we learn feature representation such that their encoding serves as suitable priors for a global solution of the graph-matching objective. This approach is distinct from prior contrastive learning methods that focus on merely optimizing pair-wise distances [13, 14] between transformed images or learning contrastive distances with positive and negative samples [15-19]. It is worthwhile noting that previous contrastive learning methods are special instances of our general framework (Figure (1), right).
LVM-Med has several advantages over existing approaches. First, it integrates advanced pair-wise image similarity taken from prior SSL methods into vertex affinities, resulting in both global and local information that can be efficiently fused. Second, it uncovers underlying structures of feature embeddings by utilizing edge constraints, enhancing robustness in the presence of similar entities in medical datasets. Third, though combinatorial problems are typically non-differentiable, LVM-Med can efficiently calculate gradients through the discrete combinatorial loss function using modern implicit maximum likelihood estimation techniques. Consequently, LVM-Med can scale successfully on large-scale data. In a wide range of 15 medical experiments, LVM-Med sets a new state-of-the-art in fully fine-tuning or prompt-based segmentation, linear and fully fine-tuning image classification, and domain generalization, outperforming several vision-language models trained on a hundred million image-text instances.
We summarize major contributions in this work, including:
(i) We present a collection of large-scale medical datasets, serving as a resource for exploring and evaluating self-supervised algorithms.
(ii) We propose LVM-Med, a novel SSL approach based on second-order graph matching. The proposed method is flexible in terms of integrating advanced pair-wise image distance and being able to capture structural feature embedding through the effective utilization of second-order constraints within a global optimization framework.
(iii) On both ResNet-50 and ViT architectures, LVM-Med consistently outperforms multiple existing self-supervised learning techniques and foundation models across a wide range of downstream tasks.
2.1 Self-supervised learning in medical image analysis
The latest approaches of global feature SSL rely on shared embedding architecture representations that remain invariant to different viewpoints. The variation lies in how these methods prevent collapsing solutions. Clustering methods [20-22] constrain a balanced partition of the samples within a set of cluster assignments. Contrastive methods [15-18] uses negative samples to push far away dissimilar samples from each other through contrastive loss, which can be constructed through memory bank [23], momentum encoder [24], or graph neural network [19]. Unlike contrastive learning, instant-based learning depends on maintaining the informational context of the feature representations by either explicit regularization [13, 25] or architectural design [26, 27]. Our work relates to contrastive and instance-based learning, where a simplified graph-matching version of 1-N or 1-1 reverts to these approaches.
In contrast to global methods, local methods specifically concentrate on acquiring a collection of local features that depict small portions of an image. A contrastive loss function can be used on those feature patches at different criteria such as image region levels [28], or feature maps [29, 14]. These strategies are also widely applied in the medical context, thereby pre-text tasks based on 3D volume's properties, such as reconstructing the spatial context [30], random permutation prediction [31] and self-restoration [32, 33], are proposed. Our LVM-Med model on this aspect can flexible unifying both global and local information by adding them to the affinities matrixes representing the proximity of two graphs, enhancing expressive feature representations.
2.2 Vision-language foundation models
In order to comprehend the multi-modal world using machines, it is necessary to create foundational models that can operate across diverse modalities and domains [34]. CLIP [3] and ALIGN [4] are recognized as groundbreaking explorations in foundation model development. These models demonstrate exceptional proficiency in tasks such as cross-modal alignment and zero-shot classification by learning contrastive pretraining on extensive image-text pairs from the web, despite the presence of noise. To further support multi-modal generation tasks such as visual question answering or video captioning, recent works such as FLAVA [5] and OmniVL [35] are designed to learn cross-modal alignment as well as image-video language models. Conversely, the SAM model [6] utilized a supervised learning strategy with over 1 billion masks on 11 million user-prompt interactions and achieved impressive zero-shot segmentation performance on unseen images. While many efforts have been proposed for natural image domains, limited research has been conducted on large-scale vision models for medical imaging. This motivated us to develop the LVM-Med model.
2.3 Graph matching in visual computing
Graph matching is a fundamental problem in computer vision, which aims to find correspondences between elements of two discrete sets, such as key points in images or vertices of 3D meshes, and used in numerous vision tasks, including 3D vision [36, 37], tracking [38, 39], shape model learning [40, 41], and many others [42-45]. In this framework, the vertices of the matched graphs correspond to the elements of the discrete sets to be matched. Graph edges define the cost structure of the problem, namely, second order, where pairs of matched vertices are penalized in addition to the vertex-to-vertex matchings. This allows us to integrate the underlying geometrical relationship between vertices into account but also makes the optimization problem NP-hard. Therefore, many approximate approaches have been proposed to seek acceptable suboptimal solutions by relaxing discrete constraints [46, 47]. In other directions, gradient estimation techniques for black-box solvers are employed to make the hybrid discrete-continuous matching framework be differentially end-to-end [48-50]. Our LVM-Med follows the latter direction and, for the first time, presents the formulation of contrastive learning as a second-order graph-matching problem.
Figure 2: LVM-Med Overview. Avg is the average pooling layer, MPN denotes for message passing network, GM indicates the combinatorial solver, and  $(c^v, c^e)$  represents vertex and edge affinity matrices. For each image  $x_i$  in batch size, we generated two distorted versions and fed them into the feature representation  $f_\theta$  and another projector  $h_\theta$ . The obtained embeddings  $z_i^\ell$ ,  $\ell \in (s, t)$  are used to build two graphs  $G^s$ ,  $G^t$ . We further design a message passing network  $g_\epsilon$  that aggregate feature per node by their neighbor information. Then we compute vertex and edge affinities  $c^v$ ,  $c^e$  and use them to solve the graph matching. The output afterward is compared with pairs of ground truth  $(x_i^s, x_i^t)$ ,  $i \in (1,.., N)$  representing distorted images generated from the same sample. In the backward pass, we use modern gradient-estimation techniques to approximate  $\frac{\partial L}{\partial c^v}$  and  $\frac{\partial L}{\partial c^e}$ .
3.1 Dataset construction
We provide detailed information about the collected datasets in the Appendix. The data was collected from publicly available resources, which include a diverse set of modalities and body organs as illustrated in Figure 1 (left). The data format is a combination of 2D images and 3D volumes as well as X-ray, MRI, CT, Ultrasounds, etc. For datasets whose data dimensions are 3D volumes, we slice them into 2D images. To avoid potential test data leaking for downstream tasks, we use the default training partition in each dataset; otherwise, we randomly sample with  $20\%$  total images. In total, we obtain approximately 1.3 million images. More statistics on the dataset are presented in the Appendix.
3.2 Contrastive learning as graph matching
Figure 2 provides an illustration of our LVM-Med method, which learns the feature representation  $f_{\theta}$  by matching two distorted views derived from the same input image through a graph-matching formulation. Below we describe in detail each component.
3.2.1 Graph construction on feature embedding
Given a batch of  $N$  images  $\pmb{B} = \{\pmb{x}_1, \pmb{x}_2, \dots, \pmb{x}_N\}$  sampled from a dataset, we generate for each image  $\pmb{x}_i \in B$  two transformed images  $\pmb{x}_i^s$  and  $\pmb{x}_i^t$  by using two transformations  $s, t \sim T$  sampled from  $T$ , a set of pre-defined image transformations. After the transformations, each image is of shape  $(C \times H \times W)$ , where  $C$  is the number of channels and  $(H, W)$  the original spatial dimensions. These distorted images are fed into an encoder  $f_\theta: \mathbb{R}^{C \times H \times W} \to \mathbb{R}^{D \times R \times S}$  to produce two representations  $\pmb{y}_i^s = f_\theta(\pmb{x}_i^s)$  and  $\pmb{y}_i^t = f_\theta(\pmb{x}_i^t)$  where  $D$  is the number of feature channels and  $(R, S)$  are the spatial dimensions of the feature map. On each such representation, we perform an average pooling
operation  $\operatorname{Avg}:\mathbb{R}^{D\times R\times S}\to \mathbb{R}^D$  followed by another projection  $h_\phi :\mathbb{R}^D\to \mathbb{R}^F$  to form two feature embeddings  $\pmb {z}_i^s = h_\phi (\mathrm{Avg}(\pmb {y}_i^s))$ , and  $\pmb {z}_i^t = h_\phi (\mathrm{Avg}(\pmb {y}_i^t))\in \mathbb{R}^F$  with  $F < D$ .
Given a set of embeddings for a batch  $B$ , we construct two graphs  $G^{s}$  and  $G^{t}$  where, for each pair  $(\pmb{x}_i^s,\pmb{x}_i^t)$  of corresponding distorted images, we add a node representing  $\pmb{x}_i^s$  to  $G^{s}$  and a node representing  $\pmb{x}_i^t$  to  $G^{t}$ . Hence, for each  $\ell \in \{s,t\}$ , we construct a graph  $G^{\ell} = (V^{\ell},E^{\ell})$  with  $V^{\ell} = \{x_{1}^{\ell},\dots,x_{N}^{\ell}\}$  the set of vertices and  $E^{\ell}$  the set of edges  $e_{ij}^{\ell} = (x_i^\ell ,x_j^\ell)$ . The node-level feature matrix is given by  $X^{\ell} = [z_1^{\ell};\ldots ;z_N^{\ell}]\in \mathbb{R}^{N\times F}$  which associates each vertex  $x_{i}^{\ell}$  with its feature embedding  $z_{i}^{\ell}$ . We create edges for each graph  $G^{\ell}$  through a  $k$ -nearest neighbors algorithm using the feature matrix  $X^{\ell}$ . The adjacency matrix  $A^{\ell}\in \mathbb{R}^{N\times N}$  is defined as  $A_{ij}^{\ell} = 1$  if  $e_{ij}^{\ell}\in E^{\ell}$  and  $A_{ij} = 0$  otherwise. With the two graph structures given, we obtain a node-attributed graph  $G^{\ell} = (V^{\ell},A^{\ell},X^{\ell})$  on which a graph neural network  $g_{\varepsilon}$  is used to aggregate the nodes' features. In particular,  $g_{\varepsilon}$  computes an embedding  $\hat{Z}^{\ell} = g_{\varepsilon}(X^{\ell},A^{\ell})$  by performing message passing operations. We set  $g_{\varepsilon}$  to be a graph convolutional network [51, 52] consisting of  $l + 1$  layers  $g_{\varepsilon} = \{g_l,g_{l - 1},\dots,g_0\}$  where the output of layer  $l$  is computed as
$$
H _ {l} ^ {\ell} = \sigma \left(\tilde {D} ^ {- \frac {1}{2}} \left(\boldsymbol {A} ^ {\ell} + \boldsymbol {I} _ {N}\right) \tilde {D} ^ {- \frac {1}{2}} H _ {l - 1} ^ {\ell} g _ {l - 1}\right), \tag {1}
$$
where  $\pmb{I}_N$  is the identity matrix modeling self-connections;  $\tilde{D}$  is a diagonal matrix with  $\tilde{D}_{ii} = \sum_{j}A_{ij}^{\ell}$ ;  $g^{l - 1}$  are the trainable parameters for each layer;  $\sigma (\cdot)$  is an activation function; and  $H_0^\ell = X^\ell$ . We use the outputs of the last layer as embeddings for the nodes, that is,  $\hat{\pmb{Z}}^{\ell} = H_{l}^{\ell}\in \mathbb{R}^{N\times F}$  given the shared graph network  $g_{\varepsilon}$ .
We now have two graphs  $G^{s}, G^{t}$  with node attribute matrices  $\hat{Z}^{s}, \hat{Z}^{t}$ , the outputs of the graph neural networks. Next, a graph-matching problem is constructed and solved where the gold matching is given by the pairs  $(\pmb{x}_i^s, \pmb{x}_i^t) \forall i \in \{1,..,N\}$ .
3.2.2 Learning affinities with global and local context
To represent potential connections for a pair of node  $(\pmb{x}_i^s,\pmb{x}_a^t)$  where  $\pmb{x}_i^s\in G^s$ $\pmb{x}_a^t\in G^t$  , we design a vertex affinity matrix  $c^{v}\in \mathbb{R}^{|V^{s}||V^{t}|}$  where  $c_{ia}^{v}$  is the prior (feature-based) similarity between  $\pmb{x}_i^s$  and  $\pmb{x}_a^t$  . An advantage of our formulation is its ability to integrate advanced pair-wise distance can be smoothly integrated to  $c_{ia}^{v}$  , resulting in more expressive proximity representation. In particular, we leverage both global and local consistency derived from feature embeddings of distorted images. The global distance used in several prior works can be computed as  $c_{ia}^{glo}(\pmb{x}_i^s,\pmb{x}_a^t) = \cos (\hat{\pmb{z}}_i^s,\hat{\pmb{z}}_a^t)$  where  $\cos (\cdot)$  denotes cosine similarity;  $\hat{z}_m^\ell$  is the embedding of  $\pmb{x}_m^\ell$ $(\ell \in \{s,t\}$ $m\in \{i,a\})$  obtained after message passing in Eq. (1).
Compared to global methods that implicitly learn features for the entire image, local methods concentrate on explicitly learning a specific group of features that characterize small regions of the image. As a result, they are more effective for dense prediction tasks such as segmentation [29, 14, 53]. While recent works applied these tactics as a part of pair-wise minimization conditions [54, 28] Instead, we integrate them as a part of vertex costs  $c_{ia}^{v}$  and use it to solve the graph matching problem. Indeed, we adapt both location- and feature-based local affinity computed as:
$$
c _ {i a} ^ {l o} \left(\boldsymbol {x} _ {i} ^ {s}, \boldsymbol {x} _ {a} ^ {t}\right) = \mathbb {E} _ {p \in \boldsymbol {P}} \cos \left(\boldsymbol {q} _ {p} ^ {s}, \boldsymbol {q} _ {\mathrm {m} (p)} ^ {t}\right) + \mathbb {E} _ {p \in \boldsymbol {P}} \cos \left(\boldsymbol {q} _ {p} ^ {s}, \boldsymbol {q} _ {\mathrm {m} ^ {\prime} (p)} ^ {t}\right) \tag {2}
$$
where  $\pmb{P} = \{(r,s) | (r,s) \in [1, \dots, R] \times [1, \dots, S]\}$  be the set of coordinates in the feature map  $\pmb{y}_i^s \in \mathbb{R}^{D \times R \times S}$  of  $\pmb{x}_i^s$ ;  $\pmb{q}_p^\ell$  ( $\ell \in \{s,t\}$ ) be the feature vector at position  $p$ ;  $\mathrm{m}(p)$  denote the spatial closest coordinate to  $p$  in coordinates of feature map  $\pmb{y}_a^t$  estimated through transformations on original image  $\pmb{x}_i$ ; finally  $\mathrm{m}'(p)$  represents the closest feature vector to  $p$  in  $\pmb{y}_a^t$  using  $l^2$  distance. Intuitively, the local cost in Eq. (2) enforces invariance on both spatial location and between embedding space at a local scale. Our final affinity cost is computed as:
$$
c _ {i a} ^ {v} \left(\boldsymbol {x} _ {i} ^ {s}, \boldsymbol {x} _ {a} ^ {t}\right) = \alpha \left(c _ {i a} ^ {g l o} \left(\boldsymbol {x} _ {i} ^ {s}, \boldsymbol {x} _ {a} ^ {t}\right)\right) + (1 - \alpha) \left(c _ {i a} ^ {l o} \left(\boldsymbol {x} _ {i} ^ {s}, \boldsymbol {x} _ {a} ^ {t}\right) + c _ {i a} ^ {l o} \left(\boldsymbol {x} _ {a} ^ {t}, \boldsymbol {x} _ {i} ^ {s}\right)\right) \tag {3}
$$
3.2.3 Self-supervision through second-order graph matching
While the standard graph matching problem for vertex-to-vertex correspondences can be used in our setting (LAP), it fails to capture the similarity between edges. If there are duplicated entities
represented by distinct nodes in the same graph, the LAP will consider them identical and skip their neighboring relations. For instance, during the image sampling, two consecutive image slides were sampled from a 3D volume, resulting in their appearances have a small difference. In such cases, it is complicated to correctly identify those augmented images generated from the same one without using information from the relations among connected nodes in the constructed graph. To address this problem, we introduce additional edge costs  $c^e \in \mathbb{R}^{|E^s||E^t|}$  where  $c_{ia,jb}^e$  represents the similarity between an edge  $v_{ij}^s = (\pmb{x}_i^s, \pmb{x}_j^s) \in E^s$  and  $v_{ab}^t = (\pmb{x}_a^t, \pmb{x}_b^t) \in E^t$ . These edge costs (second-order) are computed as  $c_{ia,jb}^e = \cos((\hat{\pmb{z}}_i^s - \hat{\pmb{z}}_j^s), (\hat{\pmb{z}}_a^t - \hat{\pmb{z}}_b^t))$ .
We now establish the second-order graph-matching problem. Denoting  $\pmb{v} = \{0,1\}^{|V^s||V^t|}$  be indicator vector of matched vertices, i.e.,  $v_{ia} = 1$  if the vertex  $\pmb{x}_i^s \in V^s$  is matched with  $\pmb{x}_a^t \in V^t$  and  $v_{ia} = 0$  otherwise. The node correspondence between two graphs  $G^s$  and  $G^t$  that minimizes the global condition stated as:
$$
\mathbb {M} \left(\boldsymbol {c} ^ {v}, \boldsymbol {c} ^ {e}\right) = \underset {\boldsymbol {v} \in U (\mathbf {1}, \mathbf {1})} {\arg \min } - \sum_ {i, a} c _ {i a} ^ {v} v _ {i a} - \sum_ {i, j, a, b} c _ {i a, j b} ^ {e} v _ {i a} v _ {j b}
$$
where  $U(\mathbf{1},\mathbf{1}) = \{\pmb {v}\in \{0,1\}^{N\times N}|\pmb {v}\pmb{1}_N = \mathbf{1},\pmb {v}^T\pmb {1}_N = \mathbf{1}\}$
and  $\mathbf{1}_N$  be a  $n$ -dimension one-value vector. The constraint  $U(\mathbf{1},\mathbf{1})$  restricts  $\pmb{v}$  satisfying the one-to-one matching. Essentially, the Eq. (4) solves the vertex-to-vertex correspondence problem using both node and edges affinities, which can be seen as a form of structural matching (Figure (1), right) and generally can be integrated with higher-order graph constraints as triangle connections or circles. In the experiment, we found out that Eq. (4) significantly improved downstream task performance compared to the pure linear matching approach (Table (6)). Since the Eq. 4 in general is an NP-Hard problem [55] due to its combinatorial nature, we thus use efficient heuristic solvers based on Lagrange decomposition techniques [56].
3.2.4 Backpropagating through a graph matching formulation
With  $\hat{\pmb{v}} = \mathsf{GM}(\pmb{c}^v, \pmb{c}^e)$  a solution obtained from the solver, we use the Hamming distance and an optimal solution  $\pmb{v}^*$  to define the following loss function
$$
L (\hat {\boldsymbol {v}}, \boldsymbol {v} ^ {*}) = \hat {\boldsymbol {v}}. (1 - \boldsymbol {v} ^ {*}) + \boldsymbol {v} ^ {*}. (1 - \hat {\boldsymbol {v}}). \tag {5}
$$
The proposed approach aims to learn the feature representation function  $f_{\theta}$  such that its output minimizes Eq. (5). However, this is a difficult problem because the partial derivatives of the loss function w.r.t vector costs  $c^v, c^e$ , i.e.,  $\frac{\partial L}{\partial c^v}$  and  $\frac{\partial L}{c^e}$ , are zero almost everywhere [48, 57] due to the objective function in Eq. (4) being piece-wise constant, preventing direct gradient-based optimization.