capbtabboxtable[][]

# Structure Inducing Pre-Training

Matthew B. A. McDermott<EMAIL_ADDRESS>
CSAIL, MIT Brendan Yap<EMAIL_ADDRESS>
CSAIL, MIT Peter Szolovits<EMAIL_ADDRESS>
CSAIL, MIT Marinka Zitnik<EMAIL_ADDRESS>
DBMI, Harvard University

###### Abstract

We present a theoretical analysis from first principles that establishes a
novel connection between relational inductive bias of pre-training and fine-
tuning performance while providing an extended view on general pre-training
models. We further explore how existing pre-training methods impose relational
inductive biases, finding that the vast majority of existing approaches focus
almost exclusively on modelling relationships in an intra-sample manner,
rather than a per-sample manner. We build upon these findings with simulations
and empirical studies on standard benchmarks spanning 3 data modalities and 10
downstream tasks. These investigations validate our theoretical analyses, and
provides a recipe to produce new pre-training methods which incorporate
provably richer inductive biases than do existing methods in line with user
specified relational graphs.

##  1 Introduction

The pre-training (PT)/fine-tuning (FT) learning paradigm has had tremendous
impact on machine learning (Devlin et al., 2019; Deng et al., 2009; Brown et
al., 2020a). In PT/FT, we pre-train an encoder
$f_{\boldsymbol{\theta}}:\mathcal{X}\to\mathcal{Z}$ which maps our domain of
interest $\mathcal{X}$ into a latent space $\mathcal{Z}$. This encoder
$f_{\boldsymbol{\theta}}$ is then transferred for use in various downstream
tasks (which are not known at pre-training time). We evaluate PT/FT systems
via the transfer performance of $f_{\boldsymbol{\theta}}$. This transfer
performance is inherently a question of geometry; in particular whether or not
the geometry of $\mathcal{Z}$ is constrained by the pre-training process to
reflect relationships of interest for the fine-tuning tasks.

Despite the importance of the latent space geometry, we posit that many recent
large-scale PT/FT models do not explicitly consider the latent space geometry
during pre-training. While we discuss this argument fully in Section 2,
consider as a motivating example the RoBERTa PT model (Liu et al., 2019b).
This model pre-trains a transformer neural network via a _masked language
modelling_ (MLM) task, in which the model is fed input sequences of tokens,
some of which are masked, and tasked to recover the identity of the masked
tokens. While masked language modelling imposes strong constraints on how the
representations of individual _tokens_ relate to one another, nothing in this
task explicitly constrains how the representations of _samples_ (e.g.,
sentences) relate to one-another. Concretely, consider trying to use this
model for sentiment analysis (Maas et al., 2011). This task would be
implemented via binary classification leveraging the pre-trained, per-sample
(e.g., the [CLS]) embeddings; however, we can make no guarantees about whether
or not the PT process will have ensured these sample embeddings meaningfully
reflect positive vs. negative sentiment (Figure 1).111Note that in NLP, we can
leverage the flexibility of natural language to sidestep this problem;
however, the same flexibility is not possible in other domains, so the lack of
explicit PT objectives focusing on the latent space geometry at a per-sample
granularity is still a marked problem. See Section 2 for further commentary.

Present work. Here we propose a framework that enables us to examine existing
PT methods and establish a new model class for structure inducing PT,
specifically with regards to per-sample latent space geometry. In particular,
we introduce a new perspective on the traditional PT paradigm, in which the PT
objective is subdivided into two components: first, a pure language-model
inspired imputation/denoising objective that leverages the intra-sample/per-
token relationships, and, second, a loss term explicitly driven to regularize
the geometry of the per-sample latent space $\mathcal{Z}$ and reflect the
connectivity patterns of a user-specified PT graph, $G_{\mathrm{PT}}$. By
relying on graphs to capture the structure we wish to induce in $\mathcal{Z}$,
this PT paradigm can capture diverse relationships, such as those motivated by
external knowledge (e.g., Zitnik et al. (2019)), self-supervised constraints
extracted from the data (e.g., Wang et al. (2019a); Hu et al. (2020a)), or
graphs summarizing distances between samples in an alternate modality.
Moreover, this new PT framework is simultaneously specific enough to allow us
to make theoretical guarantees about how different PT graphs impact the FT
performance and general enough to encompass a variety of existing PT systems,
including BERT (Devlin et al., 2019), ALBERT (Lan et al., 2019), RoBERTa (Liu
et al., 2019b), MT-DNN (Liu et al., 2019a), PLUS (Min et al., 2020), and many
others.

Our work advances PT/FT research through three major contributions. First, we
establish theoretical results quantifying how the structure of the graph
$G_{\mathrm{PT}}$ relates to FT task performance. Crucially, this
formalization in our new PT paradigm offers insight into when PT will or won’t
add value over supervised learning alone. Second, we show that our new
framework for PT provides a recipe for pre-training with rich, per-sample
latent space geometry constraints. Specifically, we devise three new PT
methods leveraging different kinds of PT graphs and demonstrate their utility
on three data modalities and 10 FT tasks, showing in all cases that these new
methods perform at or above the level of comparable baseline methods. Finally,
third, we include a discussion on how this loss subdivision creates
opportunities for new research directions that are of interest to the PT/FT
research community.

The rest of our work proceeds as follows. In Section 2, we discuss the extent
to which existing PT methods (in particular language model derived PT methods)
do or do not regularize per-sample latent space geometry. We show that this
research direction remains under-explored or neglected by comparing the
existing research on incorporating _per-token_ knowledge into PT systems to
our complementary goal of regularizing _per-sample_ structure. Next, in
Section 3, we formally define and examine our new PT paradigm, detailing
precisely how we decompose the PT loss and what constraints we place on the
different components of our new PT formulation. In Section 4, we illustrate
theoretically how we can formally relate PT graph structure to FT performance
under our PT paradigm. In Section 5 we further validate these theoretical
findings in semi-synthetic experiments. In Section 6, we further empirically
examine our theoretical model on real-world datasets spanning 3 modalities and
10 FT tasks. Our analysis shows across all three data modalities, including
protein sequences, natural language, and networks, that leveraging per-sample
latent space constraints yields PT methods that match or outperform existing
PT methods. In Section 7 we discuss the significance of out PT paradigm, and
outline the new research directions it prompts. Finally, in Sections 8 and 9,
we discuss related literature and conclude the work.

Figure 1:  Language model pre-training methods typically constrain the _per-
token_ latent space through the language modelling objective. However, they
impose no or only much weaker constraints on the _per-sample_ latent space.

##  2 Pre-training & Latent-space Geometry

In this section, we offer more details on our motivating hypothesis introduced
in Section 1. First, we define both the per-token and per-sample latent
spaces, and offer commentary on why this distinction matters particularly
outside of the NLP domain. Second we present an analysis of over 27 different
language model (LM) or LM-derived PT methods exploring how their objective
functions impose structure across both types of latent space. This analysis
crystallizes our claim that methods to regularize the per-sample latent space
are under-explored, and motivates our new PT paradigm in Section 3.

###  2.1 Per-sample vs. Per-token Latent Space

Consider a PT model $f_{\boldsymbol{\theta}}$ trained using a masked language
modelling objective (MLM), under which certain tokens in the sentence are
masked, and the encoder is trained to recover the masked tokens. Given an
input sentence
$\boldsymbol{s}_{j}=w_{1}^{(j)},w_{2}^{(j)},\ldots,w_{n}^{(j)}$,
$f_{\boldsymbol{\theta}}$ will output both per-token embeddings of the
individual tokens $w_{i}^{(j)}$ and a whole-sample embedding of the
sentence/sample $\boldsymbol{s}$. At a minor risk of ambiguous notation, we
will denote these via
$f_{\boldsymbol{\theta}}(w_{i}^{(j)}|\boldsymbol{s}_{j})$ denoting the per-
token embedding of token $i$ and $f_{\boldsymbol{\theta}}(\boldsymbol{s}_{j})$
denoting the whole-sample embedding of $\boldsymbol{s}_{j}$. For the BERT
model (Devlin et al., 2019), for example,
$f_{\boldsymbol{\theta}}(\boldsymbol{s}_{j})$ will be given by the output
embedding of the [CLS] token and
$f_{\boldsymbol{\theta}}(w_{i}|\boldsymbol{s}_{j})$ will be given by the
output embedding of the $i$-th token.

In this setting, there are two ways to interpret the “latent space”
$\mathcal{Z}$ of $f_{\boldsymbol{\theta}}$. First, we can take the space
induced by the embeddings of free-text tokens individually,
$\mathcal{Z}^{(\text{T})}=\\{f_{\boldsymbol{\theta}}(w_{i}|\boldsymbol{s}_{j})\forall
i,j\\}$. This is the _per-token latent space_. Second, we can take the space
induced by the embedding of the entire sentence/sample,
$\mathcal{Z}^{(\text{S})}=\\{f_{\boldsymbol{\theta}}(\boldsymbol{s}_{j})\forall
j\\}$, which is the _per-sample latent space_. In the rest of this paper, we
will use $\mathcal{Z}$ in general to refer to either of these latent space, as
indicated based on the surrounding text.

Both of these spaces are very different and are useful in different contexts;
for a task like named entity recognition, where the unit of classification is
single or short span of tokens, analyzing the per-token latent space will be
more informative, whereas for a task like sentiment analysis, where the unit
of classification is an entire sample, the per-sample latent space would be
preferred (Devlin et al., 2019).

###  2.2 NLP vs. other domains

In Figure 1, we illustrated that even while a PT objective like masked
language modelling will directly enrich the per-token space
$\mathcal{Z}^{(\text{T})}$, it will not necessarily similarly constrain the
per-sample latent space $\mathcal{Z}^{(\text{S})}$, and that this lack of per-
sample latent space constraints could yield weaker models at FT time. One
seeming contradiction to this is that methods like RoBERTa (Liu et al., 2019b)
have offered significant performance improvements within NLP for both per-
token and per-sample tasks.

In fact, this observation does not contradict our hypothesis, but instead
reflects a unique advantage of the natural language modality that does not
apply in other contexts. In particular, in the NLP domain, we can leverage the
flexibility of language to sidestep any deficit in the per-sample latent space
by re-framing per-sample tasks as per-token, language modelling tasks.
Significant literature exists documenting this phenomena, through the lenses
of both prompting, cloze-filling models, and general theory (Brown et al.,
2020a; Saunshi et al., 2021; Sanh et al., 2021). LM-inspired models applied in
non-NLP domains lack this ability to re-frame per-sample tasks as per-token
tasks. Accordingly, existing work on leveraging in non-NLP domains often
augment LM with auxiliary, per-sample level supervised tasks to improve
performance (Hu et al., 2020b; Yoon et al., 2020; McDermott et al., 2020; Min
et al., 2020).

###  2.3 Categorization of Language Model PT Methods

In Table 1, we survey 27 distinct PT methods and categorize how their
objective functions constrained the respective latent spaces. Our analysis
indicates that constraining the per-sample latent space geometry is under-
explored relative to the per-token latent space. A variety of methods for
constraining the per-token latent space have been explored, including direct
LM, entity and relation masking, KG alignment methods, joint token and entity
embeddings, attention over external knowledge bases, and embedding aggregation
via known relations. Despite this breadth of methods focused on improving
modelling at the per-token level, across all methods save for KEPLER (Wang et
al., 2019c), only single- or multi-task classification is used to regularize
the per-sample latent space. Further details can be found in Appendix Section
B.

Method |  Masked or standard language modelling |  Named entity masking |  Relation masking$\dagger$ |  Per-token knowledge graph alignment |  Named entity recognition and linking |  (Unconstrained) attention over a KG |  Joint token and entity embeddings |  Aggregating embedding across relations$\dagger$ |  Single-task classification |  Multi-task classification |  Whole-sample knowledge graph alignment
---|---|---|---|---|---|---|---|---|---|---|---
Per-token | Per-sample
GPT-series (Radford et al., 2018, 2019; Brown et al., 2020a) | ✓ |  |  |  |  |  |  |  |  |  |
RoBERTa (Liu et al., 2019b) | ✓ |  |  |  |  |  |  |  |  |  |
BERT (Devlin et al., 2019) | ✓ |  |  |  |  |  |  |  | ✓ |  |
ALBERT (Lan et al., 2019) | ✓ |  |  |  |  |  |  |  | ✓ |  |
MT-DNN (Liu et al., 2019a) | ✓ |  |  |  |  |  |  |  |  | ✓ |
ERICA (Qin et al., 2020) | ✓ |  |  | ✓ |  |  |  |  |  |  |
COLAKE (Sun et al., 2020a) | ✓ | ✓ | ✓ |  |  |  |  |  |  |  |
LUKE (Yamada et al., 2020) | ✓ | ✓ |  |  |  |  |  |  |  |  |
ERNIE (Sun et al., 2019) | ✓ | ✓ |  |  |  |  |  |  |  |  |
ERNIE 2.0 (Sun et al., 2020b) | ✓ | ✓ |  |  |  |  |  |  |  | ✓ |
ERNIE 3.0 (Sun et al., 2021) | ✓ | ✓ | ✓ |  |  |  |  |  |  | ✓ |
ERNIE (Zhang et al., 2019) | ✓ | ✓ |  |  |  |  | ✓ |  | ✓ |  |
KnowBERT (Peters et al., 2019) | ✓ |  |  |  | ✓ | ✓ | ✓ |  |  |  |
KgPLM (He et al., 2020) | ✓ | ✓ |  |  |  |  |  |  | ✓ |  |
KEPLER (Wang et al., 2019c) | ✓ |  |  |  |  |  |  |  |  |  | ✓
JAKET (Yu et al., 2020) | ✓ |  |  |  |  |  | ✓ | ✓ |  |  |
SMedBERT (Zhang et al., 2021b) | ✓ |  |  | ✓ |  |  | ✓ | ✓ | ✓ |  |
KeBioLM (Yuan et al., 2021) | ✓ |  |  |  | ✓ | ✓ | ✓ |  |  |  |
Coke (Su et al., 2021) | ✓ | ✓ |  |  |  |  | ✓ | ✓ | ✓ |  |
MG-BERT (BehnamGhader et al., 2021) | ✓ |  |  |  |  |  |  | ✓ |  |  |
EHR PT (McDermott et al., 2020) | ✓ |  |  |  |  |  |  |  |  | ✓ |
Graph PT (Hu et al., 2020b) | ✓ |  |  |  |  |  |  |  |  | ✓ |
MG-BERT (Zhang et al., 2021c) | ✓ |  |  |  |  |  |  |  |  |  |
PLUS (Min et al., 2020) | ✓ |  |  |  |  |  |  |  | ✓ |  |
TAPE (Rao et al., 2019) | ✓ |  |  |  |  |  |  |  |  |  |
Table 1:  A subset of existing pre-training methods, broken down by how they
constrain per-token and per-sample latent space geometries. This list excludes
a number of methods that augment LMPT with additional, fine-tuning time
procedures to incorporate knowledge, as they do not change the PT step but
only the FT step. † These models rely on entity identification from source
text and on relation incorporation. As such, they struggle to generalize to FT
tasks where relations are unknown.

##  3 Re-framing the Pre-training Paradigm

In Section 2, we saw that there has been comparably little study into how to
regularize the per-sample latent space for LM and LM-inspired PT methods. In
this section, we re-cast the PT paradigm in order to explicitly separate per-
token and per-sample latent space regularization. We begin by outlining the PT
problem formally, then introduce our new PT paradigm, clarifying explicitly
how methods within our framework are constrained and how their PT objectives
must operate, then conclude with example methods that can be realized within
our framing.

