# Full quantitative near-field characterization of strongly coupled exciton-
plasmon polaritons in thin-layered WSe2 on a monocrystalline gold platelet

Laura N. Casses  [  Binbin Zhou  Qiaoling Lin  Annie Tan  Diane-Pernille
Bendixen-Fernex de Mongex  Korbinian J. Kaltenecker  Sanshui Xiao  Martijn
Wubs  Nicolas Stenger<EMAIL_ADDRESS>[

###### Abstract

Exciton-plasmon polaritons (EPPs) are attractive both for the exploration of
fundamental phenomena and applications in nanophotonics. Previous studies of
EPPs mainly relied on far-field characterization. Here, using near-field
optical microscopy, we quantitatively characterize the dispersion of EPPs
existing in 13-nm-thick tungsten diselenide (WSe2) deposited on a
monocrystalline gold platelet. We extract from our experimental data a Rabi
splitting of 81 meV, and an experimental effective polariton loss of 55 meV,
demonstrating that our system is in the strong-coupling regime. Furthermore,
we measure for the first time at visible wavelengths the propagation length of
these EPPs for each excitation energy of the dispersion relation. To
demonstrate the quantitative nature of our near-field method to obtain the
full complex-valued wavevector of EPPs, we use our near-field measurements to
predict, via the transfer matrix method, the far-field reflectivities across
the exciton resonance. These predictions are in excellent agreement with our
experimental far-field measurements. Our findings open the door towards the
full near-field study of light-manipulating devices at the nanoscale.

###### keywords:

near-field microscopy, s-SNOM, exciton-plasmon polaritons, WSe2,
monocrystalline gold, strong coupling

DTU] Department of Electrical and Photonics Engineering, Technical University
of Denmark, 2800 Kongens Lyngby, Denmark [Nanophoton] Centre for
Nanophotonics, Technical University of Denmark, 2800 Kongens Lyngby, Denmark
[CNG] Centre for Nanostructured Graphene, Technical University of Denmark,
2800 Kongens Lyngby, Denmark [Nanophoton] Centre for Nanophotonics, Technical
University of Denmark, 2800 Kongens Lyngby, Denmark [CNG] Centre for
Nanostructured Graphene, Technical University of Denmark, 2800 Kongens Lyngby,
Denmark Current address: Attocube Systems AG, Eglfinger Weg 2, 85540 Haar,
Germany [Nanophoton] Centre for Nanophotonics, Technical University of
Denmark, 2800 Kongens Lyngby, Denmark [CNG] Centre for Nanostructured
Graphene, Technical University of Denmark, 2800 Kongens Lyngby, Denmark
[Nanophoton] Centre for Nanophotonics, Technical University of Denmark, 2800
Kongens Lyngby, Denmark [CNG] Centre for Nanostructured Graphene, Technical
University of Denmark, 2800 Kongens Lyngby, Denmark DTU] Department of
Electrical and Photonics Engineering, Technical University of Denmark, 2800
Kongens Lyngby, Denmark [Nanophoton] Centre for Nanophotonics, Technical
University of Denmark, 2800 Kongens Lyngby, Denmark [CNG] Centre for
Nanostructured Graphene, Technical University of Denmark, 2800 Kongens Lyngby,
Denmark

##  1 Introduction

Near-field microscopy is a powerful and versatile tool that has been used for
the direct observation of polaritons in different structures across a wide
range of excitation energies at the nanoscale 1, 2, 3, 4, 5, as well as the
study of the local dielectric function of various materials 6, 7, 8, 9, 10.
Many studies using a scattering-type scanning near-field optical microscope
(s-SNOM) 11 have focused on the quantitative near-field characterization of
polaritons in the mid-infrared wavelength range12, 13, 14, 15, 16, 17. In
contrast, the full s-SNOM characterization of polaritons in the visible to
near-infrared range remains unexplored. The confinement of polaritons in this
range is indeed moderate compared to other polaritons, such as plasmons in
graphene 18. Thus, in a reflection configuration, the polaritons can be
launched both by the s-SNOM tip and by the sharp edge of metallic samples. As
shown in previous works 19, 20, 21, this leads to three main excitation
channels for the polaritons, namely the tip-launched, the edge-launched, and
the tip-reflection edge-launched polaritons. The plurality of possible
excitation paths leads to complex interference patterns, which complicates the
analysis of the measurements.

As many excitonic excitations have energies in the visible and near-infrared
range 22, 3, 23, 24, the study and full quantitative characterization of
polaritons in this range are highly relevant for the future development of the
field. In particular, exciton-plasmon polaritons (EPPs) are hybrid states
emerging from the coupling between excitons 25 and surface plasmon polaritons
(SPPs) 26. These polaritons have the potential to become a building block for
applications 27, 28 such as nanoscale polaritonic lasers 29, 30, and new
quantum emitters for quantum information processing 31. A pioneering work
measured EPPs resulting from the interaction between excitons in tungsten
diselenide (WSe2) and SPPs on gold, for thicknesses of WSe2 ranging from 35 to
110 nm 32. While showing impressive confinement of the EPPs at higher
thicknesses, the high absorption losses associated with this thick excitonic
material hindered the retrieval of the upper polariton branch. Thus, the Rabi
splitting energy was obtained from numerical calculations. In addition, the
propagation length of the EPPs, directly related to the amount of radiative
and non-radiative losses present in the polaritons 33, was not determined.
This is however a crucial information to be retrieved experimentally in order
to assess the performances of a polaritonic device.

In our previous works 20, 21, we showed that by carefully choosing the
azimuthal angle of incidence relative to the edge of monocrystalline gold
platelets, we could isolate the excitation path of interest, in this case the
tip-launched SPP excitation path. Using Fourier analysis, we could retrieve
both the wavelength and the propagation length of moderately confined SPPs on
monocrystalline gold platelets.

In this work, we study EPPs propagating in a 13-nm flake of WSe2 deposited on
a monocrystalline gold platelet. We use a s-SNOM in the reflection
configuration 21 to measure the edge-launched EPP excitation path for
excitation energies around the A-exciton energy of WSe2 (1.63 eV). The choice
of perpendicular incidence and a defocused beam allows us to measure a clear
signal coming from the EPPs. With this information, we are able to retrieve
with good accuracy the experimental dispersion relation, for the lower as well
as for the upper polariton branches. As we can resolve both branches, we are
able to determine a Rabi splitting based on experimental values only, without
using numerical calculations or simulations. Furthermore, we determine the
propagation lengths of the EPPs for all the measured energies. Thereby, we
reconstruct the full complex-valued polariton wavevector, giving us access to
the fundamental properties of EPPs, such as the effective losses of the
hybridized light-matter excitation. The experimental determination of the
polariton losses shows us that our system is in the strong-coupling regime. In
addition, the near-field experimental values allow us to determine the in-
plane dielectric function of WSe2. These values, obtained from the near-field
measurements, are subsequently used to predict the far-field reflectivity of
our sample. We find good agreement with the reflectivity measured directly
with a far-field microscope.

This work reports the first quantitative full near-field characterization of
the nanoscale properties of strongly coupled light-matter states in the
visible to near-infrared energy range, the relevant range for most EPPs found
in nature. The methodology presented here pushes forward the analytical
capabilities of state-of-the-art s-SNOM allowing the full retrieval of the
coupling strength, propagating losses and effective losses of EPPs.
Furthermore, our results show that it is possible to relate near-field
information with far-field optical properties. This link between the near- and
far-field will be crucial in the design of new devices to manipulate light and
excitons on the nanoscale for future low-threshold nano-lasers, new fast
quantum gates, quantum emitters integrated in nanocavities, as well as the
development of technology-relevant metasurfaces.

##  2 Results and discussion

###  2.1 Near-field measurements of the EPPs

A sketch of the near-field experiment is presented in Figure 1a. To
investigate the properties of EPPs in the near-field, we use a 100-nm thick
monocrystalline gold platelet on silicon dioxide, protected by a 2-nm layer of
aluminum oxide. These platelets are similar to the ones studied in previous
works 34, 20, 21. A 13-nm thick flake of WSe2 is mechanically exfoliated and
transferred 35 (see Methods) onto the gold platelet such that it covers the
edge of the platelet. The EPPs are measured using a s-SNOM in the reflection
configuration 20 (see Methods). A tunable laser is focused on the s-SNOM tip
such that the laser beam is perpendicular to the edge of the platelet in the
plane of the sample (see red arrow in Figure 1b). This configuration allows to
excite EPPs of wavelength $\Lambda_{\text{ep}}$ and propagation length
L${}_{\text{p}}^{\text{ep}}$ at the sharp edge of the monocrystalline
platelet, for excitation energies from 1.5 to 1.7 eV.

The experiments are performed such that the sample is scanned under the fixed
combination of the oscillating tip and excitation light always focused on the
tip apex. With a highly focused beam, the excitation of the edge-launched EPPs
would therefore depend on the tip-edge position and thus influence the
determination of L${}_{\text{p}}^{\text{ep}}$. To overcome this challenge, we
slightly defocus the beam in front of the s-SNOM (see Methods), as in Ref. 2.
The resulting beam diameter at the tip apex is estimated to be about 50 µm.

Figure 1: Near-field measurement of exciton-plasmon polaritons (EPPs)
propagating on a thin layer of WSe2 deposited on a monocrystalline gold
platelet. (a) Illustration of the experiment. Here, $\Lambda_{\text{ep}}$ is
the wavelength of the edge-launched EPPs, and $L_{\text{p}}^{\text{ep}}$ is
the propagation length of the edge-launched EPPs. The 2-nm layer of aluminum
oxide between the gold and WSe2 is not represented here. (b) Near-field
intensity at three different excitation energies. The origin of the $x$-axis
is taken at the edge of the gold platelet, and the red arrow represents the
direction of the incident light.

Figure 1b shows three examples of the near-field intensity obtained with our
measurements, below (1.54 eV), around (1.62 eV), and above (1.69 eV) the
A-exciton energy. The maps from all the measurements can be found in the
Supplementary Information. The edge of the platelet corresponds to the origin
of the $x$-axis. The interference patterns seen on the maps and parallel to
the edge correspond to the interference between the edge-launched EPPs and the
background light scattered directly from the sample 19, 20. The wavevector of
the interference fringes $K_{\text{el}}$ can be expressed as 19

$K_{\text{el}}(\theta,\varphi)=\frac{2\pi}{\lambda_{0}}\left(-\sin(\theta)\sin(\varphi)+\sqrt{\sin^{2}(\theta)\sin^{2}(\varphi)-\sin^{2}(\theta)+n^{2}}\right),$
(1)

