# Refcat: The Internet Archive Scholar Citation Graph

Martin Czygan

Internet Archive
San Francisco, CA, USA
<EMAIL_ADDRESS>
Helge Holzmann

Internet Archive
San Francisco, CA, USA
<EMAIL_ADDRESS>
Bryan Newbold

Internet Archive
San Francisco, CA, USA
<EMAIL_ADDRESS>

###### Abstract

As part of its scholarly data efforts, the Internet Archive (IA) releases a
first version of a citation graph dataset, named _refcat_ , derived from
scholarly publications and additional data sources. It is composed of data
gathered by the fatcat cataloging project111https://fatcat.wiki (the catalog
that underpins IA Scholar), related web-scale crawls targeting primary and
secondary scholarly outputs, as well as metadata from the Open
Library222https://openlibrary.org project and
Wikipedia333https://wikipedia.org. This first version of the graph consists of
over 1.3B citations. We release this dataset under a CC0 Public Domain
Dedication, accessible through Internet
Archive444https://archive.org/details/refcat_2021-07-28. The source code used
for the derivation process, including exact and fuzzy citation matching, is
released under an MIT license555https://gitlab.com/internetarchive/refcat. The
goal of this report is to describe briefly the current contents and the
derivation of the dataset.

Index terms— Citation Graph, Web Archiving

##  1 Introduction

The Internet Archive released a first version of a citation graph dataset
derived from a corpus of about 2.5B raw references666Number of raw references:
2,507,793,772 gathered from 63,296,308 metadata records (which are collected
from various sources or based on data obtained by PDF extraction and
annotation tools such as GROBID (Lopez, 2009)). Additionally, we consider
integration with metadata from Open Library and Wikipedia. We expect this
dataset to be iterated upon, with changes both in content and processing.

According to (Jinha, 2010) over 50M scholarly articles have been published
(from 1726) up to 2009, with the rate of publications on the rise (Landhuis,
2016). In 2014, a study based on academic search engines estimated that at
least 114M English-language scholarly documents are accessible on the web
(Khabsa and Giles, 2014).

Modern citation indexes can be traced back to the early computing age, when
projects like the Science Citation Index (1955) (Garfield, 2007) were first
devised, living on in commercial knowledge bases today. Open alternatives were
started such as the Open Citations Corpus (OCC) in 2010 - the first version of
which contained 6,325,178 individual references (Shotton, 2013). Other notable
projects include CiteSeer (Giles et al., 1998), CiteSeerX (Wu et al., 2019)
and CitEc777https://citec.repec.org. The last decade has seen the emergence of
more openly available, large scale citation projects like Microsoft Academic
(Sinha et al., 2015) and the Initiative for Open Citations888https://i4oc.org
(Shotton, 2018). In 2021, over one billion citations are publicly available,
marking a “tipping point” for this category of data (Hutchins, 2021).

While a paper will often cite other papers, more citable entities exist such
as books or web links and within links a variety of targets, such as web
pages, reference entries, protocols or datasets. References can be extracted
manually or through more automated methods, by accessing relevant metadata or
structured data extraction from full text documents. Automated methods offer
the benefits of scalability. The completeness of bibliographic metadata in
references ranges from documents with one or more persistent identifiers to
raw, potentially unclean strings partially describing a scholarly artifact.

##  2 Related Work

Two typical problems in citation graph development are related to data
acquisition and citation matching. Data acquisition itself can take different
forms: bibliographic metadata can contain explicit reference data as provided
by publishers and aggregators; this data can be relatively consistent when
looked at per source, but may vary in style and comprehensiveness when looked
at as a whole. Another way of acquiring bibliographic metadata is to analyze a
source document, such as a PDF (or its text), directly. Tools in this category
are often based on conditional random fields (Lafferty et al., 2001) and have
been implemented in projects such as ParsCit (Councill et al., 2008), Cermine
(Tkaczyk et al., 2014), EXCITE (Hosseini et al., 2019) or GROBID (Lopez,
2009).

The problem of citation matching is relatively simple when common, persistent
identifiers are present in the data. Complications mount, when there is
_Identity Uncertainty_ , that is “objects are not labeled with unique
identifiers or when those identifiers may not be perceived perfectly” (Pasula
et al., 2003). CiteSeer has been an early project concerned with citation
matching (Giles et al., 1998). A taxonomy of potential issues common in the
matching process has been compiled by (Olensky et al., 2016). Additional care
is required, when the citation matching process is done at scale (Fedoryszak
et al., 2013). The problem of heterogenity has been discussed in the context
of datasets by (Mathiak and Boland, 2015).

Projects and datasets centered around citations or containing citation data as
a core component are COCI, the “OpenCitations Index of Crossref open DOI-to-
DOI citations”, which was first released
2018-07-29999https://opencitations.net/download and has been regularly updated
(Peroni and Shotton, 2020). The
WikiCite101010https://meta.wikimedia.org/wiki/WikiCite project, “a Wikimedia
initiative to develop open citations and linked bibliographic data to serve
free knowledge” continuously adds citations to its
database111111http://wikicite.org/statistics.html. Microsoft Academic Graph
(Sinha et al., 2015) is comprised of a number of
entities121212https://docs.microsoft.com/en-us/academic-
services/graph/reference-data-schema with _PaperReferences_ being one relation
among many others.

##  3 Dataset

We released the first version of the _refcat_ dataset in a format used
internally for storage and to serve queries (and which we call _biblioref_ or
_bref_ for short). The dataset includes metadata from fatcat (the catalog
underpinning IA Scholar), the Open Library project and inbound links from the
English Wikipedia. The dataset is integrated into the fatcat.wiki website and
allows users to explore inbound and outbound
references131313https://guide.fatcat.wiki/reference_graph.html.

The format records source and target identifiers, a few metadata attributes
(such as year or release stage, i.e. preprint, version of record, etc) as well
as information about the match status and provenance.

The dataset currently contains 1,323,423,672 citations across 76,327,662
entities (55,123,635 unique source and 60,244,206 unique target work
identifiers; for 1,303,424,212 - or 98.49% of all citations - we do have a DOI
for both source and target). The majority of matches - 1,250,523,321 - is
established through identifier based matching (DOI, PMIC, PMCID, ARXIV, ISBN).
72,900,351 citations are established through fuzzy matching techniques, where
references did not contain identifiers141414This not necessary mean that the
records in question do not have an identifier; however if an identifier
existed, it was not part of the raw reference.. Citations from the Open
Citations’ COCI corpus151515Reference dataset COCI v11, released 2021-09-04,
http://opencitations.net/index/coci and _refcat_ overlap to the most part, as
can be seen in Table 1.

Set |  | Count
---|---|---
COCIv11 (C) |  | 1,186,958,897
_refcat-doi_ (R) |  | 1,303,424,212
C $\cap$ R | overlap | 1,046,438,515
C $\setminus$ R | COCIv11 only | 140,520,382
R $\setminus$ C |  _refcat-doi_ only | 256,985,697
Table 1: Comparison between Open Citations COCI corpus (v11, 2021-09-04) and _refcat-doi_ , a subset of _refcat_ where entities have a known DOI. At least 150,727,673 (58.7%) of the 256,985,697 references in _refcat-doi_ only record links within a specific dataset provider; here GBIF with DOI prefix: 10.15468. Edge type | Count
---|---
doi-doi | 1,303,424,212
target-open-library | 20,307,064
source-wikipedia | 1,386,941
Table 2: Counts of classic DOI to DOI references as well as outbound
references matched against Open Library as well as inbound references from the
English Wikipedia.