###  3.1 Current Pre-training Problem Formulation

Given a pre-training dataset
$\boldsymbol{X}_{\mathrm{PT}}\in\mathcal{X}^{N_{\mathrm{PT}}}$, a pre-training
method aims to learn an encoder $f_{\theta}:\mathcal{X}\to\mathcal{Z}$ such
that $f_{\theta}$ can be transferred to downstream tasks that are typically
unknown at pre-training time. Note that we can leverage additional data at PT
time to inform the training of $f_{\theta}$ (e.g., PT specific labels
$\boldsymbol{Y}_{\mathrm{PT}}$), but $f_{\theta}$ itself _must_ take only
samples from $\mathcal{X}$ as input so it can be effectively transferred to
the fine-tuning setting. Pre-training methods attempt to solve this problem by
training $f_{\theta}$ via a pre-training loss $\mathcal{L}_{\mathrm{PT}}$ over
the dataset $\boldsymbol{X}_{\mathrm{PT}}$. In language model (LM) or LM
derived PT system, this loss $\mathcal{L}_{\mathrm{PT}}$ is typically some
variant of a masked imputation loss, which we note in Section 2 only
explicitly regularizes the per-token latent space, not the per-sample latent
space.

###  3.2 Re-casting Pre-training Problem Formulation

To investigate pre-training methods, we re-cast the above mentioned pre-
training formulation as follows. As in the traditional setting, our input here
is a dataset $\boldsymbol{X}_{\mathrm{PT}}\in\mathcal{X}^{N_{\mathrm{PT}}}$,
we aim to learn an encoder $f_{\theta}:\mathcal{X}\to\mathcal{Z}$ that enables
effective transfer performance to FT tasks, and $f_{\theta}$ will be learned
via a PT task defined with loss $\mathcal{L}_{\mathrm{PT}}$.

We differ from traditional PT setups in two regards:

  1. 1.

We assume that as an input to the PT problem we also have a graph
$G_{\mathrm{PT}}=(V,E)$ that contains as vertices the individual pre-training
samples within $\boldsymbol{X}_{\mathrm{PT}}$:
$\\{\boldsymbol{x}_{\mathrm{PT}}^{(i)}|1\leq i\leq N_{\mathrm{PT}}\\}\subseteq
V$.

  2. 2.

We decompose the loss $\mathcal{L}_{\mathrm{PT}}$ into two components,
weighted with hyperparameter $0<\lambda_{\mathrm{SI}}<1$:
$\mathcal{L}=(1-\lambda_{\mathrm{SI}})\mathcal{L}_{\mathrm{M}}+\lambda_{\mathrm{SI}}\mathcal{L}_{\mathrm{SI}}$.
These components are (1) a traditional, per-token objective with loss
$\mathcal{L}_{\mathrm{M}}$ (e.g., a language model) and (2) a structure-
inducing objective with loss $\mathcal{L}_{\mathrm{SI}}$ designed to
regularize the per-sample latent space geometry in accordance with the
relationships encoded in $G_{\mathrm{PT}}$.

Beyond standard constraints (e.g., continuity, differentiability), under our
PT paradigm, we further constrain the form of the new loss term
$\mathcal{L}_{\mathrm{SI}}$ via the following constraints:

  1. 1.

There must exist a distance function
$d:\mathcal{Z}\times\mathcal{Z}\to\mathbb{R}$, radius $r\in\mathbb{R}$, and
loss value $\ell^{*}\in\mathbb{R}$ such that at any solution
$\boldsymbol{\theta}^{*}$ such that
$\mathcal{L}_{\mathrm{SI}}(\boldsymbol{\theta}^{*})<\ell^{*}$ we must have the
property that the induced embeddings
$\boldsymbol{z}_{i}=f_{\boldsymbol{\theta}^{*}}(\boldsymbol{x}_{i})$ must
recover the graph $G_{\mathrm{PT}}$ under a radius nearest neighbor
connectivity algorithm via distance function $d$ and radius $r$. More
formally, it must be the case that $(\boldsymbol{x}_{i},\boldsymbol{x}_{j})\in
E$ if and only if
$d(f_{\boldsymbol{\theta}^{*}}(\boldsymbol{x}_{i}),f_{\boldsymbol{\theta}^{*}}(\boldsymbol{x}_{j}))<r$.
Note that by the assumption of continuity in $\mathcal{L}_{\mathrm{SI}}$ and
$f$, this correspondence will hold in some (potentially very small) region
about $\boldsymbol{\theta}^{*}$ as well.

  2. 2.

Such a solution $\boldsymbol{\theta}^{*}$ from part (1) must be possible to
obtain (though we make no requirements about feasibility of convergence to
this $\boldsymbol{\theta}^{*}$). Assuming that all input $\boldsymbol{x}_{i}$
are distinct, this requirement is more about the geometry of the graph
$G_{\mathrm{PT}}$ and the feasibility of realizing embeddings for
$G_{\mathrm{PT}}$ within the latent space $\mathcal{Z}$—e.g., if
$\mathcal{Z}=\mathbb{R}$, then it will often be impossible to produce
embeddings that fully reflect the geometry of the graph $G_{\mathrm{PT}}$, and
we wish to exclude such cases from our framework.

Note that while we also take the graph $G_{\mathrm{PT}}$ an input to the PT
problem, _we still cannot use it as a direct input to the function
$f_{\boldsymbol{\theta}}$_. Just like in traditional pre-training,
$f_{\boldsymbol{\theta}}$ must take as input only samples from $\mathcal{X}$
_because otherwise we could not apply it to the same, general class of FT
problems over domain $\mathcal{X}$_. Thus, while there are methods that
operate over graphs directly—e.g., GCNN graph embedding methods (Kipf and
Welling, 2017)—these are outside our scope and _cannot_ be used in place of a
standard PT system.

Figure 2: We re-cast the PT paradigm as taking in a graph $G_{\mathrm{PT}}$ as
an auxiliary input. $G_{\mathrm{PT}}$ is used to define a new structure-
inducing objective $\mathcal{L}_{\mathrm{SI}}$, which pushes a pre-training
encoder $f_{\boldsymbol{\theta}}$ to embed samples such that closeness in the
latent space corresponds to edges in $G_{\mathrm{PT}}$.

This re-cast PT setup is outlined visually in Figure 2. Note that this process
does not constrain the network architecture used as the PT encoder
$f_{\boldsymbol{\theta}}$, nor does it constrain the form of the per-token
objective $\mathcal{L}_{\mathrm{M}}$. As such, this perspective on PT is
sufficiently general so as to encompass many existing PT methods. However,
through the flexibility in the form of the graph $G_{\mathrm{PT}}$ and the
structure of the loss $\mathcal{L}_{\mathrm{SI}}$ (subject to the above
constraint), we are still able to naturally realize numerous novel PT methods.

###  3.3 Examining Existing Methods via This Re-Framing

We can realize many existing PT methods naturally within our re-framed PT
setup. As we can see from Table 1, their are 2 primary vehicles by which
existing pre-training methods constrain the per-sample latent space geometry:
Single-task and multi-task classification. In each case, we will show that
there exist natural graphs $G_{\mathrm{PT}}$ and losses
$\mathcal{L}_{\mathrm{SI}}$ which realize these methods under our revised PT
problem.222 Note this omits the unique case of KEPLER (Wang et al., 2019c),
which leverages whole-sample knowledge graph alignment which can only be
realized by our loss subdivision when restricted to graphs with a single edge
type.  After outlining simple single and multi-task classification, we will
then show how this re-framed PT problem poses clear opportunities to design
new, richer PT objectives by discussing two new $\mathcal{L}_{\mathrm{SI}}$
objectives which have not been explored to date.

MLM and per-sample classification. The simplest manner of augmenting PT to
induce per-sample structure would be to add a single-task classification task
at the level of the whole sample. This is the manner used by a number of
existing methods, including BERT Devlin et al. (2019), PLUS (Min et al.,
2020), and many others. These methods augment pre-training by adding a
classification head that ingests the per-sample latent space $\mathcal{Z}$ and
uses a single linear layer to produce a predicted label
$\mathcal{Y}_{\mathrm{PT}}$, training this process via, e.g., a standard cross
entropy loss. In this setting, we can view the matrix used in the linear layer
of the classification head as a set of embeddings for each class, and the
cross entropy loss as a function which seeks to maximize the inner product
between the embedding of a sample $\boldsymbol{z}\in\mathcal{Z}$ and the
embedding of its associated class $y\in\mathcal{Y}$, while minimizing the
inner product between non-matching pairs of samples and classes. Approaching
optimality, the inner product of a matching pair will be large while those of
any non-matching pairs low, with a large margin separating the two.

If we construct a graph $G_{\mathrm{PT}}$ whose nodes are given by all samples
in $\boldsymbol{X}_{\mathrm{PT}}$ augmented with auxiliary nodes for the
classes $y\in\mathcal{Y}$, and edges $(\boldsymbol{x}_{i},y_{i})$ connecting
samples to their appropriate class labels, we can see then that the standard
cross entropy loss will produce an embedding that will recover this graph as
we approach optimality using the cosine distance function. Thus, a simple
classification objective per-sample is a valid example of a method under our
restrictions, albeit one using a highly constrained graph $G_{\mathrm{PT}}$.

MLM and Multi-task Classification. Much like the single-task case, if we
construct a graph $G_{\mathrm{PT}}$ whose nodes are given by all samples in
$\boldsymbol{X}_{\mathrm{PT}}$ augmented with auxiliary nodes for the classes
$y\in\mathcal{Y}$, and edges $(\boldsymbol{x}_{i},y_{i})$ connecting samples
to their appropriate class labels (with now each node $\boldsymbol{x}_{i}$
being connected to potentially many different labels $y_{j}$ corresponding to
each of its tasks), we can see then that the standard cross entropy loss will
again produce an embedding that will recover this graph as we approach
optimality using the cosine distance function. Thus, a multi-task
classification objective per-sample is a valid example of a method under our
loss subdivision, albeit one using a highly constrained graph
$G_{\mathrm{PT}}$.

###  3.4 Defining New Methods via This Re-Framing

In both the single and multi-task classification settings, our graphs had very
rigid structural constraints based on the supervised labeled pairs
$(\boldsymbol{x}_{i},y_{i})$. In the context where our graph $G_{\mathrm{PT}}$
is not so constrained, then the methods described above will not generalize.

However, our loss subdivision itself, and the notion that
$\mathcal{L}_{\mathrm{SI}}$ must recover the graph $G_{\mathrm{PT}}$ as it
approaches optimality immediately suggests new possibilities: In particular,
if we design a new loss function that pushes embeddings of samples to be close
to their $G_{\mathrm{PT}}$ neighbors, then this loss function will satisfy our
constraints and help to induce the desired structure in $\mathcal{Z}$. There
are many existing lines of research we could leverage here, including graph
embeddings (Gao and Huang, 2018), structure-preserving metric learning
(Hadsell et al., 2006), and topological data analysis (Moor et al., 2020),
among others.

In this example (and for the concrete methods we profile in Section 6), we
investigate two new losses drawn from metric learning: a contrastive loss
(Hadsell et al., 2006), and the multi-similarity loss (Wang et al., 2019d),
both of which we adapt to the setting where relationships are determined via a
graph rather than supervised labels.

The multi-similarity loss, parametrized by $w_{+}$, $w_{-}$, and $t$, is given
below:

$\displaystyle\mathcal{L}_{\mathrm{SI}}$
$\displaystyle=\frac{1}{Nw_{+}}\log\left(1+\\!\\!\\!\\!\sum_{(i,j)\in
E}e^{-w_{+}\left(\left\langle
f(\boldsymbol{x}_{i}),f(\boldsymbol{x}_{j})\right\rangle-t\right)}\right)+\frac{1}{Nw_{-}}\log\left(1+\\!\\!\\!\\!\sum_{(i,j)\not\in
E}e^{w_{-}\left(\left\langle
f(\boldsymbol{x}_{i}),f(\boldsymbol{x}_{j})\right\rangle-t\right)}\right),$

Our contrastive loss is modeled after Hadsell et al.’s version, adapted for
the SI setting over a graph $G_{\mathrm{PT}}$. For this loss, we assume we are
given the following mappings: ‘$\mathrm{pos}$’, which maps $\boldsymbol{x}$
into a positive node (i.e., linked to $\boldsymbol{x}$ in $G_{\mathrm{PT}}$),
and ‘$\mathrm{neg}$’, which maps $\boldsymbol{x}$ into a negative node (i.e.,
not linked to $\boldsymbol{x}$ in $G_{\mathrm{PT}}$). The union of a set of
points $\boldsymbol{x}$ and its images under $\mathrm{pos}$ and $\mathrm{neg}$
mappings then form our full minibatch. This loss is then specified by the
positive and negative margin parameters $\mu_{+}$ and $\mu_{-}$ as:

$\displaystyle\mathcal{L}_{\mathrm{SI}}^{(\text{CL})}$
$\displaystyle=\frac{1}{N}\sum_{\boldsymbol{x}_{i}\in\boldsymbol{X}}\max(\mathcal{D}(\boldsymbol{x}_{i},\mathrm{pos}(\boldsymbol{x}_{i}))-\mu_{+},0)+\frac{1}{N}\sum_{\boldsymbol{x}_{i}\in\boldsymbol{X}}\max(\mu_{-}-\mathcal{D}(\boldsymbol{x}_{i},\mathrm{neg}(\boldsymbol{x}_{i})),0).$

In addition to a loss term, we can also use negative sampling strategies to
improve efficiency. Using the full graph $G_{\mathrm{PT}}$, which is not
available in many contexts where negative sampling is employed, we can
leverage the _$G_{\mathrm{PT}}$ distance_ between samples, which provides a
complementary source of information beyond embedding space distance alone. In
practice, one could use this, for example, to limit negative samples to occur
within the same connected component, but more complex strategies based on
graph sampling (Zeng et al., 2020) could also be used.

In our analyses in Sections 5 and 6, we refer to methods that leverage these
losses and more general graphs $G_{\mathrm{PT}}$ than the simple graphs used
for classification tasks as “Structure-inducing Pre-training” (SIPT) methods,
because these losses are designed specifically to induce structure in the per-
sample latent space.

##  4 Theoretical Analysis of Pre-training Methods

Next, we examine theoretical implications of our analysis from Section 3. In
particular, we examine how we can relate PT graph structure to FT task
performance under our paradigm. Our central claim towards this end is provided
in Theorem 1, below. In informal language, this theorem states that as it
approaches optimality, a fully pre-trained embedder $f$ over graph
$G_{\mathrm{PT}}$ will produce an embedding space such that nearest-neighbor
performance for any downstream task will be lower bounded by the performance
that could be obtained via a nearest neighbor algorithm over the graph
$G_{\mathrm{PT}}$.

###### Theorem 1

Let $\boldsymbol{X}_{\mathrm{PT}}$, $G_{\mathrm{PT}}$ be a PT dataset and
graph, respectively, and let $f_{\boldsymbol{\theta}^{*}}$ be an encoder pre-
trained under a PT objective permissible under our framing which realizes a
$\mathcal{L}_{\mathrm{SI}}$ value no more than $\ell^{*}$. Then, under
embedding $f$ and given sufficient data, the nearest-neighbor accuracy for a
FT task $y$ will converge as dataset size increases to at least the local
consistency (see Definition 2) of $y$ over $G_{\mathrm{PT}}$.

The full proof of Theorem 1 is given in the Appendix. However, we do offer a
sketch of the proof of this claim here, starting with a formal definition of
“Local consistency”, in Definition 2.

######  Definition 2 (Local Consistency)