where $\lambda_{0}$ is the wavelength of the incident light,
$\theta=60^{\circ}$ is the polar angle between the incident light and the
surface of the sample, $\varphi=-90^{\circ}$ is the azimuthal angle between
the incident light and the edge of the sample, and $n$ is the real part of the
effective refractive index of the confined mode of the multilayer structure.
At perpendicular incidence, Equation 1 becomes equivalent to the expression
for the wavelength of the interference fringes in Ref. 3.

It should be highlighted that $K_{\text{el}}$ is the wavevector of the
interference fringes obtained with the near-field measurements. Due to the
interference with the background scattered light, $K_{\text{el}}$ is different
from the wavevector of the confined EPPs mode, $K_{\text{ep}}~{}=~{}2\pi
n/\lambda_{0}$. The wavevector of the EPPs $K_{\text{ep}}$ can be calculated
as

$K_{\text{ep}}=K_{\text{el}}\frac{n}{-\sin(\theta)\sin(\varphi)+\sqrt{\sin^{2}(\theta)\sin^{2}(\varphi)-\sin^{2}(\theta)+n^{2}}}.$
(2)

The energy dependence of the propagation length is already apparent from the
near-field intensity maps, with a shorter propagation length above than below
the exciton energy, due to the additional absorption losses of WSe2. As the
incident light has an azimutal angle $\varphi$ perpendicular to the edge, the
propagation length of the interference fringes is equivalent to the
propagation length of the EPPs:
$L_{\text{p}}^{\text{el}}=L_{\text{p}}^{\text{ep}}$.

###  2.2 Real-space and Fourier analysis of the EPPs

Figure 2a presents the profiles made by averaging the experimental maps in the
direction parallel to the edge, for all the measured excitation energies.
Above the A-exciton energy - corresponding to the upper polariton branch -
exponentially decaying interference patterns can be seen up to a distance of
about 2 to 4 µm, depending on the excitation energy. Below the exciton energy,
the propagation lengths are longer, and the interference patterns are more
complex. The fringes indeed display a revival behavior, hinting to the
contribution from two different excitation paths.

Figure 2: (a) Fringe profiles on the WSe2/Al2O3/gold structure at the
different excitation energies measured with the s-SNOM. The red sections of
the profiles are used for the Fourier transform. (b) Fourier transforms of the
profiles (red crosses), and their fits (black curves). The red arrows
highlight the position of the edge-launched polaritons peaks of wavevector
$K_{\text{el}}$. The black dotted line is a guide to the eye.

For a better analysis of the wavevector components at different energies, we
apply the Fourier transform on the profiles, using the same procedure as our
previous work 21. The first 1.4 µm of the profiles (see the black dotted lines
in Figure 2a) induce a background interfering with the rest of the signal,
possibly due to multiple scattering between the tip and the sample at the edge
of the platelet. Thus, the sections from $x=0$ µm to $x=1.4$ µm are removed
for the Fourier transform. The resulting spectra are shown in Figure 2b. The
peak corresponding to the first type of edge-launched EPPs 20 is indicated by
the red arrows and the dotted line. As the energy varies, the peak positions
present a sizeable back-bending indicating coupling between the SPPs and the
A-exciton. Above the exciton energy, mainly one peak can be seen in each
spectrum. A second peak can sometimes be seen around 12 µm-1, and has a
wavevector corresponding to the tip-reflected edge-launched polaritons
excitation path 19, 20. Below the exciton energy, a small peak around 17-19
µm-1 can also be observed. This peak corresponds to the tip-launched EPPs
excitation path. The revival of the interference fringes at lower energies is
thus due to the interference between tip-launched and edge-launched EPPs. The
peak due to the tip-launched EPPs cannot be seen at higher energies, because
the tip-launched EPPs travel twice the distance compared to the edge-launched
EPPs. Hence, their signal becomes too weak due to higher losses in the
material.

To extract the value of the EPP wavevector as a function of the energy, we fit
the spectra as the sum of two Lorentzian functions. We use Lorentzian
functions since edge-launched polaritons can be well approximated by plane
waves 20, 21. The results of the fitting are presented as the black lines in
Figure 2b. The fit of the spectra allows to retrieve the wavevector
$K_{\text{el}}$ and the propagation length $L_{\text{p}}^{\text{el}}$ of the
edge-launched EPP interference patterns. However, the interference profiles
have a finite interval length, which implies for the Fourier transform that
our signal is inherently multiplied with a square window. Thus, each spectral
shape is the convolution between a Lorentzian function and a sinc function 36.
Therefore, while the Fourier transform gives an accurate value for the
wavevectors, the accurate determination of the propagation length is hindered.
To fit the propagation length with a higher accuracy, we directly fit the
interference profiles from Figure 2a. To help the fitting procedure, the value
of the wavevectors are fixed and taken from the fitting results of the Fourier
transformed spectra. Our analysis thus makes use of the advantages of both the
Fourier and the real spaces.

###  2.3 Reconstruction of the full complex-valued wavevector

The wavevectors of the EPPs are calculated from the retrieved values of
$K_{\text{el}}$ using Equation 2. The experimental dispersion relation
resulting from this process is depicted as the white line in Figure 3a. The
vertical error bars correspond to the linewidth of the incident laser, and the
horizontal ones are the fitting uncertainty on the peak position. To compare
our experimental results with theory, we use the transfer matrix method (TMM)
adapted to the case of an anisotropic material. Since WSe2 is a uniaxial
anisotropic crystal, we use the $2\times 2$ matrices presented in Ref. 37 in
the case of p-polarization. For our calculations, we use values of the
dielectric function of gold from McPeak et al. 38, as it has been shown to
have good agreement with the dielectric function of monocrystalline gold 21,
39. The dielectric function of aluminum oxide is set at 2.56. To obtain the
dielectric function of WSe2, we first fit the WSe2 dielectric function data
from the tabulated values of Munkhbat et al. 40. Then, we measure with a far-
field microscope (see Methods) the reflectivity of our multilayer structure,
and compare it with the reflectivity calculated using the values from Ref. 40.
The experimental and calculated reflectivities differ, as can be seen in
Figure 3d. This difference in reflectivities hints at a difference in the in-
plane dielectric function of WSe2. There could be several reasons for this
difference, as the measurement of the dielectric function of WSe2 in Ref. 40
was performed on a thicker WSe2 flake deposited on silicon dioxide. To assess
the influence of the substrate, we also measure the reflectivity of our WSe2
flake on the silicon dioxide substrate next to the gold platelet. The
comparison between the reflectivities of WSe2 on the monocrystalline gold
platelet and on silicon dioxide can be found in the Supplementary Information.
This comparison shows that the reflectivity of our WSe2 flake on silicon
dioxide follows much better the reflectivity calculated using the values from
Ref. 40. Thus, we attribute the difference in reflectivity of WSe2 mainly to a
change in the in-plane dielectric constant of WSe2 due to the presence of a
metallic substrate. To take this change into account, we adapt the parameters
from the in-plane WSe2 dielectric function fitted based on Ref. 40 to
correspond better to our experimental reflectivity. The dielectric function of
WSe2 so obtained can be found in the Supplementary Information.

The theoretical dispersion relation, plotted as the orange line in Figure 3a,
is calculated by finding the poles of the reflection coefficient $r_{p}$ of
our multilayer structure 41. The numerical solution to this first approach is
a complex-valued wavevector, from which the real part gives the dispersion
relation. The colour plot shows, as a second approach, the imaginary part of
the reflection coefficient Im($r_{p}$) - which has been linked to the optical
density of states (DOS)3, 4, 32. The two approaches both give consistent
results and show a clear agreement with the experimental results.

Figure 3: Retrieval of the EPP properties. (a) Dispersion relation of the
exciton-plasmon polaritons for a 13-nm-thick WSe2 deposited on gold. The
colormap displays the imaginary part of the reflection coefficient, calculated
with the TMM. The white line with error bars corresponds to the experimental
wavevector extracted from the data. The vertical error bars correspond to the
linewidth of the incident laser, and the horizontal ones are the fitting
uncertainty on the peak position. The orange line to the theoretical
dispersion relation calculated using the TMM. The red dashed line represents
the light line in air, the horizontal dashed line the energy of the A-exciton
of WSe2, and the blue dashed line the dispersion of the sample without the
A-exciton. (b) Dispersion relation of the EPPs compared to the coupled-
oscillator model (COM). The two polariton branches are plotted in purple, and
the experimental wavevector is in black. (c) Experimental (black curve) and
theoretical (orange curve) propagation lengths of the polaritons. The
horizontal error bars correspond to the uncertainty of the fit. (d) Comparison
between the reflectivity of 13-nm-thick WSe2 on gold, calculated using the
WSe2 dielectric function from Munkhbat et al. 40 (blue dotted line), measured
directly with a conventional far-field microscope (green line), calculated
using a dielectric function adapted to fit the far-field reflectivity (red
line), and extracted from the near-field measurements (purple squares).

To characterize the Rabi splitting, we first consider the difference between
the upper and lower polariton branches at zero detuning, with zero detuning
defined as the intersection of the exciton mode and the surface plasmon mode
without coupling to the A-exciton of WSe2 42. The dispersion of these two
modes is shown as the yellow and blue dashed lines in Figure 3a, respectively.
The vertical cross-section of the theoretical DOS at the zero detuning
($q_{0}=9.4$ µm-1) gives a value for the Rabi splitting energy of
$E_{\text{Rabi}_{\text{0}}}=89$ meV. This method has, however, the
disadvantage of being correlated with the experiment in an indirect manner, as
the Rabi splitting energy is measured from the theoretical calculations and
the comparison between the theory and the experimental data is here
qualitative. In our case, the center of the experimental and theoretical back-
bendings are considerably shifted compared to the zero detuning determined by
the method explained above. This effect has also been observed in a previous
theoretical study 43. Adding or removing the A-exciton resonance for a 13-nm
thick WSe2 has indeed a stronger influence on the WSe2 dielectric function
overall, and thus on the confinement of the surface plasmons, compared to the
monolayer case. Thus, the determination of the zero detuning condition is
difficult since it depends strongly on the thickness of the excitonic material
used and on the retrieval of the A-exciton. A method relying only on
experimental data is therefore preferable to find the amplitude of the Rabi
splitting. At the center of the theoretical back-bending ($q_{1}=9.0$ µm-1),
we find a value for the Rabi splitting of $E_{\text{Rabi}_{\text{1}}}=77$ meV,
while $E_{\text{Rabi}_{\text{2}}}=83$ meV at the center of the experimental
back-bending ($q_{2}=8.9$ µm-1). By comparison, fitting our experimental
dispersion data with the coupled-oscillator model (COM) 43 (see Figure 3b)
gives a value of $E_{\text{Rabi}_{\text{4}}}=81$ meV. In the following we will
use the more conservative value of the Rabi splitting determined
experimentally, i.e. 81 meV. The experimentally and theoretically determined
values of the Rabi splitting fall all within a 10 percent interval around this
value, showing the excellent agreement between experiment and theory.