We started to include non-traditional citations into the graph, such as links
to books included in Open Library and links from the English Wikipedia to
scholarly works. For links between Open Library we employ both identifier
based and fuzzy matching; for Wikipedia references we used a published dataset
(Singh et al., 2020) and we are contributing to upstream projects related to
wikipedia citation extraction, such as _wikiciteparser_
161616https://github.com/dissemin/wikiciteparser to generate updates from
recent Wikipedia dumps171717Wikipedia dumps are available on a monthly basis
from https://dumps.wikimedia.org/.. Table 2 lists the counts for these links.
Additionally, we are examining web links appearing in references: after an
initial cleaning procedure we currently find 25,405,592 web links181818The
cleaning process is necessary because OCR artifacts and other metadata issues
exist in the data. Unfortunately, even after cleaning not all links will be in
the form as originally intended by the authors. in the reference corpus, of
which 4,828,283 (19%) have been preserved as of August 2021 with an HTTP 200
status code in the Wayback Machine191919https://archive.org/web/ of the
Internet Archive. As an upper bound - if we additionally include all
redirection (HTTP 3XX) and server error status codes (HTTP 5XX) - we find a
total of 14,306,019 (56.3%) links preserved.

We ran a live URL check202020All links accessed on 2021-10-04 and 2021-10-05.
over a sample of 364415 links which appear in the reference corpus _and_ have
a HTTP 200 status code archival copy in the Wayback Machine. Of the 364415
links we find 305476 (83.8%) responding with an HTTP 200 OK, whereas the rest
of the links yield a variety of HTTP status codes, like 404, 403, 500 and
others - resulting in about 16% of the links in the reference corpus preserved
at the Internet Archive being currently inaccessible on the web212121We used
the https://github.com/miku/clinker command line link checking tool. \- making
targeted web crawling and preservation of scholarly references a key activity
for maintaining citation integrity.

##  4 System Design

###  4.1 Constraints

The constraints for the system design are informed by the volume and the
variety of the data. The capability to run the whole graph derivation on a
single machine222222We used a shared virtual server with 24 cores and 48G of
main memory. The most memory-intensive part of the processing currently are
the buffers set aside for _GNU sort_. was a minor goal as well. In total, the
raw inputs amount to a few terabytes of textual content, mostly newline
delimited JSON. More importantly, while the number of data fields is low,
certain documents are very partial with hundreds of different combinations of
available field values found in the raw reference data. This is most likely
caused by aggregators passing on reference data coming from hundreds of
sources, each of which not necessarily agreeing on a common granularity for
citation data and from artifacts of machine learning based structured data
extraction tools.

Each combination of fields may require a slightly different processing path.
For example, references with an Arxiv identifier can be processed differently
from references with only a title.

###  4.2 Data Sources

Reference data comes from two main sources: explicit bibliographic metadata
and PDF extraction. The bibliographic metadata is taken from fatcat, which
itself harvests and imports web accessible sources such as Crossref, Pubmed,
Arxiv, Datacite, DOAJ, dblp and others into its catalog (as the source
permits, data is processed continuously or in batches). Reference data from
PDF documents has been extracted with GROBID232323GROBID v0.5.5, with the TEI-
XML results being cached locally in a key-value store accessible with an S3
API242424Currently, https://github.com/chrislusf/seaweedfs is used. Archived
PDF documents result from dedicated web-scale crawls of scholarly domains
conducted with multiple open-source crawler technologies created by the
Internet Archive and a variety of seed lists targeting journal homepages,
repositories, dataset providers, aggregators, web archives and other venues. A
processing pipeline merges catalog data from the primary database and cached
data from the key-value store and generates the set of about 2.5B references
records, which currently serve as an input for the citation graph derivation
pipeline.

###  4.3 Methodology

Overall, a map-reduce style (Dean and Ghemawat, 2010) approach is
followed252525While the operations are similar, the processing is not
distributed but runs on a single machine. For space efficiency, zstd (Collet
and Kucherawy, 2018) is used to compress raw data and derivations., which
allows for some uniformity in the processing. We extract _(key, document)_
tuples (as TSV) from the raw JSON data and sort by key. We then group
documents with the same key and apply a function on each group in order to
generate our target schema or perform additional operations such as
deduplication or fusion of matched and unmatched references for indexing.

The key derivation can be exact (via an identifier like DOI, PMID, etc) or
based on a value normalization, like “slugifying” a title string. For
identifier based matches we can generate the target schema directly. For fuzzy
matching candidates, we pass possible match pairs through a verification
procedure, which is implemented for _release entity_
262626https://guide.fatcat.wiki/entity_release.html. pairs. This procedure is
a domain dependent rule based verification, able to identify different
versions of a publication, preprint-published pairs and documents, which are
are similar by various metrics calculated over title and author fields. The
fuzzy matching approach is applied on all reference documents without any
identifier (a title is currently required).

We currently implement performance sensitive parts in the Go programming
language272727https://golang.org/, with various processing stages (e.g.
conversion, map, reduce, …) represented by separate command line tools. A thin
task orchestration layer using the luigi
framework282828https://github.com/spotify/luigi (Bernhardsson and Freider,
2018), which has been used in various scientific pipeline application, like
(Schulz et al., 2016), (Erdmann et al., 2017), (Lampa et al., 2019), (Czygan,
2014) and others. allows for experimentation in the pipeline and for single
command derivations, as data dependencies are encoded with the help of the
orchestrator. Within the tasks, we also utilize classic platform tools such as
GNU _sort_ (McIlroy, 1971).

During a last processing step, we fuse reference matches and unmatched items
into a single, indexable file. This step includes deduplication of different
matching methods (e.g. prefer exact matches over fuzzy matches). This file is
indexed into a search index and serves both matched and unmatched references
for the web application, allowing for further collection of feedback on match
quality and possible improvements.

With a few schema conversions, fuzzy matching has been be applied to Wikipedia
articles and Open Library (edition) records as well. The aspect of precision
and recall are represented by the two stages: we are generous in the match
candidate generation phase in order to improve recall, but we are strict
during verification, in order to control precision. Quality assurance for
verification is implemented through a growing list of test cases of real
examples from the catalog and their expected or desired match status292929The
list can be found under:
https://gitlab.com/internetarchive/refcat/-/blob/master/skate/testdata/verify.csv.
It is helpful to keep this test suite independent of any specific programming
language..

##  5 Limitations and Future Work

As with other datasets in this field we expect this dataset to be iterated
upon.

  * •

The fatcat catalog updates its metadata continuously303030A changelog can
currently be followed here: https://fatcat.wiki/changelog. and web crawls are
conducted regularly. Current processing pipelines cover raw reference snapshot
creation and derivation of the graph structure, which allows to rerun the
processing pipeline based on updated data as it becomes available.

  * •

Metadata extraction from PDFs depends on supervised machine learning models,
which in turn depend on available training datasets. With additional crawls
and metadata available we hope to improve models used for metadata extraction,
improving yield and reducing data extraction artifacts in the process.

  * •

As of this version, a number of raw reference docs remain unmatched, which
means that neither exact nor fuzzy matching has detected a link to a known
entity. Metadata might be missing. However, parts of the data will contain a
reference to a catalogued entity, but in a specific, dense and harder to
recover form.

  * •

The reference dataset contains millions of URLs and their integration into the
graph has been implemented as a prototype. A full implementation requires a
few data cleanup and normalization steps.

##  6 Acknowledgements

This work is partially supported by grants from the _Andrew W. Mellon
Foundation_ , especially ”Ensuring the Persistent Access of Open Access
Journal Literature: Phase II“ (1910-07256, Jefferson Bailey, Principal
Investigator).

## Appendix: Reference Relations

Figure 1 shows the schematic reference relations.

Figure 1: Schematics of the main reference entities; green: included in the
corpus; orange: currently in development.

## Appendix: Data Quality

A note on data quality: While we implement various data quality measures,
real-world data, especially coming from many different sources will contain
issues. Among other measures, we keep track of match reasons, especially for
fuzzy matching to be able to zoom in on systematic errors more easily (see
Table 3).