Let $y:X\to\mathcal{Y}$ be a task over a domain $X$, and let
$G=(V,E),X\subseteq V$ be a graph. Then the _local consistency_
$\mathrm{LC}_{G}(y)$ is the probability that a node’s label $y(x)$ will agree
with the majority of the labels of $x$’s neighbors in $G$:

$\mathrm{LC}_{G}(y)=\mathbb{P}\left(y(x)=\operatorname*{argmax}_{\ell\in\mathcal{Y}}\sum_{x^{\prime}\in
X|(x,x^{\prime})\in E}\mathds{1}_{y(x^{\prime})=\ell}\right).$

Note this is closely related to _homophily_ (Zhu et al., 2020; Huang and
Zitnik, 2020; Zhang et al., 2016).

With local consistency defined, proving Theorem 1 takes several steps: _(1)_
Quantifying in what conditions we can have a valid $\mathcal{L}_{\mathrm{SI}}$
and $\ell^{*}$ (e.g., when approximately-optimal embeddings will exist), _(2)_
, Noting that by definition of a “valid” $\mathcal{L}_{\mathrm{SI}}$ term,
$G_{\mathrm{PT}}$ is recoverable under a radius-nearest-neighbor algorithm in
the latent space as we approach optimality, and finally, _(3)_ concluding that
local consistency will be reflected in final embeddings via nearest-neighbor
classifiers, which can be seen by noting that if we use the same distance
function and radius $d$ and $r$ associated with loss value $\ell^{*}$ in our
nearest neighbor classifier, our sets of neighbors will exactly match the
neighbors in $G_{\mathrm{PT}}$, and therefore the local consistency will
translate directly into realizable accuracy, establishing the guarantee.

###  4.1 Consequences of This Theory

We can see immediately several consequences of this theoretical result, such
as the following, both of which are stated and formally proved in the Appendix
Section D.3

  1. 1.

Graphs that are independent cliques or stars make minimal guarantees about FT
task consistency, even as dataset sizes increases.

  2. 2.

Graphs that are defined as nearest-neighbor graphs on manifolds over which the
task label is continuous almost everywhere approach unity as dataset size
approaches infinity.

These two properties show why the _richness_ of the graph $G_{\mathrm{PT}}$
matters. In particular, as discussed in Section 3.3, most pre-training methods
today can be realized under our framework as PT methods with $G_{\mathrm{PT}}$
given by independent cliques or stars; therefore, the extent to which they
provide significant theoretical guarantees about the per-sample latent space
structure is minimal, and moving to new methods with richer graphs has the
potential to markedly improve performance, a fact we will see realized on
synthetic datasets in Section 5 and on real datasets in Section 6.

##  5 Semi-synthetic Validation of Theoretical Properties

We can further validate the theoretical analyses of Section 4 via semi-
synthetic experiments. particular, we create several datasets of natural
language sentences augmented with synthetic graphs with known relationships to
certain FT tasks (e.g., low or high local consistency, low or high rates of
noise). We then use these datasets to validate three important properties of
PT methods: First, do PT methods trained with a $\mathcal{L}_{\mathrm{SI}}$
and $G_{\mathrm{PT}}$ yield Nearest-neighbor FT performance in accordance with
our theory? In particular, do (a) FT tasks with high local consistency over
the PT graph offer better performance, and (b) those with very low local
consistency offer worse performance? Second, do PT methods trained with a
$\mathcal{L}_{\mathrm{SI}}$ and $G_{\mathrm{PT}}$ suffer significantly when
pre-training graphs are polluted with noise? Finally, third, do the latent
space geometry regularizing properties of $\mathcal{L}_{\mathrm{SI}}$ yield
methods whose embeddings more clearly cluster than embeddings produced by
traditional pre-training alone?

###  5.1 Datasets & Experimental Setup

While full experimental details are provided in the Appendix, we will first
offer a brief commentary on how we create the synthetic datasets in use here.

Pre-training & fine-tuning datasets. Across all experiments, our synthetic
datasets consist of free-text sentences333Source:
https://www.kaggle.com/mikeortman/wikipedia-sentences labeled with LDA-
produced topics used as our FT labels. To test across various kinds of graphs,
we produce a number of pre-training graphs per experiment, as detailed below.

Pre-training graphs. We use graphs spanning 3 categories. (1) A graph
(cliques) consisting of disconnected cliques, where sentences are linked in
the graph if they share the same topic label. (2) Graphs composed of nearest-
neighbor graphs defined over simplicial manifolds built using topic
probabilities to localize sentences onto simplices. We explore manifolds with
a range of topological complexity, including: Plane, Möbius, Sphere, and
Torus. Finally, (3) we define several graphs according to a mechanistic
process which allows us to control how topic labels relate to graph structure.
We build graphs in this setting that have the property that topic labels
associate with them in one of three ways: so that topics are maximally
conserved within local neighborhoods (Neighborhood), by assigning sentences to
nodes in the graph such that each graph motif corresponds to a unique topic
(Motif), and such that node topics are driven by non-local graph structural
features, on the basis of graphlet degree vectors (Structural). Of all these
graphs, we expect that topics will display a low local consistency over the
Structural graph, a moderately high local consistency over the Motif graph (as
graph motifs are all connected components), and high local consistency
everywhere else.

Experimental Setup. To answer our three questions, we will pre-train models
under both traditional LM pre-training alone and a new, structure-inducing PT
(SIPT) method within our paradigm which augments the loss with a contrastive
learning loss over $G_{\mathrm{PT}}$, with $\lambda_{\mathrm{SI}}=0.1$. Both
models use a shallow transformer encoder for $f_{\boldsymbol{\theta}}$ and a
character-level tokenization scheme. Final results are reported via the AUROC
of 3-nearest-neighbor classifiers over the latent space, per-sample
embeddings. In line with our theoretical predictions, we expect to see higher
NN FT performance in all settings where the FT task (topic prediction) has
high local consistency over the graph $G_{\mathrm{PT}}$ (all graphs except
Structural) and worse performance in the case where the local consistency is
very low (Structural).

We also assess the stability of our method as the graph $G_{\mathrm{PT}}$ is
noised using the Cliques graph by randomly adding additional edges with
varying rates.

###  5.2 Results

Across all expected settings, we see that our SIPT model, which augments LM PT
with the $\mathcal{L}_{\mathrm{SI}}$ objective and graph $G_{\mathrm{PT}}$ as
motivated by section 3.2 offers significant improvements, is robust to noise
in the PT graph, and performs in agreement with our theoretical analyses.

SIPT improves performance over LM PT by $0.26\pm 0.13$ AUROC on graphs where
the topic task has a high local consistency. As can be seen in Figure 3a, SIPT
offers significant improvements over LM PT in nearest-neighbor FT AUROC across
all graph types with strong topic local consistency.

SIPT’s empirical results are in agreement with theoretical analyses. In line
with our analyses in Section 4, SIPT only under-performs LM PT on the
Structural graph where the topic task (by design) does not have strong local
consistency. This validates our theoretical results by showing that local
consistency is a strong predictor of Nearest-neighbor FT performance.

SIPT is robust to incomplete and noisy pre-training graphs. Figure 3b shows
Nearest-neighbor FT AUROC as a function of noise rate on the Cliques graph.
For up to 15% noise, SIPT shows improvements over LM PT, and even at 50%
noise, the two approaches perform comparably.

SIPT pre-trained embeddings show stronger clustering than LM PT embeddings.
Figure 3c-d shows embeddings produced under the Möbius graph either by LM PT
or SIPT, clustered via UMAP into 2 dimensions. It is clearly visually from
this figures that SIPT embeddings show clear clusters strongly associated with
the topic-modelling FT task, whereas LM PT embeddings do not.

Figure 3: (a) Comparisons between nearest-neighbor FT AUROC (higher is better)
of LM PT models and SIPT models over various graphs with various forms of
structural alignment. $LC$ indicates the label consistency between FT task and
$G_{\mathrm{PT}}$ (see Section 4). (b) nearest-neighbor FT AUROC vs. noise
rate. Up to 10% noise SIPT dramatically outperforms LM PT, and at 50% noise,
the two approaches are equal. (c-d) Embedding space of MPT and SIPT models on
the Möbius dataset. Point colors indicate topic labels. SIPT’s embedding space
reflects the structure of pre-training graph whereas MPT does not.

###  5.3 Conclusions

From these analyses, we see that both theoretically and empirically that
augmenting PT with per-sample structure-inducing objectives can both (1) offer
significant advantages over existing PT architectures and (2) permit
analytical reasoning about over which FT tasks PT will offer improvements.
Next, in Section 6, we will validate these ideas further in real-world
settings across diverse data modalities, PT graphs, loss functions, and FT
tasks.

##  6 Real-world Experiments with SIPT Methods

In Section 4, we showed that our new PT framework permits meaningful
theoretical analyses of pre-training methods, and that those theoretical
findings are validated in semi-synthetic experiments. In this section, we
instead demonstrate that we can use our PT framework to design PT methods with
rich, per-sample geometric constraints which offer real-world performance
benefits on standard benchmarks and meaningful tasks.

In particular, we examine the efficacy of incorporating structural information
across three domains: Proteins, containing protein sequences; Abstracts
containing free-text biomedical abstracts; and Networks, containing sub-graphs
of protein-protein interaction (PPI) networks. In each data modality, we use
different pre-training datasets, test on different downstream FT tasks,
leverage distinct graphs $G_{\mathrm{PT}}$, and compare to baseline per-sample
and per-token methods. In each case, we show that methods leveraging rich per-
sample constraints via a graph $G_{\mathrm{PT}}$ as outlined in Section 3 will
match or exceed the performance of current per-token & per-sample methods.

In the remainder of this section, we do the following. First, in Section 6.1,
we detail how we train models under our framework, and describe the datasets,
pre-training graphs, and per-token / per-sample baselines against which we
compare. Then, in Section 6.2 we will present results of these analyses,
including comparisons between new methods and existing per-token and per-
sample baselines, qualitative analyses into where structure is most useful,
and ablation studies.

###  6.1 Data and Experimental Setup

We explore three different domains containing different data modalities in
this work. We describe each domain fully in Appendix Section E.2, but all
details are also summarized in Table 2. Recall that our goal in these
experiments is to probe whether or not training PT systems augmented with rich
per-sample latent space geometry constraints offers improvements over
traditional per-token or per-sample methods. To that end, we perform the
following steps.

Baselines. First, for each domain, we identify an appropriate, published per-
token and per-sample PT baseline. For example, in our Proteins domain, we use
the TAPE Transformer model (Rao et al., 2019) (trained via LM) as our per-
token baseline, and the PLUS Transformer model (Min et al., 2020) (trained via
LM plus a supervised, per-sample task) as our per-sample baseline.

PT Graph. Next, we identify a source dataset and graph $G_{\mathrm{PT}}$ to
train a new structure-inducing pre-training (SIPT) method based on our PT
paradigm. We want these graphs to be simultaneously (1) rich sources of per-
sample relationships, (2) readily extractable from public data, and (3) of
diverse types and forms across data modalities to probe our framework in
various circumstances. In our Proteins experiments, we use the Tree-of-life
dataset (Zitnik et al., 2019) and the associated multi-species protein-protein
interaction graph as our dataset $\boldsymbol{X}_{\mathrm{PT}}$ and graph
$G_{\mathrm{PT}}$, respectively. In this modality, then, our PT system is
regularizing the per-sample latent space such that proteins that interact with
one another should yield similar embeddings.

Training Structure-inducing Models. Then, we need to train models leveraging
the structure inducing models directly. For our Proteins and Abstracts
datasets, in order to minimize computational burden, we do not pre-train a
model from scratch under our new framework. Instead, we initialize a model
from the per-token baseline directly, then perform additional pre-training for
only a small number of epochs under the new loss subdivision proposed in
Section 3. In these settings, we assess both the multi-similarity and
contrastive $\mathcal{L}_{\mathrm{SI}}$ variants. On the Networks dataset, we
pre-train all models (including baselines) from scratch, and based on early
experimental results we only assess the contrastive loss variant.
$\lambda_{\mathrm{SI}}$ is chosen on the Proteins and Abstracts datasets to
maximize the performance on an internal link-retrieval task over
$G_{\mathrm{PT}}$. This process suggests that $\lambda_{\mathrm{SI}}$ of $0.1$
is a robust setting, and as such $0.1$ was used directly for the Networks
task. Finally, note that our warm-start procedure allows a powerful ablation
study: by additionally training a PT model from the per-token baseline with
$\lambda_{\mathrm{SI}}=0$, we can uniquely assess the impact of the new loss
term, rather than simply additional training or the different PT dataset. We
perform this ablation study for all applicable datasets.

FT Tasks and Evaluation. Lastly, we need to evaluate these pre-training models
across FT task benchmarks. In all cases, we use established PT/FT benchmarks,
and compare to 10 tasks in total across all 3 modalities. FT models are,
naturally, initialized from the pre-trained encoders, and FT training
procedures and hyperparameters match those in the published literature, or are
fit via minimal hyperparameter searches performed across both baselines and
novel methods on validation datasets. In the Proteins dataset, for example, we
use the TAPE benchmark of 5 FT tasks (Rao et al., 2019). All fine-tuning tasks
explored across all three of our domains are detailed in Table 3. On all FT
tasks save TAPE’s contact prediction task, which has a pairwise output space
and is thus very computationally expensive, FT evaluations were performed over
3 random seeds to assess statistical significance of results.

|  Proteins |  Abstracts |  Networks
---|---|---|---
Data Modality ($\boldsymbol{x}_{i}$ is a…) |  Protein Sequence |  Biomedical Paper Abstract |  PPI Network Ego-graph
PT Dataset |  Tree-of-life (Zitnik et al., 2019) |  Microsoft Academic Graph (Wang et al., 2019a; Hu et al., 2020a) |  Hu et al. (2020b)
FT Dataset |  TAPE (Rao et al., 2019) |  SciBERT (Beltagy et al., 2019) |  Hu et al. (2020b)
Per-token baseline |  TAPE (Rao et al., 2019) |  SciBERT (Beltagy et al., 2019) |  Attribute Masking Hu et al. (2020b)
Per-sample baseline |  PLUS (Min et al., 2020) |  None |  Multi-task learning Hu et al. (2020b)
$G_{\mathrm{PT}}$: $(\boldsymbol{x}_{i},\boldsymbol{x}_{j})\in E$ iff |  $\boldsymbol{x}_{i}$ interacts with $\boldsymbol{x}_{j}$ |  $\boldsymbol{x}_{i}$’s paper cites $\boldsymbol{x}_{j}$’s paper |  $\boldsymbol{x}_{i}$’s central protein shares at least 9 Gene Ontology (GO) labels with $\boldsymbol{x}_{j}$’s central protein.
Table 2: A summary of our datasets, tasks, and benchmarks. FT Dataset | FT Task |  Description | Metric
---|---|---|---
|  Name | Abbr. |  |
TAPE Rao et al. (2019) |  Remote Homology | RH |  A per-sequence classification task to predict protein fold category. | Accuracy
Secondary structure | SS |  A per-token classification task to predict amino acid structural properties. | Accuracy
Stability | ST |  A per-sequence, regression task to predict stability. |  Spearman’s $\rho$
Fluorescence | FL |  A per-sequence, regression task to predict fluorescence. |  Spearman’s $\rho$
Contact Prediction | CP |  An intra-sequence classification task to predict which pairs of amino acids are in contact in the protein’s 3D conformation. |  Precision @ $L/5$
SciBERT Beltagy et al. (2019) |  Paper Field | PF |  A per-sentence classification problem to predict a paper’s area of study from its title. | Macro-F1
SciCite | SC |  A per-sentence classification problem to predict citation intent | Macro-F1
ACL-ARC | AA |  A per-sentence classification problem to predict citation intent | Macro-F1
SciERC | SRE |  A per-sentence Relation Extraction task. | Macro-F1
Networks Hu et al. (2020b) |  |  |  The multi-label binary classification of 40 gene-ontology term annotations | Macro-AUROC
Table 3: Fine-tuning tasks