As the solutions to finding the poles of the reflection coefficients are
complex-valued wavevectors, the theoretical propagation length can also be
deduced from the imaginary part of these wavevectors. The theoretical and
experimental propagation lengths are plotted as a function of the energy in
Figure 3c. The experimental propagation length is found by fitting the
profiles in real space and clearly follows the theoretical calculations.

With the characterized dispersion relation and the propagation lengths, we can
reconstruct the energy dependence of the complex-valued wavevector from near-
field measurements of EPPs in the optical energy range, for the first time to
our knowledge. This allows us to calculate the effective dielectric function
of the EPPs, relying only on experimental measurements. The EPP effective
dielectric function is plotted in the Supplementary Information, as a function
of the excitation energy. The imaginary part of this dielectric function is
associated with the polariton losses. By fitting the polariton dielectric
function as the sum of two Lorentz oscillators, the polariton losses are found
to be equal to $\Gamma=55$ meV. Thus, we can infer that our system is in the
strong-coupling regime, relying only on near-field experimental results,
without inference via theoretical calculations or simulations.

As a benchmark of the quality of our data, we compare the reflectivity
measured with a far-field microscope with the reflectivity predicted from the
near-field measurements. We use the TMM to find the values of the in-plane
WSe2 dielectric function from the experimental values of the complex-valued
wavevector $\tilde{K}_{\text{ep}}=K_{\text{ep}}+i/L_{\text{p}}^{\text{ep}}$.
The new values of the WSe2 dielectric function are then used to calculate the
reflectivity. The reflectivity recalculated from the near-field data shows a
clear dip due to the A-exciton absorption, and agrees well with the direct
far-field measurement of the reflectivity (see Figure 3d). Interestingly,
almost all the data points from the near-field experiments lie in between the
reflectivity calculated using directly the data from Ref. 40 and the
reflectivity calculated using the adapted in-plane WSe2 dielectric function
(see Supplementary Information). Our procedure highlights the correspondence
between material parameters obtained from polaritonic near-field measurements
and the ones obtained with far-field methods. Compared to other s-SNOM studies
of the dielectric function 6, 7, 8, 9, 10 our method does not extract the
dielectric function at the nanoscale, but gives the averaged dielectric
function felt by the EPPs. Our method allows us to retrieve the complex-valued
dielectric function of the excitonic material in the strong coupling regime
and makes the reconstruction of far-field optical properties of nanometer-thin
materials possible.

##  3 Conclusion

In this work, we used a s-SNOM in the reflection configuration to measure EPPs
on a multilayer structure made of gold and 13-nm thick WSe2, at different
energies around the A-exciton energy. By expanding the incident laser beam and
carefully choosing the direction of the incident light, we could resolve both
the higher and lower polariton branches of the edge-launched EPP excitation
path. By using a combination of Fourier and real-space analysis, we
characterized the wavevectors and, for the first time to our knowledge, the
propagation lengths of the EPPs. We observed a clear back-bending of the
dispersion relation, hinting at strong coupling. Using the experimental data
only, we retrieved the Rabi splitting energy of $E_{\text{Rabi}}=81$ meV and
polaritonic losses of $\Gamma=55$ meV, thus demonstrating that our structure
is indeed in the strong-coupling regime. Furthermore, our analysis allowed us
to highlight the importance of the substrate for the accurate determination of
the in-plane dielectric constant of thin-layered WSe2. The quality of the data
was then demonstrated by predicting the far-field reflectivity of our sample
based on the values obtained through the near-field measurements, with
excellent agreement with the direct measurement of far-field reflectivity. The
methods developed in this work can be applied to the quantitative
characterization of the complex-valued wavevectors of different types of
polaritons, for a s-SNOM in a reflection configuration. It can thus be used
for the investigation of quantum related phenomena in polariton physics 24,
44, at the nanometer scale and at visible frequencies. Moreover, relating the
far-field and the near-field information is crucial in the design of new
polaritonic devices, in the understanding of material properties in the
strong-coupling regime, as well as in the design of new metasurfaces.

##  4 Experimental methods

Sample preparation: A WSe2 crystal (HQ graphene) was mechanically exfoliated
with a light blue tape (1007R, SPS Ltd). A 13-nm-thick WSe2 flake resulting
from the exfoliation was then transferred onto the monocrystalline gold
platelets (NanoStruct GmbH), using the dry-transfer method 35.

Near-field measurements: The measurements were performed with a commercial
s-SNOM in reflection mode (neaSNOM, Attocube Systems AG). The polaritons were
excited with a p-polarized femtosecond Ti:sapphire laser (Tsunami Model 3960,
Spectra-Physics Inc.) and the near-field signal was recorded using a 10-MHz
adjustable photodiode (2051, Newport Corp.). The laser was directed towards a
platinum-coated AFM tip (Arrow-NCPt, NanoAndMore GmbH) with a nominal apex
radius of 25 nm. The tip is oscillating above the sample at a frequency of
$f_{0}\sim 300$ kHz and with an oscillation amplitude of about 40 nm. The
signal recorded by the photodiode is demodulated at $3f_{0}$.

To defocus the beam spot at the tip apex, a combination of a divergent lens of
focal length $f_{d}=-200$ mm and a convergent lens $f_{c}=300$ mm is placed -
5 cm apart - in front of the s-SNOM setup. With this effective divergent lens,
the diameter of the spot at the tip is calculated to be 53 µm.

Far-field measurements: The reflection measurements were performed with a
customised microscopy setup based on an inverted microscope (Eclipse Ti-U,
Nikon). A white light source was focused onto the sample with a 10X objective.
The collected light was extracted with the same objective and sent to a
spectrometer (Andor Solis 303I).

LNC thanks K. R. Spiegelhauer for his help with preliminary s-SNOM
experiments, B. Munkhbat for providing the data of the WSe2 dielectric
constant now published in Ref. 40, and V. Zenin, P. Bøggild and S. Raza for
valuable discussions. LNC, SX, MW, and NS acknowledge support from Danmarks
Grundforskningsfond (Grants No. DNRF147 - NanoPhoton and No. DNRF103 - CNG);
Parts of our optical setup are financed through the IDUN Center of Excellence
founded by the Danish National Research Foundation (project number DNRF122)
and VILLUM FONDEN (project number 9301); NS acknowledges Villum Fonden
(00028233); NS and MW acknowledge support from Det Frie Forskningsråd, Natur
og Univers (0135-00403B); QL and SX acknowledge support from Det Frie
Forskningsråd, Teknologi og Produktion (9041-00333B); LNC and NS acknowledge
support from Det Frie Forskningsråd, Teknologi og Produktion (1032-00496B).

Supplementary information is available for this article. It includes the near-
field intensity maps at all measured energies, the real-space fit of the near-
field profiles, and additional details on the analysis of the dielectric
function of WSe2 profiles and the study of the EPP dielectric function.