Count | Provenance | Status | Reason
---|---|---|---
934932865 | crossref | exact | doi
151366108 | fatcat-datacite | exact | doi
65345275 | fatcat-pubmed | exact | pmid
48778607 | fuzzy | strong | jaccardauthors
42465250 | grobid | exact | doi
29197902 | fatcat-pubmed | exact | doi
19996327 | fatcat-crossref | exact | doi
11996694 | fuzzy | strong | slugtitleauthormatch
9157498 | fuzzy | strong | tokenizedauthors
3547594 | grobid | exact | arxiv
2310025 | fuzzy | exact | titleauthormatch
1496515 | grobid | exact | pmid
680722 | crossref | strong | jaccardauthors
476331 | fuzzy | strong | versioneddoi
449271 | grobid | exact | isbn
230645 | fatcat-crossref | strong | jaccardauthors
190578 | grobid | strong | jaccardauthors
156657 | crossref | exact | isbn
123681 | fatcat-pubmed | strong | jaccardauthors
79328 | crossref | exact | arxiv
57414 | crossref | strong | tokenizedauthors
53480 | fuzzy | strong | pmiddoipair
52453 | fuzzy | strong | dataciterelatedid
47119 | grobid | strong | slugtitleauthormatch
36774 | fuzzy | strong | arxivversion
Table 3: Table of match counts (top 25), reference provenance, match status
and match reason. Provenance currently can name the raw origin (e.g.
_crossref_) or the method (e.g. _fuzzy_). The match reason identifier encodes
a specific rule in the domain dependent verification process and is included
for completeness - we do not include the details of each rule in this report.

## References

  * Bernhardsson and Freider [2018] E Bernhardsson and E Freider.  Rouhani a. spotify/luigi-github, 2018.
  * Collet and Kucherawy [2018] Yann Collet and Murray Kucherawy.  Zstandard compression and the application/zstd media type.  _RFC 8478_ , 2018.
  * Councill et al. [2008] Isaac G Councill, C Lee Giles, and Min-Yen Kan.  Parscit: an open-source crf reference string parsing package.  In _LREC_ , volume 8, pages 661–667, 2008.
  * Czygan [2014] Martin Czygan.  Design and implementation of a library metadata management framework and its application in fuzzy data deduplication and data reconciliation with authority data.  _Informatik 2014_ , 2014.
  * Dean and Ghemawat [2010] Jeffrey Dean and Sanjay Ghemawat.  Mapreduce: a flexible data processing tool.  _Communications of the ACM_ , 53(1):72–77, 2010\.
  * Erdmann et al. [2017] M Erdmann, B Fischer, R Fischer, and M Rieger.  Design and execution of make-like, distributed analyses based on spotify’s pipelining package luigi.  In _Journal of Physics: Conference Series_ , volume 898, page 072047\. IOP Publishing, 2017.
  * Fedoryszak et al. [2013] Mateusz Fedoryszak, Dominika Tkaczyk, and Łukasz Bolikowski.  Large scale citation matching using apache hadoop.  In _International Conference on Theory and Practice of Digital Libraries_ , pages 362–365. Springer, 2013.
  * Garfield [2007] Eugene Garfield.  The evolution of the science citation index.  _International microbiology_ , 10(1):65, 2007\.
  * Giles et al. [1998] C Lee Giles, Kurt D Bollacker, and Steve Lawrence.  Citeseer: An automatic citation indexing system.  In _Proceedings of the third ACM conference on Digital libraries_ , pages 89–98, 1998.
  * Hosseini et al. [2019] Azam Hosseini, Behnam Ghavimi, Zeyd Boukhers, and Philipp Mayr.  Excite–a toolchain to extract, match and publish open literature references.  In _2019 ACM/IEEE Joint Conference on Digital Libraries (JCDL)_ , pages 432–433. IEEE, 2019.
  * Hutchins [2021] B Ian Hutchins.  A tipping point for open citation data.  _Quantitative Science Studies_ , pages 1–5, 2021.
  * Jinha [2010] Jinha.  Article 50 million: an estimate of the number of scholarly articles in existence.  23, Jul 2010.  doi:10.1087/20100308.
  * Khabsa and Giles [2014] Khabsa and Giles.  The number of scholarly documents on the public web.  May 2014.  doi:10.1371/journal.pone.0093949.
  * Lafferty et al. [2001] John Lafferty, Andrew McCallum, and Fernando CN Pereira.  Conditional random fields: Probabilistic models for segmenting and labeling sequence data.  2001\.
  * Lampa et al. [2019] Samuel Lampa, Martin Dahlö, Jonathan Alvarsson, and Ola Spjuth.  Scipipe: A workflow library for agile development of complex and dynamic bioinformatics pipelines.  _GigaScience_ , 8(5):giz044, 2019.
  * Landhuis [2016] Landhuis.  Scientific literature: Information overload.  535(7612), Jul 2016.  doi:10.1038/nj7612-457a.
  * Lopez [2009] Patrice Lopez.  Grobid: Combining automatic bibliographic data recognition and term extraction for scholarship publications.  In _International conference on theory and practice of digital libraries_ , pages 473–474. Springer, 2009.
  * Mathiak and Boland [2015] Brigitte Mathiak and Katarina Boland.  Challenges in matching dataset citation strings to datasets in social science.  _D-Lib Magazine_ , 21(1/2):23–28, 2015.
  * McIlroy [1971] M Douglas McIlroy.  A research unix reader: annotated excerpts from the programmer’s manual.  1971\.
  * Olensky et al. [2016] Marlies Olensky, Marion Schmidt, and Nees Jan van Eck.  Evaluation of the citation matching algorithms of cwts and i fq in comparison to the w eb of science.  _Journal of the Association for Information Science and Technology_ , 67(10):2550–2564, 2016.
  * Pasula et al. [2003] Hanna Pasula, Bhaskara Marthi, Brian Milch, Stuart J Russell, and Ilya Shpitser.  Identity uncertainty and citation matching.  In _Advances in neural information processing systems_ , pages 1425–1432, 2003.
  * Peroni and Shotton [2020] Silvio Peroni and David Shotton.  Opencitations, an infrastructure organization for open scholarship.  _Quantitative Science Studies_ , 1(1):428–444, 2020.
  * Schulz et al. [2016] Wade L Schulz, Thomas JS Durant, Alexa J Siddon, and Richard Torres.  Use of application containers and workflows for genomic data analysis.  _Journal of pathology informatics_ , 7, 2016.
  * Shotton [2013] David Shotton.  Publishing: open citations.  _Nature News_ , 502(7471):295, 2013.
  * Shotton [2018] David Shotton.  Funders should mandate open citations.  _Nature_ , 553(7686):129–130, 2018.
  * Singh et al. [2020] Harshdeep Singh, Robert West, and Giovanni Colavizza.  Wikipedia Citations: A comprehensive dataset of citations with identifiers extracted from English Wikipedia, July 2020.  URL https://doi.org/10.5281/zenodo.3940692.
  * Sinha et al. [2015] Arnab Sinha, Zhihong Shen, Yang Song, Hao Ma, Darrin Eide, Bo-June Hsu, and Kuansan Wang.  An overview of microsoft academic service (mas) and applications.  In _Proceedings of the 24th international conference on world wide web_ , pages 243–246, 2015.
  * Tkaczyk et al. [2014] Dominika Tkaczyk, Pawel Szostek, Piotr Jan Dendek, Mateusz Fedoryszak, and Lukasz Bolikowski.  Cermine–automatic extraction of metadata and references from scientific literature.  In _2014 11th IAPR International Workshop on Document Analysis Systems_ , pages 217–221. IEEE, 2014.
  * Wu et al. [2019] Jian Wu, Kunho Kim, and C Lee Giles.  Citeseerx: 20 years of service to scholarly big data.  In _Proceedings of the Conference on Artificial Intelligence for Data Discovery and Reuse_ , pages 1–4, 2019.


– arXiv Preprint 2024 Nicholas Konz1<EMAIL_ADDRESS>
Yuwen Chen1<EMAIL_ADDRESS>
Haoyu Dong1<EMAIL_ADDRESS>
Maciej A. Mazurowski1,2,3,4<EMAIL_ADDRESS>
1Dept. of Electrical and Computer Engineering, 2Dept. of Radiology, 3Dept. of
Computer Science,
4Dept. of Biostatistics & Bioinformatics, Duke University, NC, USA

# Anatomically-Controllable Medical Image Generation with Segmentation-Guided
Diffusion Models

