SlowGuess/NeurIPS_2023_OCR huggingface dataset:
for X in loader: # load a batch  $\mathrm{X} = [\mathrm{x}_1, \mathrm{x}_2, \dots, \mathrm{x}_N]$  with N samples
# apply two transformations s and t
 $\mathrm{X^s}$ ,  $\mathrm{Pos^s} = \mathrm{s(X)}$  #  $\mathrm{X^k} = [\mathrm{x_1^k}, \mathrm{x_2^k}, \dots, \mathrm{x_N^k}]$ ,  $\mathrm{Pos^k} = [\mathrm{pos_1^k}, \mathrm{pos_2^k}, \dots, \mathrm{pos_N^k}]$ ,  $\mathrm{k} \in \{\mathrm{s}, \mathrm{t}\}$ $\mathrm{X^t}$ ,  $\mathrm{Pos^t} = \mathrm{t(X)}$ 
# compute feature representations
 $Y^s = f_{\theta}(X^s)$ ;  $Y^t = f_{\theta}(X^t)$  # feature dimensions:NxDxRxs
# applying projection
 $Z^s = h_{\phi}(Avg(Y^s))$ ;  $Z^t = h_{\phi}(Avg(Y^t))$  # dimensions:NxF
# build graph structures and message passing
 $G^s = k$ -nearest-neighbor  $(Z^s, k\_connects)$ $G^t = k$ -nearest-neighbor  $(Z^t, k\_connects)$ $\hat{Z}^s = g_{\epsilon}(G^s, Z^s)$ ;  $\hat{Z}^t = g_{\epsilon}(G^t, Z^t)$ 
# compute vertex and edge affinity matrices
 $c_{ia}^{v} = \alpha * \cos(\hat{z}_{i}^{s}, \hat{z}_{a}^{t}) + (1 - \alpha) * local\_cost(y_{i}^{s}, y_{a}^{t}, pos_{i}^{s}, pos_{a}^{t})$  # affinity  $x_i^s \& x_a^t$ $c_{ia,jb}^{e} = \cos((\hat{z}_{i}^{s} - \hat{z}_{j}^{s}), (\hat{z}_{a}^{t} - \hat{z}_{b}^{t}))$  # affinity between edges  $v_{ij}^{s}, v_{ab}^{t}$ $c_v^v = \{c_{ij}^v\} \in R^{N\times N}$ ;  $c_e^e = \{c_{ia,jb}^e\} \in R^{|\mathbb{E}^s||\mathbb{E}^t|}$  #  $E^k$  be a set of edges in  $G^k, k \in \{s, t\}$ 
# perturbed costs with Gumbel noise
 $\epsilon, \epsilon' \sim Gumbel(0, 1)$ $c_v^v = c_v^v + \epsilon; c_e^e = c_e^e + \epsilon'$ 
# solving graph matching and compute loss
 $\hat{v} = GM(c_v^v, c_e^e)$ $L(\hat{v}, v^*) = \hat{v}. (1 - v^*) + v^*. (1 - \hat{v})$  # compute hamming loss
# update network
L-backward() # approximate  $(\partial L / \partial c_v^v, \partial L / \partial c_e^e)$  by Algorithm 1.
Update  $(g_{\epsilon}.params)$ , Update  $(h_{\phi}.params)$ , Update  $(f_{\theta}.params)$ 
# define local_cost
def local_cost  $(y_i^s, y_a^t, pos_i^s, pos_a^t)$ :
# location-based local cost
 $y_{i,\text{nn}}^s = \text{torch.zeros\_like}(y_i^s)$ 
for r, s in R, S:
    r', s' = argmin((L2(pos_i^s[r, s], pos_a^t[r', s'])))
     $y_{i,\text{nn}}^s [r, s] = y_a^t [r', s']$
$\mathrm{y_{i\_fil}^{s},y_{i,nn\_fil}^{s} =}$  select_top  $(\mathrm{y_i^s,y_i^s}_{\mathrm{n,n}},\gamma)$    
location_cost  $=$  cos(ysi,fil,ysi,nn_fil)   
# featured-based local cost   
 $\mathrm{y_{i,nn}^{s} =}$  torch.zeros_like(yi)   
