Auto-Encoding Knowledge Graph for Unsupervised Medical Report Generation

1 Introduction . Medical images , such as radiology and pathology images , and their corresponding reports are widely used for clinical diagnosis and treatment [ 8 , 11 ] . A medical report is usually a paragraph of multiple sentences which describes both the normal and abnormal findings in the medical image . In clinical practice , writing a report can be time-consuming and tedious for experienced radiologists , and error-prone for inexperienced radiologists [ 4 ] . Therefore , given the large volume of medical images , automatically generating reports can improve current clinical practice in diagnostic radiology and assist radiologists in clinical decision-making [ 15 , 25 ] . Specifically , it can relieve radiologists from such heavy workload and alert radiologists of the abnormalities to avoid misdiagnosis and missed diagnosis . Therefore , automatic medical report generation attracts remarkable attention in both artificial intelligence and clinical medicine . Recently , inspired by the great success of neural machine translation [ 2 , 38 , 49 , 50 , 48 ] , image captioning [ 44 , 39 , 46 , 28 , 30 , 29 ] and medical imaging analysis [ 47 , 51 , 52 ] , the data-driven deep ∗Corresponding authors . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . neural models , particularly those based on the encoder-decoder frameworks [ 15 , 16 , 25 , 22 , 41 , 45 , 53 , 54 , 6 ] , have achieved great success in advancing the state-of-the-art of medical report generation . However , these models are trained in a supervised learning manner and heavily rely on labeled paired image-report datasets [ 9 , 17 ] , which are not easy to acquire in the real world . Specifically , the medical-related data can only be manually labeled by professional radiologists , and also involves privacy issues . Therefore , the medical report generation datasets are particularly labor-intensive and expensive to obtain . As a result , the scales of existing widely-used datasets for medical report generation models [ 16 , 25 , 22 , 23 ] , i.e. , MIMIC-CXR ( 0.22M samples ) [ 17 ] and IU X-ray ( 4K samples ) [ 9 ] , are relatively small compared to image recognition datasets , e.g. , ImageNet ( 14M samples ) [ 10 ] , and image captioning datasets , e.g. , Conceptual Captions ( 3.3M samples ) [ 36 ] . In addition , the MIMIC-CXR and IU X-Ray datasets only include Chest X-Ray images , for other types of medical images ( MRI , Dermoscopy , Retinal , etc . ) of other body parts ( brain , skin , eye , etc . ) , the image-report pairs could be much less or even unavailable . Therefore , to relax the reliance on the paired data sets , making use of all available data , like independent image or report sets , is becoming increasingly important . In this paper , we propose an unsupervised model Knowledge Graph Auto-Encoder ( KGAE ) , which utilizes independent sets of images and reports in training ( the image and report set are separate and have no overlap ) . KGAE consists of a pre-constructed knowledge graph , a knowledge-driven encoder and a knowledge-driven decoder . As shown in Figure 1 , the knowledge graph works as the shared latent space of images and reports . The knowledge-driven encoder can take either the image I or the report R as queries and project them to corresponding coordinates GI and GR in the latent space . In this manner , since GI and GR share the same latent space , we can use positions in latent space to measure the relationship between images and reports which narrows the gap between visual and textual domains . In brief , to bridge the gap between vision and language domains without training on the pairs of images and reports , we adopt the knowledge graph to create a latent space and propose a knowledge-driven encoder , which includes a common mapping function to project images and reports to the same latent space . As a result , our encoder can extract the image and report knowledge representations , i.e. , the knowledge related to the image and report , they ( image , report knowledge ) share the common latent space , which allows our model to bridge the gap between vision and language domains without the training on the pairs of image and report . Next , we introduce the knowledge-driven decoder to exploit GI and GR to generate the report . In the training stage , we estimate the parameters of the decoder by reconstructing the input report R based on GR , i.e. , R→ GR → R auto-encoding pipeline ; In the prediction stage , we directly input GI into the trained decoder to generate the report . In this way , our approach can produce desirable reports without any labeled image-report pairs . Overall , the contributions of this paper are as follows : • In this paper , we make the first attempt to conduct unsupervised medical report generation where the image-report pairs are not available . To this end , we propose the Knowledge Graph Auto-Encoder ( KGAE ) . By leveraging a pre-constructed knowledge graph , we introduce the knowledge-driven encoder and decoder which are trained with independent sets of images and reports . According to the experimental results , the unsupervised KGAE can even outperform several supervised approaches . • In addition to the unsupervised mode , KGAE can also be applied in a semi-supervised or supervised manner . Under the semi-supervised setting , by using only 60 % of paired dataset , KGAE is able to achieve competitive results with current state-of-art models ; Under the supervised setting , by training on fully paired datasets as in existing works , KGAE can set new state-of-the-art performances on the IU X-ray and MIMIC-CXR , respectively . • The analysis , both quantitative and qualitative , as well as a human evaluation conducted by professional radiologists , further proves the effectiveness of our approach . 2 Related Works . Medical report generation aims to generate a relatively long paragraph to describe a given medical image . It is similar to the image captioning task [ 5 ] , which aims to generate a sentence to describe a given image . In image captioning , the encoder-decoder framework [ 39 ] , where the encoder [ 20 , 12 ] computes visual representations for the image and the decoder [ 13 , 38 ] generates a target sentence based on the visual representations , has achieved great success [ 1 , 31 , 35 , 44 ] . However , instead of only generating one single sentence , medical report generation aims to generate a long paragraph including multiple structured sentences that describe both the normal and abnormal parts [ 25 , 15 ] . To this end , given the success of encoder-decoder framework on image captioning , most existing medical report generation models attempt to exploit the hierarchical LSTM ( HLSTM ) [ 15 , 16 , 19 ] or Transformer [ 6 ] to generate an accurate , long and coherent report . However , existing models require the paired image-report datasets , which are time-consuming and expensive to collect . In this paper , we propose the unsupervised model Knowledge Graph Auto-Encoder ( KGAE ) which doesn ’ t need paired images and reports . Note that although the knowledge graph has been integrated in existing medical report generation models [ 54 , 22 ] , these approaches are supervised and require paired images and reports . Thus , their objectives and motivations of using knowledge graph are different from our work . In detail , existing knowledge-graph based medical report generation methods aim to adopt the knowledge graph to boost the performance of supervised models . However , in our work , we aim to generate a medical report without using any coupled image-report training pairs , i.e. , unsupervised medical report generation . A key challenge of unsupervised medical report generation is to bridge the gap between vision and language domains . To this end , we adopt the knowledge graph to create a latent space and propose a knowledge-driven encoder to project image and report to the same latent space . 3 Approach . We first formulate the conventional supervised medical report generation problems ; Then , we describe the proposed Knowledge Graph Auto-Encoder for unsupervised medical report generation in detail . 3.1 Conventional Supervised Medical Report Generation Models . Given a medical image I , the goal is to generate a descriptive report R. Most models [ 15 , 16 , 6 ] normally include an image encoder and a report decoder , which can be formulated as : Image Encoder : I → I ′ ; Report Decoder : I ′ → R , ( 1 ) where I ′ ∈ RNI×d denotes the image embeddings extracted by the image encoder , e.g. , ResNet-50 [ 12 ] . Then , I ′ is used to guide the generation of the target report R in the report decoder , e.g. , LSTM [ 13 ] and Transformer [ 38 ] . During training , given the ground truth report for the input image , we can train the encoder-decoder model by minimizing a supervised training loss , e.g. , cross-entropy loss . However , paired image-report ( I-R ) datasets are particularly labor-intensive and expensive to obtain . Therefore , in this paper , we aim to relax the reliance on the paired datasets and make use of the independent sets of images and reports instead . 3.2 Knowledge Graph Auto-Encoder . As shown in Figure 1 , the proposed KGAE includes a knowledge graph , a knowledge-driven encoder and a knowledge-driven decoder , which will be described in detail in the following sections . Knowledge Graph The motivation of using the knowledge graph is that writing medical reports usually requires particular domain knowledge [ 11 , 22 ] . Therefore , we can utilize the knowledge graph , which models the domain-specific knowledge structure , to serve as a bridge to correlate the knowledge representations of the visual and textual domains . In particular , we construct an off-the-shelf global medical knowledge graph G = ( V , E ) covering the common abnormalities and normalities , where V = { vi } i=1 : NKG ∈ RNKG×d is a set of nodes and E = { ei , j } i , j=1 : NKG is a set of edges . In detail , based on the report corpus , i.e. , MIMIC-CXR [ 17 ] , we consider the NKG frequent clinical abnormalities ( e.g. , `` enlarged heart size '' , `` pleural effusion '' , and `` bibasilar consolidation '' ) and normalities ( e.g. , `` heart size is normal '' and `` lungs are clear '' ) as nodes . The edge weights are calculated by the normalized co-occurrence of different nodes computed from report corpus . Please note that this knowledge graph is built purely from the training set of MIMIC-CXR , thus there is no label leakage . After that , the knowledge graph is embedded by a graph embedding module , i.e. , graph convolution network [ 32 , 24 ] , which is implemented as follows : v′i = vi + ReLU ( NKG∑ j=1 ei , jWvvj ) , ( 2 ) where ReLU ( · ) represents the ReLU activation function and Wv ∈ Rd×d is the learnable matrix . As a result , we can acquire a set of node embeddings V ′ = { v′1 , v′2 , . . . , v′NKG } ∈ R NKG×d . It is worth noting that more complex graph structures could be constructed by using more large-scale medical textbooks . Therefore , our model is not limited to the currently constructed graph and could provide a good basis for the research of unsupervised medical report generation for other domains . Knowledge-driven Encoder The knowledge-driven encoder ( KE ) is designed to utilize the knowledge graph G to extract knowledge representations ( GI ∈ RNI×d and GR ∈ RNR×d ) of both image I and report R to bridge the vision and the language domains , which can be formulated as : GI = KEI ( I , G ) ; GR = KER ( R , G ) . ( 3 ) In implementations , as shown in Figure 1 , we first adopt the ResNet-50 [ 12 ] and the Transformer [ 38 ] as the image embedding module and report embedding module to embed the image I and report R , acquiring the image embeddings I ′ ∈ RNI×d and report embeddings R′ ∈ RNR×d , respectively . Next , to extract the knowledge representations from knowledge graph G , we adopt the attention mechanism [ 38 ] to implement the KE . The motivation stems from that the attention mechanism can compute the association weights between different features and allows probabilistic many-to-many relations instead of monotonic relations , as in [ 38 , 44 , 29 ] . As a result , we take I ′ and R′ as the queries , and take the knowledge graph G ( i.e. , V ′ = { v′1 , v′2 , . . . , v′NKG } ∈ R NKG×d ) as the lookup matrix : GI = KEI ( I , G ) =F ( AttentionI ( I ′ , V ′ ) ) ; GR = KER ( R , G ) = F ( AttentionR ( R′ , V ′ ) ) Attention ( x , y ) = softmax ( xWq ( yWk ) > √ d ) yWv . ( 4 ) where x ∈ RNx×dx , y ∈ RNy×dy ; Wq ∈ Rdx×dx , Wk ∈ Rdy×dy and Wv ∈ Rdy×dy ( dx = dy ) are learnable parameters . F is implemented as two fully-connected ( FC ) layers with a ReLU in between , i.e. , FC-ReLU-FC . The resulted GI ∈ RNI×d and GR ∈ RNR×d turn out to be a set of attended ( i.e. , extracted ) knowledge related to the image and report . The knowledge representations GI and GR share the common latent space which bridges the vision and the language domains . Note that the F used in KEI and KER shares the same parameters and is thus introduced to further boost the bridging capabilities of different modalities . It means that we adopt two KEs , i.e. , KEI and KER , and each KE includes an attention model as well as a common mapping function F . The only shared weights of KEI and KER are the parameters of F , and the parameters of attention models are independent . Knowledge-driven Decoder The decoder is designed to generate the reports based on the graph representations GI or GR . For clarity , we use Gk to represent the GR and GI during the training and testing stages , respectively . In implementations , since medical report generation requires generating a long paragraph , we choose the ( three-layer ) Transformer [ 38 ] as the basic module of our decoder and incorporate the proposed Knowledge-driven Attention ( KA ) to effectively model the long sequences . At each decoding step t , the decoder takes the embedding of current input word xt = wt + et ∈ R2d as input , where wt and et denote the word embedding and fixed position embedding , respectively , and generate each word rt in report R = { r1 , r2 , . . . , rT } , which can be defined as follows : ht = Attention ( xt , x1 : t ) h′t = KA ( ht , Gk , B ) , where Gk = { GR , Training GI , Testing ; rt ∼ pt = softmax ( FFN ( h ′ t ) Wp ) , ( 5 ) where FFN stands for Feed-Forward Network in the original Transformer [ 38 ] ; Wp ∈ R2d×|D| is the learnable parameter ( |D| denotes the vocabulary size ) ; B represents the knowledge bank in the KA . In particular , the knowledge-driven attention ( KA ) is inspired by the success of many works that attempt to augment a working memory into the network , e.g. , memory network . Memory network preserves a dynamic knowledge base for subsequent inference [ 37 , 21 , 42 , 43 , 40 , 46 , 33 ] . For example , [ 37 , 21 ] and [ 42 , 43 ] proposed to preserve the context of a document and visual knowledge to solve language modeling and visual question answering , respectively . In this paper , since the knowledge representations Gk are extracted from the off-the-shelf global knowledge graph G , we incorporate a knowledge memory mechanism to distill and preserve the fine-grained medical knowledge related to report generation task , which encourages our decoder to generate more accurate and desirable reports . Specifically , we introduce the knowledge bank B = { b1 , b2 , . . . , bNB } ∈ RNB×d , where NB stands for the total number of the knowledge corresponding to report generation . As a result , given the graph representations Gk ∈ RNk×d , the knowledge memory mechanism is defined as follows : Bk = softmax ( GkB > ) B . ( 6 ) Through above operation , each feature/vector in the Gk ∈ RNk×d can distill the fine-grained knowledge preserved in B to acquire Bk ∈ RNk×d for accurate report generation . Based on the above mechanism , the knowledge-driven attention in Eq . ( 5 ) is defined as : h′t = KA ( ht , Gk , B ) = Attention ( ht , [ Gk ; Bk ] ) = Attention ( ht , [ Gk ; softmax ( GkB > ) B ] ) . ( 7 ) where [ · ; · ] denotes the concatenation operation . It is worth noting that these operations are all differentiable , thus the bank B can be learned in an end-to-end fashion . In our subsequent analysis , we will show that the introduced knowledge memory mechanism indeed distills and preserves the desired medical knowledge , and thus boost the generation of reports .

Summary: The authors propose an unsupervised method for the task of medical report generation. The model can be trained using unpaired images and reports. To bridge the image and language representations, the model uses a pre-constructed knowledge graph to create a latent space and their encoder could map images and reports to the same latent space.   Specifically, 1) in the knowledge graph, the nodes are N_{KG} (= 200) frequent clinical abnormalities and normalities. The edge weights are calculated by the normalized co-occurrence of different nodes computed from the MIMIC-CXR training set. The KG embedding is trained from a graph neural network. 2) The encoder is fine-tuned on the CheXpert dataset from ResNet-50 pre-trained on ImageNet. The encoder uses the attention mechanism where queries are image or report embeddings and keys/values are KG entity embeddings. The encoder is trained with multi-label classification with binary cross entropy loss for common radiographic observations classification. 3) The decoder is a transformer-based model augmented with a memory network. The decoder is trained from scratch by re-constructing a given report. During inference time, the encoder maps an image into the latent space, and the decoder generates a report from its latent representation.   The authors evaluated their method on two datasets: IU X-ray and MIMIC-CXR. The model achieved decent performance at the unsupervised setting. Furthermore, fine-tuned using 60% paired data (60% semi-supervised), the model reaches the previous SOTA performance. Fine-tuned using 100% paired data (supervised), the model sets the new SOTA with significant gains.