###### Abstract

Diffusion models have enabled remarkably high-quality medical image
generation, which can help mitigate the expenses of acquiring and annotating
new images by supplementing small or imbalanced datasets, along with other
applications. However, these are hampered by the challenge of enforcing global
anatomical realism in generated images. To this end, we propose a diffusion
model for anatomically-controlled medical image generation. Our model follows
a multi-class anatomical segmentation mask at each sampling step and
incorporates a random mask ablation training algorithm, to enable conditioning
on a selected combination of anatomical constraints while allowing flexibility
in other anatomical areas. This also improves the network’s learning of
anatomical realism for the completely unconditional (unconstrained generation)
case. Comparative evaluation on breast MRI and abdominal/neck-to-pelvis CT
datasets demonstrates superior anatomical realism and input mask faithfulness
over state-of-the-art models. We also offer an accessible codebase and release
a dataset of generated paired breast MRIs. Our approach facilitates diverse
applications, including pre-registered image generation, counterfactual
scenarios, and others.

###### keywords:

Diffusion models, image generation, breast MRI, CT

## Introduction

Denoising diffusion probabilistic models [Ho et al.(2020)Ho, Jain, and Abbeel]
(DDPMs, or just “diffusion models”) have shown extensive promise in medical
image analysis and beyond [Kazerouni et al.(2023)Kazerouni, Aghdam, Heidari,
Azad, Fayyaz, Hacihaliloglu, and Merhof] due to their ability to generate
high-quality, high-resolution images. However, standard generative models like
DDPMs can still fail to create anatomically plausible tissue (Fig
LABEL:fig:eg_imgs). Our proposed solution is to incorporate anatomical
information as a prior for image generation via a segmentation mask for
different types of tissue, organs, etc., providing the network with a more
direct learning signal for anatomical realism. Image-space-based diffusion
models are especially amenable for strict segmentation guidance because a
conditioning mask can be used as-is for every small step of the denoising
process, with no conversion to some abstract latent space as in latent
diffusion models (or GANs) where complex spatial conditions may be lost. In
this paper, we introduce an image-space diffusion model for segmentation-
guided image generation: Segmentation-Guided Diffusion.

fig:eg_imgs

Figure 1: Standard diffusion models (right) can fail to create realistic
tissue, even if the overall image appears high-quality. From left to right:
(1) randomly sampled real breast MRI images, (2) synthetic images created by
our segmentation-guided diffusion model, and (3) by the same model but without
mask guidance; insets show areas with fibroglandular tissue.

Segmentation-guided generation would be even more flexible if only certain
object classes could be constrained in an input mask, while others are free to
be inferred/“filled in” by the model (also allowing the model to be applied to
incomplete segmentations [Tian et al.(2023)Tian, Yu, and Fu]). This also opens
up further applications such as the generation of anatomically
paired/registered data, “counterfactual” analysis of existing annotated data,
and others. To solve this, we propose a mask-ablated training strategy to
provide the model with all possible combinations of missing classes in masks
during training, teaching it to make such inferences when generating new
images. Notably, we also show that this training strategy helps the model
learn to generate more realistic tissue in the case of completely
unconditional (empty mask) inference (Sec. 3.2), allowing for the “best of
both worlds” of both anatomical accuracy and diverse and flexible
unconditional generation.

We evaluate our model’s ability to generate images conditioned on a range of
anatomical objects of interest for breast MRI and neck-to-pelvis CT datasets
and compare it to state-of-the-art mask-conditional and unconditional
generative models in terms of anatomical realism and faithfulness to input
masks of generated images. Notably, our approach far outperforms the recent
mask-conditional latent diffusion model ControlNet [Zhang et al.(2023)Zhang,
Rao, and Agrawala] in both metrics. We also publicly release our code at
https://github.com/mazurowski-lab/segmentation-guided-diffusion with a focus
on usability on any dataset of choice, along with a dataset of paired “pre-
registered” generated breast MRIs.

## Related Work

#### Diffusion Models for Medical Image Analysis

Since their inception and eventual advancement over generative adversarial
networks (GANs) for high-quality image generation, diffusion models have seen
extensive application in medical image analysis in recent years [Kazerouni et
al.(2023)Kazerouni, Aghdam, Heidari, Azad, Fayyaz, Hacihaliloglu, and Merhof].
Applications have included direct image generation [Pinaya et
al.(2022b)Pinaya, Tudosiu, Dafflon, Da Costa, Fernandez, Nachev, Ourselin, and
Cardoso, Khader et al.(2023)Khader, Müller-Franzes, Tayebi Arasteh, Han,
Haarburger, Schulze-Hagen, Schad, Engelhardt, Baeßler, Foersch, et al.], image
segmentation [Wolleb et al.(2022b)Wolleb, Sandkühler, Bieder, Valmaggia, and
Cattin], anomaly detection [Pinaya et al.(2022a)Pinaya, Graham, Gray, Da
Costa, Tudosiu, Wright, Mah, MacKinnon, Teo, Jager, et al., Wolleb et
al.(2022a)Wolleb, Bieder, Sandkühler, and Cattin], cross-modality image
translation [Lyu and Wang(2022)], and image denoising [Gong et al.(2023)Gong,
Johnson, El Fakhri, Li, and Pan], just to name a few.

#### Image Generation from Spatial Conditions

Generating an image from a mask can be considered as an image-to-image
translation or conditional generation task, where an input image is typically
converted to some latent representation from which the image in the desired
output space is decoded. Examples include GAN-based models [Yang et
al.(2019)Yang, Dvornek, Zhang, Chapiro, Lin, and Duncan, Choi et
al.(2018)Choi, Choi, Kim, Ha, Kim, and Choo, Wang et al.(2018)Wang, Liu, Zhu,
Tao, Kautz, and Catanzaro], or diffusion-based models [Wolleb et
al.(2022a)Wolleb, Bieder, Sandkühler, and Cattin] which partially noise the
image, then guide the denoising process towards the target domain via various
forms of image class guidance [Dhariwal and Nichol(2021), Ho and
Salimans(2021)]. As opposed to our approach, these methods do not directly
incorporate pixel-wise spatial guidance in the image generation process.

ControlNet [Zhang et al.(2023)Zhang, Rao, and Agrawala] is a recent method for
adding various fine-tuned spatial conditioning controls to Stable Diffusion
[Rombach et al.(2022)Rombach, Blattmann, Lorenz, Esser, and Ommer], a large
pre-trained text-to-image latent diffusion model. While ControlNet relies on
first encoding an input mask to some latent space before feeding it to the
diffusion model, our model always receives mask signals from image space
directly. We find that this allows for more transparency and faithfulness with
how the network can follow the mask conditions, and better anatomical realism
(Sec. 3.1), and that ControlNet is generally challenging to adapt to medical
images for various reasons (Appendix C).

##  1 Method

###  1.1 A Brief Review of Diffusion Models

Denoising diffusion probabilistic models [Ho et al.(2020)Ho, Jain, and Abbeel]
(DDPMs, or diffusion models for short) are a type of generative latent
variable model (LVM) that learns to sample from some data distribution
$p(x_{0})$ ($x_{0}\in\mathbb{R}^{n}$) by defining a stochastic process that
gradually converts the data to noise—the forward process $q(x_{t}|x_{t-1})$,
constructed of conditional Gaussians—and learning to reverse this process via
a learned denoising process $p_{\theta}(x_{t-1}|x_{t})$, where $\theta$ is the
model parameters. Data is generated by iteratively sampling from
$p_{\theta}(x_{t-1}|x_{t})$, beginning with a Gaussian noise sample $x_{T}\sim
p(x_{T})$, for $t=T-1,\ldots,0$ until an image $x_{0}$ is recovered.