## References

  * Chen et al. 2012 Chen, J. et al. Optical nano-imaging of gate-tunable graphene plasmons. _Nature_ 2012, _487_ , 77–81
  * Andryieuski et al. 2014 Andryieuski, A.; Zenin, V. A.; Malureanu, R.; Volkov, V. S.; Bozhevolnyi, S. I.; Lavrinenko, A. V. Direct characterization of plasmonic slot waveguides and nanocouplers. _Nano Lett._ 2014, _14_ , 3925–3929
  * Hu et al. 2017 Hu, F.; Luan, Y.; Scott, M.; Yan, J.; Mandrus, D.; Xu, X.; Fei, Z. Imaging exciton–polariton transport in MoSe2 waveguides. _Nature Photonics_ 2017, _11_ , 356–360
  * Hu et al. 2019 Hu, F.; Luan, Y.; Speltz, J.; Zhong, D.; Liu, C. H.; Yan, J.; Mandrus, D. G.; Xu, X.; Fei, Z. Imaging propagative exciton polaritons in atomically thin ${\mathrm{WSe}}_{2}$ waveguides. _Phys. Rev. B_ 2019, _100_ , 121301
  * Menabde et al. 2022 Menabde, S. G.; Boroviks, S.; Ahn, J.; Heiden, J. T.; Watanabe, K.; Taniguchi, T.; Low, T.; Hwang, D. K.; Mortensen, N. A.; Jang, M. S. Near-field probing of image phonon-polaritons in hexagonal boron nitride on gold crystals. _Science Advances_ 2022, _8_ , eabn0627
  * Govyadinov et al. 2013 Govyadinov, A. A.; Amenabar, I.; Huth, F.; Carney, P. S.; Hillenbrand, R. Quantitative Measurement of Local Infrared Absorption and Dielectric Function with Tip-Enhanced Near-Field Microscopy. _The Journal of Physical Chemistry Letters_ 2013, _4_ , 1526–1531
  * Tranca et al. 2015 Tranca, D.; Stanciu, S.; Hristu, R.; Stoichita, C.; Tofail, S.; Stanciu, G. High-resolution quantitative determination of dielectric function by using scattering scanning near-field optical microscopy. _Scientific Reports_ 2015, _5_ , 1–9
  * Ruta et al. 2020 Ruta, F. L.; Sternbach, A. J.; Dieng, A. B.; McLeod, A. S.; Basov, D. Quantitative nanoinfrared spectroscopy of anisotropic van der Waals materials. _Nano Letters_ 2020, _20_ , 7933–7940
  * Álvarez-Pérez et al. 2020 Álvarez-Pérez, G. et al. Infrared permittivity of the biaxial van der Waals semiconductor $\alpha$-MoO3 from near-and far-field correlative studies. _Advanced Materials_ 2020, _32_ , 1908176
  * Zhang et al. 2022 Zhang, S. et al. Nano-spectroscopy of excitons in atomically thin transition metal dichalcogenides. _Nature Comm._ 2022, _13_ , 1–8
  * Keilmann and Hillenbrand 2004 Keilmann, F.; Hillenbrand, R. Near-field microscopy by elastic light scattering from a tip. _Phil. Trans. R. Soc. A_ 2004, _362_ , 787–805
  * Fei et al. 2012 Fei, Z.; Rodin, A.; Andreev, G. O.; Bao, W.; McLeod, A.; Wagner, M.; Zhang, L.; Zhao, Z.; Thiemens, M.; Dominguez, G.; Fogler, M. M.; Neto, A. H. C.; Lau, C. N.; Keilmann, F.; Basov, D. N. Gate-tuning of graphene plasmons revealed by infrared nano-imaging. _Nature_ 2012, _487_ , 82–85
  * Dai et al. 2014 Dai, S. et al. Tunable phonon polaritons in atomically thin van der Waals crystals of boron nitride. _Science_ 2014, _343_ , 1125–1129
  * Yoxall et al. 2015 Yoxall, E.; Schnell, M.; Nikitin, A. Y.; Txoperena, O.; Woessner, A.; Lundeberg, M. B.; Casanova, F.; Hueso, L. E.; Koppens, F. H.; Hillenbrand, R. Direct observation of ultraslow hyperbolic polariton propagation with negative phase velocity. _Nature Photonics_ 2015, _9_ , 674–678
  * Ma et al. 2018 Ma, W. et al. In-plane anisotropic and ultra-low-loss polaritons in a natural van der Waals crystal. _Nature_ 2018, _562_ , 557–562
  * Taboada-Gutiérrez et al. 2020 Taboada-Gutiérrez, J. et al. Broad spectral tuning of ultra-low-loss polaritons in a van der Waals crystal by intercalation. _Nat. Mater._ 2020, _19_ , 964–968
  * Bylinkin et al. 2021 Bylinkin, A.; Schnell, M.; Autore, M.; Calavalle, F.; Li, P.; Taboada-Gutièrrez, J.; Liu, S.; Edgar, J. H.; Casanova, F.; Hueso, L. E.; Alonso-Gonzalez, P.; Nikitin, A. Y.; Hillenbrand, R. Real-space observation of vibrational strong coupling between propagating phonon polaritons and organic molecules. _Nat. Photon._ 2021, _15_ , 197–202
  * Woessner et al. 2015 Woessner, A.; Lundeberg, M. B.; Gao, Y.; Principi, A.; Alonso-González, P.; Carrega, M.; Watanabe, K.; Taniguchi, T.; Vignale, G.; Polini, M.; Hone, J.; Hillenbrand, R.; Koppens, F. H. L. Highly confined low-loss plasmons in graphene–boron nitride heterostructures. _Nat. Mater._ 2015, _14_ , 421–425
  * Walla et al. 2018 Walla, F.; Wiecha, M. M.; Mecklenbeck, N.; Beldi, S.; Keilmann, F.; Thomson, M. D.; Roskos, H. G. Anisotropic excitation of surface plasmon polaritons on a metal film by a scattering-type scanning near-field microscope with a non-rotationally-symmetric probe tip. _Nanophotonics_ 2018, _7_ , 269–276
  * Kaltenecker et al. 2020 Kaltenecker, K. J.; Krauss, E.; Casses, L.; Geisler, M.; Hecht, B.; Mortensen, N. A.; Jepsen, P. U.; Stenger, N. Mono-crystalline gold platelets: a high-quality platform for surface plasmon polaritons. _Nanophotonics_ 2020, _9_ , 509–522
  * Casses et al. 2022 Casses, L. N.; Kaltenecker, K. J.; Xiao, S.; Wubs, M.; Stenger, N. Quantitative near-field characterization of surface plasmon polaritons on monocrystalline gold platelets. _Opt. Express_ 2022, _30_ , 11181–11191
  * Basov et al. 2016 Basov, D.; Fogler, M.; De Abajo, F. G. Polaritons in van der Waals materials. _Science_ 2016, _354_ , aag1992
  * Low et al. 2017 Low, T.; Chaves, A.; Caldwell, J. D.; Kumar, A.; Fang, N. X.; Avouris, P.; Heinz, T. F.; Guinea, F.; Martin-Moreno, L.; Koppens, F. Polaritons in layered two-dimensional materials. _Nat. Mater._ 2017, _16_ , 182–194
  * Basov et al. 2021 Basov, D.; Asenjo-Garcia, A.; Schuck, P. J.; Zhu, X.; Rubio, A. Polariton panorama. _Nanophotonics_ 2021, _10_ , 549–577
  * Mueller and Malic 2018 Mueller, T.; Malic, E. Exciton physics and device application of two-dimensional transition metal dichalcogenide semiconductors. _npj 2D Materials and Applications_ 2018, _2_ , 1–12
  * Maier 2007 Maier, S. A. _Plasmonics: Fundamentals and Applications_ ; Springer Science & Business Media: New York, 2007
  * Xia et al. 2014 Xia, F.; Wang, H.; Xiao, D.; Dubey, M.; Ramasubramaniam, A. Two-dimensional material nanophotonics. _Nature Photonics_ 2014, _8_ , 899–907
  * Baranov et al. 2018 Baranov, D. G.; Wersäll, M.; Cuadra, J.; Antosiewicz, T. J.; Shegai, T. Novel Nanostructures and Materials for Strong Light–Matter Interactions. _ACS Photonics_ 2018, _5_ , 24–42
  * Byrnes et al. 2014 Byrnes, T.; Kim, N. Y.; Yamamoto, Y. Exciton–polariton condensates. _Nature Physics_ 2014, _10_ , 803–813
  * Sanvitto and Kéna-Cohen 2016 Sanvitto, D.; Kéna-Cohen, S. The road towards polaritonic devices. _Nature materials_ 2016, _15_ , 1061–1073
  * Denning et al. 2022 Denning, E. V.; Wubs, M.; Stenger, N.; Mørk, J.; Kristensen, P. T. Cavity-induced exciton localization and polariton blockade in two-dimensional semiconductors coupled to an electromagnetic resonator. _Phys. Rev. Res._ 2022, _4_ , L012020
  * B. Iyer et al. 2022 B. Iyer, R.; Luan, Y.; Shinar, R.; Shinar, J.; Fei, Z. Nano-optical imaging of exciton–plasmon polaritons in WSe2/Au heterostructures. _Nanoscale_ 2022, _14_ , 15663–15668
  * Jang et al. 2024 Jang, M.; Menabde, S. G.; Kiani, F.; Heiden, J. T.; Zenin, V. A.; Mortensen, N. A.; Tagliabue, G.; Jang, M. S. Fourier analysis of near-field patterns generated by propagating polaritons. _arXiv preprint arXiv:2402.17225_ 2024,
  * Krauss et al. 2018 Krauss, E.; Kullock, R.; Wu, X.; Geisler, P.; Lundt, N.; Kamp, M.; Hecht, B. Controlled Growth of High-Aspect-Ratio Single-Crystalline Gold Platelets. _Cryst. Growth Des._ 2018, _18_ , 1297–1302
  * Castellanos-Gomez et al. 2014 Castellanos-Gomez, A.; Buscema, M.; Molenaar, R.; Singh, V.; Janssen, L.; Van Der Zant, H. S.; Steele, G. A. Deterministic transfer of two-dimensional materials by all-dry viscoelastic stamping. _2D Materials_ 2014, _1_ , 011002
  * Bracewell 1986 Bracewell, R. N. _The Fourier Transform and its Applications - 3rd Ed._ ; McGraw-Hill: New York, 1986
  * Majérus et al. 2018 Majérus, B.; Dremetsika, E.; Lobet, M.; Henrard, L.; Kockaert, P. Electrodynamics of two-dimensional materials: Role of anisotropy. _Physical Review B_ 2018, _98_ , 125419
  * McPeak et al. 2015 McPeak, K. M.; Jayanti, S. V.; Kress, S. J. P.; Meyer, S.; Iotti, S.; Rossinelli, A.; Norris, D. J. Plasmonic Films Can Easily Be Better: Rules and Recipes. _ACS Photonics_ 2015, _2_ , 326–333
  * Lebsir et al. 2022 Lebsir, Y.; Boroviks, S.; Thomaschewski, M.; Bozhevolnyi, S. I.; Zenin, V. A. Ultimate Limit for Optical Losses in Gold, Revealed by Quantitative Near-Field Microscopy. _Nano Letters_ 2022, _22_ , 5759–5764
  * Munkhbat et al. 2022 Munkhbat, B.; Wróbel, P.; Antosiewicz, T. J.; Shegai, T. O. Optical Constants of Several Multilayer Transition Metal Dichalcogenides Measured by Spectroscopic Ellipsometry in the 300–1700 nm Range: High Index, Anisotropy, and Hyperbolicity. _ACS Photonics_ 2022, _9_ , 2398–2407
  * Zhan et al. 2013 Zhan, T.; Shi, X.; Dai, Y.; Liu, X.; Zi, J. Transfer matrix method for optics in graphene layers. _Journal of Physics: Condensed Matter_ 2013, _25_ , 215301
  * Geisler et al. 2019 Geisler, M.; Cui, X.; Wang, J.; Rindzevicius, T.; Gammelgaard, L.; Jessen, B. S.; Gonçalves, P. A. D.; Todisco, F.; Bøggild, P.; Boisen, A.; Wubs, M.; Mortensen, N. A.; Xiao, S.; Stenger, N. Single-crystalline gold nanodisks on WS2 mono- and multilayers for strong coupling at room temperature. _ACS Photonics_ 2019, _6_ , 994–1001
  * Gonçalves et al. 2018 Gonçalves, P.; Bertelsen, L.; Xiao, S.; Mortensen, N. A. Plasmon-exciton polaritons in two-dimensional semiconductor/metal interfaces. _Physical Review B_ 2018, _97_ , 041402
  * Boroviks et al. 2022 Boroviks, S.; Lin, Z.-H.; Zenin, V. A.; Ziegler, M.; Dellith, A.; Gonçalves, P.; Wolff, C.; Bozhevolnyi, S. I.; Huang, J.-S.; Mortensen, N. A. Extremely confined gap plasmon modes: when nonlocality matters. _Nature Communications_ 2022, _13_ , 1–8


# sEHR-CE: Language modelling of structured EHR data for efficient and
generalizable patient cohort expansion

Anna Munoz-Farre
BenevolentAI
<EMAIL_ADDRESS>
Harry Rose
BenevolentAI
<EMAIL_ADDRESS>
Sera Aylin Cakiroglu
BenevolentAI
<EMAIL_ADDRESS>

###### Abstract

Electronic health records (EHR) offer unprecedented opportunities for in-depth
clinical phenotyping and prediction of clinical outcomes. Combining multiple
data sources is crucial to generate a complete picture of disease prevalence,
incidence and trajectories. The standard approach to combining clinical data
involves collating clinical terms across different terminology systems using
curated maps, which are often inaccurate and/or incomplete. Here, we propose
sEHR-CE, a novel framework based on transformers to enable integrated
phenotyping and analyses of heterogeneous clinical datasets without relying on
these mappings. We unify clinical terminologies using textual descriptors of
concepts, and represent individuals’ EHR as sections of text. We then fine-
tune pre-trained language models to predict disease phenotypes more accurately
than non-text and single terminology approaches. We validate our approach
using primary and secondary care data from the UK Biobank, a large-scale
research study. Finally, we illustrate in a type 2 diabetes use case how sEHR-
CE identifies individuals without diagnosis that share clinical
characteristics with patients.

