summaries: The authors tackle the problem of time evolution of discrete sets through the view of continuous-time markov chains. The authors provide a theoretical justification of why these problems in general underspecified in the setting of cross-sectional data. They go on to demonstrate that the underspecification can be overcome by use of additional items / features that can help determine time order by deriving a bound on the mean and variance of the observed time for these additional items. An approximate formulation for likelihood maximization is then presented without any assumptions of prior knowledge of the additional items or the structure of the additional items. The theoretical framework is then tested using simulation studies and demonstrate that accuracy of state sequence recovery can be significantly improved with use of the additional items. The framework is also evaluated using a real world cancer dataset from TCGA by ordering of pairs of mutation and copy number aberration events.

summaries: The authors model the time evolution of discrete states (e.g. genetic mutations) using a continuous-time Markov chains. They show that the resulting learning task is generally underspecified given cross-sectional data. The authors suggest including additional independent items that can help determine time order, and resolve this underspecification.  The authors implement an approximate likelihood maximization method for learning continuous-time Markov chains, which shows good performance on high-dimensional data, is faster than competing methods, and shows promising results on synthetic and real cancer data.

summaries: This paper describes a continuous-time Markov Chain method for modeling the time evolution of events. By including independent items the authors are able to increase the limit on the number of items modeled. This research is particularly relevant as larger single-cell sequencing data sets are being collected.

summaries: This paper deals with learning a particular class of multi-component continuous time Markov chains (CTMC) when time stamps are unobserved in the data. In this class binary components change their state irreversibly, rates follow a pairwise interaction model and one component changes at a time. The authors show that modeling independent elements provides information about the posterior distribution of observation time and provide an efficient stochastic approximation to the likelihood’s gradient.

abstractText: Modeling the time evolution of discrete sets of items (e.g., genetic mutations) is a fundamental problem in many biomedical applications. We approach this problem through the lens of continuous-time Markov chains, and show that the resulting learning task is generally underspecified in the usual setting of cross-sectional data. We explore a perhaps surprising remedy: including a number of additional independent items can help determine time order, and hence resolve underspecification. This is in sharp contrast to the common practice of limiting the analysis to a small subset of relevant items, which is followed largely due to poor scaling of existing methods. To put our theoretical insight into practice, we develop an approximate likelihood maximization method for learning continuous-time Markov chains, which can scale to hundreds of items and is orders of magnitude faster than previous methods. We demonstrate the effectiveness of our approach on synthetic and real cancer data.

authors:
L.B. Alexandrov
P.H. Jones
D.C. Wedge
J.E. Sale
P.J. Campbell
S. Nik-Zainal
M.R. Stratton

title: Clock-like mutational processes in human somatic cells

venue: Nature Genetics,

authors:
L.B. Alexandrov
J. Kim
N.J. Haradhvala
M.N. Huang
A.W. Tian Ng
Y. Wu
A. Boot
K.R. Covington
D.A. Gordenin
E.N. Bergstrom
S.M.A. Islam
N. Lopez-Bigas
L.J. Klimczak
J.R. McPherson
S. Morganella
R. Sabarinathan
D.A. Wheeler
V. Mustonen
G. Getz
S.G. Rozen
M.R. Stratton

title: The repertoire of mutational signatures in human cancer

venue: Nature,

authors:
C. Archambeau
M. Opper

title: Approximate inference for continuous-time markov processes

venue: Bayesian Time Series Models

authors:
N. Beerenwinkel
R.F. Schwarz
M. Gerstung
F. Markowetz

title: Cancer evolution: Mathematical models and computational inference

venue: Systematic Biology,

title: Notes on the sum and maximum of independent exponentially distributed random variables with different scale parameters

venue: arXiv, art

authors:
E. Cerami
J. Gao
U. Dogrusoz
B.E. Gross
S.O. Sumer
B.A. Aksoy
A. Jacobsen
C.J. Byrne
M.L. Heuer
E. Larsson
Y. Antipin
B. Reva
A.P. Goldberg
C. Sander
N. Schultz

title: The cBio cancer genomics portal: An open platform for exploring multidimensional cancer genomics data

venue: Cancer Discovery,

authors:
S. Cristea
J. Kuipers
N. Beerenwinkel

title: pathTiMEx: Joint inference of mutually exclusive cancer pathways and their progression dynamics

venue: Journal of Computational Biology,

authors:
D.T. Crommelin
E. Vanden-Eijnden

title: Fitting timeseries by continuous-time markov chains: A quadratic programming approach

venue: Journal of Computational Physics,

authors:
S. Yadlowsky
T. Yun
X. Zhai
D. Sculley

title: Underspecification presents challenges for credibility in modern machine learning

authors:
J. Duchi
E. Hazan
Y. Singer

title: Adaptive subgradient methods for online learning and stochastic optimization

venue: Journal of Machine Learning Research (JMLR),

authors:
M. Gerstung
M. Baudis
H. Moch
N. Beerenwinkel

title: Quantifying cancer progression with conjunctive bayesian networks

authors:
S.C. Sahinalp
Y. Ji
M. Peifer
F. Markowetz
V. Mustonen
K. Yuan
W. Wang
Q.D. Morris
P.T. Spellman
D.C. Wedge
P.V. Loo

title: The evolutionary history of 2,658 cancers

authors:
G. Grimmett
D. Stirzaker

title: Probability and random processes

authors:
C. Jolly
P. Van Loo

title: Timing somatic events in the evolution of cancer

venue: Genome Biology,

authors:
M.D.M. Leiserson
D. Blokh
R. Sharan
B.J. Raphael

title: Simultaneous identification of multiple driver pathways in cancer

venue: PLoS Computational Biology,

authors:
Liu
Y.-Y
S. Li
F. Li
L. Song
J.M. Rehg

title: Efficient learning of continuous-time hidden markov models for disease progression

venue: In Neural Information Processing Systems (NeurIPS),

authors:
K.L. Mengersen
R.L. Tweedie

title: Rates of convergence of the Hastings and Metropolis algorithms

venue: The Annals of Statistics,

authors:
P. Miron
A. Fatima
A. Langerød
Aparicio
S.A.J. R
A. Tutt
A.M. Sieuwerts
Borg
G. Thomas
A.V. Salomon
A.L. Richardson
Børresen-Dale
A.-L
P.A. Futreal
M.R. Stratton
P.J. Campbell

title: The life history of 21 breast cancers

venue: Cell,

authors:
M. Opper
G. Sanguinetti

title: Variational inference for markov jump processes

title: Maximum likelihood trajectories for continuous-time markov chains

authors:
V. Rao
Y. Teh

title: Mcmc for continuous-time discrete-state systems

authors:
B.J. Raphael
F. Vandin

title: Simultaneous inference of cancer pathways and tumor progression from cross-sectional mutation data

authors:
F. Sanchez-Vega
M. Mina
more

title: Oncogenic signaling pathways in the cancer genome

venue: atlas. Cell,

authors:
R. Schill
S. Solbrig
T. Wettig
R. Spang

title: Modelling cancer progression using mutual hazard

authors:
Y. Song
D.P. Kingma

title: How to train your energy-based models

authors:
M.A. Suchard
R.E. Weiss
J.S. Sinsheimer

title: Bayesian selection of continuous-time markov chain evolutionary models

venue: Molecular Biology and Evolution,

authors:
S. Turajlic
N. McGranahan
C. Swanton

title: Inferring mutational timing and reconstructing tumour evolutionary histories

venue: Biochimica et Biophysica Acta (BBA) - Reviews on Cancer,

authors:
T. Watanabe
S. Nobusawa
P. Kleihues
H. Ohgaki

title: IDH1 mutations are early events in the development of astrocytomas and oligodendrogliomas

venue: The American Journal of Pathology,

authors:
J. P
P.A. Alldredge
T.N. Wylie
J.R. Walker
M.A. Watson
S.E. Heath
W.D. Shannon
N. Varghese
R. Nagarajan
J.E. Payton
J.D. Baty
S. Kulkarni
J.M. Klco
M.H. Tomasson
P. Westervelt
M.J. Walter
T.A. Graubert
J.F. DiPersio
L. Ding
E.R. Mardis
R.K. Wilson

title: The origin and evolution of mutations in acute myeloid leukemia

heading: 1 Introduction