Any step of the forward process can be written explicitly as
$x_{t}=\sqrt{\overline{\alpha}_{t}}x_{0}+\sqrt{1-\overline{\alpha}_{t}}\epsilon$
where $\epsilon\sim\mathcal{N}(0,I_{n})$, and $\alpha_{t}:=1-\beta_{t}$ given
the variance of the additive pre-scheduled noise $\beta_{t}$, and
$\overline{\alpha}_{t}:=\prod_{s=1}^{t}\alpha_{s}$. As is typical for LVMs
[Kingma and Welling(2013)], DDPMs can be trained by maximizing the evidence
lower bound (ELBO) with respect to $\theta$ given observed data $x_{0}$ as
$\mathrm{ELBO}(\theta;x_{0})=\mathbb{E}_{x_{1:T}\sim
q(x_{1:T}|x_{0})}\left[\log
p(x_{T})+\sum_{t=1}^{T}\log\frac{p_{\theta}(x_{t-1}|x_{t})}{q(x_{t}|x_{t-1})}\right]$,
where $x_{1:T}:=x_{1},\ldots x_{T}$. The ELBO can be approximately optimized
in a relatively simple form by training a network $\epsilon_{\theta}(x_{t},t)$
to predict the noise $\epsilon$ added to each datapoint $x_{0}$ for various
time steps $t$, with the loss
$L=\mathbb{E}_{x_{0},t,\epsilon}\left[||\epsilon-\epsilon_{\theta}(x_{t},t)||^{2}\right]=\mathbb{E}_{x_{0},t,\epsilon}\left[||\epsilon-\epsilon_{\theta}(\sqrt{\overline{\alpha}_{t}}x_{0}+\sqrt{1-\overline{\alpha}_{t}}\epsilon,t)||^{2}\right]$,
which has proven to be the typically superior DDPM loss in practice [Nichol
and Dhariwal(2021)].

###  1.2 Adding Segmentation Guidance to Diffusion Models

Rather than sampling from the unconditional distribution $p(x_{0})$, our goal
is to condition the generation of some image $x_{0}\in\mathbb{R}^{c\times
h\times w}$ to follow some multi-class anatomical mask
$m\in\\{0,\ldots,C-1\\}^{h\times w}$ (i.e., sample from $p(x_{0}|m)$), where
$c$ is the number of channels of $x_{0}$ and the mask has $C$ possible values
of $0,\ldots,C-1$ associated with each object class. While this does not alter
the noising process $q(x_{t}|x_{t-1})$, the data log-likelihood becomes mask-
conditioned as $\log p(x_{0}|m)$, thereby similarly modifying the reverse
process $p_{\theta}(x_{t-1}|x_{t},m)$ and noise-predicting network
$\epsilon_{\theta}$, changing the ELBO to

$\mathrm{ELBO}(\theta;(x_{0},m))=\mathbb{E}_{x_{1:T}\sim
q(x_{1:T}|x_{0})}\left[\log
p(x_{T})+\sum_{t=1}^{T}\log\frac{p_{\theta}(x_{t-1}|x_{t},m)}{q(x_{t}|x_{t-1})}\right].$
(1)

We propose to implement this with a simple approach of concatenating the mask
$m$ channel-wise to the network input at every denoising step, written as
$\epsilon_{\theta}(x_{t},t|m)$. We modify the network to take an additional
channel allocated for $m$ as input, i.e.,
$\epsilon_{\theta}:\mathbb{R}^{(c+1)\times h\times
w}\rightarrow\mathbb{R}^{c\times h\times w}$, which can be any image-to-image
model (see Sec. 3). This results in a loss of

$L_{m}=\mathbb{E}_{(x_{0},m),t,\epsilon}\left[||\epsilon-\epsilon_{\theta}(x_{t},t|m)||^{2}\right]=\mathbb{E}_{(x_{0},m),t,\epsilon}\left[||\epsilon-\epsilon_{\theta}(\sqrt{\overline{\alpha}_{t}}x_{0}+\sqrt{1-\overline{\alpha}_{t}}\epsilon,t|m)||^{2}\right]$
(2)

for training our model, where each training image $x_{0}$ has some paired
anatomical mask $m$. Sampling can be completed with the typical DDPM algorithm
[Ho et al.(2020)Ho, Jain, and Abbeel] if stochastic results are desired given
some input mask, although in practice we use the DDIM (denoising diffusion
implicit model) algorithm [Song et al.(2021)Song, Meng, and Ermon], which
provides $20\times$ faster sampling than the standard DDPM algorithm with
little reduction in visual quality.

This simple method generates images that are very faithful to input masks
(Fig. LABEL:fig:samples), because the denoising process is conditioned on the
mask at each of its many gradual steps, allowing the network to follow the
masks because they provide helpful spatial information that is directly
correlated with the optimal denoised model output that minimizes the loss.

###  1.3 Mask-Ablated Training and Sampling

Given that our model is mask-guided, the quality of these masks is important;
a generated image may be misleading if the input mask is not fully annotated,
known as the partial label problem [Tian et al.(2023)Tian, Yu, and Fu] in
medical image analysis. This is because the model may assume that un-annotated
objects should not be present in the output image whatsoever (associating the
missing/zero pixel labels as background or some other object), when in reality
we may desire for the model to simply “fill in”/infer the unprovided objects.

To alleviate this problem, we propose a mask-ablated training strategy
(Algorithm 2), which provides examples of masks with various numbers and
combinations of classes removed for the model to learn to generate images from
during training. This can be thought of as a form of self-supervised learning
(Appendix B). We set all $2^{C-1}$ of these possible combinations of classes
being removed from a given mask in training (including the cases of no classes
or all classes removed) to occur with equal probability so that the model can
handle each equally, although any other balancing of these probabilities could
be used.

Dodowhile

[H] number of mask classes $C$, dataset $p(x_{0},m)$. converged $x_{0},m\sim
p(x_{0},m)$
$c=1,\ldots,C-1$ $\delta\sim\mathrm{Bernoulli}(0.5)$
$\delta=1$ $m[m=c]=0$
$\epsilon\sim\mathcal{N}(0,I_{n})$; $t\sim\mathrm{Uniform}(\\{1,\dotsc,T\\})$
$x_{t}=\sqrt{\bar{\alpha}_{t}}x_{0}+\sqrt{1-\bar{\alpha}_{t}}\epsilon$
Update $\theta$ with
$\qquad\nabla_{\theta}\left\|\epsilon-\epsilon_{\theta}(x_{t},t|m)\right\|^{2}$

Figure 2: Segmentation-guided diffusion model training with mask ablation.

##  2 Datasets

Breast MRI: Our first dataset is a 100-patient subset of the Duke Breast
Cancer MRI dataset [Saha et al.(2018)Saha, Harowicz, Grimm, Kim, Ghate, Walsh,
and Mazurowski]. We use all 2D image slices from the fat-saturated gradient
echo T1-weighted pre-contrast sequence, with a train/test split of 70/15
patients, resulting in $\sim$12000/2500 slice images per split. We also keep a
held-out training set of 15 patients for additional experiments. All images
have full segmentation annotations for (1) the breast, (2) blood vessels, and
(3) fibroglandular/dense tissue provided at [Lew(2023)]. Notably, the
fibroglandular tissue and blood vessels have very high variability in shape,
size, and other morphological characteristics, posing a challenge for
generative models to realistically capture. CT Organ: Our second dataset is
CT-Organ [Rister et al.(2020)Rister, Yi, Shivakumar, Nobashi, and Rubin]. We
use a 40-patient subset of neck-to-pelvis and abdominal CT scans, which
include segmentation annotations for liver, bladder, lungs, kidney, and bone.
This results in a train/test split of $\sim$11000/2100 2D slice images, given
a patient-wise split of 24/8, as well as a held-out training set of 8
patients. All generative models are trained on the training sets, and the
auxiliary segmentation network, introduced next, is trained on the held-out
training sets.

##  3 Experiments

#### Training, Architecture and Implementational Details

All images are resized to $256\times 256$ and normalized to $[0,255]$. We use
a UNet architecture [Ronneberger et al.(2015)Ronneberger, Fischer, and Brox]
for the denoising model $\epsilon_{\theta}$, modified to take two channels
(image and mask) as input; see Appendix A.1 for additional training and
architecture details.

###  3.1 Comparison to Existing Image Generation Models