##  1 Introduction

Electronic health records (EHR) are collected as part of routine medical care,
and include demographic information, disease diagnoses, laboratory results,
medication prescriptions, etc., providing a patient’s clinical state over
time. Recent machine learning techniques have been used to exploit the
richness of EHR data at scale for diagnosis, prognosis, treatment and
understanding of disease [Li et al., 2020, Steinberg et al., 2021]. Many
medical terminologies are used across clinical data sets, and the standard
practice involves mapping clinical terms across different resources [Li et
al., 2020, Hassaine et al., 2020] or onto common data models [Stang et al.,
2010].

Here, we propose sEHR-CE (language modelling of Structured EHR data for
patient Cohort Expansion), a novel framework based on transformers to enable
the integrated analysis of multiple clinical resources without relying on any
manual curation and mapping. Using text descriptions of concepts as input, our
method generalises across data modalities and terminologies, i.e. text and
structured EHR. This enables us to leverage a plethora of pre-trained language
models like PubMedBERT [Gu et al., 2022] to encode each patient’s medical
record. We ask the model to learn representations of clinical histories from
diagnosed patients (cases) to predict phenotypes (such as diseases). In the
absence of a directly comparable model, we evaluate the performance of our
model against that of Li et al. [2020], as the state-of-the-art approach for
learning clinical term embeddings for future disease prediction. See Appendix
A for a more complete survey.

In our validation experiments, we will show that the model can then identify
individuals with missing diagnoses (controls) that share a similar clinical
history with cases, indicating they might have the disease or be at risk of
developing it. The contributions of our paper are (i) the presentation of a
cohort expansion method that provides phenotype predictions outperforming non-
text and single terminology approaches, and (ii) an in-depth qualitative
evaluation demonstrating that positively predicted controls share similar
clinical representations with cases, providing a high degree of evidence that
these may be previously undiagnosed or misdiagnosed individuals. Finally, we
demonstrate our method’s potential for patient stratification by disease
severity.

##  2 Methods

Input Generation. For each EHR source and associated ontology
$a\in\mathcal{A}$, we denote the set of concepts (e.g. clinical terms) as
$\Theta_{a}$, and the set of text descriptions as $\Xi_{a}$. The total
vocabulary of concepts and text descriptions across all ontologies is denoted
by $\Theta=\bigcup\limits_{a\in\mathcal{A}}\Theta_{a}$ and
$\Xi=\bigcup\limits_{a\in\mathcal{A}}\Xi_{a}$, respectively. For each patient,
we define their full clinical history through time and across sources by the
concatenation of sequences of clinical descriptions
$(\xi_{\theta_{1}},\ldots,\xi_{\theta_{t}})$, $\xi_{\theta_{i}}\in\Xi$,
$i=1,\ldots,t$, ordered in time. More details and an example can be found in
Appendix B.1, and Figure 4. To form the input, we process the raw text
sequence of descriptions into tokens (e.g. words and sub-words)
$X=W(\xi_{\theta_{1}},\ldots,\xi_{\theta_{t}})=(x_{1},\ldots,x_{n})$ under a
fixed size vocabulary $V$ with the WordPiece tokenizer $W$ [Wu et al., 2016].

Label Generation. Let $\Delta=(d_{1},\ldots,d_{D})$ denote an ordered set of
clinical outcomes or events, in our case disease phenotypes $d_{i}$. For each
phenotype $d_{i}\in\Delta$, we let $\mathbf{1}_{d_{i}}$ be an indicator
function that assigns a binary label to individuals according to the presence
or absence of $d_{i}$. To define $\mathbf{1}_{d_{i}}$, we rely on external
oracles or phenotype definitions, such as CALIBER (Appendix B.2). Figure 2
shows a schematic of input and label generation for a given patient.

Figure 1: Schematic of mapping a unified clinical history to an input text
sequence with multi-label target vector using an oracle annotation 222Any
patient data shown is simulated and does not represent data of real patients..

Model Design. Let $X^{(p)}=(x^{(p)}_{1},\ldots,x^{(p)}_{n})$ denote the
tokenized input sequence of individual $p$. It forms the input to an encoding
function $\mathbf{x}^{(p)}_{1},\ldots,\mathbf{x}^{(p)}_{n}$ =
Encoder$(X^{(p)})$, where each $\mathbf{x}_{i}$ is a fixed-length vector
representation of each input token $x_{i}$. Let
$\mathbf{y}^{(p)}=(y^{(p)}_{1},\ldots,y^{(p)}_{D})$,
$y^{(p)}_{i}\in\\{0,1\\}$, be labels denoting presence or absence of
phenotypes $d_{1},\ldots,d_{D}$. Given a learned representation over inputs,
we decode over $\mathbf{y}^{(p)}$ under the predictive model
$\mathbf{P}(\mathbf{y}^{(p)}|X^{(p)})$. We calculate the probability of each
phenotype $d_{i}$ given the input sequence encoding
$\mathbf{P}(y^{(p)}_{i}|\textbf{x}^{(p)}_{1},\dots,\textbf{x}^{(p)}_{n})$ via
a decoder module. Specifically, we decode the representation into logits per
phenotype $z^{(p)}_{1},\dots,z^{(p)}_{D}=$
Decoder$(\textbf{x}^{(p)}_{1},\dots,\textbf{x}^{(p)}_{n})$ and calculate the
probability per phenotype as
$\mathbf{P}(y^{(p)}_{d}|\textbf{x}^{(p)}_{1},\dots,\textbf{x}^{(p)}_{n})=\sigma(z^{(p)}_{d})$,
where $\sigma$ denotes the sigmoid function (Figure B.2). We will omit the
superscript $(p)$ denoting the sample index in the remainder of the text.

###  2.1 Data Augmentation

Clinical Masking. We mask input descriptions $\xi_{\theta}$ from term-
description pairs $(\theta,\xi_{\theta})$ with
$\mathbf{1}_{d}(\theta,\xi_{\theta})=1$ for $d\in\Delta$ using the following
clinical masking strategy during training and validation as in Devlin et al.
[2019], Wei and Zou [2019]: Remove $\xi_{\theta}$ with 80% probability, retain
$\xi_{\theta}$ with 10% probability, and replace $\xi_{\theta}$ with a
randomly selected description from the corpus with 10% probability. Figure 2
displays a worked example. During testing, these term-description pairs are
fully removed from the input sequence.

Figure 2: Example of clinical masking strategy for a given patient history
with Heart Failure.

Comorbidities. This masking approach is straightforward if an individual has
only one positive label, but many people have comorbidities, e.g. co-occurring
conditions. To allow for a sample with multiple positive labels
$d_{i_{1}},\ldots,d_{i_{n}}$, we create $n$ input samples by replicating both
the input sequence and target vector of phenotype labels. Consider the $j$th
replicated input sequence, and let $d=d_{i_{j}}$. We mask this replicated
input sequence with the masking strategy for phenotype $d$ described in the
previous section; e.g. $\xi_{\theta}$ is masked if
$\mathbf{1}_{d}(\theta,\xi_{\theta})=1$ (Figure 3).

###  2.2 Defining a Loss Function in the Presence of Comorbidities and varying
Prevalences

Loss weights. The data augmentation approach for comorbidities may lead to our
model overfitting when an input sequence contains the descriptions associated
with an existing phenotype $d$ that is not masked. To account for the
contribution of the target vector $\mathbf{y}$ to the loss function in these
scenario, we define a binary masking vector
$\mathbf{\gamma}^{j}=(\gamma^{j}_{1},\ldots,\gamma^{j}_{D})$ where
$\gamma^{j}_{j}=1$ and $\gamma^{j}_{k}=0$, $k=1,\ldots,D$, $k\not=j$.
Individuals with no positive phenotype labels are assigned an all-zero masking
vector (Figure 3). Then, for a given input text sequence $X$, target label
vector $\mathbf{y}=(y_{1},\ldots,y_{D})$ and masking vector
$\mathbf{\gamma}=(\gamma_{1},\ldots,\gamma_{D})$, we define the loss weights
by setting $\omega_{d}=0$ if $y_{d}=1$ and $\gamma_{d}=0$, and $\omega_{d}=1$
if $y_{d}=0$ or $y_{d}=1$ and $\gamma_{d}=1$.

Figure 3:  Example of an input sequence containing descriptions relevant to
two positive phenotype labels in the target vector $y$. The input sequence is
duplicated, and each phenotype is masked once and a loss weights vector is
defined. Clinical terms relating to a directly associated phenotype $d$ remain
in the input data ($y_{d}=1$ and $\gamma_{d}=0$) but get assigned a loss
weight $w_{d}=0$ and so do not contribute to the loss function.

Positive weights. Because some diseases are very rare, whereas others are very
common, our prediction classes are expected to be highly imbalanced in the
practical setting. For cohort expansion, we want to increase recall while
balancing a decline in precision. For $d\in\Delta$, let $TN_{d}$ and $TP_{d}$
denote the total number of negative and positive examples of $d$ in the
training set. We define the positive weight $\rho_{d}$ as
$\rho_{d}=TN_{d}/TP_{d}$, for all d $\in\Delta.$

Our loss function is defined as a mean-reduced binary cross-entropy loss
function over phenotypes, where disease prevalence and comorbidities are
handled with loss and positive weights:

$\mathcal{L}^{(p)}=-\frac{1}{|\Delta|}\sum_{d\in\Delta}\omega_{d}^{(p)}(\rho_{d}y_{d}^{(p)}\log\sigma({z_{d}^{(p)}})\\\
+(1-y_{d}^{(p)})\log\sigma(1-z_{d}^{(p)})),$ (1)

where $\sigma$ denotes the sigmoid function, $\omega_{d}$ the comorbidity-
derived loss weight, $\rho_{d}$ the positive weight, and $z_{d}^{(p)}$ the
predicted probability for phenotype $d\in\Delta$ for sample (e.g. individual)
$p$.

##  3 Experiments and Results

Data. This research has been conducted using the UK Biobank (UKBB) Resource
under Application Number 43138, a large-scale research study of around $500$k
individuals [Sudlow et al., 2015]. We restrict the data set to those that have
entries in both hospital and GP records, reducing our cohort to $155$k. To
assess the quality of our model predictions, we choose four diseases that
differ in terms of prevalence and clinical characteristics (Appendix C): Type
2 diabetes mellitus (T2DM), Heart failure (HF), Breast Cancer, and Prostate
Cancer. We use validated phenotype definitions from CALIBER Kuan et al. [2019]
to label patients with each of the diseases (Appendix C.1). We test the
performance of our model on its ability to diagnose known cases, compare it to
other methods, and evaluate associations of clinical features with predictions
on T2DM.