text:
Modeling the time evolution of physical processes with discrete states is an important machine learning problem that spans a wide range of application domains, including molecular dynamics (Crommelin & Vanden-Eijnden, 2006), phylogenetics (Suchard et al., 2001), and computational medicine (Liu et al., 2015). Continuous-time markov chains have found success in such modeling tasks, supported by extensive research into learning and inference in such models (Opper & Sanguinetti, 2008; Perkins, 2009; Archambeau & Opper, 2011; Rao & Teh, 2012).
While the majority of previous work has focused on modeling a relatively small number of discrete states, many interesting applications involve the interaction between sets of items, which results in exponentially large state spaces. For example, when modeling the time progression of cancer, we want to consider a number of genetic alterations (e.g., mutations, copy number variations, etc.) that exhibit complex time-related dependencies (Beerenwinkel et al., 2014). To model how sets of n such alterations evolve over time, we need to consider a state space of size 2n. What makes the problem even more challenging is the fact that the available data sets are typically cross-sectional; that is, they consist of (unordered) sets of items observed at unknown time points without any further information about the history of the underlying process.
In this paper, we focus on the problem of learning a particular parametric family of continuous-time Markov chains from such cross-sectional data, and show that the resulting problem is in general underspecified. The issue of underspecification is many-faceted and has been shown to permeate a
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
large number of machine learning systems, often leading to poor generalization, lack of robustness, and spurious relationships when interpreting the resulting models (D’Amour et al., 2020). In our setting, we explore a perhaps surprising remedy: including a number of additional (approximately) independent items can help determine the time order of process events, and hence resolve underspecification. We theoretically show that these extra items act as a “background clock”, since counting the number of their occurrences in a data sample can help us estimate the time at which that sample was observed. In sharp contrast to common practice, which limits the analysis to a small subset of relevant, highly-interacting items, our insight suggests that scaling up the learning procedure can be crucial for the robustness of the inferred models. We thus show that items deemed a priori unimportant to the application at hand may in fact be particularly valuable for recovering the time properties of the underlying physical process.
Existing learning approaches are, unfortunately, not well-suited for practically applying this insight; for example, the state-of-the-art method by Schill et al. (2019) for learning continuous-time Markov chains to model cancer progression, scales exponentially in the number of items considered, thus limiting the analysis to around 20 items. To alleviate this issue, we propose an approximate likelihood maximization method that relies on a fast gradient approximation. On problems of 20 items our approach runs almost 1000 times faster that the state of the art, while it can also scale to problems involving hundreds of items. In experiments on real cancer data, we demonstrate how some previous results may have been artifacts of underspecification, and how scaling up the analysis can result in more robust models of cancer progression.

heading: 2 Problem setup

text:
Given a ground set V = {1, . . . , n}, we consider a continuous-time Markov chain (CTMC) {X(t)}t≥0 on state space 2V (Grimmett & Stirzaker, 2001). Thus, the states of the Markov chain are subsets of V , and can equivalently be identified as binary vectors of size n. CTMCs are commonly represented by a generator matrix Q ∈ R2n×2n . Since the rows and columns of this matrix correspond to states in 2V , we will use qS R to denote the entry of Q at row indexed by S and column indexed by R. For S ̸= R, the entry qS R represents the infinitesimal rate of transitioning from state S to state R, that is,
qS R = lim δt→0+ P(Xt+δt = R |Xt = S) δt , for S ̸= R,
with qS R ≥ 0. The quantity q̃S := −qS S represents the total infinitesimal rate of leaving state S, with q̃S ≥ 0. Since the total rate of leaving S is equal to the sum of the rates of transitioning to any other state R, we have q̃S = ∑ R∈2V qS R, for all S ∈ 2V .
We are interested in modeling physical processes that induce an increasing accumulation of items over time, for example, the accumulation of genetic alterations when modeling cancer progression. To this end, we impose the following two constraints on the CTMC. First, we assume that the chain always starts from the empty state, that is, X(0) = ∅. Second, we only allow transitions that add a single item to the current set, that is, for S ̸= R, qS R ̸= 0 only if R = S ∪ {i} for some i ∈ V . Since items are never removed, this implies that limt→∞Xt = V .
In this paper, we focus on learning the CTMC from a given data set. Since learning directly the exponentially sized generator Q is a hopeless endeavor, we impose a particular parametric form on this matrix, recently introduced by Schill et al. (2019). We define a parameter matrix Θ ∈ Rn×n (also denoted by θ as a vector), and parameterize the off-diagonal entries of the generator matrix as
qS S∪{j}(θ) = exp ( θjj +
∑ i∈S θij
) ,
for all S ⊆ V , and j ∈ V \S. The off-diagonal elements of Θ encode positive (attractive) or negative (repulsive) pairwise interactions between items in V . The presence of item i increases (θij > 0) or decreases (θij < 0) the rate of adding item j by a multiplicative factor of wij := eθij . The diagonal elements θjj encode the intrinsic rate of adding element j when ignoring all other interactions.
Obtaining a sample from a CTMC under all aforementioned constraints can be done using the following sequential procedure. At each step, given the current state S, we draw a “holding time” h
Algorithm 1: Sampling a set from the marginal CTMC Input :Parameters θ Draw tobs ∼ Exp(1) t← 0, S ← ∅ while t < tobs do
Draw h ∼ Exp(q̃S(θ)) Compute pi = qS S∪{i}(θ)/q̃S(θ), for i ∈ V \ S Draw x ∼ Cat(V \ S, (pi)i∈V \S) t← t+ h, S ← S ∪ {x}
return S
from an exponential distribution with parameter q̃S . This represents how much time will pass until we add the next item. Then, we draw the next item x ∈ V \S according to a categorical distribution with probabilities ∝ qS S∪{x}. The result is a sequence of items (σ1, . . . , σn), σi ∈ V , together with a sequence of holding times (h1, . . . , hn), hi ∈ R. Ideally, our data set would consist of a number of such item sequences and holding times; we would then proceed to learn the parameters θ by maximizing the likelihood of our CTMC model. Unfortunately, such detailed data is typically not available in practice. In particular, it is rather common in biomedical applications to only have access to cross-sectional data, i.e., to only observe the current state (in the form of an unordered subset of V ) at a particular point in time without knowing how the process reached that state. Furthermore, it is often the case that the time point of each observation can only be roughly estimated or, worse, is completely unknown.
Similarly to previous work, we assume that we are given a data setD = {S(1), . . . , S(N)}, S(i) ⊆ V , with each S(i) representing a state observed at a potentially different unknown time Tobs. More specifically, we assume that Tobs is a random variable following an exponential distribution with density p(t) = e−t (Gerstung et al., 2009; Schill et al., 2019), and investigate the problem of maximizing the marginal likelihood,
p(D;θ) = N∏ i=1 p(S(i);θ) = N∏ i=1 ∫ ∞ 0 p(S(i) | t;θ)p(t)dt. (1)
Algorithm 1 shows how to sample a set from p(S;θ). The difference from our previous description of sampling from a CTMC lies in the fact that the procedure is only run until we reach the observation time tobs, and the result is an unordered set without any information about the holding times.

heading: 3 The value of “unimportant” items

text: Our ultimate goal is to retrieve the ordered interaction structure present in the data by learning the parameter vector θ. Before attempting to maximize the marginal likelihood of eq. (1), though, it is natural to first ask whether the limited information available in the data is sufficient to infer such ordered interactions in the first place. We start with a simple example to showcase this issue.

heading: 3.1 Warmup example

text:
Consider a ground set V = {1, 2} and true model parameters θ∗11 = θ∗21 = 0, and θ∗12 = −θ∗22 = α, for some α > 0. When α is sufficiently large, the probability of item 2 occurring before item 1 goes to zero, therefore the model encodes the fact that item 1 is a prerequisite (and therefore, always appears before) item 2. As a consequence, we will also have p({2};θ∗) ≈ 0 for the resulting marginal distribution of sets. One may intuitively think that this could be enough information to infer the true time order; the following proposition shows that this is not the case.
Proposition 1. There is a one-dimensional family of parameters θ = θ(s), and s1, s2 ∈ R, such that p(· ;θ(s)) = p(· ;θ∗), for all s1 < s < s2.
The proof follows by simple algebra, and can be found in Appendix A together with an illustration of the parameter family. Interestingly, for s → s1, the model θ(s) encodes the fact that, when we
observe both items, item 2 always appears before item 1, which is the opposite of what is encoded in our true model.

heading: 3.2 Independent items as a background clock

text:
While the above example paints a pessimistic picture for inferring the correct time order, we show that this is still possible given some additional side information. Suppose that we are given another ground set V+ whose items have no interaction with the items in V , that is, the parameter matrix Θfull of the resulting model over Vfull = V ∪ V+ has a block diagonal structure with blocks Θ ∈ R|V |×|V | and Θ+ ∈ R|V+|×|V+|. Then, it is easy to see that the distributions over V and V+ are conditionally independent given time t. That is, for any S ⊆ V , S+ ⊆ V+, we can write p(S ∪S+ | t;θfull) = p(S | t;θ) p(S+ | t;θ+). Note that the same statement does not hold when considering the marginal distributions, i.e., p(S∪S+;θfull) ̸= p(S;θ) p(S+;θ+). Using this conditional independence property, we can rewrite the marginal probability of S ∪ S+ as follows:
p(S ∪ S+;θfull) = ∫ ∞ 0 p(S ∪ S+ | t;θfull)p(t)dt
= ∫ ∞ 0 p(S | t;θ) p(S+ | t;θ+)p(t)dt (by cond. ind.)
= ∫ ∞ 0 p(S | t;θ) p(t |S+;θ+)p(S+;θ+)dt (by Bayes’ rule)
⇒ p(S |S+;θfull) = ∫ ∞ 0 p(S | t;θ) p(t |S+;θ+)dt (dividing with p(S+;θ+)). (2)
Comparing equations (1) and (2), we see that the information about S gained by observing S+ can be explained via a posterior observation time distribution p(t |S+;θ+) that refines the prior p(t) by taking into account S+.
To gain further insight into this time posterior, we analyze in more detail a simplified setup, in which V+ consists of m identically distributed independent items parameterized by θ+, that is, Θ+ = θ+Im. In this case, |S+| | t follows a binomial distribution over m variables with success parameter 1−e−w+t, wherew+ := eθ+ . Intuitively, we expect these independent items to act as a “background clock”: counting the number of items |S+| should give us an indication about the observation time t. Moreover, as we increase the number of items m we should be able to get increasingly accurate estimates of the observation time. The following theorem formalizes these ideas.
Theorem 1. Let S+ ⊆ V be randomly drawn according to the CTMC with parameter matrix Θ+, and let t∗ be the true observation time of S+. Then, for any δ ∈ (0, 1), there exists an m0, such that for all m ≥ m0, the mean and variance of the posterior time distribution p(t |S+;θ) can be bounded as follows with probability at least 1− δ,
|Mpost − t∗| ≤ C1(θ+, t∗) √ logm
m +O
( logm
m ) Vpost ≤ C2(θ+, t∗) 1
m +O
( 1
m2
) .
The quantities C1 and C2 are constant w.r.t. m and encode how suitable a particular rate (quantified through θ+) of an item is for estimating the specific observation time t∗. Intuitively, larger rates are better suited for estimating smaller times and vice versa. We further illustrate this effect in Appendix B together with the detailed proof of the theorem. Note that for m → ∞ we get p(t |S+;θ+) → δ(t− t∗) in distribution, and therefore p(S |S+;θfull)→ p(S | t∗;θ). This means that observing an arbitrarily large amount of independent items is equivalent to knowing the true observation time. In Appendix C we show that a similar asymptotic convergence behavior can be observed even when the items in V+ are not identically distributed or exhibit some dependencies.

