SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Manuel Tran<sup>1,3,4</sup>
Yashin Dicente Cid
Amal Lahiani
Fabian J. Theis<sup>3,4</sup>
Tingying Peng $^{4,*}$
Eldad Klaiman $^{1,*}$
Roche Diagnostics GmbH, Roche Diagnostics S.L.
<sup>3</sup>Technical University of Munich, <sup>4</sup>Helmholtz Munich
Training multimodal foundation models is challenging due to the limited availability of multimodal datasets. While many public datasets pair images with text, few combine images with audio or text with audio. Even rarer are datasets that align all three modalities at once. Critical domains such as healthcare, infrastructure, or transportation are particularly affected by missing modalities. This makes it difficult to integrate all modalities into a large pre-trained neural network that can be used out-of-the-box or fine-tuned for different downstream tasks. We introduce LoReTTa (Linking mOdalities with a tRansitive and commutativE pre-Training sTrAtegy) to address this understudied problem. Our self-supervised framework unifies causal modeling and masked modeling with the rules of commutativity and transitivity. This allows us to transition within and between modalities. As a result, our pre-trained models are better at exploring the true underlying joint probability distribution. Given a dataset containing only the disjoint combinations  $(A,B)$  and  $(B,C)$ , LoReTTa can model the relation  $A\leftrightarrow C$  with  $A\leftrightarrow B\leftrightarrow C$ . In particular, we show that a transformer pre-trained with LoReTTa can handle any mixture of modalities at inference time, including the never-seen pair  $(A,C)$  and the triplet  $(A,B,C)$ . We extensively evaluate our approach on a synthetic, medical, and reinforcement learning dataset. Across different domains, our universal multimodal transformer consistently outperforms strong baselines such as GPT, BERT, and CLIP on tasks involving the missing modality tuple.
Integrating multiple modalities, such as images, text, or speech, promises to provide a more comprehensive and holistic understanding of otherwise intractable problems. By leveraging the unique and complementary strengths of each data type, this approach results in a more synergistic representation of the underlying phenomena – enabling deep learning systems to excel and adapt to real-world scenarios. Thus, there has been impressive progress in the field of multimodal learning in recent years. While most research focuses on multimodal datasets where all data points have all modalities, there has also been work on missing modalities [44, 45]. However, all these models and methods rely heavily on at least a subset of samples where all modality combinations are still present (Figure 1). In practice, it is not easy to obtain such datasets when we consider more than two modalities, e.g. three modalities. In the worst case, we might end up with a dataset containing only modalities  $A$  and  $B$ , and another disjoint dataset containing modalities  $B$  and  $C$ . Thus, the combination  $A$  and  $C$
(a) All Modality Combinations Present
(b) All Modality Combinations Sometimes Present
(c) Some Modality Combinations Never Present
Figure 1: Venn diagrams showing the relationship between datasets with different modalities  $A$ ,  $B$ , and  $C$ . Overlapping datasets indicate that the dataset contains samples with aligned modalities (i.e., audio, image, and text files belonging to the same concept). While recent work has mostly focused on datasets where at least some samples have all modality combinations available (a, b), we investigate the case where some modality combinations, e.g.  $(A, C)$  and  $(A, B, C)$ , are missing entirely (c).
is unconditionally missing. So we asked ourselves, "Is it possible to pre-train a multimodal neural network that can handle any combination of modalities at inference time? Optimally, the model will achieve better downstream performance the more modality combinations are given. This includes the never-seen combinations  $(A,C)$  and  $(A,B,C)$ .
To the best of our knowledge, this scenario has never been considered in the literature. We propose a novel transitive pre-trained transformer (Figure 2) that can deal with such situations. LoReTTa (Linking mOdalities with a tRansitive and commutativE pre-Training sTrAtegy) can link modalities and transition between them using generative pre-training (predicting the next token), the commutative property  $(A,B) = (B,A)$ , and the transitive relation  $(A\to B)\wedge (B\to C)\Rightarrow (A\to C)$ . We show theoretically and experimentally that by learning the relationship within and between tokens from different data distributions, LoReTTa is well suited for training multimodal models.
This has far-reaching implications for safety-critical domains such as healthcare, infrastructure, and transportation. In medicine, for example, it is already difficult to obtain a large dataset with a single modality to train powerful foundation models – such as those used in general computer vision and natural language processing. Finding a dataset with two or more modalities is even more challenging. But combining them is important because it is common practice to look at different types of data, such as radiological images, genetic sequences, and microscopic slides, to determine the optimal treatment. Similarly, infrastructure monitoring models or automated driving systems require input from a variety of sensors. However, not all modality combinations might be present in the same training set. With LoReTTa, we demonstrate a unique way to address this challenging and under-researched problem.
With the introduction of transformers [63], the fields of natural language processing (NLP) and computer vision (CV) have experienced a significant shift in performance and capabilities. These models excel at a wide range of complex tasks, including machine translation, text summarization, question answering, sentiment analysis, natural language generation, image classification, image captioning, semantic segmentation, object detection, and face recognition - achieving state-of-the-art performance in each area [40]. Transformers have also opened up new possibilities for multimodal learning due to their lack of inductive bias. They can process and fuse different types of data, such as audio, image, video, and text, using self- and cross-attention mechanisms. As a result, transformers are considered to be universal, modality-agnostic learners [43, 22].
However, their lack of inductive bias comes at a price. Unlike other neural network architectures, transformers make few assumptions about the structure of the data [18]. This makes them highly expressive, but they also require a large amount of data to learn effectively and generalize well [28]. Consequently, self-supervised learning (SSL) on large amounts of unlabeled data has been proposed for transformers. In NLP, this strategy has been very successful, with methods such as masked
(a) Causal Modeling (Switch Pairs for Commutative Modeling)
(b) + Causal Masked Modeling
(c) + Transitive Modeling
(d) Before & After Transitive Pre-Training
Figure 2: LoReTTa consists of two novel self-supervised strategies: commutative and transitive modeling. (a) In commutative modeling, we apply causal modeling in a commutative manner to generate modality  $A$  from  $B$  and modality  $B$  from  $A -$  given the aligned input sample  $(A,B)$ . For the aligned but disjoint data point  $(B^{\prime},C)$ , we apply the same technique. (b) To ensure that the model learns bidirectional relations between the input tokens, we apply a modified variant of generative pre-training called causal masked modeling. (c, d) Next, we use transitive pre-training to learn any missing conditional joint distributions. The idea is simple. We randomly select a sample and use the linking modality  $B$  to predict the missing modality  $C$ , which is then used to reconstruct the existing modality  $A$ . The last step is crucial because it ensures that all modalities are properly aligned.
modeling and causal modeling used to train large language models (LLMs) such as BERT [23] and GPT [49] variants [42, 46].
Recently, attempts have been made to unify both approaches through frameworks like prefix modeling [53], permutation modeling [71], causal masked modeling [1], or unified language learning [59]. When applied to image classification, token-based pre-training methods [8, 15] have demonstrated their ability to outperform established computer vision approaches such as SimCLR [17] and DINO [13]. As a result, these sequence modeling techniques have been successfully adapted and extended to a wide range of domains, encompassing programming [16], biochemistry [41], and reinforcement learning [35].
Adding or combining different modalities to the training process could open up three possibilities: positive transfer, cross-modal understanding, and cross-modal generation. Positive transfer occurs when learning in one modality helps to improve performance in another modality. For example, it has been shown that aligning images with text leads to better classification performance on ImageNet [51]. Cross-modal understanding refers to the explanation of one modality in terms of the other. Models with this capability can describe an image using text [4] or infer protein structures from sequences [30]. With cross-modal generation, one can even generate one modality using the other. Examples are text-to-image generation [54] and text-to-music generation [3]. Many modern multimodal foundation models are based on either contrastive learning [26], masked modeling [66], or causal modeling [72].
In this paper, we present a novel pre-training method that combines recent advances in self-supervised learning into a unified framework. LoReTTa uses causal modeling for generative tasks and masked modeling for discriminative tasks. It extends both approaches by using commutative and transitive modeling to seamlessly integrate and impute modalities. This allows us to learn expressive features and relationships within multimodal datasets. The pre-trained models can be used for various downstream tasks such as infilling, classification, survival prediction, or cross-modal generation. The most similar to our work is the RNN-based model MCTN [47]. However, it can only be used as a classifier, does not model all combinations of modalities, and requires all modalities to be aligned and present at training time. LoReTTa, on the other hand, does not need such strong assumptions.
Model architecture: LoReTTa (Figure 2) uses the autoregressive transformer decoder architecture described in Radford et al. [50]. This model includes a modified initialization, pre-normalization, and GELU activation. In addition, we choose RMSNorm [74] over LayerNorm [7], as in Rae et al. [52] and Brown et al. [12]. We also replace the StandardAttention [63] with FlashAttention [21] to reduce memory usage and runtime - similar to Touvron et al. [61]. Following Chowdhery et al. [19], we train without dropout and do not use biases in any of the dense kernels or layer norms. While any general sequence model can work for next token prediction, we choose the transformer [63] in particular for its simplicity and scalability. Theoretically, however, our method also works with other backbones, such as the recently proposed RWKV [10] and Hyena [48].
Data representation: Since transformers take a sequence of tokens as input, we need to tokenize our data first. We use common methods such as raw byte encoding [56, 29, 15], learned lookup tables [33, 58, 57], and vector quantization [62, 55, 54, 3]. After tokenization, we use a parameterized embedding function to project all tokens into a common feature space. In addition, we add modality-specific tokens to indicate the beginning and end of a modality. We also learn modality-specific absolute position encodings that are added to the feature vectors. The representations for each modality are then concatenated, but only if they belong to the same object/subject. For example, if we have an image-audio pair or an image-text pair describing the same concept, we would concatenate their token sequences and feed them into the transformer. However, if only one modality is available, we use only that modality. In summary, our pre-processing step is described below:
- We tokenize and project all samples  $a, b, c$  into a common vector space, obtaining the token embeddings  $[a_1, \ldots, a_l], [b_1, \ldots, b_m]$ , and  $[c_1, \ldots, c_n]$ .
- If  $a$  and  $b$  are aligned, then we concatenate them as  $s = [a_1, \ldots, a_l; b_1, \ldots, b_m]$  or  $s = [b_1, \ldots, b_m; a_1, \ldots, a_l]$ , which ensures commutativity.
Pre-training strategy: We train the transformer using causal language modeling. Given a sequence of tokens  $s_{1:L}$  and parameters  $\theta$ , the data is modeled using the chain rule of probability
$$
\log p _ {\theta} \left(s _ {1}, \dots , s _ {L}\right) = \sum_ {l = 1} ^ {L} \log p _ {\theta} \left(s _ {l} \mid s _ {1}, \dots , s _ {l - 1}\right) \tag {1}
$$
Thus, the task of the model is to predict the next token given the sequence of previous tokens. As shown in GATO [56] and PALM-E [24], this objective is modality agnostic. The systems even learn rich internal world models [38, 36]. We ensure commutativity by randomly mixing the order of concatenation when appending tokens from two modalities to an input sequence, as described above. This allows the model to infer modality  $A$  from  $B$  and modality  $\bar{B}$  from  $A$ . We call this technique commutative modeling. Interestingly, this simple modification has not been used in recent multimodal foundation models.
Since generative pre-trained transformers only consider information from the left and ignore information from the right, they are known to lack bidirectional context [65], which is crucial for many downstream tasks [53]. Thus, we combine masked modeling (GPT-style models) and causal modeling (BERT-style models) into one framework. This recently developed approach is known in the literature under different names: fill-in-the-middle (FIM) [9], causal masked modeling (CM3) [1, 25], hybrid unidirectional modeling (HybUni) [5], span corruption + language modeling (SCLM) [59]. Conceptually, the idea is to randomly mask out a span (or spans) of tokens, replace them with sentinel tokens, and move the original tokens to the end. This change does not alter the training process based on Equation 1; however, it now allows the model to use the information to the right of the mask to predict the masked positions, which is not possible with traditional unidirectional causal modeling. The hybrid view not only leads to better performance in downstream tasks after fine-tuning but also retains the utility of language modeling [5, 59]. With CM3Leon [73], causal masked modeling has been scaled up to large vision-language models capable of generating and filling both text and images.
Transitive modeling: We now extend several language learning paradigms (masked, prefix, causal, causal masked) with our new technique, which we call transitive modeling. To do this, we return to the problem formulation outlined in Section 1. In particular, it requires us to combine modalities  $A$  and
$C$ , given only the disjoint combinations  $(A, B)$  and  $(B', C)$ . We propose to randomly sample data points  $(a, b), (b, a), (b', c)$ , or  $(c, b')$  during training and apply transitive modeling. For simplicity, we will consider the pair  $(a, b)$ . Since the model is trained to predict modality  $C$  from  $B$ , we can infer the pseudo data point  $\hat{c}$  from  $b$ . Now we train the transformer to predict point  $\hat{a}$  using the newly inferred sample  $\hat{c}$  in a way that minimizes the loss with respect to the original sample  $a$ , ensuring that all our modalities are aligned (Figure 2c). In this way, we model the conditional distribution of the missing pair  $(C, A)$  - which, due to commutativity, is equal to  $(A, C)$ . This idea has some conceptual similarities to cycle consistency. However, our method is much more general since it can be applied to many combinations of modalities as long as there is a linking modality, see Figure 2d.
In fact, upon closer inspection, transitive modeling is quite different from cycle consistency. The most popular version of cycle consistency was introduced in CycleGAN [75]. Given an input  $a$  in domain/modality  $A$ , one wants to generate an aligned output  $b$  in domain/modality  $B$ . The model then computes a discriminative loss  $D$  on  $b$  and a reconstruction loss  $L$  on the original  $a$  and the predicted  $\hat{a}$ . A different version of cycle consistency was used in MCTN [47]: Given aligned inputs  $(a, b)$  from modalities  $A$  and  $B$ ,  $a$  is used to predict  $\hat{b}$ , and the predicted  $\hat{b}$  is used to predict  $\hat{a}$ . Then the reconstruction loss  $L$  is applied to compare  $a$  and  $\hat{a}$  as well as  $b$  and  $\hat{b}$ . Transitive modeling, on the other hand, starts with the aligned modalities  $(a, b)$  and  $(b', c)$ , it uses  $b$  to predict  $\hat{c}$  and  $\hat{c}$  to predict  $\hat{a}$ , the reconstruction loss then compares  $a$  and  $\hat{a}$ . Since we use commutative modeling, the other direction starting from a different sample  $(b', c)$  is also learned. In summary, we have the following three approaches:
- CycleGAN: Given  $a$ , model  $a \rightarrow \hat{b} \rightarrow \hat{a}$  and calculate  $L(a, \hat{a}) + D(\hat{b})$ .
- MCTN: Given  $(a,b)$ , model  $a\rightarrow \hat{b}\rightarrow \hat{a}$ , and calculate  $L(a,\hat{a}) + L(b,\hat{b})$
- LoReTTa: Given  $(a, b)$ , model  $b \to \hat{c} \to \hat{a}$  and calculate  $L(a, \hat{a})$ .
Thus, transitive modeling does not simply add another step to cycle consistency by extending  $a \rightarrow \hat{b} \rightarrow \hat{a}$  to  $a \rightarrow \hat{b} \rightarrow \hat{c} \rightarrow \hat{a}$ . In fact, we model  $b \rightarrow \hat{c} \rightarrow \hat{a}$ . Intuitively, this ensures that the model does not "cheat" by memorizing the input  $a$  in order to reconstruct  $a$ . We also do not use a cycle loss, because  $a$  is not used as an input and output. Thus, our loss  $L(a, \hat{a})$  is not cyclic with respect to the model's inputs. However, we still make sure that the predicted modality is consistent with the data as a whole. This has a similar flavor to cycle consistency, but as can be seen above, it is not the same. In summary, LoReTTa uses masked modeling to transition within modalities, causal modeling to transition between modalities, and vice versa thanks to commutative modeling. Moreover, not all modality combinations need to be available at training time, as we use transitive modeling to predict the missing modalities and connect them to the original data. In our experiments, we show that this approach allows us to learn very rich multimodal feature representations.
We consider a dataset with modalities  $X_{1},\ldots ,X_{N}$  to be aligned if there exists a subset  $S$  sampled from a joint probability distribution  $P(X_{1},\dots,X_{N};Y)$  [37]. Here,  $Y$  represents the label space. This form of alignment, where only labels are shared across modalities, is called weak alignment [14]. Conversely, when there is a functional relationship at the level of physical objects, we consider the alignment to be strong [14]. The definition of strong and weak alignment also extends to datasets where we do not have all modality combinations at once. For example, during data collection, we may recover only the paired samples  $(x_{1},x_{2})$  and  $(x_2',x_3')$ , where  $x_{2}$  and  $x_{2}^{\prime}$  do not belong to the same object/subject. Thus, we have datasets of the form  $(X_{1},X_{2})$  and  $(X_{2},X_{3})$  instead of  $(X_{1},X_{2},X_{3})$ . This is different from the case where the dataset still contains some samples of the form  $(x_{1},x_{2},x_{3})$ , which was considered in previous work. To model the unseen relationship  $(X_{1},X_{3})$  in a self-supervised way, we cannot simply use techniques like masked modeling or causal modeling. The neural network would have to implicitly learn
$$
p \left(x _ {3} \mid x _ {1}\right) = \frac {p \left(x _ {1} , x _ {3}\right)}{p \left(x _ {1}\right)}, \tag {2}
$$
by deriving  $p(x_{1},x_{3}) = \sum_{x_{2}}p(x_{1},x_{2},x_{3})$  and  $p(x_{1},x_{2},x_{3}) = p(x_{3}|x_{2},x_{1})p(x_{2}|x_{1})p(x_{1})$  from the chain rule of probability. Empirically [6, 39, 69], both the unidirectional and bidirectional
Table 1: Perplexity on the SVL-MNIST test set. The subscript describes the pre-training datasets with image (I), text (T), or speech (A), while the blue colors highlight the unseen modality pairs.
Training	Testing
I	T	A	(I, T)	(I, A)	(T, A)
CM3(I)	1.55	-	-	-	-	-
CM3(T)	-	2.05	-	-	-	-
CM3(A)	-	-	18.63	-	-	-
C2M3(I,T)	1.62	2.28	-	1.76	-	-
C2M3(I,A)	1.61	-	20.09	-	4.42	-
C2M3(T,A)	-	2.25	19.85	-	-	10.68
C2M3(I,T), (I,A)	1.64	2.19	21.37	1.76	4.39	30.88
C2M3(T,I), (T,A)	1.65	2.18	19.69	1.76	104.27	11.26
C2M3(A,I), (A,T)	1.64	2.19	23.45	9.70	4.66	11.92
LoReTTa(I,T), (I,A)	1.56	2.17	18.77	1.72	3.14	11.21
LoReTTa(T,I), (T,A)	1.59	2.10	18.95	1.70	5.04	10.03
LoReTTa(A,I), (A,T)	1.60	2.14	20.35	3.13	4.36	10.65
C2M3(I,T,A)	1.55	2.35	19.43	1.72	4.25	10.55
approaches fail this task even within one modality because it is difficult to model  $p(x_{3}|x_{2},x_{1})$  due to the missing set  $(X_{1},X_{2},X_{3})$ . This is where our proposed transitive modeling paradigm comes in.
Suppose we have a dataset of the form  $(X_{1},X_{2}),(X_{2},X_{3}),\ldots ,(X_{N - 1},X_{N})$ , where  $X_{k}$  and  $X_{l}$  are not aligned when  $k + 1\neq l$ . By applying causal modeling to all pairs, we can learn the transition  $X_{1}\rightarrow X_{2}\rightarrow \dots \rightarrow X_{N - 1}\rightarrow X_{N}$ . Adding our proposed commutative modeling, we get an even stronger bidirectional relation:  $X_{1}\leftrightarrow X_{2}\leftrightarrow \dots \leftrightarrow X_{N - 1}\leftrightarrow X_{N}$ . Now, let us interpret each modality as a node on a high-dimensional graph. This gives us a connected graph (minimum spanning tree) where we can walk from one modality to any other modality. Thus, given a real data point  $x_{i}$ , we can use transitive modeling to sample pseudo data points  $x_{j}$  by transitioning from  $x_{i}$  to  $x_{j}$ . Having sampled many pairs  $(x_{i},x_{j})$ , we get a new dataset  $(X_{i},X_{j})$ . It can be used again in the generative pre-training process to learn the new transition  $X_{i}\leftrightarrow X_{j}$ . We apply this idea to all  $i$  and  $j$ . In the end, we learn the conditional distribution between all samples and obtain a fully connected graph. This idea holds for more complicated modality combinations as long as there is a bridging modality, i.e., we have a spanning tree (see Figure 2d). We can also sample all missing tuples  $(x_{k},\ldots ,x_{l})$ . Thus, LoReTTa allows us to approximate any missing joint probability distribution  $P(X_{k},\ldots ,X_{l})$  by sampling and training with the missing modality combinations.
Ideally, there would be no generalization error after pre-training. In practice, however, we would not be able to accurately recover the ground truth of the missing modality if it were available. For the datasets  $(X_{1},X_{2})$  and  $(X_{2},X_{3})$ , the learned generative model  $f$  would make the following errors:
$$
f \left(x _ {1} \mid x _ {2}\right) = x _ {1} + e _ {2, 1}, \quad f \left(x _ {2} \mid x _ {1}\right) = x _ {2} + e _ {1, 2}, \quad f \left(x _ {2} \mid x _ {3}\right) = x _ {2} + e _ {3, 2}, \quad f \left(x _ {3} \mid x _ {2}\right) = x _ {3} + e _ {2, 3} \tag {3}
$$
where all  $e_{i,j}$  are vector-valued error terms. Without loss of generality, we do not keep track of the sign of the error, since we can always write  $e_{i,j}' = -1 \cdot e_{i,j}$ . Assuming that  $f$  is differentiable and using the multivariate Taylor expansion, we obtain
$$
x _ {3} = f \left(x _ {3} \mid f \left(x _ {2} \mid x _ {1}\right) + e _ {1, 2}\right) + e _ {2, 3} \approx f \left(x _ {3} \mid f \left(x _ {2} \mid x _ {1}\right)\right) + J _ {f} \left(x _ {3} \mid f \left(x _ {2} \mid x _ {1}\right)\right) J _ {f} \left(x _ {2} \mid x _ {1}\right) e _ {1, 2} + e _ {2, 3}, \tag {4}
$$