sEHR-CE. We use the pre-trained language model PubMedBERT [Gu et al., 2022] as
the encoder of the tokenised input sequences of clinical term descriptions.
Since our input systematically differs from the general scientific text on
which PubMedBERT was trained, we fine-tune on the masked-language modeling
(MLM) task using the full UKBB cohort. The proposed model sEHR-CE uses the
fine-tuned encoder and a fully connected linear layer as the decoder. To train
on the multi-label classification task of outcome prediction, we split our
data set into five equally sampled folds $f_{0}$,…$f_{4}$ containing unique
patients, and mask the data according to our strategy (Section 2.1, Figure
C.1). We train a total of five models on three folds, holding back folds
$f_{i}$ for validation and $f_{(i+1)\bmod{5}}$ for testing for model $i$,
$i=1,\ldots,5$ (Figure C.2). All results presented are predictions on the
independent test set. For masking and training details, see Appendix C.2.

Baseline Models. We compare the performance of our model sEHR-CE to BEHRT [Li
et al., 2020], which takes a tokenised sequence of clinical terms, age and
position embeddings as input. Ontologies from hospital and GP records are
mapped to CALIBER definitions, removing unmapped terms (more details in
Appendix C.3). We train five such BEHRT models to predict an individual
developing the four phenotypes on the same five data splits as sEHR-CE.
Similarly, we train five sEHR-CE models restricted to CALIBER tokens (denoted
sEHR-CE-codes) for comparison. Figure C.3 shows the distributions of predicted
probabilites for all phenotypes across all methods. sEHR-CE shows the best
performance across all four phenotypes in terms of recall at $0.5$ and AUC on
the test set (Table 1, Figure C.4). BEHRT performs slightly better than sEHR-
CE-codes, indicating a benefit of adding visit position and age. Performance
varies across phenotypes, presumably due to different clinical characteristics
making some diseases easier to predict than others.

| T2D | HF | Breast Cancer | Prostate Cancer | Average
---|---|---|---|---|---
| Recall | PRC | Recall | PRC | Recall | PRC | Recall | PRC | Recall, std | PRC
sEHR-CE-codes | 0.64 | 0.70 | 0.76 | 0.60 | 0.42 | 0.45 | 0.36 | 0.26 | 0.55 $\pm$ 0.19 | 0.50
BEHRT | 0.66 | 0.70 | 0.78 | 0.61 | 0.48 | 0.53 | 0.41 | 0.29 | 0.58 $\pm$ 0.17 | 0.53
sEHR-CE | 0.74 | 0.74 | 0.85 | 0.69 | 0.55 | 0.55 | 0.57 | 0.47 | 0.68 $\pm$ 0.14 | 0.61
Table 1: Average and phenotype specific recall and AUCPR at $0.5$ on the test
sets.

###  3.1 Clinical evaluation on Type 2 Diabetes Mellitus

T2DM lends itself as a use case to qualitatively evaluate the predictions of
missed cases, as it is a well-studied, slowly developing disease with varying
disease severity. We used thresholds based on the 98th, 90th and 12th
percentiles of sEHR-CE’s predicted probabilities (Figure C.3) to define five
different groups (Table C.1): regular cases, cases predicted with high
probability, cases predicted with low probability, regular controls and
controls predicted with high probability (missed cases).

Measures of Disease Severity. Haemoglobin A1c (HbA1c) is a blood biomarker
used to diagnose and monitor diabetes. A level of 48mmol/mol or higher is
considered diabetes; while a value between 42 and 48 mmol/mol is considered
pre-diabetes. The input data did not include such biomarker data, so we use it
here for evaluation. We aggregated HbA1c measurements taken in primary care
(GP) of each individual by taking the $95$-th percentile value. Figure 4(a)
shows that cases predicted with high probability had the highest HbA1c mean
levels. Cases identified with low probability were in the prediabetic range of
HbA1c levels, indicating a less severe state. Missed cases had elevated HbA1c
levels close to the prediabetic stage, representing individuals at risk of
developing T2DM. We further investigated the association of the T2DM predicted
probabilities with other measures of disease severity, expanded in Appendix
C.4. Taken together, our results demonstrate that sEHR-CE’s predicted
probabilities of being diagnosed with T2DM are associated with disease
severity.

Polygenic risk scores. Genetic risk for complex diseases like T2DM arises from
many genetic changes that, when taken together, can increase an individual’s
risk of developing the disease, which can be defined by polygenic risk scores
(PRS). Sinnott-Armstrong et al. [2021] developed PRS based on the UK Biobank
participants for a set of diseases, including T2DM. We computed and
standardised T2DM PRS across all individuals in our cohort [Lewis and Vassos,
2017]. Figure 4(b) demonstrates that higher predicted probability of T2DM was
associated with a higher genetic risk.

(a)

(b)

Figure 4: 4(a) HbA1c distribution across groups in the test sets with diabetes
diagnosis ranges shown as dashed lines. Significant p-values are indicated
with *** (t-test, $\alpha=.001$). 4(b) Median T2DM prediction of individuals
in the test sets grouped by the percentiles of the polygenic risk score.

##  4 Conclusion

We presented a data-driven method for cohort expansion based on language
modelling. Our approach fuses primary and secondary care data via text, and we
propose a data augmentation approach to allow for comorbidities in a patient’s
history. Our method predicts disease phenotype labels more accurately than
non-text and single terminology approaches. We presented a high degree of
evidence that our model identifies previously undiagnosed individuals that can
extend the original cohort for downstream analysis. Future work will consider
methods that are less restrictive on sequence length [Kitaev et al., 2020,
Beltagy et al., 2020] and allow for irregular time steps [Shukla and Marlin,
2021] and age [Kazemi et al., 2019], as well as adding more data sources (e.g.
medications).

##  5 Acknowledgements

This research has been conducted using the UK Biobank Resource under
Application Number 43138. Using real patient data is crucial for clinical
research and to find the right treatment for the right patient. We would like
to thank all participants who are part of the UK Biobank, who volunteered to
give their primary and secondary care and genotyping data for the purpose of
research. UK Biobank is generously supported by its founding funders the
Wellcome Trust and UK Medical Research Council, as well as the British Heart
Foundation, Cancer Research UK, Department of Health, Northwest Regional
Development Agency and Scottish Government.

We are particularly grateful to Prof. Spiros Denaxas, and Drs. Aaron Sim,
Andrea Rodriguez-Martinez, Nicola Richmond, Sam Abujudeh, and Julien Fauqueur
for their feedback, insightful comments and the many inspiring conversations.