heading: 3.3 Discussion

text:
To illustrate how this result applies to our two-item example discussed before, we define Θ+ = Im[θ (1) + , . . . , θ (m) + ] T , where each θ(i)+ is drawn uniformly in [−4,−2]. This corresponds to individual
frequencies of items in S+ that roughly range from 0.01 to 0.1. We then draw N = 500 samples from the CTMC defined by θfull. Figure 1 shows the posterior time density corresponding to each sample with S = {1} (red) and S = {1, 2} (blue), for different numbers m of independent items. To avoid clutter, we do not show the densities for samples with S = ∅. As expected, whenm = 0 all samples follow the prior time density, and we have no way to distinguish their time order. However, as we keep increasing m we notice that the two colors keep separating from each other in time, and it becomes clear that the red posteriors tend to concentrate around earlier observation times than the blue ones. This is a clear indication that item 1 comes before item 2, and the additional time information obtained by the posterior time estimates should help pinpoint the true model within the family discussed in Proposition 1.
Due to computational considerations, it is common practice to constrain the ground set to a small subset of items that interact strongly and are deemed interesting to the application at hand; for example, genetic alterations that are known to have a role in cancer development (Raphael & Vandin, 2015; Schill et al., 2019). Our results in this section challenge this practice. We show that items believed to be independent of the “interesting items”, and thus considered unimportant to the analysis, can in fact be particularly valuable in estimating the observation time of data samples and recovering the time order properties of the underlying physical process.
For the remainder of this paper, we will not assume any prior knowledge about the block-diagonal structure of Θfull or the values of the parameters θ+. Furthermore, we will make no distinction between “items of interest” and independent items. If there is indeed a block-diagonal parameter structure in the data, the addition of a regularization term that promotes sparsity can help recover this structure in the presence of noise. For this reason, similarly to what was used by Schill et al. (2019), the final objective we maximize in practice is the L1-regularized marginal log-likelihood,
F (D;θ) = 1 N N∑ d=1 log p(S(d);θ)− λ ∑ i̸=j |θij |. (3)

heading: 4 Efficient approximate likelihood maximization

text: To put the previously discussed theoretical insight into practice, we need to be able to efficiently maximize the above objective for ground sets V containing potentially hundreds of items. When using a first-order method for optimization, the bulk of the required computation is devoted to obtaining the gradient of the marginal log-likelihood log p(S;θ) with respect to the parameters θ. Computing the exact gradient requires O(2n) computation; we propose here a method to compute a gradient approximation in a much more efficient manner. We start by deriving expressions for the likelihood in some simplified setups, and then use these results as building blocks for our setup of interest.

heading: 4.1 Full sequences

text:
First, we consider a setup in which we observe the arrival order of the items in the CTMC, which corresponds to building up an ordered sequence σ instead of a set S in Algorithm 1. Furthermore, we assume that tobs = +∞, which means that the returned sequence will be a permutation of the ground set V . By definition of the CTMC, the probability of observing a sequence σ = (σ1, . . . , σn) ∈ SV , where SV denotes the permutation group over V , can be written as
p(σ;θ) = n∏
i=1
P ( Xi = σ[i] | Xi−1 = σ[i−1] ) = n∏ i=1 qσ[i−1] σ[i](θ) q̃σ[i−1](θ) . (4)
We denote subsequences using the bracket notation, σ[i] := (σ1, . . . , σi), with σ[0] = (). Also, with a slight abuse of notation we directly use sequences in some places that require sets; therefore, the notation q̃σ is to be understood as q̃set(σ), where set(σ) := {σi | i ∈ {1, . . . , |σ|}}.

heading: 4.2 Partial sequences given time

text:
As a next step, instead of tobs = +∞, we assume that we are given the observation time tobs, and would like to derive the probability of observing a partial sequence σ = (σ1, . . . , σk), with k ≤ n, at time tobs. We can express this as the probability of the intersection of two events in the space of outcomes of the continuous-time Markov chain. Event A identifies the outcomes in which the first k ordered elements agree with our partial sequence σ. Event B identifies the outcomes for which the observation time tobs falls between the addition of the k-th and k + 1-th elements; that is, Tk < tobs < Tk+1. We can then express the desired probability as p(σ | tobs;θ) = P (A)P (B | A). The probability of event A can be directly written using (4), except that, in this case the product is taken from i = 1 to k, instead of n. Regarding event B, first note that the sequence of events( {Ti < tobs | A} )n i=1 is decreasing. As a result, we can write
P (B | A) = P (Tk < tobs < Tk+1 | A) = P (Tk < tobs | A)− P (Tk+1 < tobs | A) . (5)
By definition of the CTMC, each holding time Hi is an exponentially distributed random variable whose parameter depends on the i − 1 items that have been added up to that point. We can also define the jump time Tk := ∑k i=1Hi, which represents the time at which the k-th change of state occurs. When we condition on knowing the first k items of the sequence, i.e., event A, the jump times Tk and Tk+1 are distributed as sums of k and k+1 independent exponential random variables respectively. The CDF of the sum of r independent exponential variables with rates λ1, . . . , λr, also known as the hypoexponential distribution (Bibinger, 2013), is given by
FHEXP(y;λ1, . . . , λr) =
( r∏
i=1
λi
) r∑
i=1
1− e−λiy λi ∏ j ̸=i(λj − λi) .
In our case, the rates λi of each exponential distribution are q̃[σi](θ), thus we can compute the terms of (5) as P (B | A) = FHEXP(tobs; q̃[σ1], . . . , q̃[σk])− FHEXP(tobs; q̃[σ1], . . . , q̃[σk+1]).

heading: 4.3 Marginal partial sequences

text:
Next we consider the probability of observing a partial sequence σ = (σ1, . . . , σk), with k ≤ n, without knowledge of the observation time tobs. We define events A and B as before, except that in
this case the observation time Tobs is a random variable, and event B is defined by Tk < Tobs < Tk+1. In particular, we assume that Tobs is an exponential random variable with rate λobs = 1. Rather than naively integrating over time the probability derived before, we can obtain the following greatly simplifed expression by making use of the memoryless property of exponential random variables (Bertsekas & Tsitsiklis, 2008) when computing P (B | A).
Proposition 2. The marginal probability of a partial sequence σ = (σ1, . . . , σk) can be written as
p(σ;θ) =
( k∏
i=1
qσ[i−1] σ[i](θ) 1 + q̃σ[i−1](θ)
) 1
1 + q̃σ[k](θ) . (6)
The proof can be found in Appendix D, but the form of this expression allows for an intuitive interpretation. We can modify the original continuous-time Markov chain by adding an extra dummy state Xobs, such that for all X ⊆ V , qX Xobs(θ) = 1, and qXobs X(θ) = 0. The first property implies that there is a fixed rate 1 to transition from any state to Xobs, while the second property implies that Xobs is a terminal state. It is easy to see then, that eq. (6) expresses exactly the probability of observing a “full” sequence σ in this modified chain. Note that a “full” sequence in this chain is no longer a permutation of V , but rather any sequence that ends with Xobs.

heading: 4.4 Marginal sets

