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: For automatic medical report generation, this work 1) builds a knowledge graph from the report corpus of MIMIC-CXR, 2) designs a knowledge-driven encoder which uses a ResNet-50 and a BiLSTM trained in a multi-label classification task, 3) designs a knowledge-driven decoder based on Transformer and memory network. The proposed model doesn’t require any image-report training pairs, in other words, it uses independent sets of images and reports in training, but it outperforms several supervised models on MIMIC-CXR and IU X-ray datasets. This method can be further extended to semi-supervised and supervised settings, and outperforms current SOTA.