###  6.2 Results

Domain | Task | Vs. Per-Token PT | vs. Per-Sample
---|---|---|---
RRE | $\Delta$ | RRE | $\Delta$
Proteins | RH |  *7.0%,1.2[1]$\pm$[2] | $\boldsymbol{\uparrow}$ |  *8.4%,2.4[1]$\pm$[2] | $\boldsymbol{\uparrow}$
FL |  *-0.8%,1.3[1]$\pm$[2] | $\sim$ |  *12.8%,1.1[1]$\pm$[2] | $\boldsymbol{\uparrow}$
ST |  *13.1%,2.5[1]$\pm$[2] | $\boldsymbol{\uparrow}$ |  *2.2%,2.8[1]$\pm$[2] | $\sim$
SS |  *4.5%,0.2[1]$\pm$[2] | $\boldsymbol{\uparrow}$ |  *4.5%,0.2[1]$\pm$[2] | $\boldsymbol{\uparrow}$
CP |  *10.5%[1]$\pm$[2] ∗ | $\boldsymbol{\uparrow}$ | N/A
Abstracts | PF |  *0.3%,0.2[1]$\pm$[2] | $\sim$ | N/A
SC |  *2.4%,4.1[1]$\pm$[2] | $\sim$ | N/A
AA |  *17.7%,6.5[1]$\pm$[2] | $\boldsymbol{\uparrow}$ | N/A
SRE |  *6.7%,0.4[1]$\pm$[2] | $\boldsymbol{\uparrow}$ | N/A
Networks |  |  *7.8%,5.2[1]$\pm$[2] | $\sim$ |  *5.1%,2.7[1]$\pm$[2] | $\boldsymbol{\uparrow}$
Table 4:  Relative reduction of error (RRE; defined to be
$\frac{\text{[baseline error]}-\text{[$G_{\mathrm{PT}}$ model
error]}}{\text{[baseline error]}}$) of models trained under our framework vs.
published per-token or per-sample baselines. Higher numbers indicate models
under our framework reduce error more, and thus outperform baselines. The
$\Delta$ column indicates whether the model offers a statistically significant
improvement ($\uparrow$), no significant change ($\sim$), or a statistically
significant decrease ($\downarrow$). Statistical significance is assessed via
a $t$-test at significance level $p<0.1$. Per-sample analysis and variance
estimates for CP were infeasible due to the computational cost of this task.

Incorporating $\mathcal{L}_{\mathrm{SI}}$ performs comparably to or improves
over all baselines across all 3 domains and 10 FT tasks. In Table 4, we show
the relative reduction of error444
$\mathrm{RRE}=\frac{(1-\text{baseline})-(1-\text{SIPT})}{(1-\text{baseline})}$.
RRE is unitless and suitable for comparing across tasks.  of the best
performing model leveraging rich structural constraints (i.e., maximal over
the multi-similarity loss or contrastive loss) vs. the per-token or per-sample
baselines. We can see that in 10/15 cases, incorporating per-sample latent
space regularization improves over existing methods, and in no case does it do
worse than either baseline. Furthermore, In some cases, the gains in
performance are quite significant, with improvements of approximately 17%
(0.05 macro-F1 raw change) on AA, 6% on SRE (0.01 macro-F1 raw change), and 4%
on RH (2% accuracy raw change). We further establish a new SOTA on AA and RH,
and match SOTA on FL, ST, & PF.

We also see in Figure 4 how performance evolves over FT iterations for the
Networks dataset, to determine if the improvements observed at the final
converged values are present throughout training. We see that methods using
per-sample structure inducing losses converge faster to better performance
than do other methods.

Figure 4:  FT AUROC as a function of iteration over the Networks dataset. The
novel method leveraging richer per-sample constraints converges faster and
obtains better performance than LM or multi-task PT.

Note that raw results across all settings are present in Appendix Section F.1,
Tables 7,8.

These performance gains are present across diverse modalities and
$G_{\mathrm{PT}}$s. Note that these performance gains persist over all three
data modalities and all different $G_{\mathrm{PT}}$ types we use here. This
shows that explicitly regularizing the per-sample latent space geometry offers
value across NLP, non-language sequences, and networks domains, as well as
while leveraging graphs including those defined by external knowledge, those
inferred from self-supervised signals in the data directly, and those defined
by nearest-neighbor methods over multi-task label spaces.

Observed performance improvements are consistent with theory. Hyperparameter
tuning finds that far less structure-inducing is necessary on our Abstracts
dataset ($\lambda_{\mathrm{SI}}=0.01$) than on our Proteins dataset
($\lambda_{\mathrm{SI}}=0.1$). This agrees with our guiding hypothesis that
per-sample latent space regularization is much more necessary on non-NLP
domains than it is on NLP domains. Further details on these analyses are in
the Appendix, Section F.2

Observed gains are uniquely attributable to the novel loss
$\mathcal{L}_{\mathrm{SI}}$. As outlined in Section 6.1, we also perform
ablation studies to determine how much of the observed gains in Table 4 are
due to the novel loss component, as opposed to, for example, continued
training, new PT data, or the batch selection procedures used in our method
which also indirectly leverage the knowledge inherent in $G_{\mathrm{PT}}$.
Unsurprisingly, there are some gains observed due to these other factors, and
performance gains shrink as compared to continued training of per-token models
alone over the new PT dataset. However, these gains do not change the observed
relationships at all, and more generally comparing against the maximal
performance baseline overall, including the ablation study, does not effect
the statistically significant relationships at all at $p<0.1$. Full ablation
study results, along with all other raw results, can be found in Appendix
Section F.1, Tables 7,8.

##  7 Discussion

In this section, we will provide unifying commentary and emphasize several key
strengths of our work, namely that (1) we crystallize the distinction between
enforcing per-token and per-sample latent space structure, and highlight why
this is important via theory and experiment, (2) we offer a new framework for
understanding and developing PT methods, which itself permits a natural
vehicle to associate PT task structure with FT performance, which is
especially important in non-NLP domains, and (3) we expose novel directions
for future work and PT method development focused on enriching the extent to
which models capture per-sample structure. The possible utility of these new
directions for research is underscored by our empirical results in Section 6,
which show consistent improvement over existing baselines by taking into
account richer notions of per-sample structure.

Per-token vs. per-sample latent space structure. In Table 1, we clarify how a
battery of existing LMPT methods leverage and enforce structure in their
learned embeddings. From this, we can see a clear disparity in attempts at
enriching structure at a _per-sample_ level. In NLP, when many downstream
tasks can be realized as per-token tasks through methods such as prompting,
this discrepancy is not so problematic; however, in non-NLP domains, such as
protein sequences, graphs, timeseries, or tabular data prompting approaches
are unlikely to offer such advantages, and therefore methods leveraging per-
sample structure are critical.

Through our careful experiments, we can show that these differences matter in
practice. On Proteins, Abstracts, and Networks we show that continuing LMPT
methods for several additional iterations augmented with per-sample structural
regularization via $\mathcal{L}_{\mathrm{SI}}$ via a variety of different
graphs offers consistent (if modest) improvements to FT performance across a
battery of downstream tasks. We believe that with further exploration of
methods leveraging structural constraints, even larger gains can be realized,
which motivates the value of our work.

Alignment of PT task structure to FT performance. Through a relatively simple
modification of the PT paradigm, our re-cast PT problem is able to realize a
learning framework in which PT task structure can be directly and simply
related to FT task suitability (as measured through nearest-neighbor accuracy
in the PT latent space). In general, identifying when and why PT will offer
advantages for FT is a very challenging task and has prompted significant
research (Saunshi et al., 2019); while our theoretical analyses sidestep the
traditional difficulties through the addition of the per-sample graph
structure, this does not diminish the fact that we are able to form clear
links between the PT context and FT performance. This theory is validated via
our synthetic experiments, which also establish that simple, metric-learning
motivated structure-inducing losses can be resilient to graph noise.

Novel research directions highlighted by our work. Our work highlights novel
directions for research relating to how we can best pre-train models to
capture per-sample latent space structure. Are there novel losses that are
better suited than the classical metric learning losses we adapt here for the
graph-structured context—e.g., can we leverage the graph-distance alongside
the intra-batch distance to develop improved negative sampling strategies? Are
there adaptations of our restrictions that allow a PT system to capture
multiple graphs (or multiple edge types) simultaneously through different
embeddings? Can we (and do we ever need to) reflect forms of structure beyond
the simple nearest neighbor relationships we leverage in this work—e.g., such
as by leveraging higher-order topological considerations, or by matching a
full distance function rather than a discrete graph? Can we refine the simple
theoretical results here to better take into account the convergence
properties of trained models to understand when and how to effectively
converge to solutions which recover $G_{\mathrm{PT}}$? All these questions and
more are prompted by our PT framework, and this shows that our work can help
prompt new, impactful work in this space.

##  8 Related Work

This paper builds upon a wealth of previous research at the intersection of
PT/FT, neural language models, graph representation learning, and theoretical
analysis of metric learning.

###  8.1 Pre-training Strategies With Explicit Consideration of Geometry

Per-token latent space. As outlined in Section 2 and Table 1, a significant
amount of research in pre-training language models explored ways to
incorporate structure into the per-token latent space. For example, prior work
considered selectively masking named entities (Sun et al., 2020a; Yamada et
al., 2020; Sun et al., 2019, 2020b, 2021; Zhang et al., 2019; He et al., 2020;
Su et al., 2021), aligning per-token representations to KG embeddings (Qin et
al., 2020; Zhang et al., 2021b), performing named entity recognition (Peters
et al., 2019; Yuan et al., 2021), and producing joint token-entity embeddings
(Zhang et al., 2019; Peters et al., 2019; Yu et al., 2020; Zhang et al.,
2021b; Yuan et al., 2021; Su et al., 2021).

Per-sample latent space. As outlined in Section 2 and Table 1, incorporating
structure at the per-sample level in language model pre-training remains
under-explored or neglected research area. The most common approach by far is
to leverage additional classification tasks at the per-sample level (Devlin et
al., 2019; Lan et al., 2019; Zhang et al., 2019; He et al., 2020; Zhang et
al., 2021b; Su et al., 2021; Min et al., 2020; Liu et al., 2019a; Sun et al.,
2020b, 2021; McDermott et al., 2020; Hu et al., 2020b). Further, KEPLER (Wang
et al., 2019c) aligns per-sample embeddings of free-text descriptions of
entities to KG embeddings of those entities to match trans-E (Mroueh et al.,
2012).

Methods with geometrically motivated auxiliary FT/pre-FT stages. Beyond the
methods discussed above and in Section 2 and Table 1, there are also a number
of works that do not explicitly change the PT stage of language-model pre-
training. Instead, those methods add additional stages between PT and FT to
impose geometric constraints at either a per-token or per-sample level. This
includes studies that add “knowledge adapters” to incorporate KG embeddings
into a pre-trained language model at a per-token level (Wang et al., 2020;
Meng et al., 2021; Lu et al., 2021), adapting input text or embeddings to
incorporate KG information (Liu et al., 2020; Faldu et al., 2021; Yan et al.,
2021; Poerner et al., 2020), or adapting sample embeddings via metric learning
techniques based on either self-supervised spatial tasks or on FT tasks
(Reimers and Gurevych, 2019; Gunel et al., 2021).

###  8.2 Pre-training Strategies With No Explicit Consideration of Geometry

Pre-training using language models has been extensively studied in NLP. To
this end, researchers have explored autoregressive approaches (Brown et al.,
2020b; Radford et al., 2018, 2019; Brown et al., 2020a), masked language
modelling approaches (Peters et al., 2018; Liu et al., 2019b), and
discriminative modelling approaches (Clark et al., 2019). These methods have
also been adapted to non-NLP domains where masked imputation can still be
employed, including networks (Hu et al., 2020b; Wang et al., 2019b), tabular
data (Yoon et al., 2020), time series (McDermott et al., 2020), and protein
sequences (Rao et al., 2019).

###  8.3 Pre-training Theory

Prior theoretical analyses found that pre-trained language models reflect
properties consistent with domain knowledge (Coenen et al., 2019; Tenney et
al., 2019; Rao et al., 2020). Additionally, Lee et al. (2020) suggests that
masked-based PT transfers positively to downstream tasks that obey certain
conditional independence relationships with PT learning. Further, Saunshi et
al. (2019) posits that masked-based PT works because one can re-frame most NLP
tasks as language modelling tasks.

###  8.4 Graph Representation Learning

Attributed graph embedding and knowledge graphs. Our work reflects topics in
deep attributed graph embedding (DAGE), in which an embedding model $f$ is
learned for a graph $G$ equipped with node attributes $\boldsymbol{X}$ such
that the connectivity structure of $G$ is reflected in the image of $f$ (Gao
and Huang, 2018; Cui et al., 2020; Li et al., 2018). DAGE algorithms are
typically not adaptable to contexts in which the graph is not known at
inference time, and so it does not directly apply to our setting.

Graph neural networks (GNNs). There is also a wide variety of research on
graph neural networks (Li et al., 2021), including, e.g., work on semi-
supervised graph classification (Kipf and Welling, 2017), inductive
representation learning on graphs (Hamilton et al., 2017), and many others.
While relevant to our study in principle, these methods are not directly
applicable in either the original or our revised PT setting. As we state in
Section 3, although we take as input the graph $G_{\mathrm{PT}}$, this cannot
be processed with a standard GNN. This is because we still want to apply the
encoder $f_{\boldsymbol{\theta}}$ at fine-tuning time to contexts with no
associated graph, using only the input $\boldsymbol{x}$ itself (e.g., the
sentence or protein sequence).

Structure-preserving metric learning. A sub-field within metric learning that
our analyses rely on is structure-preserving metric learning (SPML) (Vert and
Yamanishi, 2004; Shaw and Jebara, 2009; Shaw et al., 2011). In SPML, rather
than using supervised labels, a graph $G$ defines when samples should be close
together, such that $G$ can be recovered in the embedding space via a
connectivity algorithm (Shaw et al., 2011). To the best of our knowledge,
there has been minimal investigation of adaptation of SPML in deep metric
learning.

##  9 Conclusion

In this work, we propose a new perspective on the traditional PT framework.
Under our paradigm, the PT loss is subdivided into two components: one which
optimizes for capturing per-token/intra-sample relationships, and one novel
component explicitly designed to regularize the per-sample latent space to
capture user-specified relationships, encoded with a pre-training graph
$G_{\mathrm{PT}}$. Through theory, semi-synthetic experiments, and real-world
experiments spanning 3 data modalities and 10 FT tasks, we have shown that
this new PT perspective offers significant value. This value presents itself
through new insights into the shortcomings of existing models, clear theory
relating PT task formulation to FT task performance, and real-world
illustrations that new methods can be designed leveraging our formulation and
diverse graphs $G_{\mathrm{PT}}$ to outperform existing per-sample and per-
token PT/FT models. Ultimately, we envision that this work can motivate
further research and have a considerable impact on pre-training methodology.