## References

  * Li et al. [2020] Yikuan Li, Shishir Rao, José Roberto Ayala Solares, Abdelaali Hassaine, Rema Ramakrishnan, Dexter Canoy, Yajie Zhu, Kazem Rahimi, and Gholamreza Salimi-Khorshidi.  BEHRT: Transformer for electronic health records.  _Sci. Rep._ , 10(1):7155, 2020.
  * Steinberg et al. [2021] Ethan Steinberg, Ken Jung, Jason A. Fries, Conor K. Corbin, Stephen R. Pfohl, and Nigam H. Shah.  Language models are an effective representation learning technique for electronic health record data.  _Journal of Biomedical Informatics_ , 113:103637, 2021.
  * Hassaine et al. [2020] Abdelaali Hassaine, Dexter Canoy, Jose Roberto Ayala Solares, Yajie Zhu, Shishir Rao, Yikuan Li, Mariagrazia Zottoli, Kazem Rahimi, and Gholamreza Salimi-Khorshidi.  Learning multimorbidity patterns from electronic health records using non-negative matrix factorisation.  _Journal of Biomedical Informatics_ , 112:103606, 2020.
  * Stang et al. [2010] Paul E Stang, Patrick B Ryan, Judith A Racoosin, J Marc Overhage, Abraham G Hartzema, Christian Reich, Emily Welebob, Thomas Scarnecchia, and Janet Woodcock.  Advancing the science for active surveillance: rationale and design for the observational medical outcomes partnership.  _Ann. Intern. Med._ , 153(9):600–606, 2010.
  * Gu et al. [2022] Yu Gu, Robert Tinn, Hao Cheng, Michael Lucas, Naoto Usuyama, Xiaodong Liu, Tristan Naumann, Jianfeng Gao, and Hoifung Poon.  Domain-specific language model pretraining for biomedical natural language processing.  _ACM Trans. Comput. Heal._ , 3(1):2:1–2:23, 2022\.
  * Wu et al. [2016] Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoc V. Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, Jeff Klingner, Apurva Shah, Melvin Johnson, Xiaobing Liu, Łukasz Kaiser, Stephan Gouws, Yoshikiyo Kato, Taku Kudo, Hideto Kazawa, Keith Stevens, George Kurian, Nishant Patil, Wei Wang, Cliff Young, Jason Smith, Jason Riesa, Alex Rudnick, Oriol Vinyals, Greg Corrado, Macduff Hughes, and Jeffrey Dean.  Google’s neural machine translation system: Bridging the gap between human and machine translation, 2016.  arxiv:1609.08144.
  * 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 _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 4171–4186. Association for Computational Linguistics, 2019.
  * Wei and Zou [2019] Jason Wei and Kai Zou.  EDA: Easy data augmentation techniques for boosting performance on text classification tasks.  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 6382–6388. Association for Computational Linguistics, 2019.
  * Sudlow et al. [2015] Cathie Sudlow, John Gallacher, Naomi Allen, Valerie Beral, Paul Burton, John Danesh, Paul Downey, Paul Elliott, Jane Green, Martin Landray, Bette Liu, Paul Matthews, Giok Ong, Jill Pell, Alan Silman, Alan Young, Tim Sprosen, Tim Peakman, and Rory Collins.  UK Biobank: An open access resource for identifying the causes of a wide range of complex diseases of middle and old age.  _PLOS Medicine_ , 12(3):1–10, 2015.
  * Kuan et al. [2019] Valerie Kuan, Spiros Denaxas, Arturo Gonzalez-Izquierdo, Kenan Direk, Osman Bhatti, Shanaz Husain, Shailen Sutaria, Melanie Hingorani, Dorothea Nitsch, Constantinos A Parisinos, et al.  A chronological map of 308 physical and mental health conditions from 4 million individuals in the English National Health Service.  _The Lancet Digital Health_ , 1(2):e63–e77, 2019\.
  * Sinnott-Armstrong et al. [2021] Nasa Sinnott-Armstrong, Yosuke Tanigawa, David Amar, Nina Mars, Christian Benner, Matthew Aguirre, Guhan Ram Venkataraman, Michael Wainberg, Hanna M Ollila, Tuomo Kiiskinen, Aki S Havulinna, James P Pirruccello, Junyang Qian, Anna Shcherbina, FinnGen, Fatima Rodriguez, Themistocles L Assimes, Vineeta Agarwala, Robert Tibshirani, Trevor Hastie, Samuli Ripatti, Jonathan K Pritchard, Mark J Daly, and Manuel A Rivas.  Genetics of 35 blood and urine biomarkers in the UK Biobank.  _Nat. Genet._ , 53(2):185–194, 2021.
  * Lewis and Vassos [2017] Cathryn M Lewis and Evangelos Vassos.  Prospects for using risk scores in polygenic medicine.  _Genome Med._ , 9(1):96, 2017.
  * Kitaev et al. [2020] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya.  Reformer: The efficient transformer.  In _International Conference on Learning Representations_ , 2020.
  * Beltagy et al. [2020] Iz Beltagy, Matthew E. Peters, and Arman Cohan.  Longformer: The long-document transformer.  2020\.  arxiv:2004.05150.
  * Shukla and Marlin [2021] Satya Narayan Shukla and Benjamin Marlin.  Multi-time attention networks for irregularly sampled time series.  In _International Conference on Learning Representations_ , 2021.
  * Kazemi et al. [2019] Seyed Mehran Kazemi, Rishab Goel, Sepehr Eghbali, Janahan Ramanan, Jaspreet Sahota, Sanjay Thakur, Stella Wu, Cathal Smyth, Pascal Poupart, and Marcus Brubaker.  Time2vec: Learning a vector representation of time, 2019.  arXiv:1907.05321.
  * Henderson et al. [2017] Jette Henderson, Joyce C. Ho, Abel N. Kho, Joshua C. Denny, Bradley A. Malin, Jimeng Sun, and Joydeep Ghosh.  Granite: Diversified, sparse tensor factorization for electronic health record-based phenotyping.  In _2017 IEEE International Conference on Healthcare Informatics (ICHI)_ , pages 214–223, 2017.
  * Miotto et al. [2016] Riccardo Miotto, Li Li, Brian A. Kidd, and Joel T. Dudley.  Deep patient: An unsupervised representation to predict the future of patients from the electronic health records.  _Scientific Reports_ , 6(1):26094, 2016.
  * Choi et al. [2016a] Edward Choi, Mohammad Taha Bahadori, Andy Schuetz, Walter F. Stewart, and Jimeng Sun.  Doctor AI: Predicting clinical events via recurrent neural networks.  In _Proceedings of the 1st Machine Learning for Healthcare Conference_ , volume 56 of _PMLR_ , pages 301–318, 2016a.
  * Pham et al. [2016] Trang Pham, Truyen Tran, Dinh Phung, and Svetha Venkatesh.  DeepCare: A deep dynamic memory model for predictive medicine.  In _Advances in Knowledge Discovery and Data Mining_ , pages 30–41. Springer International Publishing, 2016.
  * Alaa and van der Schaar [2019] Ahmed M. Alaa and Mihaela van der Schaar.  Attentive state-space modeling of disease progression.  In _Advances in Neural Information Processing Systems_ , volume 32. Curran Associates, Inc., 2019.
  * Choi et al. [2016b] Edward Choi, Mohammad Taha Bahadori, Elizabeth Searles, Catherine Coffey, Michael Thompson, James Bost, Javier Tejedor-Sojo, and Jimeng Sun.  Multi-layer representation learning for medical concepts.  In _Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining_ , KDD ’16, page 1495–1504. Association for Computing Machinery, 2016b.
  * Denaxas et al. [2018] Spiros Denaxas, Pontus Stenetorp, Sebastian Riedel, Maria Pikoula, Richard Dobson, and Harry Hemingway.  Application of clinical concept embeddings for heart failure prediction in UK EHR data.  _NIPS ML4H: Machine Learning for Health_ , 2018.
  * Cai et al. [2018] Xiangrui Cai, Jinyang Gao, Kee Yuan Ngiam, Beng Chin Ooi, Ying Zhang, and Xiaojie Yuan.  Medical concept embedding with time-aware attention.  In _Proceedings of the 27th International Joint Conference on Artificial Intelligence_ , IJCAI’18, page 3984–3990. AAAI Press, 2018.
  * Darabi et al. [2020] Sajad Darabi, Mohammad Kachuee, Shayan Fazeli, and Majid Sarrafzadeh.  TAPER: Time-aware patient EHR representation.  _IEEE journal of biomedical and health informatics_ , 24(11):3268—3275, 2020.
  * Nguyen et al. [2017] Phuoc Nguyen, Truyen Tran, Nilmini Wickramasinghe, and Svetha Venkatesh.  Deepr: A convolutional net for medical records.  _IEEE Journal of Biomedical and Health Informatics_ , 21(1):22–30, 2017.
  * Hur et al. [2022] Kyunghoon Hur, Jiyoung Lee, Jungwoo Oh, Wesley Price, Younghak Kim, and Edward Choi.  Unifying heterogeneous electronic health records systems via text-based code embedding.  In _Proceedings of the Conference on Health, Inference, and Learning_ , volume 174 of _PMLR_ , pages 183–203, 2022.
  * Chapman et al. [2021] Martin Chapman, Shahzad Mumtaz, Luke V Rasmussen, Andreas Karwath, Georgios V Gkoutos, Chuang Gao, Dan Thayer, Jennifer A Pacheco, Helen Parkinson, Rachel L Richesson, Emily Jefferson, Spiros Denaxas, and Vasa Curcin.  Desiderata for the development of next-generation electronic health record phenotype libraries.  _Gigascience_ , 10, 2021.
  * [29] WHO.  Who. diabetes.  https://www.who.int/news-room/fact-sheets/detail/diabetes, Accessed: 19-08-2022.
  * [30] Medscape.  Type 2 diabetes mellitus treatment & management.  https://emedicine.medscape.com/article/117853-treatment, Accessed: 03-08-2022.
  * Ferreira et al. [2019] João Pedro Ferreira, Sarah Kraus, Sharon Mitchell, Pablo Perel, Daniel Piñeiro, Ovidiu Chioncel, Roberto Colque, Rudolf A de Boer, Juan Esteban Gomez-Mesa, Hugo Grancelli, Carolyn S P Lam, Antoni Martinez-Rubio, John J V McMurray, Alexandre Mebazaa, Gurusher Panjrath, Ileana L Piña, Mahmoud Sani, David Sim, Mary Walsh, Clyde Yancy, Faiez Zannad, and Karen Sliwa.  World heart federation roadmap for heart failure.  _Glob. Heart_ , 14(3):197–214, 2019.
  * Ly et al. [2013] Diana Ly, David Forman, Jacques Ferlay, Louise Brinton, and Michael Cook.  An international comparison of male and female breast cancer incidence rates.  _International journal of cancer_ , 132, 2013.
  * Rawla [2019] Prashanth Rawla.  Epidemiology of prostate cancer.  _World Journal of Oncology_ , 10:63–89, 2019.
  * Sechidis et al. [2011] Konstantinos Sechidis, Grigorios Tsoumakas, and Ioannis P. Vlahavas.  On the stratification of multi-label data.  In _ECML/PKDD_ , 2011.

##  Appendix A Related Work

ML approaches have been applied to EHR data either using a cross-sectional
matrix of diagnosis terms (e.g. an ICD-10 term) [Henderson et al., 2017,
Miotto et al., 2016] or sequential data as input [Choi et al., 2016a, Pham et
al., 2016, Alaa and van der Schaar, 2019, Choi et al., 2016b, Denaxas et al.,
2018, Cai et al., 2018, Darabi et al., 2020]. In the former, methods are blind
to the order in which diagnoses occur and subsequently how a patient’s disease
profile develops. In the latter, most methods do not learn embeddings for the
full sequence of diagnoses in a patient’s medical history, and instead learn
embeddings per diagnosis term or, at most, a short sequence of terms whether
they are utilising LSTMs [Choi et al., 2016a], RNNs [Pham et al., 2016], CNNs
[Nguyen et al., 2017], or transformer models [Li et al., 2020]. To be able to
deal with heterogeneous ontologies from different EHR data sources, all of
these models rely on noisy and often lossy mappings across ontologies or on
phenotyping algorithms, eg. manually curated groupings of ontology terms, such
as CALIBER [Kuan et al., 2019].

In contrast, our method uses the textual description of terms to learn
representations across the full sequence of diagnoses in a patient’s history.
Recent work by Hur et al. [2022] aims to unify EHR records by learning
description-based embeddings from multiple data sources. Our work was
developed in parallel independently and addresses the specific use case of
cohort expansion, instead of merely providing examples of potential downstream
applications to assess improvements of predictive power.

##  Appendix B Method details and figures

###  B.1 Fusing ontologies via text

We consider all EHR data sources with their ontologies where each concept has
a textual descriptor. For example, the ontologies of GP and hospital records
are made up of clinical terms (e.g. Read version 2/ Clinical Terms Version 3
and ICD9/ICD10 codes, respectively) and their description. For each EHR source
and associated ontology $a\in\mathcal{A}$, we denote the set of concepts (e.g.
clinical terms in the case of GP or hospital records) within this ontology as
$\Theta_{a}$ and the set of text descriptions as $\Xi_{a}$. The total
vocabulary of concepts and text descriptions across all ontologies is denoted
by $\Theta=\bigcup\limits_{a\in\mathcal{A}}\Theta_{a}$ and
$\Xi=\bigcup\limits_{a\in\mathcal{A}}\Xi_{a}$, respectively. For each patient,
we define their full clinical history through time and across sources as the
sequence of time-indexed concepts as $(\theta_{1},\dots,\theta_{t})$,
$\theta_{i}\in\Theta$, $i=1,\ldots,t$.

The premise of our work relies on the assumption that for every concept
$\theta\in\Theta,$ there exists a unique text description
$\xi_{\theta}\in\Xi$. For example, under the ICD-10 terminology, the
alphanumeric code E11.9 has the associated description ‘Type 2 diabetes
mellitus without complications’. Thus the clinical history of each patient can
be uniquely represented by the concatenation of sequences of clinical
descriptions $(\xi_{\theta_{1}},\ldots,\xi_{\theta_{t}})$,
$\xi_{\theta_{i}}\in\Xi$, $i=1,\ldots,t$, ordered in time. Figure 4 shows an
example of fusing primary care and hospitalization records.

Figure B.1: Fusing primary (GP) and secondary (hospital) data into a single
paragraph444Any patient data shown is simulated and does not represent data of
real patients..

###  B.2 Label Generation: Oracle Feature Tagging for Disease Phenotyping