text:
Finally, we consider the probability of observing a set S ⊆ V , without knowledge of the order in which the individual items arrived, and still under the assumption that the observation time Tobs is an exponential random variable with rate λobs = 1. Making use of the result in the previous section, we can compute this probability by summing over all partial sequences that are permutations of S, that is, p(S;θ) = ∑ σ∈SS p(σ;θ). However, this computation is infeasible for all but very small set sizes, since it requires summing over |S|! terms. To alleviate this problem, we propose a method to approximate the gradient of log p(S;θ). This gradient can be written as
∇θ log p(S;θ) = 1
p(S;θ) ∇θp(S;θ) =
1
p(S;θ) ∑ σ∈SS ∇θp(σ;θ) = ∑ σ∈SS p(σ;θ) p(S;θ) ∇θ log p(σ;θ).
In the last expression, we observe that the gradient corresponding to each permutation σ is weighed by the probability of observing that permutation of the given set S. This suggests a stochastic approximation of the desired gradient by first sampling m permutations σ(1), . . . , σ(M) according to p(· |S;θ) := p(· ;θ)/p(S;θ), and then computing the average of the resulting gradients,
∇θ log p(S;θ) ≈ 1
M M∑ i=1 ∇θ log p(σ(i);θ). (7)
Conceptually similar approaches have been used to approximate the gradient of the log-normalizer when maximizing the likelihood of energy-based probabilistic models (Song & Kingma, 2021).
It remains to show how to obtain samples from p(· |S;θ), which we do by using a Markov chain Monte Carlo method on SS . More concretely, we employ a Metropolis-Hastings chain (Levin et al., 2009); at each time step, given the current permutation σ, the chain proposes a new permutation σnew according to proposal distribution Q(σnew |σ), and transitions to σnew with probability
paccept = min ( 1, p(σnew |S;θ)Q(σ |σnew;θ) p(σ |S;θ)Q(σnew |σ;θ) ) .
A simple choice for Q is the uniform distribution over all permutations regardless of the current state. We instead employ a more sophisticated proposal that is based on the form of the marginal probability (6), and leads to faster and more stable learning in practice. While a detailed discussion and theoretical analysis of mixing is beyond the scope of this paper, in Appendix E we present our proposal in detail, and compare it to the uniform.

heading: 5 Further related work

text:
There has been a long line of research focused on understanding and reconstructing the evolutionary history of tumors (Nik-Zainal et al., 2012; Welch et al., 2012; Turajlic et al., 2015; Mitchell
et al., 2018; Jolly & Van Loo, 2018). However, these papers do not attempt to systematically model the timed interaction of genetic alterations across a data set. More recently, Gerstung et al. (2020) showed that modern data resulting from deep sequencing procedures may contain partial order information about specific genetic events. Incorporating this additional information into an analysis like ours, could potentially help with further reducing underspecification. Finally, there is evidence that specific classes of mutational processes generate mutations at a constant rate, and are largely unaffected by other factors (Alexandrov et al., 2015; Campbell et al., 2020; Alexandrov et al., 2020). The resulting so-called mutational signatures have been likened to a biological clock, and could be candidates for explicitly defining the set V+ discussed earlier (cf. Section 3.2).

heading: 6 Experiments

text: The following experimental setup is common to all our experiments. To optimize the objective (3), we use a proximal AdaGrad method (Duchi et al., 2011). We fix the initial step size to η = 1, and the regularization weight to λ = 0.01. To initialize the parameter matrix Θ, we train a diagonal model for 50 epochs, and then draw each off-diagonal entry from Unif([−0.2, 0.2]). For the gradient approximation (7), we use M = 50 samples in addition to 10 burn-in samples that are discarded. The code used to run our experiments can be found at https://github.com/3lectrologos/time/ tree/clean.

heading: 6.1 Synthetic data

text:
We start by evaluating our approach on two synthetic data sets drawn from known true models. The first data set is identical to the one discussed in Figure 1, and contains observations of two items that tend to follow a specific time order (item 1 comes before item 2). The second data set contains observations of five items following the graphical structure shown on the upper right of Figure 2. Edges with arrows ← represent positive interactions (θij > 0), while edges with vertical lines ⊢ represent negative interactions (θij < 0). This type of structure, containing groups of items with bidirectional negative interactions, is common in genomics data; in fact, this exact structure can be found in the model learned from real cancer data in our following experiment (see Appendix F).
Our goal is to evaluate how well we can recover the true model by optimizing the regularized marginal likelihood; in particular, we want to examine the effect of adding a number m of independent items to the data. To quantify the quality of recovery, we compare the true distribution ptrue of marginal sequences, computed using eq. (6), to the distribution of marginal sequences p∗ of the learned model, approximated by drawing 106 marginal sequences from that model. In Figure 2 we depict the KL divergence dKL(p∗ ∥ ptrue) for increasing numbers m of extra items. To compare against the ideal case, we also train a model by maximizing the regularized marginal likelihood given the true observation times (cf. Section 4.2). The KL divergence between this model and the true one is shown as a dashed line in the figure. We repeated each experiment 96 times; sources of randomness include the choice of independent items, the parameter initialization, and the gradient approximation. We see that the recovery quality improves for both data sets with increasing m, and rapidly approaches the ideal. This shows that observing a large enough number of additional items is practically equivalent to knowing the true observation times, and thus validates Theorem 1.

heading: 6.2 Real cancer data