## References

  * Agarwal et al. (2021) Oshin Agarwal, Heming Ge, Siamak Shakeri, and Rami Al-Rfou.  Knowledge graph based synthetic corpus generation for knowledge-enhanced language model pre-training.  In _Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pages 3554–3565, Online, June 2021. Association for Computational Linguistics.  doi: 10.18653/v1/2021.naacl-main.278.  URL https://aclanthology.org/2021.naacl-main.278.
  * BehnamGhader et al. (2021) Parishad BehnamGhader, Hossein Zakerinia, and Mahdieh Soleymani Baghshah.  Mg-bert: Multi-graph augmented bert for masked language modeling.  In _Proceedings of the Fifteenth Workshop on Graph-Based Methods for Natural Language Processing (TextGraphs-15)_ , pages 125–131, 2021.
  * Beltagy et al. (2019) Iz Beltagy, Kyle Lo, and Arman Cohan.  SciBERT: A Pretrained Language Model for Scientific Text.  In _EMNLP_ , 2019.
  * Brown et al. (2020a) Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _arXiv preprint arXiv:2005.14165_ , 2020a.
  * Brown et al. (2020b) Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language Models are Few-Shot Learners.  _arXiv preprint arXiv: 2005.14165_ , 2020b.
  * Clark et al. (2019) Kevin Clark, Minh-Thang Luong, Quoc V. Le, and Christopher D. Manning.  ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators.  In _ICLR_ , 2019.
  * Coenen et al. (2019) Andy Coenen, Emily Reif, Ann Yuan, Been Kim, Adam Pearce, Fernanda Viégas, and Martin Wattenberg.  Visualizing and Measuring the Geometry of BERT.  In _NeurIPS_ , 2019.
  * Cohan et al. (2019) Arman Cohan, Waleed Ammar, Madeleine van Zuylen, and Field Cady.  Structural Scaffolds for Citation Intent Classification in Scientific Publications.  In _Proceedings of the 2019 Conference of the NAACL: Human Language Technologies, Volume 1 (Long and Short Papers)_ , 2019.
  * Cui et al. (2020) Ganqu Cui, Jie Zhou, Cheng Yang, and Zhiyuan Liu.  Adaptive graph encoder for attributed graph embedding.  In _Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_, pages 976–985, 2020.
  * Deng et al. (2009) Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE conference on computer vision and pattern recognition_ , pages 248–255. Ieee, 2009.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.  In _NAACL_ , 2019.
  * Faldu et al. (2021) Keyur Faldu, Amit Sheth, Prashant Kikani, and Hemang Akabari.  Ki-bert: Infusing knowledge context for better language and domain understanding.  _arXiv preprint arXiv:2104.08145_ , 2021.
  * Févry et al. (2020) Thibault Févry, Livio Baldini Soares, Nicholas FitzGerald, Eunsol Choi, and Tom Kwiatkowski.  Entities as experts: Sparse memory access with entity supervision.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 4937–4951, Online, November 2020\. Association for Computational Linguistics.  doi: 10.18653/v1/2020.emnlp-main.400.  URL https://aclanthology.org/2020.emnlp-main.400.
  * Gao and Huang (2018) Hongchang Gao and Heng Huang.  Deep attributed network embedding.  In _Twenty-Seventh International Joint Conference on Artificial Intelligence (IJCAI))_ , 2018.
  * Gunel et al. (2021) Beliz Gunel, Jingfei Du, Alexis Conneau, and Veselin Stoyanov.  Supervised contrastive learning for pre-trained language model fine-tuning.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=cu7IUiOhujH.
  * Gururangan et al. (2020) Suchin Gururangan, Ana Marasović, Swabha Swayamdipta, Kyle Lo, Iz Beltagy, Doug Downey, and Noah A. Smith.  Don’t Stop Pretraining: Adapt Language Models to Domains and Tasks.  In _ACL_ , 2020.
  * Hadsell et al. (2006) R. Hadsell, S. Chopra, and Y. LeCun.  Dimensionality Reduction by Learning an Invariant Mapping.  In _CVPR_ , 2006.
  * Hamilton et al. (2017) William L Hamilton, Rex Ying, and Jure Leskovec.  Inductive representation learning on large graphs.  In _Proceedings of the 31st International Conference on Neural Information Processing Systems_ , pages 1025–1035, 2017.
  * He et al. (2020) Bin He, Xin Jiang, Jinghui Xiao, and Qun Liu.  Kgplm: Knowledge-guided language model pre-training via generative and discriminative learning.  _arXiv preprint arXiv:2012.03551_ , 2020.
  * Hou et al. (2018) Jie Hou, Badri Adhikari, and Jianlin Cheng.  DeepSF: deep convolutional neural network for mapping protein sequences to folds.  _Bioinformatics_ , (8), 2018.
  * Hu et al. (2020a) Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec.  Open Graph Benchmark: Datasets for Machine Learning on Graphs.  _arXiv: 2005.00687_ , 2020a.
  * Hu et al. (2020b) Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec.  Strategies for Pre-training Graph Neural Networks.  In _ICLR_ , 2020b.
  * Huang and Zitnik (2020) Kexin Huang and Marinka Zitnik.  Graph meta learning via local subgraphs.  In _NeurIPS_ , 2020.
  * Jurgens et al. (2018) David Jurgens, Srijan Kumar, Raine Hoover, Dan McFarland, and Dan Jurafsky.  Measuring the Evolution of a Scientific Field through Citation Frames.  _Transactions of the ACL_ , 6, 2018.
  * Kipf and Welling (2017) Thomas N. Kipf and Max Welling.  Semi-supervised classification with graph convolutional networks.  In _5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings_. OpenReview.net, 2017.  URL https://openreview.net/forum?id=SJU4ayYgl.
  * Klausen et al. (2019) Michael Schantz Klausen, Martin Closter Jespersen, Henrik Nielsen, Kamilla Kjaergaard Jensen, Vanessa Isabell Jurtz, Casper Kaae Sønderby, Morten Otto Alexander Sommer, Ole Winther, Morten Nielsen, Bent Petersen, and Paolo Marcatili.  NetSurfP-2.0: Improved prediction of protein structural features by integrated deep learning.  _Proteins_ , (6), 2019.
  * Lan et al. (2019) Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut.  ALBERT: A Lite BERT for Self-supervised Learning of Language Representations.  In _ICLR_ , 2019.
  * Lee et al. (2020) Jason D. Lee, Qi Lei, Nikunj Saunshi, and Jiacheng Zhuo.  Predicting What You Already Know Helps: Provable Self-Supervised Learning.  _arXiv: 2008.01064_ , 2020.
  * Li et al. (2021) Michelle M Li, Kexin Huang, and Marinka Zitnik.  Representation learning for networks in biology and medicine: Advancements, challenges, and opportunities.  _arXiv:2104.04883_ , 2021.
  * Li et al. (2018) Ye Li, Chaofeng Sha, Xin Huang, and Yanchun Zhang.  Community detection in attributed graphs: An embedding approach.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 32, 2018.
  * Liu et al. (2020) Weijie Liu, Peng Zhou, Zhe Zhao, Zhiruo Wang, Qi Ju, Haotang Deng, and Ping Wang.  K-bert: Enabling language representation with knowledge graph.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 34, pages 2901–2908, 2020.
  * Liu et al. (2019a) Xiaodong Liu, Pengcheng He, Weizhu Chen, and Jianfeng Gao.  Multi-Task Deep Neural Networks for Natural Language Understanding.  In _ACL_ , 2019a.
  * Liu et al. (2021) Ye Liu, Yao Wan, Lifang He, Hao Peng, and S Yu Philip.  Kg-bart: Knowledge graph-augmented bart for generative commonsense reasoning.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 35, pages 6418–6425, 2021.
  * Liu et al. (2019b) Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov.  RoBERTa: A Robustly Optimized BERT Pretraining Approach.  _arXiv:1907.11692 [cs]_ , 2019b.  arXiv: 1907.11692.
  * Lu et al. (2021) Yinquan Lu, Haonan Lu, Guirong Fu, and Qun Liu.  Kelm: Knowledge enhanced pre-trained language representations with message passing on hierarchical relational graphs.  _arXiv preprint arXiv:2109.04223_ , 2021.
  * Maas et al. (2011) Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts.  Learning word vectors for sentiment analysis.  In _Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies_ , pages 142–150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics.  URL https://aclanthology.org/P11-1015.
  * McDermott et al. (2021) Matthew McDermott, Brendan Yap, Harry Hsu, Di Jin, and Peter Szolovits.  Adversarial contrastive pre-training for protein sequences.  _arXiv preprint arXiv:2102.00466_ , 2021.
  * McDermott et al. (2020) Matthew B. A. McDermott, Bret Nestor, Evan Kim, Wancong Zhang, Anna Goldenberg, Peter Szolovits, and Marzyeh Ghassemi.  A Comprehensive Evaluation of Multi-task Learning and Multi-task Pre-training on EHR Time-series Data.  _arXiv: 2007.10185_ , 2020.
  * Meng et al. (2021) Zaiqiao Meng, Fangyu Liu, Thomas Hikaru Clark, Ehsan Shareghi, and Nigel Collier.  Mixture-of-partitions: Infusing large biomedical knowledge graphs into bert.  _arXiv preprint arXiv:2109.04810_ , 2021.
  * Min et al. (2020) Seonwoo Min, Seunghyun Park, Siwon Kim, Hyun-Soo Choi, and Sungroh Yoon.  Pre-Training of Deep Bidirectional Protein Sequence Representations with Structural Information.  _arXiv: 1912.05625_ , 2020.
  * Moor et al. (2020) Michael Moor, Max Horn, Bastian Rieck, and Karsten Borgwardt.  Topological autoencoders.  In _International Conference on Machine Learning_ , pages 7045–7054. PMLR, 2020.
  * Mroueh et al. (2012) Youssef Mroueh, Tomaso Poggio, Lorenzo Rosasco, and Jean-Jacques E Slotine.  Multiclass learning with simplex coding.  In _Proceedings of the 25th International Conference on Neural Information Processing Systems-Volume 2_ , pages 2789–2797, 2012.
  * Mukherjee (2015) Amiya Mukherjee.  Approximation Theorems and Whitney’s Embedding.  In _Differential Topology_ , pages 43–67. Springer International Publishing, Cham, 2015.  ISBN 978-3-319-19045-7.  doi: 10.1007/978-3-319-19045-7˙2.
  * Pedregosa et al. (2011) F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay.  Scikit-learn: Machine learning in Python.  _Journal of Machine Learning Research_ , 12:2825–2830, 2011\.
  * Peters et al. (2018) Matthew Peters, Mark Neumann, Mohit Iyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer.  Deep Contextualized Word Representations.  In _NAACL_ , 2018.
  * Peters et al. (2019) Matthew E Peters, Mark Neumann, Robert Logan, Roy Schwartz, Vidur Joshi, Sameer Singh, and Noah A Smith.  Knowledge enhanced contextual word representations.  In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_ , pages 43–54, 2019.
  * Poerner et al. (2020) Nina Poerner, Ulli Waltinger, and Hinrich Schütze.  E-BERT: Efficient-yet-effective entity embeddings for BERT.  In _Findings of the Association for Computational Linguistics: EMNLP 2020_ , pages 803–818, Online, November 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.findings-emnlp.71.  URL https://aclanthology.org/2020.findings-emnlp.71.
  * Qin et al. (2020) Yujia Qin, Yankai Lin, Ryuichi Takanobu, Zhiyuan Liu, Peng Li, Heng Ji, Minlie Huang, Maosong Sun, and Jie Zhou.  Erica: improving entity and relation understanding for pre-trained language models via contrastive learning.  _arXiv preprint arXiv:2012.15022_ , 2020.
  * Radford et al. (2018) Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever.  Improving language understanding by generative pre-training.  2018\.
  * Radford et al. (2019) Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8):9, 2019.
  * Raffel et al. (2019) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu.  Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer.  _arXiv: 1910.10683_ , 2019.
  * Rao et al. (2019) Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Peter Chen, John Canny, Pieter Abbeel, and Yun Song.  Evaluating Protein Transfer Learning with TAPE.  In _NeurIPS_. Curran Associates, Inc., 2019.
  * Rao et al. (2020) Roshan Rao, Joshua Meier, Tom Sercu, Sergey Ovchinnikov, and Alexander Rives.  Transformer protein language models are unsupervised structure learners.  _bioRxiv preprint bioRxiv: 2020.12.15.422761_ , 2020.
  * Reimers and Gurevych (2019) Nils Reimers and Iryna Gurevych.  Sentence-BERT: Sentence embeddings using Siamese BERT-networks.  In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_ , pages 3982–3992, Hong Kong, China, November 2019. Association for Computational Linguistics.  doi: 10.18653/v1/D19-1410.  URL https://aclanthology.org/D19-1410.
  * Ribeiro and Forbus (2021) Danilo Neves Ribeiro and Kenneth Forbus.  Combining analogy with language models for knowledge extraction.  In _3rd Conference on Automated Knowledge Base Construction_ , 2021\.
  * Rocklin et al. (2017) Gabriel J. Rocklin, Tamuka M. Chidyausiku, Inna Goreshnik, Alex Ford, Scott Houliston, Alexander Lemak, Lauren Carter, Rashmi Ravichandran, Vikram K. Mulligan, Aaron Chevalier, Cheryl H. Arrowsmith, and David Baker.  Global analysis of protein folding using massively parallel design, synthesis, and testing.  _Science_ , 357(6347), 2017.
  * Sanh et al. (2021) Victor Sanh, Albert Webson, Colin Raffel, Stephen H Bach, Lintang Sutawika, Zaid Alyafeai, Antoine Chaffin, Arnaud Stiegler, Teven Le Scao, Arun Raja, et al.  Multitask prompted training enables zero-shot task generalization.  _arXiv preprint arXiv:2110.08207_ , 2021.
  * Sarkisyan et al. (2016) Sarkisyan et al.  Local fitness landscape of the green fluorescent protein.  _Nature_ , 533(7603), 2016.
  * Saunshi et al. (2019) Nikunj Saunshi, Orestis Plevrakis, Sanjeev Arora, Mikhail Khodak, and Hrishikesh Khandeparkar.  A Theoretical Analysis of Contrastive Unsupervised Representation Learning.  In _ICML_ , 2019.
  * Saunshi et al. (2021) Nikunj Saunshi, Sadhika Malladi, and Sanjeev Arora.  A mathematical exploration of why language models help solve downstream tasks.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=vVjIW3sEc1s.
  * Shaw and Jebara (2009) Blake Shaw and Tony Jebara.  Structure preserving embedding.  In _ICML_ , 2009.
  * Shaw et al. (2011) Blake Shaw, Bert Huang, and Tony Jebara.  Learning a Distance Metric from a Network.  In _NeurIPS_ , 2011.
  * Su et al. (2021) Yusheng Su, Xu Han, Zhengyan Zhang, Yankai Lin, Peng Li, Zhiyuan Liu, Jie Zhou, and Maosong Sun.  Cokebert: Contextual knowledge selection and embedding towards enhanced pre-trained language models.  _AI Open_ , 2:127–134, 2021.
  * Sun et al. (2020a) Tianxiang Sun, Yunfan Shao, Xipeng Qiu, Qipeng Guo, Yaru Hu, Xuanjing Huang, and Zheng Zhang.  CoLAKE: Contextualized language and knowledge embedding.  In _Proceedings of the 28th International Conference on Computational Linguistics_ , pages 3660–3670, Barcelona, Spain (Online), December 2020a. International Committee on Computational Linguistics.  doi: 10.18653/v1/2020.coling-main.327.  URL https://aclanthology.org/2020.coling-main.327.
  * Sun et al. (2019) Yu Sun, Shuohuan Wang, Yukun Li, Shikun Feng, Xuyi Chen, Han Zhang, Xin Tian, Danxiang Zhu, Hao Tian, and Hua Wu.  ERNIE: Enhanced Representation through Knowledge Integration.  _arXiv:1904.09223 [cs]_ , 2019.  arXiv: 1904.09223.
  * Sun et al. (2020b) Yu Sun, Shuohuan Wang, Yukun Li, Shikun Feng, Hao Tian, Hua Wu, and Haifeng Wang.  ERNIE 2.0: A Continual Pre-Training Framework for Language Understanding.  _Proceedings of the AAAI Conference on Artificial Intelligence_ , 34(05), 2020b.  ISSN 2374-3468.  Number: 05.
  * Sun et al. (2021) Yu Sun, Shuohuan Wang, Shikun Feng, Siyu Ding, Chao Pang, Junyuan Shang, Jiaxiang Liu, Xuyi Chen, Yanbin Zhao, Yuxiang Lu, et al.  Ernie 3.0: Large-scale knowledge enhanced pre-training for language understanding and generation.  _arXiv preprint arXiv:2107.02137_ , 2021.
  * Sun et al. (2018) Zhiqing Sun, Zhi-Hong Deng, Jian-Yun Nie, and Jian Tang.  Rotate: Knowledge graph embedding by relational rotation in complex space.  In _International Conference on Learning Representations_ , 2018.
  * Tenney et al. (2019) Ian Tenney, Dipanjan Das, and Ellie Pavlick.  BERT Rediscovers the Classical NLP Pipeline.  In _ACL_ , 2019.
  * Vert and Yamanishi (2004) Jean-Philipe Vert and Yoshihiro Yamanishi.  Supervised graph inference.  In _NeurIPS_ , 2004.
  * Wang et al. (2019a) Kuansan Wang, Zhihong Shen, Chiyuan Huang, Chieh-Han Wu, Darrin Eide, Yuxiao Dong, Junjie Qian, Anshul Kanakia, Alvin Chen, and Richard Rogahn.  A Review of Microsoft Academic Services for Science of Science Studies.  _Frontiers in Big Data_ , 2019a.
  * Wang et al. (2020) Ruize Wang, Duyu Tang, Nan Duan, Zhongyu Wei, Xuanjing Huang, Guihong Cao, Daxin Jiang, Ming Zhou, et al.  K-adapter: Infusing knowledge into pre-trained models with adapters.  _arXiv preprint arXiv:2002.01808_ , 2020.
  * Wang et al. (2019b) Sheng Wang, Yuzhi Guo, Yuhong Wang, Hongmao Sun, and Junzhou Huang.  SMILES-BERT: Large Scale Unsupervised Pre-Training for Molecular Property Prediction.  In _ACM International Conference on Bioinformatics, Computational Biology and Health Informatics_ , 2019b.
  * Wang et al. (2019c) Xiaozhi Wang, Tianyu Gao, Zhaocheng Zhu, Zhiyuan Liu, Juanzi Li, and Jian Tang.  KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation.  _arXiv:1911.06136 [cs]_ , 2019c.  arXiv: 1911.06136.
  * Wang et al. (2019d) Xun Wang, Xintong Han, Weilin Huang, Dengke Dong, and Matthew R. Scott.  Multi-Similarity Loss With General Pair Weighting for Deep Metric Learning.  In _CVPR_ , 2019d.
  * Wang et al. (2014) Zhen Wang, Jianwen Zhang, Jianlin Feng, and Zheng Chen.  Knowledge graph embedding by translating on hyperplanes.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 28, 2014.
  * Whitney et al. (1992) Hassler Whitney, James Eells, and Domingo Toledo.  _Collected Papers of Hassler Whitney_ , volume 1.  Nelson Thornes, 1992.
  * Yamada et al. (2020) Ikuya Yamada, Akari Asai, Hiroyuki Shindo, Hideaki Takeda, and Yuji Matsumoto.  Luke: Deep contextualized entity representations with entity-aware self-attention.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 6442–6454, 2020.
  * Yan et al. (2021) Ruiqing Yan, Lanchang Sun, Fang Wang, and Xiaoming Zhang.  A general method for transferring explicit knowledge into language model pretraining.  _Security and Communication Networks_ , 2021, 2021.
  * Yang et al. (2021) Junhan Yang, Zheng Liu, Shitao Xiao, Chaozhuo Li, Guangzhong Sun, and Xing Xie.  Graphformers: Gnn-nested language models for linked text representation.  _arXiv preprint arXiv:2105.02605_ , 2021.
  * Yao et al. (2019) Liang Yao, Chengsheng Mao, and Yuan Luo.  Kg-bert: Bert for knowledge graph completion.  _arXiv preprint arXiv:1909.03193_ , 2019.
  * Yoon et al. (2020) Jinsung Yoon, Yao Zhang, James Jordon, and Mihaela van der Schaar.  VIME: Extending the Success of Self- and Semi-supervised Learning to Tabular Domain.  In _NeurIPS_ , 2020.
  * Yu et al. (2020) Donghan Yu, Chenguang Zhu, Yiming Yang, and Michael Zeng.  Jaket: Joint pre-training of knowledge graph and language understanding.  _arXiv preprint arXiv:2010.00796_ , 2020.
  * Yuan et al. (2021) Zheng Yuan, Yijia Liu, Chuanqi Tan, Songfang Huang, and Fei Huang.  Improving biomedical pretrained language models with knowledge.  In _Proceedings of the 20th Workshop on Biomedical Language Processing_ , pages 180–190, Online, June 2021. Association for Computational Linguistics.  doi: 10.18653/v1/2021.bionlp-1.20.  URL https://aclanthology.org/2021.bionlp-1.20.
  * Zeng et al. (2020) Hanqing Zeng, Hongkuan Zhou, Ajitesh Srivastava, Rajgopal Kannan, and Viktor Prasanna.  Graphsaint: Graph sampling based inductive learning method.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=BJe8pkHFwS.
  * Zhang et al. (2016) Daokun Zhang, Jie Yin, Xingquan Zhu, and Chengqi Zhang.  Homophily, structure, and content augmented network representation learning.  In _ICDM_ , 2016.
  * Zhang et al. (2021a) Ningyu Zhang, Shumin Deng, Xu Cheng, Xi Chen, Yichi Zhang, Wei Zhang, Huajun Chen, and Hangzhou Innovation Center.  Drop redundant, shrink irrelevant: Selective knowledge injection for language pretraining.  In _In IJCAI_ , 2021a.
  * Zhang et al. (2021b) Taolin Zhang, Zerui Cai, Chengyu Wang, Minghui Qiu, Bite Yang, and Xiaofeng He.  Smedbert: A knowledge-enhanced pre-trained language model with structured semantics for medical text mining.  _arXiv preprint arXiv:2108.08983_ , 2021b.
  * Zhang et al. (2021c) Xiao-Chen Zhang, Cheng-Kun Wu, Zhi-Jiang Yang, Zhen-Xing Wu, Jia-Cai Yi, Chang-Yu Hsieh, Ting-Jun Hou, and Dong-Sheng Cao.  Mg-bert: leveraging unsupervised atomic representation learning for molecular property prediction.  _Briefings in Bioinformatics_ , 2021c.
  * Zhang et al. (2019) Zhengyan Zhang, Xu Han, Zhiyuan Liu, Xin Jiang, Maosong Sun, and Qun Liu.  ERNIE: Enhanced Language Representation with Informative Entities.  In _Proceedings of the 57th Annual Meeting of the ACL_. ACL, 2019\.
  * Zhu et al. (2020) Jiong Zhu, Yujun Yan, Lingxiao Zhao, Mark Heimann, Leman Akoglu, and Danai Koutra.  Generalizing graph neural networks beyond homophily.  In _NeurIPS_ , 2020.
  * Zitnik et al. (2019) Marinka Zitnik, Rok Sosič, Marcus W. Feldman, and Jure Leskovec.  Evolution of resilience in protein interactomes across the tree of life.  _Proceedings of the National Academy of Sciences_ , 116(10), 2019.