Given a set of clinical terms $\Theta$ and text descriptions $\Xi$, we rely on
external oracles to assign labels to a given set of target phenotypes
$\Delta$. We assume that for each phenotype $d\in\Delta$ there exists a
mapping
$\mathbf{1}_{d}:\Theta\times\Xi\rightarrow\\{0,1\\}$,
$(\theta,\xi_{\theta})\mapsto\delta_{d}$ indicating whether the presence of
$d$ can be inferred from the clinical term and its description. An aggregated
phenotype label of $1$ is assigned, if $\mathbf{1}_{d}(\theta,\xi_{\theta})=1$
for any of the term-description pairs $(\theta,\xi_{\theta})$ in the input
sequence, and $0$ otherwise (Figure 2). Here, $\Delta$ is a set of disease
phenotypes, that can be taken from disease-specific phenotyping algorithms
[Chapman et al., 2021], such as CALIBER [Kuan et al., 2019], which consists of
manually-curated sets of clinical terms across primary and secondary care
ontologies for defining 308 chronic and acute disease phenotypes.

###  B.3 Model Design

Figure B.2 shows a diagram of sEHR-CE.

Figure B.2: Diagram of the sEHR-CE model. The model is fine-tuned on the MLM
task. We then use the pre-trained encoder and train sEHR-CE. The input tokens
are first encoded and the hidden vector of the [CLS] token is passed to the
decoder, a fully connected linear layer. The output is passed through a
sigmoid function to generate probabilities for each phenotype.

##  Appendix C Experiment design and model predictions

To assess the quality of our model predictions, we chose four diseases that
differ in terms of prevalence and clinical characteristics. Type 2 diabetes
mellitus (T2DM) is one of the most prevalent chronic diseases worldwide [WHO,
], and patients are primarily diagnosed and managed in primary care (GP)
[Medscape, ]. Heart failure (HF) is one of the main causes of death in the
older population [Ferreira et al., 2019], and its acute manifestations are
treated in hospital care. Malignant neoplasms of the breast and of the
prostate are both less prevalent diseases almost exclusively present in only
biologically females or males, respectively [Ly et al., 2013, Rawla, 2019]. We
test the performance of our model on its ability to diagnose known cases,
compare it to other methods, and evaluate associations of clinical features
with the predictions on T2DM with available orthogonal data.

###  C.1 Data Processing

The UK Biobank (UKBB) [Sudlow et al., 2015] is a large-scale research study of
around $500$k individuals between the ages of $40$ and $54$ at the time of
recruitment. It includes rich genotyping and phenotyping data, both taken at
recruitment and during primary and secondary care encounters (GP and
hospital). We use patient records from GP and hospital visits in the form of
code ontologies Read version2/ Clinical Terms Version 3, and ICD-9/10 together
with their textual descriptions. We restrict the data set to individuals that
have entries in both hospital and GP records, which reduces our cohort to
$154,668$ individuals. Requiring individuals to have entries in their GP
records reduces bias towards acute events that usually present in hospitals,
but we note that removing individuals without any hospital records may still
bias the data towards more severe cases. We use CALIBER, previously validated
phenotype definitions from Kuan et al. [2019] to label patients with T2DM, HF,
malignant neoplasm of the breast, and malignant neoplasm of the prostate.

A patient can be admitted to the hospital for multiple days. We treat an
entire hospital admission as one point in time using the admission date, and
only keep unique ICD-10/ICD-9 codes for each visit. We aggregate visits that
are less than a week apart into one visit keeping only unique codes. This
approach removes repeated codes, thus avoiding redundancy and reducing
sequence length.

Only patients with at least $5$ clinical terms present in their clinical
history are included to allow for sufficient information for any predictions,
reducing the cohort to a final $129,932$ individuals. We use phenotype
definitions from CALIBER [Kuan et al., 2019] to label patients with T2DM, HF,
malignant neoplasm of the breast, and malignant neoplasm of the prostate. Each
phenotype definition consists of a list of ICD10 and Read2, Read3 codes and
their children.

The usage of PubMedBERT restricts the length of input sequences we can use. To
avoid excluding relevant clinical information by truncating the input
sequences, we break up patient histories longer than the limit into multiple
input sequences of smaller length with the same target vector.

###  C.2 Model Training

We use the pre-trained language model PubMedBERT [Gu et al., 2022] as the
encoder of the tokenised input sequences of clinical term descriptions. Since
our input systematically differs from the general scientific text on which
PubMedBERT was trained, we fine-tuned on the masked-language modeling (MLM)
task (Figure B.2, Figure C.1), by masking words (e.g. descriptions) at random
following the original BERT paper [Devlin et al., 2019]. The model, fine-tuned
using the full UKBB cohort of $138,079$ patients, was trained with early
stopping for $5$ epochs with a batch size of $32$ and a learning rate of
$4\times 10^{-5}$ using gradient descent with an AdamW optimizer, and weight
decay of $0.01$. The output dimension of the encoder was $768$.

The proposed model sEHR-CE is using the fine-tuned encoder and a fully
connected linear layer as the decoder. To train on the multi-label
classification task of outcome prediction, we split our data set into five
equally sampled folds ${f_{0},...f_{4}}$ containing unique patients, using a
stratified sampling method to maintain the same phenotype proportion in every
split [Sechidis et al., 2011], and mask the data according to our strategy
(Section 2.1, Figure C.1).

Figure C.1: Example of input diagnosis and different masking strategies for
MLM fine-tuning and classification training. A description is encoded with
multiple word embeddings. For MLM fine-tuning, words are masked at random; for
classification training, whole descriptions are masked using clinical masking
strategy described in Section 2.1.

We train a total of five models for $3$ epochs on three folds, holding back
folds $f_{i}$ for validation and $f_{(i+1)\bmod{5}}$ for testing for model
$i$, $i=1,\ldots,5$ (Figure C.2). This maintains a $60/20/20$ training,
validation and testing split overall while providing us with enough training
and testing examples. All results presented are predictions of each model on
its respective independent test set. We used a learning rate of $10^{-5}$, and
a warm-up proportion of $0.25$. Performance was monitored every $0.25$ epochs
on the validation fold.

Figure C.2: Five models trained using five independent folds from the total
data set using stratified sampling.

###  C.3 Model evaluation

As explained in 3, we compare the performance of our model sEHR-CE to BEHRT
[Li et al., 2020]. BEHRT takes a tokenised sequence of clinical terms, age and
position embeddings as input. Ontologies from hospital and GP records are
mapped to CALIBER definitions [Kuan et al., 2019], removing unmapped terms.
Phenotype definitions in CALIBER include different categories (for example,
‘diabetes’ contains categories ‘type 1’ and ‘type 2’), that were ignored by
the original BEHRT publication, so we expanded the token set to define a token
per CALIBER phenotype and category. A transformer model is pre-trained to
predict masked tokens before it is trained to predict a set of possible
diagnoses an individual may develop given the input sequence. Similarly, we
trained five sEHR-CE models restricted to CALIBER tokens (denoted sEHR-CE-
codes) for comparison. All results presented are predictions of each model on
its respective independent test set.

Figure C.3 shows the predicted probabilities for cases and controls across all
phenotypes and models, and figure C.4 shows the AUPRC curves for each
phenotype and method in the test sets. To avoid inclusion of too many false
positives, we defined missing cases as those controls with predicted
probability in the $98$th percentile.

Figure C.3: Predicted probabilities for cases and controls in test sets across
all phenotypes. Vertical lines indicate the $98$th percentile threshold.
Figure C.4: AUPRC curves for each phenotype in the test sets.

###  C.4 Expanded Qualitative Evaluation of T2DM Diagnosis Prediction

Five groups were used to evaluate the model’s T2DM diagnosis predictions.
Table C.1 shows the percentiles of sEHR-CE’s predicted probabilities to define
each group, along with each size.

Patient group | Size
---|---
Cases | 16431
Controls | 113501
| Controls with high probability
---
(p>=0.85, 98th percentile)
2020
| Cases with high probability
---
(p>=0.985, 90th percentile)
2072
| Cases with low probability
---
(p<=0.25, 12th percentile)
2343
Table C.1: Case and control cohorts and groups of interest based on sEHR-CE’s
predicted probability for T2DM in the test sets.

We then investigated the association of predicted probability and several
proxies of disease severity: number of GP and hospital admissions, survival
and risk of cardiovascular disease.

####  C.4.1 Number of GP visits and hospital admissions

As expected, both cases and controls with a high predicted probability of a
T2DM diagnosis, exhibit a slightly higher number of GP and hospital visits
than the other groups (Figure C.5), indicating that they are experiencing a
more severe form of T2DM requiring care. This is particularly higher in the
case of hospital visits, indicating patients experiencing acute events: both
cases and controls with a high predicted probability visit a hospital
approximately $10$ times more often than their low probability counterparts.

Although the model was not given information from which data source the input
data was coming from, this analysis indicates that it has learned to associate
acute events with disease severity.

Figure C.5: Distribution of number of hospital and GP codes per patient group.
Matched controls on sex.

####  C.4.2 Survival analysis

To compare survival across different groups of individuals, we use the Kaplan-
Meier estimator with all-cause mortality as the endpoint with right-censored
data accounting for individuals without any event occurrence since the last
follow-up. Both cases and controls with high predicted probabilities had the
lowest survival, followed by general cases, controls and finally cases with
low predicted probability (Figure 6(a)), indicating that the model’s predicted
probability is associated with survival.

####  C.4.3 Cardiovascular Risk

T2DM is a known risk factor and comorbidity of cardiovascular disease, which,
in turn, is the most prevalent cause of death in T2DM patients. The GP records
contain Framingham and QRISK3 scores; these are two scores that assess an
individual’s risk of developing cardiovascular disease within the next 10
years, based on several coronary risk factors. The Framingham score is derived
from an individual’s age, gender, total cholesterol, high-density lipoprotein
cholesterol, smoking habits, and systolic blood pressure, whereas the QRISK
score, which is almost exclusively used now, extends this score with
additional factors such as body mass index, ethnicity, measures of
deprivation, chronic kidney disease, rheumatoid arthritis, atrial
fibrillation, diabetes mellitus, and anti-hypertensive treatment. Both cases
and controls with high predicted probability of having T2DM had a higher risk
of developing cardiovascular disease compared to their low predicted
probability counterparts (Figure 6(b)) indicating that the model has learned
to associate the risk of developing both diseases at the same time.

(a)

(b)

Figure C.6: 6(a) Death survival plots for different patient groups. 6(b)
Framingham and QRISK cardiovascular scores vs T2DM probability prediction.
Each point represents the median of each percentile of cardiovascular risk.