text:
We next evaluate our approach on a data set that contains N = 378 tumor samples of glioblastoma multiforme, an aggressive type of brain cancer. The data is part of the TCGA PanCancer Atlas project (https://www.cancer.gov/tcga), and we obtained a preprocessed version via cBioPortal (Cerami et al., 2012). We followed some further filtering procedures from previous work (Leiserson et al., 2013) and ended up with a ground set of n = 410 items, each of which represents a point mutation, amplification, or deletion of a specific gene. Analyzing the interaction structure between such genetic alterations is of fundamental importance to cancer research, as it can help illuminate the processes that are involved in cancer initiation and progression.
For the remainder of this section, we assume that the items in V are ordered by decreasing frequency, and when selecting a subset we keep the most frequent items. In Table 1, we show that our learning approach runs about 1000 times faster on a ground set of size n = 20 compared to the exponentially scaling exact approach proposed by Schill et al. (2019).
In the absence of ground truth, we use the following procedure to evaluate how some of the learned parameters behave as we increase the size of the ground set. Given a pair of items a, b ∈ V , and n−2 other (most frequent) items, we first learn a CTMC model on the n total items using our approach. Next we approximate the sequence distribution of just a and b (marginally over time and over all other items) by sampling 106 marginal sequences from the learned model. Finally, we compute the proportion of sequences in which a occurred before b, given that both where observed; when this proportion is greater than 0.5, we infer that a is more likely to occur before b, and vice versa. In Figure 3 we present these computed proportions for four chosen pairs of genetic alterations that are interesting both from a biological standpoint, as well as in terms of learning behavior.
First, we discover two significant ordered interactions, shown in the top row of the figure, which were not reported in previous work analyzing glioblastoma data (Raphael & Vandin, 2015; Cristea et al., 2017; Schill et al., 2019). Our results suggest that amplification events for genes EGFR and PDGFRA tend to occur before their respective point mutation events. Interestingly, the cooccurrence of amplifications and mutations for EGFR has been observed before (Leiserson et al.,
2013; Sanchez-Vega et al., 2018), but there is little known about their time order. Note that the time order of the PDGFRA alterations can be robustly inferred only after including more than 70 additional items in the analysis.
Second, we confirm some previously reported interactions, for example, the tendency of IDH1 mutations to occur before TP53 ones, as shown in the bottom left of the figure. This was observed before by Schill et al. (2019), and is also supported by some biological evidence (Watanabe et al., 2009). Again, we can see that our result becomes robust only after n = 50. This indicates that the previous observation of this time order may have been in part due to a fortuitous choice of the optimization setup, e.g., fixed initialization (see Appendix G). The final result in the bottom right of the figure seems to support this claim: the amplifications of MDM2 and CDK4 do not seem to have a clear time order, even though previous work has inferred both that CDK4(A) occurs before MDM2(A) (Cristea et al., 2017), as well as the opposite order (Schill et al., 2019).
Finally, the parameter matrix Θ learned from this data set (see Appendix G) has indeed an approximately block-diagonal structure, which confirms the practical relevance of Theorem 1.

heading: 7 Conclusion

text:
Our experiments demonstrated that, in the presence of underspecification, special care is required when interpreting results that may be significantly affected by particular choices in the optimization setup. In this paper, we proposed an effective way to mitigate such effects, namely scaling up the problem and taking into account additional items that are commonly discarded as insignificant. We hope that our approach will inspire further work into combating underspecification and learning robust time evolution models.
Acknowledgments and Disclosure of Funding
This work was partially supported by a Swiss National Science Foundation Early Postdoc.Mobility fellowship, by NSF CAREER award 1553284, NSF BIGDATA IIS-1741341, and NSF award 1717610, as well as a grant from the US National Cancer Institute (1R35CA220523).

title: Scaling up Continuous-Time Markov Chains Helps Resolve Underspecification


summaries: This paper proposes a self-supervised pre-training framework for document understanding called UniDoc. UniDoc takes multimodal data (image feature and text feature from document) as input and uses gated cross attention for learning cross-modal correlation. The superiority of UniDoc has been empirically shown on several downstream tasks.

summaries: This paper proposes a new unified framework for pre-training image/text encoders for capturing both visual and textual features in documents. The main Transformer model is pre-trained with multimodal embeddings as input features (visual & textual features) through three pre-training tasks: Masked Sentence Modeling, Visual Contrastive Learning, and Vision-Language Alignment. The pre-trained model can be fine-tuned on downstream tasks on which the proposed UniDoc model outperforms baselines.

summaries: This paper proposes a multi-modal pretraining method for document understanding. It involves a new hierarchical model architecture and three pretraining tasks. It compares with previous SOTA methods on three tasks and achieves stronger performance on two of them.   The major contribution is the newly proposed multi-modal pretraining tasks and the architecture is a combinations of models from both domains.

summaries: This paper presents UniDoc as a multimodal pre-trained model for document understanding. UniDoc consists of 3 layers: a feature extraction layer, which extracts sentence embeddings and region embeddings from the original document; a fusion layer, which fuses textual and visual information based on a gated cross-attention mechanism; a task layer, which uses 3 tasks as pre-training objectives to optimize model parameters. Evaluations are conducted on 4 downstream tasks, including Form Understanding, Receipt Understanding, Document Classification and Document Object Detection. Comparing to several SOTA baselines, UniDoc performs better on Form Understanding and Receipt Understanding, but worse on Document Classification. The paper also includes some ablation studies to investigate the effects of V/L/V+L, pre-training tasks and visual backbones, and some interesting future directions.

abstractText: Document intelligence automates the extraction of information from documents and supports many business applications. Recent self-supervised learning methods on large-scale unlabeled document datasets have opened up promising directions towards reducing annotation efforts by training models with self-supervised objectives. However, most of the existing document pretraining methods are still language-dominated. We present UDoc, a new unified pretraining framework for document understanding. UDoc is designed to support most document understanding tasks, extending the Transformer to take multimodal embeddings as input. Each input element is composed of words and visual features from a semantic region of the input document image. An important feature of UDoc is that it learns a generic representation by making use of three self-supervised losses, encouraging the representation to model sentences, learn similarities, and align modalities. Extensive empirical analysis demonstrates that the pretraining procedure learns better joint representations and leads to improvements in downstream tasks.

authors:
Zichao Yang
Diyi Yang
Chris Dyer
Xiaodong He
Alex Smola
Eduard Hovy

title: Hierarchical attention networks for document classification

authors:
Jade Goldstein
Vibhu O Mittal
Jaime G Carbonell
Mark Kantrowitz

title: Multi-document summarization by sentence extraction

venue: In NAACL-ANLP Workshop,

authors:
Zihang Dai
Zhilin Yang
Yiming Yang
William W Cohen
Jaime Carbonell
Quoc V Le
Ruslan Salakhutdinov

title: Transformer-xl: Attentive language models beyond a fixed-length context

authors:
Yiheng Xu
Minghao Li
Lei Cui
Shaohan Huang
Furu Wei
Ming Zhou

title: Layoutlm: Pre-training of text and layout for document image understanding

venue: In SIGKDD,

authors:
Yang Xu
Yiheng Xu
Tengchao Lv
Lei Cui
Furu Wei
Guoxin Wang
Yijuan Lu
Dinei Florencio
Cha Zhang
Wanxiang Che

title: Layoutlmv2: Multi-modal pre-training for visuallyrich document understanding

venue: In ACL,

authors:
Peizhao Li
Jiuxiang Gu
Jason Kuen
Vlad Morariu
Handong Zhao
Rajiv Jain
Varun Manjunatha
Hongfu Liu

title: Selfdoc: Self-supervised document representation learning

authors:
Seunghyun Park
Seung Shin
Bado Lee
Junyeop Lee
Jaeheung Surh
Minjoon Seo
Hwalsuk Lee

title: Cord: A consolidated receipt dataset for post-ocr parsing

venue: In NeurIPS Workshop,

authors:
Adam W Harley
Alex Ufkes
Konstantinos G Derpanis

title: Evaluation of deep convolutional nets for document image classification and retrieval

venue: In ICDAR,

authors:
Jiuxiang Gu
Zhenhua Wang
Jason Kuen
Lianyang Ma
Amir Shahroudy
Bing Shuai
Ting Liu
Xingxing Wang
Gang Wang
Jianfei Cai

title: Recent advances in convolutional neural networks

venue: Pattern Recognition,

authors:
Jiasen Lu
Dhruv Batra
Devi Parikh
Stefan Lee

title: Vilbert: Pretraining task-agnostic visiolinguistic representations for vision-and-language tasks

venue: In NeurIPS,

authors:
Jiuxiang Gu
Jason Kuen
Shafiq Joty
Jianfei Cai
Vlad Morariu
Handong Zhao
Tong Sun

title: Self-supervised relationship probing

authors:
Jia Deng
Wei Dong
Richard Socher
Li-Jia Li
Kai Li
Li Fei-Fei

title: Imagenet: A large-scale hierarchical image database

venue: In CVPR,

authors:
Weijie Su
Xizhou Zhu
Yue Cao
Bin Li
Lewei Lu
Furu Wei
Jifeng Dai

title: Vl-bert: Pre-training of generic visual-linguistic representations

venue: In ICLR,

authors:
Gen Li
Nan Duan
Yuejian Fang
Daxin Jiang
Ming Zhou

title: Unicoder-vl: A universal encoder for vision and language by cross-modal pre-training

venue: In AAAI,

authors:
Akshay Gadi Patil
Omri Ben-Eliezer
Or Perel
Hadar Averbuch-Elor

title: Read: Recursive autoencoders for document layout generation

venue: In CVPR Workshop,

title: Fast r-cnn

authors:
Herve Jegou
Matthijs Douze
Cordelia Schmid

title: Product quantization for nearest neighbor search

authors:
Aaron van den Oord
Oriol Vinyals

title: Neural discrete representation learning

authors:
Alexei Baevski
Henry Zhou
Abdelrahman Mohamed
Michael Auli

title: wav2vec 2.0: A framework for self-supervised learning of speech representations

authors:
Eric Jang
Shixiang Gu
Ben Poole

title: Categorical reparameterization with gumbel-softmax

authors:
Frederick Tung
Greg Mori

title: Similarity-preserving knowledge distillation

venue: In ICCV,

authors:
D. Lewis
G. Agam
S. Argamon
O. Frieder
D. Grossman
J. Heard

title: Building a test collection for complex document information processing

venue: In SIGIR,

authors:
G. Jaume
H. Kemal Ekenel
J. Thiran

title: Funsd: A dataset for form understanding in noisy scanned documents

venue: In ICDAR Workshop,

authors:
Xu Zhong
Jianbin Tang
Antonio Jimeno Yepes

title: Publaynet: largest dataset ever for document layout analysis

authors:
Nils Reimers
Iryna Gurevych

title: Sentence-bert: Sentence embeddings using siamese bertnetworks

venue: In EMNLP,

authors:
Adina Williams
Nikita Nangia
Samuel R Bowman

title: A broad-coverage challenge corpus for sentence understanding through inference

venue: In NAACL,

authors:
Daniel Cer
Mona Diab
Eneko Agirre
Inigo Lopez-Gazpio
Lucia Specia

title: Semeval-2017 task 1: Semantic textual similarity-multilingual and cross-lingual focused evaluation

authors:
Diederik P Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

authors:
Rafał Powalski
Łukasz Borchmann
Dawid Jurkiewicz
Tomasz Dwojak
Michał Pietruszka
Gabriela Pałka

title: Going full-tilt boogie on document understanding with text-image-layout transformer

title: Tesseract: An open-source optical character recognition

venue: engine. Linux J.,

heading: 1 Introduction

text:
Document intelligence is a broad research area that includes techniques for information extraction and understanding. Unlike plain-text documents in natural language processing (NLP) [1, 2], a physical document can be composed of multiple elements: tables, figures, charts, etc. In addition, a document usually includes rich visual information, and can be one of various types of documents (scientific paper, form, resume, etc.), with various combinations of multiple elements and layouts. Complex content and layout, noisy data, font and style variations make automatic document understanding very challenging. For example, to understand text-rich documents such as letters, a system needs to focus almost exclusively on text content, paying attention to a long sequential context, while processing semi-structured documents such as forms requires the system to analyze spatially distributed short words, paying particular attention to the spatial arrangement of the words. Following the success of BERT [3] on NLP tasks, there has been growing interest in developing pretraining methods for document understanding [4, 5, 6]. Pretrained models have achieved state-of-the-art (SoTA) performance across diverse document understanding tasks [7, 8].
Huge training datasets help pretraining models to learn a good representation for downstream tasks. However, we observe three major problems with the current pretraining setup: (1) documents are composed of semantic regions. Most of the recent document pretraining works follow BERT and split documents into words. However, unlike the sequence-to-sequence learning in NLP, documents have a hierarchical structure (words form sentences, sentences form a semantic region, and semantic regions form a document). Also, the importance of words and sentences are highly context-dependent, i.e., the same word or sentence may have different importance in a different context. Moreover, current transformer-based document pretraining models suffer from input length constraints. Also, input
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
length becomes a problem for text-rich documents or multi-page documents. (2) documents are more than words. The semantic structure of the document is not only determined by the text within it but also the visual features such as table, font size and style, and figure, etc. Moreover, the visual appearance of the text within a block are often overlooked. Most of recent BERT-based pretraining works only take the words as input without considering multimodal content and alignment of multimodal information within semantic regions. (3) documents have spatial layout. Visual and layout information is critical for document understanding. Recent works encode spatial information via 2D position encoding and model spatial relationships with self-attention, which computes attention weights for long inputs [4, 5]. However, for semi-structured documents, such as forms and receipts, words are more related to their local surroundings. This corresponds strongly with human intuition – when we look at magazines or newspapers, the receptive fields are modulated by our reading order and attention. Based on the above observations, we ask the following question: Can unified document pretraining benefit all of these different kinds of documents?
We propose a unified pretraining framework for document understanding, shown in Fig. 1. Our model integrates image information in the pretraining stage by taking advantage of the transformer architecture to learn cross-modal interactions between visual and textual information. To handle textual information, we encode sentences using a hierarchical transformer encoder. The first level of the hierarchical encoder models the formation of the sentences from words. The second level models the formation of the document from sentences. With the help of the hierarchical structure, UDoc learns how words form sentences and how sentences form documents. Meanwhile, it reduces model computation complexity exponentially and increases the number of input words. This also mimics human reading behaviors since the sentence/paragraph is a reasonable unit for people to read and understand—people rarely check the interactions between arbitrary words across different regions in order to understand an article. Convolution has been very successful in the extraction of local features that encode visual and spatial information [9], so we use convolution layers as a more efficient complement to self-attention for addressing local intra-region dependencies in a document image. Meanwhile, self-attention uses all input tokens to generate attention weights for capturing global dependencies. Thus, we combine convolution with self-attention to form a mixed attention mechanism that combines the advantages of the two operations.
We depart from previous vision-language pretraining [10, 11] by extracting both the textual and visual features for each semantic region. We propose a novel gated cross-attentional transformer that enables information exchange between modalities. A visually-rich region (figure, chart, etc) may have stronger visual information than textual information. Instead of treating outputs from both modalities identically, we design a gating mechanism that can dynamically control the influence of textual and visual features. This approach enables cross-modal connections and allows for variable highlight the relevant information in visual and textual modality and enables cross-modal connections. During pretraining, the CNN-based visual backbone and multi-layer gated cross-attention encoder are jointly trained in both pretraining and fine-tuning phases.
Our contributions are summarized as follows: (1) We introduce UDoc, a powerful pretraining framework for document understanding. UDoc is capable of learning contextual textual and visual information and cross-modal correlations within a single framework, which leads to better performance. (2) We present Masked Sentence Modeling for language modeling, Visual Contrastive Learning for vision modeling, and Vision-Language Alignment for pretraining. (3) We present extensive experiments and analyses to validate the effectiveness of the proposed UDoc. Extensive experiments and analysis provide useful insights on the effectiveness of the pretraining tasks and show outstanding performance on various downstream tasks.

heading: 2 Related Work

text:
Self-supervised learning has shown great success in producing generic representations that learn from large-scale unlabeled corpora [3]. Like the development of pretraining in computer vision [12] and NLP [3], there has been a surging interest in self-supervised learning for Vision-Language (VL) tasks [10, 13, 14, 11]. Transformers [3] are the key technology that enables learning contextualized representations from large-scale unlabeled training data. The unique characteristics of document images (spatial layout and multiple elements) distinguish document image pretraining from pretraining works in NLP and VL domains. In the NLP domain, the inputs are pure texts without spatial layouts (bounding boxes). In the VL domain, the inputs are the visual objects and captions. While for
document images, the input elements are spatially distributed, and the visual and textual information co-occur within the semantic regions.
Several recent works have explored pretraining on document images [4, 5, 15]. LayoutLM [4] extends BERT to learn contextualized word representations for document images through multi-task learning. It takes a sequence of OCR words as input during pretraining and incorporates the 2D position embedding as input for each token. However, LayoutLM only considers textual information during pretraining without modeling the alignment between visual and textual information–visual information is only incorporated into the model during the fine-tuning stage. The most recent version, LayoutLMv2 [5], improves on this by incorporating the image encoder into pretraining and jointly training the image encoder along with the BERT model. LayoutLMv2 splits the document image into several parts and concatenates the visual embeddings and text embeddings into a single sequence. Apart from masked language learning (MLM), LayoutLMv2 also considers image-text alignment and image-text matching during pretraining. The most related work to ours is SelfDoc [6], which proposes a multimodal document pretraining framework. It first extracts the document object proposals from pre-trained Faster R-CNN [16] and then applies OCR for each proposal to get the words. It takes the pre-extracted RoI features and sentence embeddings as input, and models the perform learning over the textual and visual information using the cross-modality encoder.
There is a noticeable difference between our proposed method, UDoc, and other concurrent works in document image pretraining. UDoc is a multimodal end-to-end pretraining framework for document images. Unlike the fixed document object detector in [6], the parameters of the image encoder with RoI align, which derive the visual features for semantic regions, are also updated in UDoc. In contrast to [5], our visual features come from the semantic regions instead of splitting the image into fixed regions. Like the object-level semantic elements in natural images, for document images, we represent the typical document layout elements such as paragraph, title, figure, and table as semantic regions. Moreover, to learn the contextualized visual representations, UDoc masks visual information in the latent space and learns contextualized representations by solving a contrastive learning task defined over a quantization of the latent visual embeddings.

heading: 3 Method

heading: 3.1 Model Architecture

text:
Fig. 1 illustrates our approach, UDoc, which consists of four components: feature extraction, feature embedding, multi-layer gated cross-attention encoder, and pretraining tasks. Given a document image and the locations of document elements (sentence or RoI), UDoc takes image regions and words that correspond to each document elements as inputs, and extracts their respective embeddings through a visual feature extractor and a sentence encoder. These embeddings are then fed into a transformer-based encoder to learn the cross-modal contextualized embeddings that integrate both visual features and textual features.
In the feature extraction step, we first employ an off-the-shelf OCR tool [17] to extract text from a document image I, where the words are grouped into sentences S = {s1, . . . , sN} whose corresponding bounding boxes are P = {p1, . . . , pN}. For each sentence bounding box pi, we use a ConvNet-based backbone fImEnc and RoI Align [18] fRoIAlign to extract the pooled RoI features vi. To obtain a feature embedding, we extract the sentence embedding si for each sentence si via a pretrained sentence encoder fSentEnc. Each region’s RoI feature vi is discretized into a finite set of visual representations vQi ∈ VQ via product quantization [19]. The multi-layer Gated Cross-Attention encoder takes the position information, masked visual features Ṽ and masked textual features S̃ as inputs, and then it generates the contextualized multimodal representations (HlV and H l S , l ∈ [1, L]) and outputs the predicted features (V̂ and Ŝ), where L is the number of stacked transformer blocks.
More formally, the pretraining procedure can be decomposed into the following steps:
I OCR−−→ ( P S ) fImEnc+fRoIAlign−−−−−−−−→ fSentEnc ( V,VQ S ) −−−→ fMask ( Ṽ S̃ ) −→ ( HlV HlS ) −→ ( V̂ Ŝ ) −→ LPretraining (1)
where fMask denotes the masking function that randomly masks RoI features and sentence embeddings with the respective probabilities pvMask and p s Mask. LPretraining is composed of three pretraining tasks: Masked Sentence Modeling (MSM), Visual Contrastive Learning (VCL), and Vision-Language Alignment (VLA). Next, we provide details mentioned in Eq. 1.
Feature Extraction and Embedding. Formally, a document image I ∈ RW×H consists of N regions, where each region’s bounding box is characterized by a 6-d vector, as pi = {xLTW , yLT H , xRB W , yRB H , w W , h H }, where w and h are of the width and height the region, W and H are the width and height of I, while (xLT, yLT) and (xRB, yRB) denote the coordinates of the top-left and bottom-right corners respectively. The 6-d vector is mapped into a high-dimensional representation via a linear mapping function.
The visual embedding is the sum of the mapped RoI feature and position embedding. Likewise, textual embedding is the sum of sentence embedding and position embedding. We also have different types of segments to distinguish different modalities. The input sequence to the transformer-based encoder starts with a special start element ([CLS] and full visual features), then it is followed by multimodal elements, and it ends with a special ending element ([SEP]+full visual features). For the special elements ([CLS] and [SEP]), the corresponding full visual features are features extracted from the whole input image, by applying fImEnc to an RoI covering the whole input image.
Quantization Module. Unlike the fixed image encoder in [6], we jointly learn the image encoder in an end-to-end fashion alongside the multimodal model. A visual representation can be learned by predicting the visual features of the masked regions, but it is challenging to predict such features exactly, since they are unconstrained and of continuous representation. To constrain the representation space of the visual features and facilitate the end-to-end learning of image encoder (see Task #2 in Sec. 3.2), we follow [20, 21] and use vector quantization to discretize the visual features V = {v1, . . . ,vN} into a finite set of representations VQ = {vQ1 , . . . ,v Q N}. Specifically, we define latent embedding spaces e ∈ RC×E , where C is the number of codebooks, and E is the number of entries for each codebook. For each vi, we first map it to logits v`i ∈ RC×E , and calculate the probability for the j-th codebook entry in i-th group as pc,e = exp((v`c,e + ge)/τ)/ ∑E k=1 exp((v ` c,k + gk)/τ), where τ is a non-negative temperature, g1:E are i.i.d samples drawn from Gumbel(0,1) distribution. During the forward pass, we choose one entry vector from each codebook by ẽi ∼ argmaxepc,e and generate the quantized representation vQi by a concatenation of {ẽ1, . . . , ẽG} which is then followed by a linear transformation. During the backward pass, the gradients are computed through a Gumbel-Softmax estimator [22].
Gated Cross-Attention. To model the interactions among multimodal inputs, we introduce a multimodal transformer with gated cross-attention to model the cross-modality relationships. Let Hl+1m be output features at the l-th layer for one modality m, and let n be another modality (m,n ∈ {V, S}). We obtain the features at (l + 1)-th layer as:
Hl+1m = fLN ( fLN ( Hlm + f l Cross-Att(H l m|Hln) ) + f lFF ( fLN(H l m + f l Cross-Att(H l m|Hln)) )) (2)
where fLN denotes layer normalization [23]. The feed-forward sub-layer fFF in Eq. 2 is further composed of two fully-connected sub-layers, both wrapped in residual adds and fLN.
The core part of Eq. 2 is the cross-attention fCross-Att(·). Given the intermediate representations Hlm and Hln, the cross-attention output for modality m is computed as:
fCross-Att(H l m|Hln) = [Cross-Att1(Hlm|Hln); . . . ; Cross-Atth(Hlm|Hln)]U (3) Cross-Atti(Hlm|Hln) = softmax ( f iq(H l m)f i k(H l n) T / √ d ) f iv(H l n) (4)
where f iq(H l m), f i k(H l n), and f i v(H l n) are the query, key, and value calculated by linear mapping layers for the i-th head. d is the model dimension, h is the number of heads, and U is the weight matrix that combines the outputs of the heads.
Considering the substantial diversity of document images and the different information needs of differing document types, we use a gating mechanism [24] to dynamically weight the outputs of the visual and textual branches. Specifically, we feed the concatenated the visual and textual features to a non-linear network fGate([Hl+1m ;H l+1 n ]), which generates the modality-specific attention weights α l m and αln, and returns the weights separately to their respective modality-specific branches to perform element-wise products. We multiply the features for modality m with its modality-specific attention weight, and compute the updated feature as: Hl+1m = H l+1 m (1 + α l m), same that for modality n.

heading: 3.2 Training Tasks and Objectives

text:
The full pretraining objective of UDoc (right block in Fig. 1) is defined as: LPretraining = LMSM + LVCL + LVLA. In the rest of this section, we describe each task in detail.
Task #1 : Masked Sentence Modeling. This task is similar to the MLM task utilized in BERT. The key difference is that we mask sentences instead of tokens. During pretraining, each sentence and RoI of the input document is randomly and independently masked. For the masked sentence, its token is replaced with a special sentence of [MASK]. The model is trained to predict the masked sentence feature, based on the unmasked words and the visual features. The goal is to predict the masked sentence embeddings based on the contextual information from the surrounding sentences and image regions, by minimizing the smooth L1 loss [16]:
LMSM(Θ) = ∑ i smoothL1(si − fUDoc(si|s\i, Ṽ)) (5)
where Θ is the trainable parameters and fUDoc(.) outputs the unmasked textual feature, s\i is the surrounding features for the i-th input, Ṽ are the image features with random masking.
Task #2 : Visual Contrastive Learning. We learn visual feature representations by solving a visual contrastive learning task which requires estimating the true quantized latent RoI representation. Given a prediction v̂i ∈ V̂ for the masked RoI ṽi ∈ Ṽ, the model needs to estimate the positive quantized representation vQi in a set of quantized candidate representations V
Q. Good representations are learned by maximizing the agreement between output representation and quantized representation of the same RoIs as follows:
LVCL(Θ) = − ∑ ṽi∈Ṽ ( log exp(sim(v̂i,v Q i )/κ)∑ vQj exp(sim(v̂i,v Q j )/κ) ) + λ 1 CE C∑ c=1 E∑ e=1 pc,e log pc,e (6)
where sim(·, ·) computes the cosine similarity between two vectors, λ is a hyperparameter, and κ is a temperature scalar. The second term encourages the model to use the codebook entries more equally.
Task #3 : Vision-Language Alignment. To enforce the alignment among different modalities, we explicitly encourage alignment between words and image regions via similarity-preserving knowledge distillation [25]. Note that, unlike the text-image alignment in LayoutLMv2 [5] which splits the image into four regions and predicts whether the given word is covered or not on the image side, we align the image and text belonging to the same region. The goal is to minimize the differences
between the pairwise similarities of sentence embeddings and the pairwise similarities of image region features:
LVLA(Θ) = 1
N ×N ||fNorm(S · S>)− fNorm(HLV ·HL>V )||2F (7)
where S is the unmasked input sentence embeddings, HLV is the mapped visual representations of the final layer, || · ||F is the Frobenius norm, and fNorm performs L2 normalization.

heading: 4 Experiment

heading: 4.1 Pretraining UDoc

text:
Pretraining corpus. We build our pretraining corpus based on IIT-CDIP Test Collection 1.0 [26], which contains more than 11M scanned document images. To differentiate pretraining from finetuning, we filter out the document images of RVL-CDIP [8] from IIT-CDIP since it is a subset of IIT-CDIP, and sample 1M document images as our pretraining corpus.
Table 1 shows the dataset statistics. IIT-CDIP only provides the OCR texts in XML format. We extract words and their locations by applying EasyOCR [17] on document images. As shown in Fig. 3 (a), EasyOCR provides two kinds of output modes: non-paragraph and paragraph. The paragraph mode groups the non-paragraph results into text regions. We think document image pretraining should be treated differently than sequence-based pretraining in NLP, since the words in the document (2D) are arranged according to spatial layouts, while the words in
NLP corpora are sequential (1D). Considering the special characteristics of documents (complex layout, multi-pages) and the limited input length of BERT models, it is not intuitive to formulate the input at the word level. Hence, we adopt the paragraph-level outputs as the basic input elements since textual regions provide semantically more meaningful information than independent words.
There are some advantages to our design: (1) the region-level design hierarchically encodes document elements and this facilitates the modeling of latent relationships at the region level which has higher-level semantics than the word level. (2) the hierarchical encoding also overcomes the input size limitation of word-level BERT-based models [4, 5]. Fig. 2 shows the distribution of words per region on RVL-CDIP. It can be seen that even though we consider regionlevel input, for some semi-structured documents, single-words dominate the inputs; this somehow forces UDoc to pay attention to word-level in-
puts. Unlike MLM that predicts the masked word, UDoc predicts the textual embedding of the masked input with MSM.
Pretraining setting. We initialize the sentence encoder fSentEnc with BERT-NLI-STSb-base [29] pretrained for NLI [30] and STS-B [31]. The ResNet-50 backbone in the image encoder is pretrained on the PubLayNet training set [28]. All the parameters (except fSentEnc and fImEnc) are randomly initialized. During pretraining, we freeze the parameters of fSentEnc and jointly train the visual encoder and multi-modal UDoc model in an end-to-end fashion. Such an end-to-end training allows the ConvNet and Transformer to realize their full potentials in spatial and sequence modeling for pretraining. UDoc contains 12 layers of gated cross-attention transformer blocks. We set the hidden size to 768 and the number of heads to 12, the maximum number of regions N to 64, and the maximum input sequence length for fSentEnc to 512. The pretraining is conducted on 8 NVIDIA Tesla
V100 32GB GPUs with a batch size of 64. It is trained with Adam optimizer [32], with an initial learning rate of 10−5, weight decay of 10−4, and learning rate warmup in the first 20% iterations.
To learn a useful multimodal representation, random masking is applied to both textual and visual inputs. For MSM, we set the mask probability psMask for input sentences to 15%. 80% among the masked sentences are replaced by special sentence [CLS, MASK, SEP], while 10% sentences are replaced by random sentences sampled from other documents, and 10% remains unchanged. For VCL, the λ is set to 0.1, κ is set to 0.1, the mask probability pvMask is set to 7.5% and the masked RoI features are filled with zeros. The temperature τ is annealed from 2.0 to 0.5 by a factor of of 0.999995 at every iteration. We select the pretraining checkpoint with the lowest LPretraining for finetuning stage.

heading: 4.2 Finetuning Tasks

text:
Form Understanding. Form understanding requires the model to predict the label for each semantic entity. We use FUNSD [27] as the evaluation dataset. It contains 149/50 training/testing images. Fig. 3 (b) shows a sample from FUNSD. Each semantic entity comprises a list of words, a label, and a bounding box. The officially-provided OCR texts and bounding boxes are used during training and testing. We take the semantic entities as input and feed the concatenated visual and textual output representations to a classifier. We apply cross-entropy loss for finetuning. The model is finetuned for 100 epochs with a learning rate of 10−5 and batch size of 16. All the parameters except fSentEnc are trained. One of question, answer, header or other is predicted for each semantic entity. We use entity-level F1 score as the evaluation metric.
Receipt Understanding. Receipt understanding requires the model to recognize a list of text lines with bounding boxes. The performance on this task is evaluated on CORD [7] dataset. It contains 626/247 receipts for training/testing. The receipts are labeled with 30 types of entities under 4 categories: company, date, address, and total. Like FUNSD, we feed the concatenated visual and textual output representations to the classifier. The model is finetuned for 200 epochs with a batch size of 16 and a learning rate of 10−5. The evaluation metric is entity-level F1 score.
Document Classification. Document classification involves predicting the category for each document image. We use RVL-CDIP [8] as the target dataset. It consists of 320K/40K/40K training/validation/testing images under 16 categories. The OCR words and bounding boxes are extracted by EacyOCR. To fine-tune UDoc on RVL-CDIP, we compute the overall representation as an elementwise product between the visual and textual representations averaged from all sentences/regions, and learn a classifier on top of the overall representation with cross-entropy loss. We fine-tune the model for 30 epochs with a batch size of 64 and a learning rate of 10−5. Classification accuracy over 16 categories is used to measure model performance.
Document Object Detection. Document object detection involves decomposing a document image into semantic units. We evaluate the effectiveness of our pretrained visual backbone on PubLayNet [28]. As shown Fig. 3 (d), the documents in PubLayNet are scientific articles. PubLayNet consists of 336K/11K training/validation images with six category labels (text, title, list, figure, and
table). We train Faster-RCNN (F-RCNN) using Detectron2 [33] and initialize the visual backbone with the pretrained ResNet-50 from UDoc. The model is trained for 180k iterations with a base learning rate of 0.01 and a batch size of 8. Mean average precision (MAP) @ intersection over union (IOU) [0.50:0.95] of bounding boxes is used to measure the performance.

heading: 4.3 Results and Discussion

text:
The importance of multimodal learning. To study the effect of multimodal learning, we experiment in three different settings (1) Vision only (V): this setting omits the textual components of UDoc and adopts multilayer self-attention transformer to learn the visual representation. (2) Language only (L): this setting omits the visual encoder and keeps only the textual components. (3) Vision-Language (V+L): this setting considers both vision and language information. We first train three settings without pretraining. Table 2 shows consistent improvement across tasks for V+L over the single-stream baselines (V or L). This demonstrates that our UDoc model is able to learn important visual-linguistic relationships that benefit downstream tasks even without pretraining.
In Table 2, we find that visual information dominates the performance of document classification, while language information contributes a lot to form understanding and receipt understanding. The results also indicate that different document tasks rely on different information. For document entity recognition tasks, language information is more important than visual features. As can be seen in Fig. 3 (b) and (c), entity recognition is more word-oriented. On the other hand, document classification is more focused on global-level understanding. As a result, visual and layout information contribute a lot to the final prediction of the document classification model. This matches well with the innate abilities of humans to distinguish between document types without fully understanding the words. We also observe that gated cross-attention (V+L) achieves a better performance than the non-gated version (V+L]), as its gating mechanism can learn to adaptively determine how much each modality contributes to the output features.
Effect of pretraining tasks. We analyze the effectiveness of different pretraining settings through ablation studies over FUNSD, CORD, and RVL-CDIP, which are representative document benchmarks. Table 2 ablates the key design choices in pretraining UDoc. For experimental efficiency, UDoc models evaluated here are trained with 5 epochs on 300k training corpus. Overall, the pretraining of UDoc consistently improves the performance over all three downstream tasks. The improvement gains vary among different tasks.
We first establish two baselines: MSM+MVM in Table 2 indicates the combination of masked sentence learning and masked visual feature prediction. Similar to MSM, for MVM, we freeze the visual backbone and perform masked visual feature prediction via RoI-feature regression. MSM+VCL jointly trains the visual backbone end-to-end with contrastive learning. As shown in Table 2, MSM+MVM achieves better results than the model without pretraining. Furthermore, when combining VCL together with MSM, consistent performance gains are observed across all the benchmarks. Among the three finetuning tasks, the improvements on FUNSD and CORD are bigger than on RVL-CDIP. We think the local context modeling capability of the ConvNet-based image encoder brings more benefits to entity recognition, since entities are heavily linked and correlated to their local surroundings. When MSM, VCL, and VLA are jointly trained, we observe further performance gains across all the benchmarks. For VCL, instead of sampling the negatives from the same input document, we also try including the negative samples from other document images of the same batch. However, we find that sampling negatives from the entire batch of document images hurt the performance. This is likely because the negatives from other document images are easy to distinguish from each other. We
also consider the image-text matching task (Rel) [5], and combine Rel with MSM+VCL+VLA. It hurts the performance of all three downstream tasks. We conjecture that the image-text matching task introduces mismatched pairs of image and OCR texts as negative examples that potentially hamper the training of other tasks.
What if Masked Language Modeling is included? To study the feasibility of that, we consider MLM during pretraining. Since the number of words may be very large, we select the tokens by randomly applying a sliding window (window size 128) across all sequenced OCR words. Each word is formulated as a single-word sentence ([CLS] [Token] [SEP]). We randomly mask 15% of those sampled words ([CLS] [MASK] [SEP]) and concatenate them along with the region-based inputs. During pretraining, we add the word prediction head on top of UDoc and predict the masked words. We conduct finetuning experiments on entity recognition tasks (FUNSD and CORD), and find that such a direct combination hurts the performance: FUNSD: 87.38 (UDoc) vs. 83.76↓ (UDoc+MLM), CORD: 98.75 (UDoc) vs. 98.63↓ (UDoc+MLM). There are consistent performance drops from adding MLM. One possible reason is that the RoI features extracted by token bounding boxes might not be discriminative enough due to the tiny word-level bounding boxes.
Performance Comparison with SoTA. We further pretrain UDoc on 1M document images with 5 epochs and report the finetuning results in Table 3. UDoc outperforms previous models on FUNSD and CORD, by a significantly large margin, demonstrating that our proposed approach is highly effective, partially due to the end-to-end training of the image encoder that improves the semantic alignments between images and texts. Note that UDoc is pretrained on a subset of IIT-CDIP (1M document images), which is considerably less than the 11M document images used in LayoutLM [4] and LayoutLMv2 [5]. TILT [34] builds a 1.1M pretraining corpus by combining RVL-CDIP, UCSF Industry Documents Library, and Common Crawl. UDoc also achieves promising results on document classification. Note that both LayoutLM v2 and TILT use Microsoft OCR, which is a commercial service with a stronger OCR performance than EasyOCR, which is used in our experiments. We find that OCR plays a key role in document classification performance. As shown in Fig. 4, UDoc
performs the best on the ‘email’ category but worst on the ‘form’ category. We also report the results with different OCR engines: 93.42 (Tesseract [35]) vs. 93.96 (EasyOCR [17]) vs. 94.10 (Google OCR [36]). UDoc with EasyOCR achieves a better performance than with Tesseract since EasyOCR is powered by an advanced neural network, while Tesseract is based on less sophisticated techniques. Since different tasks require task-specific input embeddings to perform well, instead of finetuning the sentence encoder during pretraining, we explore unfreezing the sentence encoder during the finetuning stage (named as UDoc∗) and report the results in Table 3. Unsurprisingly, we see performance improvements on several downstream applications. E.g., RVL-CDIP: 93.96 (UDoc) vs. 95.05↑ (UDoc∗). However, this also makes the training more challenging in terms of computational resources and training time.
Effect of visual backbone. Additionally, we apply the trained visual backbone to document object detection on PubLayNet. The performance of the F-RCNN on the validation set is depicted in Table 4. To better compare, we establish two F-RCNN models with: (1) backbone initialized with
ResNet-50 pretrained on ImageNet; (2) backbone initialized from UDoc’s pretrained visual backbone. It can be seen that our pretrained backbone outperforms ImageNet-pretrained backbones. By leveraging UDoc, we can train different variants of the visual backbone and apply them to document-specific downstream applications, without relying on incompatible pretrained backbones from other domains (e.g., natural image). Moreover, the visual backbone of UDoc does not require any custom layers, and thus any ConvNet architecture can be used in place of ResNet.

heading: 5 Conclusion, Limitations, and Future Works

text:
We develop UDoc, a unified pretraining framework for document understanding. Our model introduces a novel joint training framework that effectively exploits the visual and textual information during pretraining and finetuning. We evaluate the UDoc comprehensively on three downstream tasks: form understanding, receipt understanding, and document image classification. Extensive empirical analysis demonstrates that the pretraining procedure can take advantage of multimodal inputs. Also, it can effectively aggregate and align visual and textual information of document images with the proxy tasks. This work has a broader impact on document applications. By finetuning the pretrained UDoc on task-specific data, document processing systems can provide better results and reduce the expensive data annotations costs. In terms of negative social impact, the document images used for pretraining may contain sensitive information and therefore the models trained on such data may inappropriately leak some private information. To address the privacy leakage, it is worthwhile to explore the combination of privacy-preserving learning and self-supervised learning.
There are interesting short- and long-term research directions for UDoc: (1) we freeze the sentence encoder during pretraining and fine-tuning phases due to computational constraints. A better document representation can be learned by jointly training the sentence encoder, visual backbone and cross-attention encoder in a completely end-to-end fashion. (2) Although impressive performance has been achieved in document entity recognition tasks such as form and receipt understanding, the classification accuracy on semi-structured documents such as forms is still inferior to that of rich-text documents. It is possible to devise a better method to model the spatial relationship among words. (3) An interesting direction is to extend UDoc to multipage/multilingual document pretraining. Additionally, there exist many text-based labeled document datasets in the NLP domain, such as document summarization. Can we transfer the knowledge learned from the text-based document domain to the image-based document domain? How to unify the pretraining of the pure-text document (1D) and image-based document (2D) in a single framework is also worth to try. Lastly, the use of different OCR tools is one of the major sources of inconsistency among the existing document pretraining works. It is worthwhile and essential to build standardized pretraining document image datasets with preprovided OCR results. In addition to scanned documents, using digital PDF as part of the pretraining data is a direction worth exploring since it provides rich metadata which could be beneficial for multimodal learning.

title: Unified Pretraining Framework for Document Understanding