##  A Code and Data Availability

Our code, synthetic datasets, and pointers to real-world datasets are
available during here: https://github.com/mmcdermott/structure_preserving_pre-
training

##  B Existing Language Model Pre-training Methods - Free-text Descriptions

In this section, we describe all of the 27 models featured in Table 1,
highlighting on which domain they operate and other key details of their
approach. We will break these models down into categories, and also include
some additional models in this section not included in Table 1 that do not fit
into the paradigm we explore here.

Language modelling alone.

Liu et al. (2019b)

General domain NLP; RoBERTa includes only a masked language modelling
objective.

Radford et al. (2018, 2019); Brown et al. (2020a)

General domain NLP; The GPT series of models use autoregressive language
modelling alone, and focus on generative language tasks, not general PT/FT,
though GPT-III does show that by reframing many classical NLP fine-tuning
tasks as generative language tasks, GPT-III can still offer a compelling zero
and few-shot solution to these tasks using only the pre-trained embedder
(Brown et al., 2020a).

Rao et al. (2019)

Protein sequences; the TAPE benchmark profiles various language-model based PT
models against a benchmark of PT/FT tasks.

Zhang et al. (2021c)

Molecular Graphs; Molecular Graph BERT (MG-BERT; no relation to MG-BERT
(BehnamGhader et al., 2021)) uses masked atom prediction to pre-train a GNN
over molecular graphs.

Language modelling & Per-token KG Integration.

Sun et al. (2019)

General domain NLP; ERNIE 1 augments traditional MLM with entity-specific
masking (e.g., masking the word “Mozart” from the sentence “Mozart was a
musician”) to force the model to recover common-sense knowledge about named
entities.

He et al. (2020)

General domain NLP; KgPLM adapts the discriminative training ideas of ELECTRA
(Clark et al., 2019) alongside the idea of entity masking explored previously.
They perform both entity masking and a discriminative loss identifying which
tokens were replaced focused on entity replacements.

Qin et al. (2020)

General domain NLP; ERICA presents a mechanism for leveraging contrastive
learning and distant supervision to incorporate external knowledge into a PLM
for improving language understanding. They augment MLM with two per-token
tasks aimed at ensuring the per-token representations within a document
reflect the structure of the KG; first by ensuring that the pooled
representations of head and tail entities are similar conditioned on a
relation (which is prepended to the document prior to embedding) and second
that relation embeddings (defined as concatenated head, tail per-token entity
embeddings) are similar within and across documents. As both tasks are done on
per-token embeddings, and never at a per-sample level, this approach induces
minimal constraints on the per-sample latent space.

Yuan et al. (2021)

Biomedical domain NLP; KeBioLM integrates a per-token KG into a biomedical
language model by augmenting token entity representations with attention look-
ups into a biomedical KG (regardless of whether the attended to entities
actually match a given entity mention in the source text, though they do only
apply this on recognized entities). In order to ensure this attention is
meaningful, they further perform named entity linking and recognition
auxiliary PT objectives, leveraging the same KG embeddings used during the
attention calculation (thereby incentivizing per-token representations to be
similar to their associated entities representations thus ensuring that those
entities are reflected in the attention over the KG, modulo the fact that the
loss is applied after the final layer, rather than the layer directly used to
do the search). KG embeddings are initialized using Trans-E (Mroueh et al.,
2012). Their usage of automatically attending over entities within their
language model (without explicit constraints on those matches) is motivated by
Févry et al.’s work in Févry et al. (2020) and has similarities to Know-BERT
(Peters et al., 2019).

Yamada et al. (2020)

General domain NLP; LUKE performs pre-training using MLM and an entity-
specific masking/recognition scheme that is a slight variation on the
traditional entity-specific masking Sun et al. proposed. At FT time, they have
other knowledge-specific integrations, including specialized query matrices in
KQV attention based on attending to either traditional tokens or entities, but
at PT time their only modulation over a ROBERTA (Liu et al., 2019b) baseline
is their entity masking task.

Peters et al. (2019)

General domain NLP; Know-BERT integrates per-token entity information into an
MLM pre-training scheme by performing unconstrained attention over a per-
entity knowledge graph (only on pre-identified candidate entity spans),
alongside any available entity linking supervision information via direct
Named Entity Linking. This has similarities with (Yuan et al., 2021) and
(Févry et al., 2020).

Sun et al. (2020a)

General domain NLP; COLAKE performs a priori entity linking on the source
text, then replaces per-token mentions with entity embeddings, and appends to
the input text sub-graphs from a (relational) knowledge graph, including both
neighboring mentions and relations in the augmented input text block. This
input is then encoded via a variant of a transformer which limits attention
flow between tokens of different types and trains the entire ensemble with
masked language, entity, and relation modelling.

Language modelling, Per-token KG Integration, & Supervised Classification.

Sun et al. (2020b, 2021)

General domain NLP; ERNIE 2.0 & 3.0 augments traditional MLM with entity-
specific masking (e.g., masking the word “Mozart” from the sentence “Mozart
was a musician”) as well as a multi-task per-sample task, largely motivated at
classifying a block of text on the basis of internal text cohesion (predict
the true order of the sentences within an input sample & identify whether the
sentences within the input sample are spatial neighbors, come from the same
document, or come from different documents). ERNIE 3.0 additionally augments
pre-training with a per-token relation-embedding task using cloze-filling as a
vehicle to perform relation extraction on pre-specified per-token KGs.

Zhang et al. (2019)

General domain NLP; ERNIE (no relation to (Sun et al., 2019, 2020b)) uses both
architectural and objective-function changes to inject per-token knowledge
into PT. Specifically, they separately embed all named entities in a sample,
use a specific architecture to join contextualized entity embeddings alongside
the embeddings of tokens that realize that entity in the span, and perform
entity-specific masking (through both entity replacement and entity-token
relationship masking). In addition, they simultaneously perform standard MLM
and next-sentence prediction in the manner of BERT (Devlin et al., 2019).

BehnamGhader et al. (2021)

General domain NLP; MG-BERT introduces a GCNN layer after BERT token which
aggregates token embeddings together over a unified graph consisting both of
co-occurrence relationships and knowledge graph relationships. Note that as
the KG layer is used during pre-training as well, this model may have
difficulty adapting to settings at FT time where the KG is unavailable, which
limits flexibility.

Yu et al. (2020)

General domain NLP; JAKET embeds entities by extracting per-token
representations of entity texts inside per-entity descriptions, then produces
updated KG embeddings via a graph convolutional neural network. Those
embeddings are then fed back into a language model alongside per-token
embeddings corresponding to those entities in the raw text. The entire system
is trained according to an MLM objective, plus entity category prediction and
relation prediction (only on the entity embeddings extracted from entity
descriptions and fed through the GCNN— _not_ on the raw entities within the
contextualized text). Notably, unlike many other methods, this approach does
leverage relations in the formation of individual per-sample and per-token
embeddings, so can’t be used at inference time without that level of
information as well.

Su et al. (2021)

General domain NLP; Coke is similar to ERNIE (Zhang et al., 2019) and JAKET
(Yu et al., 2020) in that it aggregates entity information by leveraging a
GCNN over a restricted dynamic context KG based on token-entity mentions then
integrates those augmented embeddings into the per-token embeddings of a BERT-
style pretrained model (similar to JAKET), but also leverages the denoising
entity autoencoder task of ERNIE (Zhang et al., 2019). By leveraging a GCNN
over entity relationships to directly inform per-token embeddings, relations
must be known about per-token entities at FT time as well.

Zhang et al. (2021b)

Medical domain NLP; SMedBERT leverages a complex, multi-faceted loss including
MLM, Sentence-order prediction SOP (as introduced in, e.g., ALBERT (Lan et
al., 2019)), and includes per-token KG information by aggregating token
embeddings across KG embeddings (produced via trans-H (Wang et al., 2014))
corresponding to matching entities and the neighbors of matching entities in
the KG. They also include variations on relation and entity masking to ensure
the PT model learns per-token information corresponding to the KG. This method
bares similarity to Coke (Su et al., 2021) and JAKET (Yu et al., 2020).