for r,s in R,S:  $\mathbf{r}'$  ,s'  $=$  argmin((L2(yis[r,s],ya[ra',s'])   
 $\mathrm{y_{i,nn}^{s}[r,s] = y_{a}^{t}[r',s']}$ $\mathrm{y_{i,nn}^{s},y_{i,nn\_fil}^{s} =}$  select_top  $(\mathrm{y_i^s,y_i^s}_{\mathrm{i,n,n}},\gamma)$    
feature_cost  $=$  cos(ysi,fil,ysi,nn_fil)   
return 0.5*(location_cost + feature_cost)
We trained LVM-Med with graph size of 16 nodes, each node connected to the top 5 nearest neighbors after using kNN,  $\lambda$  value in Algorithm 1 is 80, and  $\alpha = 0.8$  for associating global- and local-based similarities when computing  $c_{ij}^v$ . The size of projector  $h_\phi$  is  $2048 \times 128$  for ResNet-50, and  $768 \times 128$  for ViT. We configure the message passing network  $g_\theta$  with two convolutional layers of size 128. For the user-based prompt version, because the SAM model [9] requires an input of shape  $256 \times 14 \times 14$  for the mask decoder part, we add two additional convolutional layers with a kernel size of 1 and 3 at the end of ViT backbone to convert from shape  $768 \times 14 \times 14$  to the target shape.
B Downstream task setups
B.1 Downstream tasks
Segmentation tasks On 2D-based segmentation tasks, we employ U-Net architecture [91] and load ResNet-50 [60] trained by self-supervised learning algorithms as network backbones. With foundation models, we use TransUnet [76] and take pre-trained ViT models as the backbones. For the prompt-based segmentation, we follow the architecture of SAM [6] consisting of encoder, prompt, and mask decoder layers. We also fine-tune SAM where encoder and prompt networks are frozen, only learning decoder layers [9]. Our LVM-Med for prompt-based setting is similar to [9] except that we substitute SAM's encoders with our weights. We utilize Adam optimizer for all experiments and train architectures with Dice and Cross-Entropy loss [92]. We also normalize the norm-2 of gradient values to stabilize the training step to maximize 1. Table 7 summarizes each dataset's learning rate, number of epochs, and image resolution.
On 3D-based segmentations, we reformulate these tasks as 2D segmentation problems and make predictions on 2D slices taken from 3D volumes. Furthermore, we apply balance sampling to select equally 2D slices covering target regions and other 2D slices, not including the ground truth. Table 8 presents configurations used for 3D datasets; other settings are identical to 2D cases.
Table 7: Configurations for training 2D segmentation tasks
ISIC-2018 (Skin Lesion)	JSRT (Lung X-ray)	KvaSir (Polyp)	Drive (Vessel)	BUD (Breast Cancer)
ResNet-50.	lr = 10-4, epochs 35 shape 512 × 512 batch size 16	lr = 10-3, epochs 50 shape 224 × 224 batch size 32	lr = 10-3, epochs 35 shape 224 × 224 batch size 64	lr = 10-3, epochs 50 shape 224 × 224 batch size 16	lr = 10-4, epochs 50 shape 256 × 256 batch size 8
Foundation Model	lr = 10-4, epochs 100 shape 512 × 512 batch size 16	lr = 10-3, epochs 200 shape 224 × 224 batch size 32	lr = 10-3, epochs 200 shape 224 × 224 batch size 64	lr = 10-3, epochs 200 shape 224 × 224 batch size 16	lr = 10-4, epochs 200 shape 256 × 256 batch size 8
Prompt-based Seg.	lr = 10-4, epochs 50 shape 1024 × 1024 batch size 16	lr = 3 × 10-4, epochs 50 shape 1024 × 1024 batch size 16	lr = 3 × 10-4, epochs 20 shape 1024 × 1024 batch size 16	lr = 3 × 10-4, epochs 100 shape 1024 × 1024 batch size 16	lr = 10-4, epochs 20 shape 1024 × 1024 batch size 16
Table 8: Configurations for 3D-based-segmentation tasks
BraTS	MMWHS-CT	MMWHS-MRI	BMC
ResNet50	lr = 15 × 10-4, epochs 20 shape 224 × 224 batch size 128	lr = 10-3, epochs 20 shape 224 × 224 batch size 64	lr = 15 × 10-4, epochs 30 shape 224 × 224 batch size 64	lr = 10-3, epochs 30 shape 224 × 224 batch size 64
Foundation Model	lr = 10-4, epochs 100 shape 224 × 224 batch size 16	lr = 10-4, epochs 100 shape 224 × 224 batch size 16	lr = 10-4, epochs 100 shape 224 × 224 batch size 16	lr = 10-4, epochs 100 shape 224 × 224 batch size 16
Prompt-based Seg.	lr = 3 × 10-5, epochs 30 shape 1024 × 1024 batch size 16	lr = 5 × 10-5, epochs 30 shape 1024 × 1024 batch size 16	lr = 3 × 10-5, epochs 30 shape 1024 × 1024 batch size 16	lr = 3 × 10-4, epochs 50 shape 1024 × 1024 batch size 16
Image classification tasks We take the feature embedding outputs of each architecture and build one fully connected layer to produce desired classes for image classification tasks. We freeze the encoder layers for the linear evaluation and only train the fully connected layer. For the fully-finetuning, the whole network is trained. The Adam optimizer [63] with cross-entropy loss function and learning rates  $\{5 \times 10^{-4}, 10^{-3}\}$  are used for Brain Tumor and FGADR, respectively. To benchmark LVMMed with other state-of-the-art methods on FGADR (Figure 3 in paper), we follow the settings of DRG-Net [90] and change their encoder layers by our networks.
Object detection We use Faster-RCNN [93] for object detection tasks. The ResNet-50 of Faster-RCNN is replaced by pre-trained weights. In the Vin-Dr dataset, there is a total of 14 objects for, e.g., Aortic enlargement, Atelectasis, Calcification, etc. We use image resolutions of  $512 \times 512$ , Adam solver, and learning rate  $10^{-4}$  in 40 epochs. In the Kvasir dataset for polyp detection, we also resize images to a fixed size of  $512 \times 512$ , employ the Adam optimizer with learning rate  $2.5 \times 10^{-4}$  and batch size 8.
C LVM-Med ablation studies
C.1 Graph sizes and  $\lambda$  in backpropagation
We provide in Figure 5 and Figure 6 LVM-Med performance when changing the number of nodes in graph construction steps  $G^{s}$ ,  $G^{t}$  and  $\lambda = 80$  used in Algorithm 1 in the backpropagation step. The results are reported on the average Dice score of five 2D segmentation tasks and the average accuracy of two linear classifications on FGADR and Brain Tumor Classification. Figure 5 indicates that 16 is the best value for both classification and segmentation. Increasing the graph's nodes tends to decrease classification performance.
Figure 6 compared different values for  $\lambda \in \{70,80,90,100\}$ . We observe that  $\lambda = \{80,90\}$  achieve good results for linear classification tasks though  $\lambda = \{90,100\}$  decreases segmentation performance.
C.2 Performance on large- and small-scale
We investigate LVM-Med performance when reducing the number of datasets in the pre-training step. Especially, we trained LVM-Med on a small-scale with four datasets: LUNA2016 [94], LiTS2017 [95], BraTS2018 [64], and MSD (Heart) [96]. We compare this version with our default settings trained on 55 datasets (Section F). Two models are evaluated on dice scores of five 2D segmentation tasks, the accuracy metric of two linear image classifications, and mAP50 of two object detection tasks on VinDr and Kvasir detection. Table 9 shows that LMV-Med full leads to better performance overall, especially with the classification settings; the improvement gap is around  $3.6\%$ . In summary, we conclude that LVM-Med is beneficial when training in large-scale medical settings.
Figure 5: LVM-Med performance when varying the number of nodes in graph construction.
Figure 6: LVM-Med performance when varying the  $\lambda$  in backpropagation step.
C.3 Performance on weighting global and local similarities
We test with different  $\alpha = \{0.7, 0.8, 0.9\}$  which used to fuse global- and local-based similarities  $c_{ij}^{v}$ . Table 9 demonstrates that  $\alpha = 0.8$  is generally the best value in average across segmentation, classification, and object detection tasks.
Table 9: LVM-Med ablation studies trained with full data, small-scale, and different hyper-parameter  $\alpha$  fusing global- and local-based similarities. Results are reported on an average of five 2D segmentation, two linear classifications, and two object detection tasks. The most impacted factors are highlighted.
Method	Cls.(Acc)	Seg. (Dice)	Detect. (mAP50)
LVM-Med (full, α = 0.8)	67.47	83.05	57.1
LVM-Med (small-scale, α = 0.8)	63.83	81.97	56.03
LVM-Med (full, α = 0.7)	65.89	82.20	56.49
LVM-Med (full, α = 0.9)	65.03	81.09	57.14
C.4 Computational complexity
We present a parameter comparison of LVM-Med with other foundation models in Table 10. Our LVM-Med model, based on ResNet-50, has significantly fewer parameters, approximately 3-4 times smaller than models such as Flava or SAM, while still maintaining competitive performance. When utilizing the ViT encoder pre-trained by the SAM method, LVM-Med's parameters are comparable to the Flava model and slightly higher than Clip and Align by 1.03 and 1.43 times, respectively. However, it is important to note that both LVM-Med and SAM outperform these models by a significant margin in various settings.
Table 10: Computational complexity of our approaches and other foundation models.
Method	LVM-Med (R50)	LVM-Med (ViT)	Clip [3]	Flava [5]	Align [4]	SAM (Encoder) [6]
#Param	25.55 M	88.88 M	85.80 M	86.39 M	62.14 M	88.88 M
D Prompt-based segmentation on 3D datasets and classification tasks
We provide additional results for LVM-Med on 3D-based prompt segmentation and image classification tasks with several fully connected layers.
Table 12: Comparing SSL approaches and Foundation models on classification tasks with two evaluation protocols, Linear evaluation and full Fine-tuning. Settings used with several fully connected layers are in cyan. The best results in 2D-SSL and foundation models (two fully connected layers) are in bold; the best results overall are in bold and underlined.
Method	Linear Evaluation (Frozen)	Fine-tuning
FGADR (DR Grading)	Brain Tumor Class.	FGADR (DR Grading)	Brain Tumor Class.
2D-SSL on medical	Twin-Barlon [13]	66.86 ± 0.41	63.03 ± 0.32	66.37 ± 0.77	74.20 ± 1.38
Dino [79]	65.98 ± 1.91	62.27 ± 0.32	67.35 ± 1.36	71.91 ± 1.55
SimCLR [15]	65.30 ± 1.70	62.52 ± 1.67	67.55 ± 0.28	73.52 ± 3.56
Moco-v2 [17]	65.98 ± 1.04	62.35 ± 1.92	67.55 ± 1.79	74.53 ± 0.43
Deepcluster [20]	65.34 ± 1.93	64.47 ± 0.55	67.94 ± 1.78	73.10 ± 0.55
VicRegl [14]	64.71 ± 0.60	59.64 ± 1.36	65.69 ± 1.46	73.18 ± 2.03
LVM-Med (R50)	68.33 ± 0.48	66.33 ± 0.31	68.32 ± 0.48	76.82 ± 2.23
66.67 ± 0.84	74.70 ± 0.84	70.58 ± 0.36	78.77 ± 0.78
Foundation Model	Clip [3]	57.87 ± 0.50	57.87 ± 0.71	57.48 ± 0.86	34.86 ± 2.27
62.66 ± 0.36	67.85 ± 0.23	56.21 ± 1.86	21.74 ± 1.14
Flava [5]	31.87 ± 0.69	35.19 ± 0.43	57.18 ± 0.96	34.01 ± 5.97
32.84 ± 0.12	24.45 ± 4.30	56.01 ± 0.86	33.67 ± 8.11
Algin [4]	36.95 ± 1.04	30.71 ± 2.35	57.28 ± 0.97	63.96 ± 0.04
38.12 ± 1.45	30.34 ± 1.35	57.87 ± 0.90	61.42 ± 0.25
SAM [6]	55.13 ± 0.41	31.81 ± 4.26	58.75 ± 1.32	60.66 ± 1.36
57.48 ± 0.24	36.89 ± 1.61	58.75 ± 0.99	60.07 ± 0.31
LVM-Med (SAM's ViT)	62.46 ± 0.86	59.31 ± 0.48	63.44 ± 0.73	67.34 ± 2.08
63.83 ± 1.36	64.13 ± 1.14	59.04 ± 0.14	64.97 ± 2.71
D.1 Promt-based Segmentation on 3D datasets
We perform experiments on three 3D datasets in Table 11, including BraTS, MMWHS-MRI, and MMWHS-CT. The setup for box prompts follows 2D segmentation cases. We discover that the LMV-Med in 3D cases consistently improves the performance of fine-tuned SAM [9] as in 2D settings and attains a large margin compared with SAM without training [6]. This evidence thus confirms that LVM-Med is also effective under prompt-based scenarios.
Table 11: Prompt-based segmentation on 3D datasets.
Method	BraTS	MMWHS-MRI	MMWHS-CT
Prompt-based Seg.	SAM (fixed encoder) [9]	85.37 ± 0.07	77.64 ± 1.14	76.61 ± 1.91
SAM with Prompt (no-train) [6]	38.97 ± 0.21	59.74 ± 0.76	50.25 ± 0.33
LVM-Med (SAM's ViT)	85.76 ±0.07	78.91 ± 0.80	78.03 ± 0.93
D.2 Image classification
We aim to inspect whether foundation models improve their performance given more fully connected layers for image classification tasks with both frozen encoders or fully fine-tuning. For each method in this category and our LVM-Med (ResNet-50 and ViT), we configure two fully connected layers with sizes 512 - 256 and 512 - 128 for the Brain and FGADR respectively that map from the output dimension of each network to a number of desired classes. Table 12 presents obtained results where new settings are highlighted in color. We notice the following points. (i) Firstly, using more fully connected layers tends to improve the performance of foundation models, especially on linear evaluation. For e.g., the Clip increases from  $4.79\% - 9.98\%$  on FGADR and Brain Tumor classification tasks, respectively. Similarly, our LVM-Med with SAM's ViT also achieves better results by approximately  $1.37\%$  and  $4.82\%$  on those tasks. (ii) Secondly, LVM-Med overall attains the best results in four settings using linear or several fully connected layers with ResNet-50. LVM-Med with ViT architecture also delivers the best records on three of four test cases compared with foundation models.
E Visualizing results
We provide qualitative results for prompt-based segmentation in Figure 7. We compare three approaches, including (i) the standard SAM without fine-tuning [6] (second column), (ii) SAM
with encoders and prompt networks are frozen, and only decoder layers are trained as [7] (third column), and (iii) a similar setting as (ii) but encoders taken from LVM-Med version with SAM's ViT architecture (fourth column). For all methods, we simulate box-based prompts using the ground-truth masks and define boxes covering those target regions perturbed by offset values.
Figure 7 demonstrates that the original SAM is prone to generate useless predictions (top and bottom rows) or less precise boundaries. In contrast, updated SAM and LVM-Med produce more accurate results, confirming the importance of fine-tuning to achieve adequate results. Figures in the third and fourth columns also illustrate that SAM tends to over-segment or lacks structures on an object's edges in several cases, while LVM-Med is more stable in those situations (red arrows).
SAM (no fine-tuning) [6]
SAM (fine-tuning) [9]
LVM-Med (SAM's ViT)
Figure 7: Visualizing prompt-based predictions on three datasets: BUID, Kvasir, and ISIC. Red arrows show differences between SAM (fine-tuning) and LVM-Med using SAM's ViT architecture. Best viewed in color with zoomm.
F Dataset overviews
Table 13 overviews the dataset used in our study. For each dataset, we provide its modality, data dimension, and the total of samples. If the training/testing rate is available (column Train/Test Rate), we utilize all training data; otherwise, we sample  $20\%$  total samples to avoid potential test data leaking for downstream tasks used in the pre-training step. For datasets whose data dimensions are 3D volumes, we sample 2D slices from those formats. Some datasets, such as MSD or ADNI, comprise different sub-datasets inside; we consider these sub-sets as independent ones to avoid confusion during the training steps. In summary, a total of 55 datasets are used with approximately  $40\%$  in 3D datasets and  $60\%$  in 2D images as presented in Figure 8. Moreover, we also outline ratios between distinct data modalities such as MRI, CT, X
ray, grayscale types such as Ultrasound, OCT, and finally, color images depicted in Figure 9.
Figure 8: Pie chart illustrating the ratio of the number of datasets with the 3D and 2D dimension.
Figure 9: Pie chart illustrating the ratio of different data modalities in our collected dataset.