Reviewing Labels: Label Graph Network with Top-k Prediction Set for Relation Extraction

Bo Li, Wei Ye, Jinglei Zhang, Shikun Zhang

Abstract: The typical way for relation extraction is fine-tuning large pre-trained language models on task-specific datasets, then selecting the label with the highest probability of the output distribution as the final prediction. However, the usage of the \textit{Top-k} prediction set for a given sample is commonly overlooked. In this paper, we first reveal that the \textit{Top-k} prediction set of a given sample contains useful information for predicting the correct label. To effectively utilizes the \textit{Top-k} prediction set, we propose \textbf{L}abel \textbf{G}raph Network with \textit{\textbf{Top-k}} Prediction Set, termed as \textbf{KLG}. Specifically, for a given sample, we build a label graph to review candidate labels in the \textit{Top-k} prediction set and learn the connections between them. We also design a dynamic $k$-selection mechanism to learn more powerful and discriminative relation representation. Our experiments show that \textbf{KLG} achieves the best performances on three relation extraction datasets. Moreover, we observe that \textbf{KLG} is more effective in dealing with long-tailed classes.