Language modelling & Single-task Classification.

Devlin et al. (2019)

General domain NLP; Masked language model plus binary classification of
whether text block is sequentially consistent, with samples chosen via true
positive pairs vs. randomly joined sentences.

Lan et al. (2019)

General domain NLP; Masked language model plus binary classification of
whether text block is sequentially consistent, with samples chosen via true
positive pairs vs. re-ordered positive sentence pairs.

Min et al. (2020)

Protein sequences; Masked language model plus multi-class classification of to
which protein family an input sequence belongs. Uses non-standard whole-
sequence embedding procedure (no [CLS] token).

Language modelling & Multi-task Classification.

Liu et al. (2019a)

General domain NLP; Masked language model plus multi-task classification
across a variety of NLP tasks.

Hu et al. (2020b)

Graph data; This model uses a masked imputation task similar to a masked
langauge model and a highly multi-task supervised whole-graph level
prediction. On this non-NLP domain, Hu et al. finds that the multi-task whole-
graph level task is essential for performance.

McDermott et al. (2020)

EHR Timeseries data; This model uses a masked imputation task similar to a
masked language model over timeseries data and a multi-task supervised whole-
sequence prediction task. On this non-NLP domain, McDermott et al. finds that
the multi-task whole-sequence level task is essential for performance.

Language modelling & Whole-sample KG Constraints via Entity-descriptions.

Wang et al. (2019c)

General domain NLP; KEPLER leverages the most complex inter-sample geometric
constraint observed in existing methods. It augments traditional MLM on
generic text samples with a constraint ensuring the (per-sample) embeddings of
entity descriptions pulled from pre-specified knowledge graphs (KGs) reflect
geometric constraints in accordance with the underlying relations in the KG
(in particular leveraging the Sun et al. (2018) geometric constraints). As we
will see in our theoretical analyses, these constraints are much more
restrictive on the latent space geometry, and thus imply a greater encoding of
domain knowledge in the model. Note that JAKET (Yu et al., 2020) also
leverages entity descriptions in its per-token encoding, but these
descriptions are (1) extracted via per-token embeddings, using the first
mention of the token, not whole-sample embeddings, and (2) integrated back
into the original text in a per-token manner, not optimized over directly via
geometric constraints as in KEPLER.

Methods not natively captured in our framework as they leverage relation
information during encoding directly.

Yu et al. (2020)

As outlined above, JAKET isn’t natively captured in our framework as it
actively leverages the relationships in a KG to produce KG-contextualized per-
token entity embeddings during processing. Thus, this system cannot be fine-
tuned on KG-free data without significant domain shift, unlike the PT methods
we focus on.

Su et al. (2021)

CokeBERT leverages both entities and relatonships, summarizing KG-
contextualized embeddings at a per-token level. Similar to JAKET and
Graphformer, the use of relationships here means that at FT time, if the per-
token knowledge cannot be embedded similar to at PT time, the model will face
a domain shift burden.

Zhang et al. (2021b)

Given SMedBERT leverages (per-token) relational information in producing
unified token embeddings, it is not natively realizable within our framework
which assumes that the input to the model is the raw data domain.

Sun et al. (2021)

By leveraging relational data via their relational masking task, ERNIE3 may
also have limited generalizability to other contexts.

Sun et al. (2020a)

Given COLAKE leverages (per-token) relational information to construct the
augmented inputs to the model, it requires per-token relational data about the
source inputs at both PT and FT time.

Methods orthogonal to our framework.

Liu et al. (2021)

KG-BART is a text-generation model that leverages per-token knowledge after a
text-encoder to enrich generated text with information from a textual
knowledge graph (in a per-token manner). It is neither used for general pre-
training nor does it leverage any additional per-sample constraints.

Yang et al. (2021)

Text-based Knowledge Graphs; This work produces embeddings of nodes in KGs by
combining transformer based text encodings with graph convolutional network KG
embedding methods, leveraging link prediction as the pre-training task. The
individual nodes are represented by entity descriptions / textual features,
and link prediction can be seen as inducing a geometric constraint via the
connectivity of the knowledge graph on whole-sample embeddings. However, given
that relationships are used in encoding the data as well, GraphFormer can’t be
used in a context where KG links may not be observed at FT time and should be
seen not as a general text PT method but instead simply as an advanced KG
embedding mechanism, so it does not directly fall under our framework.

Agarwal et al. (2021)

KeLM (unrelated to KELM (Lu et al., 2021)) is a method for converting a free-
text KG into textual nodes so language modelling can be used over that corpus,
and is orthogonal to the methods of pre-training.

Ribeiro and Forbus (2021)

This paper is a method for populating a KG from free-text via BERT; it has no
bearing on incorporating structure or knowledge into PT itself and as such is
not relevant to our system.

Zhang et al. (2021a)

This paper presents a method to drop redundant triples from a knowledge graph
and a regularization technique to limit the impact of added irrelevant
knowledge to per-token knowledge-enhanced PT methods such as ERNIE (Zhang et
al., 2019).

Yao et al. (2019)

Text-based Knowledge Grpahs; KG-BERT is a method for which knowledge graph
completion in which textual representations of entities and relations in KGs
are embedded by fine-tuning a pre-trained BERT style transformer for link
prediction over a given KG. As this is only for knowledge graph completion, it
is orthogonal to our study of pre-trained models in general.

Methods that only change things at FT time.

Meng et al. (2021)

Biomedical domain NLP; MOP doesn’t change anything at PT time, but further
trains sub-KG adapters on entity recognition tasks prior to FT to infuse
entity knowledge into the PT system. This is also still per-token. This is
similar in spirit to the K-Adapters work which introduced the notion of KG
adapters (Wang et al., 2020).

Liu et al. (2020)

General domain NLP; K-BERT, at PT time, is actually equivalent to BERT (Devlin
et al., 2019). However, it does do other, interesting things at FT time,
including augmenting the sentence flow with injected per-token knowledge
graphs and limiting self-attention to only flow along links supported either
by the original sentence or the injected knowledge. But, as this is only true
at FT time, at PT time it is equivalent to BERT.

Reimers and Gurevych (2019)

General domain NLP; This model, at PT time, is also actually equivalent to
BERT (Devlin et al., 2019). Like Liu et al. (2020), however, it specializes a
fine-tuning procedure for sentence information retrieval tasks, which is
actually similar to how PT is adapted in this framework.

Lu et al. (2021)

General domain NLP; KELM doesn’t modify PT objective explicitly, but insteads
enhances a model at FT time by injecting per-token knowledge via a GNN module
atop the pre-trained LM embeddings via a unified text-entity graph. It is
similar to KBERT (Liu et al., 2020) in this way, but resolves other issues
with that approach relating to knowledge ambiguity and by supporting multi-hop
reasoning, again over the per-token embeddings.

Faldu et al. (2021)

General domain NLP; KI-BERT augments BERT with KG specific information via
joint token-entity embeddings and information fusion, but does this only at FT
time.

Yan et al. (2021)

General domain NLP; K-XLNet introduces a secondary FT stage in which knowledge
injectors throughout an XL-Net architecture are further trained to leverage
knowledge (encoded via free-text entity descriptions) that is injected into
input sentences alongside matched tokens. It does not modify the XL-Net PT
stage at all.

Wang et al. (2020)

General domain NLP; K-Adapter proposes to also pre-train various knowledge
adapters that can be used alongside a pre-trained language model at fine-
tuning time. Thus, while there is a pre-training process for the adapters,
this process does not modulate the original pre-trained language model at all.
In addition, both adapters pre-trained in this work are based on per-token
knowledge graphs; one leverages concatenated entity embeddings to perform
relation classification, another predicts which token in the sentence is the
“head” in a dependency parse tree, so no per-sample constraints are applied.

Poerner et al. (2020)

General domain NLP; E-BERT injects per-token knowledge into BERT by first
aligning embeddings of a knowledge grpah with the input wordpiece embedding
space of a (fixed, pre-trained) BERT model, then using various strategies to
input them alongside their source mentions in FT text. They do no additional
pre-training of this system, and so this model only affects the model at FT
time.

Gunel et al. (2021)

General domain NLP; Gunel et al. augment LMPT systems with an additional, pre-
FT procedure in which the model is further trained using a supervised, per-
sample metric learning task leveraging FT labels directly to form the classes
used for metric learning. They do not materially change the task-independent
PT procedure at all, though their FT metric learning procedure does induce
some structure at the per-sample level.

##  C Language Model Pre-trianing on Natural Language vs. General Domains

In this section, we explore the hypothesis introduced in Section 2.2 more
fully; namely, that language model (LM) Pre-training (PT) methods have limited
ability to reflect inter-sample relationships in domains outside of NLP. We
break this hypothesis down into 3 claims:

  1. 1.

LM PT natively captures _intra-sample_ relationships.

  2. 2.

LM PT is _not_ guaranteed to capture inter-sample information in general
domains.

  3. 3.

LM PT (in particular, language modelling or masked language modelling at a
per-word level) does capture inter-sample information in NLP, via a mechanism
that is unlikely to generalize to other domains.

In the remainder of this section, we first define “intra-” vs. “inter-” sample
relationships formally, then expand each of these 3 claims to argue our
overall hypothesis. Establishing each of these claims will ultimately suggest
that just because LM PT methods work well in NLP does not mean they do not
need modification to better capture inter-sample associations on more general
domains.

Intra- vs. inter-sample relationships. We use the terms _sample_ to refer to
an individual datapoint in a dataset. _Intra-sample_ relationships thus refer
to information that helps model one part of a single sample vs. other parts of
that same sample. _Inter-sample_ , on the other hand, reflects information
that models an entire sample vs. other samples in the same dataset. For
example, if we have a dataset of sentences such as in a NLP context, then (1)
a _sample_ refers to a single sentence, (2) masked language modelling, where
we occlude a word in the sentence and try to infer it from the rest of the
words in the sentence is an _intra-sample_ modelling task, as it relates the
occluded word to the observed words within a single sample, and (3) sentence
classification, natural language inference, or question answering tasks would
all refer to _inter-sample_ tasks, as they relate entire sentences either to
each other either directly or indirectly.

(1) LM PT-derived PT natively captures _intra-sample_ relationships. This
claim is clear from the definitions above—any form of intra-sample imputation
is directly formulating an intra-sample task. While we do not claim that LM PT
will necessarily capture every intra-sample association (e.g., most LM PT
methods will have some limit on how much of the input they are allowed to mask
at once, which precludes the model from learning to impute given very sparse
input features), it is clear that it naturally reflects a rich subset of
intra-sample associations natively.

(2) LM PT-derived PT does not need to capture _inter-sample_ relationships..
Though it is clear that LM PT-derived PT natively captures intra-sample
relationships, this does not preclude it in general from capturing inter-
sample relationships. However, for any style of LM PT-derived PT, we can
quickly design a synthetic dataset that shows that inter-sample relationships
need not be reflected. For example given a dataset $\boldsymbol{X}$, simply
appending an extra, categorical feature to each sample in $\boldsymbol{X}$
that has no dependence on the existing features in $\boldsymbol{X}$ induces an
inter-sequence relation between elements of $\boldsymbol{X}$ (sharing the same
appended feature) that is _not_ capture-able by imputation PT (because the
appended feature, having no dependence on the other features, is neither
impute-able nor informative in imputing other tokens, and thus is not relevant
to the imputation task).

LM PT does capture inter-sample information in NLP, but is unlikely to in
other real-world domains. The fact that LM PT-derived PT is not guaranteed to
capture inter-sequence information in general, yet simultaneously is very
effective as a PT strategy on NLP suggests that in either NLP-specifically or
real-world domains in general LM PT-derived PT actually will regardless
reflect inter-sample information. Here, we argue that this property holds only
for NLP-specifically, due to unique properties of NLP that allow one to
consistently re-formulate inter-sample tasks as intra-sample tasks, which is
not possible in other domains. This reformulation argument has received
significant attention in the literature. In particular, Saunshi et al. (2021)
shows through both theoretical and empirical analysis how LM PT performance on
inter-sample tasks can be directly related to their intra-sample performance
through this reformulation argument presented above. More indirectly, Brown et
al. (2020b) has shown that language models are effective few-shot learners,
even in inter-sample tasks. Concretely, they show that GPT-3 achieves strong
performance on question-answering tasks even in the _zero-shot_ setting—thus
showing clearly that LM PT (here, in particular, autoregressive language-model
pre-training) explicitly captures the inter-sample task of question-answering
simply through language modeling alone. Raffel et al. (2019) has also
leveraged this reformulation trick explicitly to form a unified PT objective
to great effect.