We will now compare our segmentation-guided diffusion model (“Seg-Diff” for
short) to various state-of-the-art image generation models trained on each
dataset: (1) the segmentation-conditional model ControlNet [Zhang et
al.(2023)Zhang, Rao, and Agrawala], and the unconditional models of (2) a
diffusion model identical to ours but trained without mask conditioning, and
(3) StyleGAN2-ADA [Karras et al.(2020)Karras, Aittala, Hellsten, Laine,
Lehtinen, and Aila], a state-of-the-art GAN model. We abbreviate these as
CtrlNet, Unc-Diff and GAN, respectively, and training details for each are in
Appendix A.2. We show exemplar generated images from all models in Fig.
LABEL:fig:samples (more in Appendix D); for conditional models, input masks
were randomly sampled from the test set.

fig:samples

Figure 3: Breast MRI (top) and CT organ (bottom) images generated with
different segmentation-conditional and unconditional models. Samples from
conditional models (blue) are below their input mask, and unconditional model
(red) samples are random. For breast MRI, the breast, blood vessel, and
fibroglandular tissue segmentations are shown as white, red, and blue,
respectively, while for CT, the liver, bladder, lungs, kidneys, and bone are
in maroon, orange, pink, red, and white, respectively. More samples are in
Appendix D.

#### Evaluating Generated Image Quality

We first attempted to use the common Fréchet Inception Distance (FID) [Heusel
et al.(2017)Heusel, Ramsauer, Unterthiner, Nessler, and Hochreiter] as a
metric for generated image quality/realism compared to real data. Since the
original FID utilizes a CNN image encoder pre-trained on natural images, which
may fail to capture the characteristics of medical images, we propose to
instead use an encoder trained on the corresponding medical image dataset. We
observe that samples generated from both the segmentation-guided and
unconditional diffusion models achieve promising results, (e.g.,
$\mathrm{FID}\simeq 0.5$ for breast MRI), yet the CNN-based FID fails to
capture the global feature of anatomical realism that can differ in images
generated by these models (e.g., fibroglandular tissue as shown in Fig.
LABEL:fig:eg_imgs).

tab:qual

Table 1: Quality of generated images. Real test set performance (Dice coeff.)
of a segmentation network for the objects of interest, trained on real data
vs. synthetic data generated by different models.

Training set | Breast MRI | CT Organ
---|---|---
Real | 0.8376 | 0.9075
Synthetic: |  |
Our model | 0.7934 | 0.8981
ControlNet | 0.7570 | 0.0000

We propose to better measure anatomical realism by instead determining how
well an auxiliary segmentation model for the objects of interest can be
trained solely on these synthetic images to be able to generalize to real
data, using the input masks as targets (training details in Appendix A.3). We
compare the performance of the segmentation model trained on (a) the real
held-out training set (Sec. 2) and (b) the set of images generated from all
masks corresponding to these images. We split the real test set in half (by
patient) into a validation set and a test set to use for these models. The
results for this are in Table LABEL:tab:qual; we see that for both datasets,
the segmentation network trained on our model’s synthetic data barely performs
worse ($<0.04-0.01$ Dice) than the network trained on real data, implying that
our generated images are both highly realistic and faithful to the input masks
(especially considering certain objects’ segmentation difficulty), while
outperforming ControlNet.

tab:faithful

Table 2: Faithfulness of generated images to input masks. $m$ denotes
generative model input masks, and $m^{pred}_{gen}$ and $m^{pred}_{real}$
denote the masks predicted for (a) the generated images or (b) the real images
corresponding to the input masks, respectively, by an auxiliary segmentation
model.

| Breast MRI | CT Organ
---|---|---
Model | $\mathrm{Dice}(m^{pred}_{gen},m)$ | $\mathrm{Dice}(m^{pred}_{gen},m^{pred}_{real})$ | $\mathrm{Dice}(m^{pred}_{gen},m)$ | $\mathrm{Dice}(m^{pred}_{gen},m^{pred}_{real})$
| Segmentation-guided
---
diffusion (ours)
0.9027 | 0.8593 | 0.8980 | 0.8797
ControlNet | 0.3636 | 0.3604 | 0.1132 | 0.1126

#### Evaluating Faithfulness of Generated Images to Input Masks

To measure how well our model follows an input mask for image generation, we
use the aforementioned segmentation network trained on the real training set,
to predict segmentations $m^{pred}_{gen}$ on images that were generated from
the masks $m$ in the test set, and measure their overlap with (a) $m$ and (b)
the model’s predicted segmentations $m^{pred}_{real}$ for the input masks’
original corresponding real images. We see our model’s generated images have
high overlap with both ($>0.85$ Dice coeff., Table LABEL:tab:faithful),
indicating that our model closely followed the input masks when generating the
images, also far outperforming ControlNet’s generated images.

#### Understanding ControlNet’s Performance

ControlNet performed much worse for each metric because it failed to generate
images that follow input masks closely (if at all). We suspect that this is
because (1) their instructions for adapting their model to out-of-distribution
(e.g., medical) images was untested in their paper, and (2) their segmentation
guidance training requires joint image+text embeddings, the latter of which is
not immediately applicable to our dataset. We discuss this in more detail in
Appendix C.

fig:samples_ablated

Figure 4: Generating images from masks with classes removed, shown for breast
MRI. Rows 1 and 3: input segmentations, Rows 2 and 4: corresponding output
images.

###  3.2 Sampling from Ablated Masks

We will next explore our model’s ability to be conditioned on a range of
combinations of present or missing anatomical classes in input masks,
following our mask-ablated training algorithm (Sec. 1.3). In Fig.
LABEL:fig:samples_ablated we demonstrate the effect on generated images of
ablating certain classes from an input mask for breast MRI (more examples,
including for CT Organ, in Appendix D). For example, we see that constraining
the blood vessels and fibroglandular tissue (BV+FGT) in breast MRI, yet
keeping the breast free, results in images that have the latter two classes
pre-registered while the breast shape varies, and vice-versa.

Notably, we also see that when the model is provided with an empty mask
(unconditioned), it typically generates more anatomically realistic images
than the unconditionally-trained model does (e.g., compare the bottom right
image in Fig. LABEL:fig:samples_ablated to the rightmost image in Fig.
LABEL:fig:eg_imgs), due to the mask-ablated training acting as self-supervised
learning of more realistic object representations (somewhat analogous to MAE
[He et al.(2022)He, Chen, Xie, Li, Dollár, and Girshick]; see Appendix B).
This is important because the model has both the anatomical realism of a model
trained with full mask guidance and the ability to generate new diverse
samples without needing input masks.

## Conclusion

Our segmentation-guided diffusion model enables more realistic medical image
generation, as well as applications such as the synthesis of paired/pre-
registered data. We also hope that our public codebase and generated dataset
of paired breast MRIs provide helpful resources for the community. For future
work, we are interested in incorporating image-level class guidance [Ho and
Salimans(2021)] either for pathological or domain-related variables [Konz and
Mazurowski(2023)], extending our model to segmentation-guided image
translation, and extending to 3D imaging.

