Retrieval-Augmented Parsing for Complex Graphs by Exploiting Structure and Uncertainty (2023)

Abstract: Retrieval augmentation enhances generative language models by retrieving informative exemplars relevant for output prediction. However, in realistic graph parsing problems where the output space is large and complex, classic retrieval methods based on input-sentence similarity can fail to identify the most informative exemplars that target graph elements the model is most struggling about, leading to suboptimal retrieval and compromised prediction under limited retrieval budget. In this work, we improve retrieval-augmented parsing for complex graph problems by exploiting two unique sources of information (1) structural similarity and (2) model uncertainty. We propose \textit{\textbf{S}tructure-aware and \textbf{U}ncertainty-\textbf{G}uided \textbf{A}daptive \textbf{R}etrieval} \textbf{(SUGAR)} that first quantify the model uncertainty in graph prediction and identify its most uncertain subgraphs, and then retrieve exemplars based on their structural similarity with the identified uncertain subgraphs. On a suite of real-world parsing benchmarks with non-trivial graph structure (SMCalflow and E-commerce), SUGAR exhibits a strong advantage over its classic counterparts that do not leverage structure or model uncertainty.

## introduction large language models (llms) have demonstrated remarkable capabilities as effective few-shot learners @xcite . recently, a new learning paradigm called in-context learning has been developed. under this paradigm, the model is given a prompt including test inputs and a few related exemplars, and can generate outputs directly without updating its parameters. a typical approach to obtaining these exemplars is to retrieve training examples similar to the test input @xcite . however, for realistic parsing tasks with large output graphs and non-trivial structure (e.g., dialogue-oriented semantic parsing), the † co-senior authors. ‡ work done at google. input similarity alone may not be effective in identifying the most informative exemplars for aiding graph prediction. as an example mentioned in @xcite , for the test input "schedule a meeting with my manager", it is more similar to example "schedule a meeting with alex" than "who is my manager", while the latter one contains an important action for searching an org chart which is also required by the test input. in this paper, we explore effective approaches to improve the generalization performance of retrieval-augmented llms for parsing complex graphs. specifically, we consider exploiting two sources of information uniquely available to this problem: (1) structural similarity between output subgraphs, and (2) model uncertainty in graph component prediction. the motivations behind our approach are two empirical investigations on llm graph parsing (presented in section 3): (a) inadequacy of sequence-similarity retrieval (section 3.1). when output graphs exhibit non-trivial structure, the exemplars retrieved based on sequence similarity is less effective than those based on graph similarity, even when the similarity is computed with respect to the gold output graphs foot_0 . (b) llm uncertainty correlates with performance (section 3.2). we conduct a exploratory study of the quality of llm uncertainty as an indicator of its generalization performance in graph prediction, and identified a monotonic between model uncertainty v.s. accuracy for node and edge prediction (figure 2 ). this implies that model uncertainty can serve as an effective signal for identifying the subgraphs that the model is struggling to predict, thereby helping the retrieval algorithm to efficiency identify the most effective examplars for aiding model prediction, especially when the output graph is large. based on the above observations, we propose (3) structure-aware retrieval that retrieves structurally similar exemplars targeting at ŷi k (section 4.1.3). note that at iteration 0, the exemplars are retrieved by input similarity since there is no model prediction yet. structure-aware and uncertainty-guided adaptive retrieval (sugar), a retrieval-augmented llm inference framework for complex graph parsing that incorporates both structural similarity and model uncertainty into retrieval procedure (section 4). operating in an iterative manner, sugar first identifies uncertainty regions in the model's graph prediction from the previous iteration, and adaptively retrieves exemplars based on their graph similarity with the identified uncertainty subgraphs (figure 1 , section 4). in this way, sugar is able to better target model weaknesses in structural prediction by retrieving the most informative exemplars, which appears to be especially valuable in the setting of large and complex output graphs given limited retrieval budget. on a suite of real-world complex graph parsing benchmarks (i.e., smcalflow and ecommerce), sugar exhibits distinct strength over its classic counterparts that do not leverage uncertainty or structure, bringing clear performance improvement over the base model across iterations even when other retrieval-augmentation methods become counterproductive (section 5). further indepth analysis revealed the interesting role of exemplar quality on model uncertainty, the effectiveness of uncertainty as an early-stopping signal for retrieval iterations, and verifies the effectiveness of structural retrieval in improving model confidence (section 6) 2 . ## related work retrieval-augmented parsing. sequence-tosequence (seq2seq) models have achieved stateof-the-art performance on many natural language processing tasks including complex parsing, e.g., dialogue-oriented semantic parsing and meaning representation parsing @xcite @xcite . the general approach is to treat the output structure as a sequence and fine-tune a seq2seq model to learn the mapping between input sentences and output structures. to reduce the reliance on largescaled annotated data, several work augment the input with retrieved exemplars from the training data, with different strategy to select exemplars. for unsupervised retrieval, @xcite and @xcite retrieve exemplars with similar input encodings from a pre-trained neural encoder. @xcite retrieves exemplars for which the input and output (from the preliminary prediction) has high tf-idf similarity with the sample input. the above work mainly focused on fine-tuning settings. for supervised retrieval, @xcite suggest to use language models themselves to label examples that can serve as good prompts, and train a prompt retriever from this signal. in this work, we focus on unsupervised retrievers that do not rely on additional training data beyond the candidate pool they retrieve from. iterative retrieval. while llms can generate coherent outputs via single-time retrieval-based augmentation, they often fall short in addressing more complex tasks. to address this, there have been various attempts to retrieve exemplars more than one time. most of the work focused on ad- ## motivations in this section, we present two ablation studies which serve as motivations for our methods. ## sugar: structure-aware and uncertainty-guided adaptive retrieval this section describes details of sugar for parsing complex structures. typically, the output structure is a semantic graph that is rooted, directed, acyclic and labeled @xcite . problem formulation. we aim to solve a graph parsing problem that maps from a natural language utterance x to a target graph representation @xmath0, e⟩, where n is the set of nodes and @xmath1. for seq2seq models, g is represented as a linearized graph string y. following @xcite , we adopt pen-man annotation @xcite to linearize all graph structures in this work, which is a serialization format for the directed, rooted graphs used to encode semantic dependencies (details for graph linearization can be found in appendix c). figure 1 shows an overview of sugar and algorithm 1 summarizes the detailed process. in section 4, we illustrate the retrieval process based algorithm 1 retrieval-augmented inference with sugar precompute: test input x output: final graph prediction ŷi initialize: @xmath2. ŷ0 = llm(x, e0) ▷ initial prediction. # iterative retrieval with early stopping. on preliminary predictions ŷi at step i (the retrieval process in figure 1 ), which includes three steps: ## analysis: the role of uncertainty 6.1 uncertainty quality matters we notice that sugar does not work very well on the random base retriever, and this might be due to the relatively poor uncertainty quality of the random base retriever in comparison to the other two base retrievers. to validate our hypothesis, we further evaluate the uncertainty quality of the three base retrievers. we can find that the calibration for random base retriever is consistently worse than bm 25 and casper use , indicating that the relativelyconfident neighbors in uncertain subgraphs for random retriever might fail to capture accurate contexts for iterative retrieval. this also indicates that improving llms' calibration can be a fruitful venue for improving retrieval-based data augmentation. ## conclusions in this work, we present structural-aware and uncertainty-guided adaptive retrieval (sugar), a new retrieval-augmented parsing framework us- our future work includes considering more advantage graph similarity metric beyond smatch (e.g., incorporates additional similarity metric between the space of node and edge properties), and also larger scale and more fine-grained evaluation on graph parsing benchmarks with distinct properties (e.g., long-tail generalization, graph of specific families) to further elucidate in what setting is the graph similarity most effective. furthermore, it is also of interest to study the generalization of this approach to a broader class of modern program synthesis problems, e.g., code generation.