Critical to this argument is the fact that not only is it possible to perform
this reformulation, but that these reformulated examples qualify both as (1)
valid NLP samples themselves and (2) sufficiently reasonable NLP samples that
they may be observed in large PT datasets. For example, the reformulation
trick leveraged for the “review sentiment analysis” task in Saunshi et al.
(2019) is to append the tokens “:)” after the end of positive-sentiment
reviews and “:(” after negative-sentiment reviews. This addition of an emoji
indicating sentence sentiment is very plausible in real-world datasets.

Other domains do not permit this inter- to intra-sample reformulation, at
least not while maintaining validity and plausibility. For example, protein
sequences do not permit arbitrary modifications to amino acid sequence to
still remain valid, plausible proteins, and even if they did there is no way
to naturally encode the outcome to all inter-sample tasks of interest. Thus,
this argument suggests strongly that the efficacy of LM PT for inter-sample
tasks on NLP is not guaranteed, and possibly not even likely, to extend to
other domains of interest.

##  D Further Theoretical Analysis of SIPT

We proceed with additional theoretical analysis of SIPT and formal proofs of
the claims made in the main paper. Note that throughout all results, we will
assume that the PT and FT datasets are iid, that FT tasks, though they may be
unobserved over PT samples, are well defined over the entire PT and FT domain
and thus true labels (if unknown) do exist for PT samples, and that the
sampling distribution of the PT/FT data has full support over the label-space
of any considered task.

###  D.1 Preliminaries

####  D.1.1 How do we define “optimal” embeddings?

Let $G=(V,E)$ be an arbitrary graph, $d_{e}\in\mathbb{N}$ be an embedding
dimension, $\mathcal{D}$ be a distance function.

######  Definition 3 (SI-optimal embedding)

We will say that there exists a SI-optimal embedding of $G$ if there exists a
positive _radius_ $r\in\mathbb{R}^{+}$ and mapping $f:V\to\mathbb{R}^{d_{e}}$
such that $\mathcal{D}(f(v_{1}),f(v_{2}))<c$ if and only if $(v_{1},v_{2})\in
E$.

Note that this is a specialization of the idea that we can leverage a
connectivity algorithm to recover $G$ from an optimal embedding; in
particular, here we take the connectivity algorithm used to be a radius
nearest neighbor algorithm with distance $\mathcal{D}$. Note that any SI-
optimal embedding also yields a _margin_ parameter
$m=\min_{(v_{i},v_{j})\not\in E}(D(f(v_{i}),f(v_{j}))-\max_{(v_{i},v_{j})\in
E}(D(f(v_{i}),f(v_{j}))$ where $m\in\mathbb{R}^{+}$ given $G$ is finite and
$f$ is an optimal embedding.

####  D.1.2 SI-optimal embeddings determine SIPT-optimal embeddings

Suppose we have PT dataset $\boldsymbol{X}_{\mathrm{PT}}$, $G_{\mathrm{PT}}$.
We will consider some encoder $f$ optimal (and its embeddings optimal as well)
provided that $f$ produces optimal embeddings both according to the LM PT
objective and the SI objective. Note that this definition is not
contradictory—the LM PT objective does not constrain the global embedding
space geometry of $f$ at all, and the SI objective is focused exclusively on
the overall geometry. Given our primary interest in this work is on the
properties of the SI loss, we will operate from the perspective of only
considering that loss.

####  D.1.3 When do optimal embeddings exist?

The study of when optimal embeddings exist for SIPT problems covers many
topics of existing interest within graph embedding, metric learning, and other
fields of machine learning. Various results are already known about limits on
optimal embeddings will exist for certain kinds of graphs; however, in full
generality typically $d_{e}\in\Theta(|V|)$ may be required. We will highlight
several results here that will be of particular relevance for our purposes.
Note that, throughout, we will assume that $G_{\mathrm{PT}}$ is _fully
inferrable_ from $\boldsymbol{X}_{\mathrm{PT}}$—in other words, there exists a
function $f:\mathcal{X}\times\mathcal{X}\to\\{0,1\\}$ that correctly predicts
whether or not there will be an edge between two nodes $\boldsymbol{x}$ and
$\boldsymbol{x}^{\prime}$. In the case that this is not possible, then there
will not exist optimal embeddings. Our focus in the rest of this section
focuses around when, given that relationship exists, we will be able to
efficiently (relative to the necessary embedding dimension $d_{e}$) produce
such embeddings given various graph structures.

Disconnected Cliques. If the graph $G_{\mathrm{PT}}$ contains a serious of
disconnected cliques, then it is trivial to see that optimal embeddings do
exist, as we can simply assign each clique a single point and map all elements
within that clique to that single point. This embedding schema works in
arbitrary dimensions, and enforces minimal geometry over the latent space.

Manifolds. Let $\mathcal{M}$ be a manifold of dimension $N$ embedded in
$\mathbb{R}^{2N}$, where such an embedding is guaranteed to exist by the
Whitney Embedding Theorem Whitney et al. (1992); Mukherjee (2015). Note that
for any point $\boldsymbol{x}\in\mathcal{M}$, with its corresponding embedding
$\pi(\boldsymbol{x})\in\mathbb{R}^{2N}$, there exists a radius
$r_{\boldsymbol{x}}\in\mathbb{R}^{+}$ such that the ball of radius $r$ around
$\pi(\boldsymbol{x})$ in $\mathbb{R}^{2N}$ intersection
$B_{r_{\boldsymbol{x}}}(\pi(\boldsymbol{x}))\cap\pi(\mathcal{M})$ is
isomorphic both to some ball around $\boldsymbol{x}$ in $\mathcal{M}$ and to a
ball of radius $r^{\prime}$ in the copy of $\mathbb{R}^{N}$ isomorphic to the
neighborhood of $\boldsymbol{x}$ on $\mathcal{M}$. Let us call the manifold
$\mathcal{M}$ valid if
$\mathrm{inf}_{\boldsymbol{x}\in\mathcal{M}}(r_{\boldsymbol{x}})>0$.

###### Theorem 4

Let $X_{\mathrm{PT}}\in\mathcal{X}^{N_{\mathrm{PT}}}$ be our PT dataset as
usual, and let $\mathcal{M}$ be a valid manifold. Suppose there exists some
(potentially unknown) isomorphism $f:\mathcal{X}\to\mathcal{M}$. If we then
set $G_{\mathrm{PT}}$ to be a radius nearest-neighbor graph over
$\mathcal{M}$, then there exist $SIPT$-optimal embeddings with as few as
$d_{e}=2d_{m}$ dimensions.

Proof  This a direct consequence of the definition (above) of a _valid_
manifold. As
$r^{*}=\mathrm{inf}_{\boldsymbol{x}\in\mathcal{M}}r_{\boldsymbol{x}}>0$, we
know that we can select some strictly positive $c\leq r^{*}$ such that all
neighborhoods within $c$ of any point in the embedding of the manifold only
intersect the manifold in a manner that preserves geodesic distance. Then, on
the basis of the Whitney embedding theorem, SIPT-optimal embeddings do exist.

###  D.2 Proofs for Section 4

###### Theorem 5

Let $\boldsymbol{X}_{\mathrm{PT}}$, $G_{\mathrm{PT}}$ be a PT dataset and
graph, respectively, such that SIPT-optimal embeddings exist and are realized
by an embedder $f$. Then, under embedding $f$ and given sufficient data, the
nearest-neighbor accuracy for a FT task $y$ will converge as dataset size
increases to _at least_ the local consistency of $y$ over $G_{\mathrm{PT}}$.

Proof  Given $f$ realizes SIPT-optimal embeddings, we know that if we define a
$r$-NN predictor via the same radius $r^{*}$ at which $f$ achieves optimality,
then this predictor will be correct exactly as often as the label of a given
node in the graph $G_{\mathrm{PT}}$ agrees with the labels of its
neighbors—which is precisely $\mathrm{LC}_{G_{\mathrm{PT}}}(y)$. This
classifier may not be well defined for small FT dataset sizes, however, as if
data is not sufficiently dense there may be no data-points within radius $r$
of a given query. Similarly, without sufficient PT data then the $\mathrm{LC}$
computed over the empirical distribution of the graph $G_{\mathrm{PT}}$ may be
a poor proxy for the true distribution. As PT and FT dataset sizes increase,
however, we know that we can always achieve at least this performance. We may
be able to achieve even higher performance if other effects motivate stronger
performance at radii smaller than $r^{*}$, but this is not guaranteed.

###  D.3 Proofs for Section 4.1

In this section, we give formal theorem statements and proofs for the 3
findings reported in Section 4.1.

Finding 1. Multi-class or supervised metric learning PT methods correspond to
cliques.
Supervised metric learning PT is when we are given a PT dataset
$\boldsymbol{X}_{\mathrm{PT}}$,
$\boldsymbol{y}_{\mathrm{PT}}\in\mathcal{Y}^{N}$ and attempt to pre-train an
encoder $f$ such that (potentially in addition to other constraints),
$f(\boldsymbol{x}_{i})$ is similar to $f(\boldsymbol{x}_{j})$ if
$y_{i}=y_{j}$. From that definition, it is clear that this is identical to a
SI objective if $G_{\mathrm{PT}}=(\boldsymbol{X}_{\mathrm{PT}},E)$ such that
$(\boldsymbol{x}_{i},\boldsymbol{x}_{j})\in E$ if and only if $y_{i}=y_{j}$.
This corresponds to a graph of disconnected cliques, one for each label
$\ell\in\mathcal{Y}$.

In supervised classification PT, we have a PT dataset
$\boldsymbol{X}_{\mathrm{PT}}$ and labels
$\boldsymbol{y}_{\mathrm{PT}}\in\mathcal{Y}^{N}$, and attempt to pre-train an
encoder $f$ and predictor $h$ such that (potentially in addition to other
constraints), $h(f(\boldsymbol{x}_{i}))=y_{i}$. We can realize an analogous SI
objective with an identical construction as that for supervised metric
learning. If $G_{\mathrm{PT}}$ has edges
$(\boldsymbol{x}_{i},\boldsymbol{x}_{j})$ such that $y_{i}=y_{j}$, then
finding SIPT-optimal embeddings under $G_{\mathrm{PT}}$ will also solve the
supervised classification objective natively, through nearest-neighbor
methods. Thus, SIPT with a graph of disconnected cliques encompasses the
supervised classification objective.

Finding 2. The local consistency of any downstream task w.r.t. such a graph
rapidly tends to a fixed, sub-optimal constant as PT dataset size increases.
Finding 2 is formally described in Theorem 6.

###### Theorem 6

Let $\boldsymbol{X}_{\mathrm{PT}}\in\mathcal{X}^{N}$, be a PT dataset with
corresponding labels $\boldsymbol{y}\in\mathcal{Y}_{\mathrm{PT}}^{N}$. Define
$G_{\mathrm{PT}}=(\boldsymbol{X}_{\mathrm{PT}},E)$ such that
$(\boldsymbol{x}_{i},\boldsymbol{x}_{j})\in E$ if and only if $y_{i}=y_{j}$.
Let $y_{\mathrm{FT}}:\mathcal{X}\to\mathcal{Y}_{\mathrm{FT}}$ correspond to a
FT downstream task. Further, define
$\mathrm{MC}(\boldsymbol{x}_{i},y_{\mathrm{FT}})=\operatorname*{argmax}_{\ell_{\mathrm{FT}}\in\mathcal{Y}_{\mathrm{FT}}}\sum_{\boldsymbol{x}_{j}\in\boldsymbol{X}_{\mathrm{PT}}|y_{i}=y_{j}}\mathds{1}_{y_{\mathrm{FT}}(\boldsymbol{x}_{i})=\ell_{\mathrm{FT}}}$
to be the majority class label for the clique containing $\boldsymbol{x}_{i}$
in $G_{\mathrm{PT}}$. Then,

$\mathrm{LC}_{G_{\mathrm{PT}}}(y_{\mathrm{FT}})=\sum_{\ell_{\mathrm{PT}}\in\mathcal{Y}_{\mathrm{PT}}}\mathbb{P}(y_{i}=\ell_{\mathrm{PT}})\mathbb{P}(y_{\mathrm{FT}}(\boldsymbol{x}_{i})=\mathrm{MC}(\boldsymbol{x}_{i},y_{\mathrm{FT}})|y_{i}=\ell).$

Proof  This follows directly from the definition of Local Consistency,
$G_{\mathrm{PT}}$, and the law of total probability. In particular,

$\displaystyle\mathrm{LC}_{G_{\mathrm{PT}}}(y_{\mathrm{FT}})$
$\displaystyle=\mathbb{P}\left(y_{\mathrm{FT}}(\boldsymbol{x}_{i})=\operatorname*{argmax}_{\ell\in\mathcal{Y}_{\mathrm{FT}}}\sum_{\boldsymbol{x}_{j}\in\boldsymbol{X}_{\mathrm{PT}}|(\boldsymbol{x}_{i},\boldsymbol{x}_{j})\in
E(G_{\mathrm{PT}})}\mathds{1}_{y_{\mathrm{FT}}(\boldsymbol{x}_{i})=\ell}\right)$
$\displaystyle=\mathbb{P}\left(y_{\mathrm{FT}}(\boldsymbol{x}_{i})=\mathrm{MC}(\boldsymbol{x}_{i},y_{\mathrm{FT}})\right)$
$\displaystyle=\sum_{\ell_{\mathrm{PT}}\in\mathcal{Y}_{\mathrm{PT}}}\mathbb{P}(y_{i}=\ell_{\mathrm{PT}})\mathbb{P}(y_{\mathrm{FT}}(\boldsymbol{x}_{i})=\mathrm{MC}(\boldsymbol{x}_{i},y_{\mathrm{FT}})|y_{i}=\ell),$

as desired.

Note that this has dependence on the PT dataset size as the probabilities
$\mathbb{P}$ are taken over the empirical distribution induced by the dataset
$\boldsymbol{X}_{\mathrm{PT}}$ and graph $G_{\mathrm{PT}}$ inherent in local
consistency — if $\boldsymbol{X}_{\mathrm{PT}}$ is too small, these empirical
distributions will be poor proxies for the true distribution and this property
will not hold as cleanly. However, once saturation is reached it will not
improve beyond this fixed upper bound relating to task correlation.

Finding 3. In contrast, manifold nearest neighbor representations yields
nearest-neighbor FT performance that converges to optimality.
Finding 3 is formally described in Theorem 7.

###### Theorem 7

Let $\boldsymbol{X}_{\mathrm{PT}}$ be a PT dataset which can be realized over
a valid manifold $\mathcal{M}$. Assume $\boldsymbol{X}_{\mathrm{PT}}$ is
sampled with full support over $\mathcal{M}$. Let
$G_{\mathrm{PT}}(\boldsymbol{X}_{\mathrm{PT}},E)$ be an $r$-nearest-neighbor
graph over $\mathcal{M}$ (e.g., $(\boldsymbol{x}_{1},\boldsymbol{x}_{2})\in E$
if and only if the geodesic distance between the two points on $\mathcal{M}$
is less than $r$:
$\mathcal{D}_{\mathcal{M}}(\boldsymbol{x}_{1},\boldsymbol{x}_{2})<r$).

Let $y_{\mathrm{FT}}$ be a FT classification task that is almost everywhere
smooth on the manifold (as $y_{\mathrm{FT}}$ is a classification task, this
means that almost everywhere
$y_{\mathrm{FT}}(\boldsymbol{x}_{1})=y_{\mathrm{FT}}(\boldsymbol{x}_{1}+\boldsymbol{\varepsilon})$
where $\boldsymbol{\varepsilon}$ is a sufficiently small transformation along
$\mathcal{M}$. Note this requirement is merely stating more formally that
$y_{\mathrm{FT}}$ is “piecewise constant” along $\mathcal{M}$, such that
compact regions of $\mathcal{M}$ correspond to individual labels of
$y_{\mathrm{FT}}$ and only on points at the boundary between two labels does
$y_{\mathrm{FT}}$ show non-constant behavior.

Then, as PT dataset size (and thus size of $G_{\mathrm{PT}}$) tends to
$\infty$ and $r$ tends to zero, the local consistency of $y_{\mathrm{FT}}$
over $G_{\mathrm{PT}}$ will likewise tend to 1.

Proof  As $r\to 0$, provided PT dataset size increases at a sufficient
associated rate so as to maintain constant minimum degree of $G$, we have the
property that the total diameter over $\mathcal{M}$ contained in a node’s
local neighborhood within $G_{\mathrm{PT}}$ likewise decreases. Given some
fixed node $\boldsymbol{x}\in\mathcal{M}$ that is within the interior of a set
of constant $y_{\mathrm{FT}}$ label, this implies that, eventually, it will
grow sufficiently small that all of $\boldsymbol{x}$’s neighbors share the
same label as $\boldsymbol{x}$ under $y_{\mathrm{FT}}$.

More concretely, it is clear that this point will occur exactly when $r$ is
the geodesic distance between $\boldsymbol{x}$ and the boundary of the
surrounding constant-label patch containing $\boldsymbol{x}$. But, it is clear
that this implies that the only sections of $\mathcal{M}$ will not have the
property that neighborhoods around points will be constant w.r.t.
$y_{\mathrm{FT}}$ labels will almost everywhere be patches within distance $r$
of the points where $y_{\mathrm{FT}}$ changes.

This implies that as $r\to 0$, then almost everywhere will the neighborhoods
around a node $\boldsymbol{x}$ be constant w.r.t. $y_{\mathrm{FT}}$. However,
this implies that almost everywhere would $y_{\mathrm{FT}}$ display perfect
local consistency, as desired.

##  E Full Experimental Details

###  E.1 Full Details for Synthetic Experiments

####  E.1.1 Details on Dataset & Graph Construction

General Setup. All our synthetic datasets leverage free-text sentences from
https://www.kaggle.com/mikeortman/wikipedia-sentences (CC BY-SA 4.0 License).
Topics were assigned to these sentences by running Latent Dirichlet Allocation
via Scikit-learn Pedregosa et al. (2011) over a Bag-of-words representation to
100 topics, with otherwise default parameters.

Given the probabilities over all 100 topics, we treated the prediction of the
most probable topic as a 100-class multi-class classification problem for our
FT task in these experiments.

Cliques Graph Setup. To construct the Cliques graph setting, we choose a
random subset of sentences as $\boldsymbol{X}_{\mathrm{PT}}$ and define
$G_{\mathrm{PT}}=(\boldsymbol{X}_{\mathrm{PT}},E)$ such that