## References

  * [Chen et al.(2020)Chen, Kornblith, Norouzi, and Hinton] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.  A simple framework for contrastive learning of visual representations.  In _International conference on machine learning_ , pages 1597–1607. PMLR, 2020.
  * [Choi et al.(2018)Choi, Choi, Kim, Ha, Kim, and Choo] Yunjey Choi, Minje Choi, Munyoung Kim, Jung-Woo Ha, Sunghun Kim, and Jaegul Choo.  Stargan: Unified generative adversarial networks for multi-domain image-to-image translation.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 8789–8797, 2018.
  * [Dhariwal and Nichol(2021)] Prafulla Dhariwal and Alexander Nichol.  Diffusion models beat gans on image synthesis.  _Advances in neural information processing systems_ , 34:8780–8794, 2021.
  * [Gong et al.(2023)Gong, Johnson, El Fakhri, Li, and Pan] Kuang Gong, Keith Johnson, Georges El Fakhri, Quanzheng Li, and Tinsu Pan.  Pet image denoising based on denoising diffusion probabilistic model.  _European Journal of Nuclear Medicine and Molecular Imaging_ , pages 1–11, 2023.
  * [He et al.(2022)He, Chen, Xie, Li, Dollár, and Girshick] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick.  Masked autoencoders are scalable vision learners.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 16000–16009, 2022.
  * [Heusel et al.(2017)Heusel, Ramsauer, Unterthiner, Nessler, and Hochreiter] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.  Gans trained by a two time-scale update rule converge to a local nash equilibrium.  _Advances in neural information processing systems_ , 30, 2017.
  * [Ho and Salimans(2021)] Jonathan Ho and Tim Salimans.  Classifier-free diffusion guidance.  In _NeurIPS 2021 Workshop on Deep Generative Models and Downstream Applications_ , 2021.
  * [Ho et al.(2020)Ho, Jain, and Abbeel] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in neural information processing systems_ , 33:6840–6851, 2020.
  * [Karras et al.(2020)Karras, Aittala, Hellsten, Laine, Lehtinen, and Aila] Tero Karras, Miika Aittala, Janne Hellsten, Samuli Laine, Jaakko Lehtinen, and Timo Aila.  Training generative adversarial networks with limited data.  _Advances in neural information processing systems_ , 33:12104–12114, 2020.
  * [Kazerouni et al.(2023)Kazerouni, Aghdam, Heidari, Azad, Fayyaz, Hacihaliloglu, and Merhof] Amirhossein Kazerouni, Ehsan Khodapanah Aghdam, Moein Heidari, Reza Azad, Mohsen Fayyaz, Ilker Hacihaliloglu, and Dorit Merhof.  Diffusion models in medical imaging: A comprehensive survey.  _Medical Image Analysis_ , page 102846, 2023.
  * [Khader et al.(2023)Khader, Müller-Franzes, Tayebi Arasteh, Han, Haarburger, Schulze-Hagen, Schad, Engelhardt, Baeßler, Foersch, et al.] Firas Khader, Gustav Müller-Franzes, Soroosh Tayebi Arasteh, Tianyu Han, Christoph Haarburger, Maximilian Schulze-Hagen, Philipp Schad, Sandy Engelhardt, Bettina Baeßler, Sebastian Foersch, et al.  Denoising diffusion probabilistic models for 3d medical image generation.  _Scientific Reports_ , 13(1):7303, 2023.
  * [Kingma and Welling(2013)] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * [Konz and Mazurowski(2023)] Nicholas Konz and Maciej A Mazurowski.  Reverse engineering breast mris: Predicting acquisition parameters directly from images.  In _Medical Imaging with Deep Learning_ , 2023.
  * [Lew(2023)] Christopher Lew.  Segmentation of breast and fibroglandular tissue in MRI: a publicly available dataset and deep learning model, 2023.  URL https://github.com/mazurowski-lab/3D-Breast-FGT-and-Blood-Vessel-Segmentation.
  * [Loshchilov and Hutter(2016)] Ilya Loshchilov and Frank Hutter.  Sgdr: Stochastic gradient descent with warm restarts.  In _International Conference on Learning Representations_ , 2016.
  * [Loshchilov and Hutter(2018)] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  In _International Conference on Learning Representations_ , 2018.
  * [Lyu and Wang(2022)] Qing Lyu and Ge Wang.  Conversion between ct and mri images using diffusion and score-matching models.  _arXiv preprint arXiv:2209.12104_ , 2022.
  * [Nichol and Dhariwal(2021)] Alexander Quinn Nichol and Prafulla Dhariwal.  Improved denoising diffusion probabilistic models.  In _International Conference on Machine Learning_ , pages 8162–8171. PMLR, 2021.
  * [Pinaya et al.(2022a)Pinaya, Graham, Gray, Da Costa, Tudosiu, Wright, Mah, MacKinnon, Teo, Jager, et al.] Walter HL Pinaya, Mark S Graham, Robert Gray, Pedro F Da Costa, Petru-Daniel Tudosiu, Paul Wright, Yee H Mah, Andrew D MacKinnon, James T Teo, Rolf Jager, et al.  Fast unsupervised brain anomaly detection and segmentation with diffusion models.  In _International Conference on Medical Image Computing and Computer-Assisted Intervention_ , pages 705–714. Springer, 2022a.
  * [Pinaya et al.(2022b)Pinaya, Tudosiu, Dafflon, Da Costa, Fernandez, Nachev, Ourselin, and Cardoso] Walter HL Pinaya, Petru-Daniel Tudosiu, Jessica Dafflon, Pedro F Da Costa, Virginia Fernandez, Parashkev Nachev, Sebastien Ourselin, and M Jorge Cardoso.  Brain imaging generation with latent diffusion models.  In _MICCAI Workshop on Deep Generative Models_ , pages 117–126. Springer, 2022b.
  * [Rister et al.(2020)Rister, Yi, Shivakumar, Nobashi, and Rubin] Blaine Rister, Darvin Yi, Kaushik Shivakumar, Tomomi Nobashi, and Daniel L Rubin.  Ct-org, a new dataset for multiple organ segmentation in computed tomography.  _Scientific Data_ , 7(1):381, 2020.
  * [Rombach et al.(2022)Rombach, Blattmann, Lorenz, Esser, and Ommer] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 10684–10695, June 2022.
  * [Ronneberger et al.(2015)Ronneberger, Fischer, and Brox] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In _Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18_ , pages 234–241. Springer, 2015.
  * [Saha et al.(2018)Saha, Harowicz, Grimm, Kim, Ghate, Walsh, and Mazurowski] Ashirbani Saha, Michael R Harowicz, Lars J Grimm, Connie E Kim, Sujata V Ghate, Ruth Walsh, and Maciej A Mazurowski.  A machine learning approach to radiogenomics of breast cancer: a study of 922 subjects and 529 dce-mri features.  _British journal of cancer_ , 119(4):508–516, 2018.
  * [Song et al.(2021)Song, Meng, and Ermon] Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=St1giarCHLP.
  * [Tian et al.(2023)Tian, Yu, and Fu] Yingjie Tian, Xiaotong Yu, and Saiji Fu.  Partial label learning: Taxonomy, analysis and outlook.  _Neural Networks_ , 2023.
  * [von Platen et al.()von Platen, Patil, Lozhkov, Cuenca, Lambert, Rasul, Davaadorj, and Wolf] Patrick von Platen, Suraj Patil, Anton Lozhkov, Pedro Cuenca, Nathan Lambert, Kashif Rasul, Mishig Davaadorj, and Thomas Wolf.  Diffusers: State-of-the-art diffusion models.  URL https://github.com/huggingface/diffusers.
  * [Wang et al.(2018)Wang, Liu, Zhu, Tao, Kautz, and Catanzaro] Ting-Chun Wang, Ming-Yu Liu, Jun-Yan Zhu, Andrew Tao, Jan Kautz, and Bryan Catanzaro.  High-resolution image synthesis and semantic manipulation with conditional gans.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 8798–8807, 2018.
  * [Wolleb et al.(2022a)Wolleb, Bieder, Sandkühler, and Cattin] Julia Wolleb, Florentin Bieder, Robin Sandkühler, and Philippe C Cattin.  Diffusion models for medical anomaly detection.  In _International Conference on Medical image computing and computer-assisted intervention_ , pages 35–45. Springer, 2022a.
  * [Wolleb et al.(2022b)Wolleb, Sandkühler, Bieder, Valmaggia, and Cattin] Julia Wolleb, Robin Sandkühler, Florentin Bieder, Philippe Valmaggia, and Philippe C Cattin.  Diffusion models for implicit image segmentation ensembles.  In _International Conference on Medical Imaging with Deep Learning_ , pages 1336–1348. PMLR, 2022b.
  * [Yang et al.(2019)Yang, Dvornek, Zhang, Chapiro, Lin, and Duncan] Junlin Yang, Nicha C Dvornek, Fan Zhang, Julius Chapiro, MingDe Lin, and James S Duncan.  Unsupervised domain adaptation via disentangled representations: Application to cross-modality liver segmentation.  In _Medical Image Computing and Computer Assisted Intervention–MICCAI 2019: 22nd International Conference, Shenzhen, China, October 13–17, 2019, Proceedings, Part II 22_ , pages 255–263. Springer, 2019\.
  * [Zhang et al.(2023)Zhang, Rao, and Agrawala] Lvmin Zhang, Anyi Rao, and Maneesh Agrawala.  Adding conditional control to text-to-image diffusion models.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 3836–3847, 2023.

##  Appendix A Additional Training and Architectural Details

###  A.1 Our model (segmentation-guided diffusion)

The denoising model (UNet)’s encoder is constructed from six standard ResNet
down-sampling blocks, with the fifth block also having spatial self-attention,
with $(128,128,256,256,512,512)$ output channels, respectively. The decoder is
simply the up-sampling reverse of the encoder. We use $T=1000$ timesteps for
the denoising and noising processes, with a standard forward process variance
schedule that linearly increases from $\beta_{1}=10^{-4}$ to $\beta_{T}=0.02$
[Ho et al.(2020)Ho, Jain, and Abbeel]. For training, we use the AdamW
optimizer [Loshchilov and Hutter(2018)] and a cosine learning rate scheduler
[Loshchilov and Hutter(2016)] with an initial learning rate of $10^{-4}$, with
$500$ linear warm-up steps. We train for 400 epochs with a batch size of 64,
and we perform all training and evaluation on four 48 GB NVIDIA A6000 GPUs. We
use the Diffusers library as a backbone [von Platen et al.()von Platen, Patil,
Lozhkov, Cuenca, Lambert, Rasul, Davaadorj, and Wolf].

###  A.2 Comparison models

#### Unconditional DDPM.

The unconditional DDPM that we compare to has identical training and
architecture to our segmentation-guided model, just without taking
segmentations as input during training.

#### StyleGAN2-ADA.

We train StyleGAN2-ADA for $2$ million image iterations (approximately 200
epochs), a batch size of 64, and all other settings set to the auto-generated
default of https://github.com/NVlabs/stylegan2-ada-pytorch (except with the
network set for one-channel images.

#### ControlNet.

We adapted ControlNet to each of our medical image datasets as was instructed
at their official tutorial for use with datasets that are out-of-distribution
(e.g., medical images) from the Stable Diffusion (SD) model’s very large
natural image training set
(https://github.com/lllyasviel/ControlNet/blob/main/docs/train.md#sd_locked).
This involved first finetuning the autoencoder for 200 epochs, then finetuning
the Stable Diffusion (SD) model for 400 epochs using the respective breast MRI
or CT organ training set images. We then finetuned the ControlNet with the
images and their corresponding masks for segmentation guidance, for 200
epochs, with empty text prompts. The pretrained (pre-finetuning) models are
from the SD v1.5 checkpoints available on Hugging Face at
https://huggingface.co/runwayml/stable-diffusion-v1-5. For all training, we
set the batch size to 128, the initial learning rate to $10^{-4}$, and adopted
cosine annealing learning rate schedulers rate with 500 steps of warm-up.

We note that our model being image-space-based, instead of a latent diffusion
model, also allows for faster and simpler training-from-scratch on a new
dataset; for example, on breast MRI we found that our model took 4-5 days to
train on one 48 GB A6000, while ControlNet took 8-9 days (including the
autoencoder, SD/latent diffusion model, and the ControlNet itself).

###  A.3 Auxiliary segmentation model

We used the UNet segmentation model from MONAI
(https://docs.monai.io/en/stable/networks.html) with single channel input and
the number of output channels set to the number of target object classes + 1
(background). The sequence of intermediate UNet channels was set to (16, 32,
64, 128, 256). For training, we set the initial learning rate to $10^{-3}$ and
adopted a cosine annealing scheduler. The input image resolution was
$256\times 256$. We trained each model for 100 epochs with a batch size of 8
and selected the models with the lowest validation loss.

##  Appendix B Mask-Ablated Training as Self-Supervised Learning

In this section, we will explain how mask-ablated training (Section 1.3) is a
form of self-supervised learning of representations of objects of interest.
During training, the model will see a given image once per epoch, but each
time, the mask paired with the image that the model learns to be conditioned
on will have a different combination of object classes missing and present. In
other words, the model sees contrasting views of the same image+varying mask
input each epoch, analogous to contrastive self-supervised learning [Chen et
al.(2020)Chen, Kornblith, Norouzi, and Hinton]. As such, the prototypical
internal representation that the denoising model uses for a given type of
object will be similar for either that object being provided or missing in an
input mask for a given image. We already found in Sec. 1.2 that providing a
real mask to guide the generation of a given type of object helps the
generated object be more realistic. Therefore, because of this representation
similarity, even if a given object class is missing from an input mask, if the
model decides that the object should be present given the other classes in the
mask (or simply randomly if the mask is empty), the representation for that
object will be decoded into a more realistic version that what would have been
learning by a model trained completely unconditionally.

This idea also has parallels with masked-autoencoding self-supervised learning
of images [He et al.(2022)He, Chen, Xie, Li, Dollár, and Girshick] as we mask
out certain object classes randomly during training, so further investigation
and exploration in this concept and training strategy is certainly warranted.

##  Appendix C Challenges of Adapting ControlNet for Medical Images

As shown in our main results, ControlNet [Zhang et al.(2023)Zhang, Rao, and
Agrawala] failed to achieve solid segmentation-guided performance when adapted
to our datasets, which we believe to be due to several reasons.

First, as mentioned, we adapted ControlNet to each of our medical image
datasets as was instructed at their official tutorial for use with datasets
that are out-of-distribution (e.g., medical images) from the Stable Diffusion
(SD) model’s very large natural image training set
(https://github.com/lllyasviel/ControlNet/blob/main/docs/train.md#sd_locked).
This involved (1) fine-tuning the latent space VAE, (2) fine-tuning the SD
model and (3) training the ControlNet module to learn segmentation guidance,
on the new given dataset. However, we note that this procedure of finetuning
the SD model was completely untested in their paper, and so was unproven.
ControlNet works by first copying the SD model weights to create the
ControlNet module itself, and then fine-tuning that module for spatial
conditioning. However, this was not evaluated (or proposed) for a custom
dataset finetuned SD model, so the performance being unreliable is
understandable.

Second, the ControlNet training algorithm requires using non-empty text
prompts for half of the images seen in training, while in our case, no text
exists for our data, such that we had to train with all text prompts empty; as
the diffusion model acts completely on the joint image-text latent space, the
effect of this difference in training on the latent space is unclear and
possibly violates the assumptions needed for ControlNet to work.

We also note that we performed each of these trainings for many epochs
(Appendix C), ruling out the possible issue of the “sudden convergence
phenomenon” of proper segmentation conditioning that was observed in the
ControlNet paper. Finally, we could foreseeable train some latent diffusion
model from scratch on each of our datasets (without text conditioning), but
this is out-of-scope from our paper and would form a completely new model, no
longer simply adapting ControlNet as-is for our datasets.

##  Appendix D Additional Sample Images

We show additional samples for both datasets for all models in Fig.
LABEL:fig:samples_many (to supplement Fig. LABEL:fig:samples). Similarly, we
show more samples with various classes removed from input masks to illustrate
the effect of mask-ablated training (Sec. 1.3) in Figures
LABEL:fig:samples_many_abla_breast and LABEL:fig:samples_many_abla_ct for
breast MRI and CT organ, respectively, to supplement Fig.
LABEL:fig:samples_ablated.

fig:samples_many

Figure 5: Additional samples from all segmentation-conditional and
unconditional models; breast MRI on the left, CT organ on the right. Please
see the caption of Fig. LABEL:fig:samples for more details.

fig:samples_many_abla_breast

Figure 6: Additional samples from our model with various classes removed from
given input segmentations for breast MRI.

fig:samples_many_abla_ct

Figure 7: Additional samples from our model with various classes removed from
given input segmentations for CT organ.