UNIREX : A Unified Learning Framework for Language Model Rationale Extraction (2022)

Abstract: An extractive rationale explains a language model’s (LM’s) prediction on a given task instance by highlighting the text inputs that most influenced the prediction. Ideally, rationale extraction should be faithful (reflective of LM’s actual behavior) and plausible (convincing to humans), without compromising the LM’s (i.e., task model’s) task performance. Although attribution algorithms and select-predict pipelines are commonly used in rationale extraction, they both rely on certain heuristics that hinder them from satisfying all three desiderata. In light of this, we propose UNIREX, a flexible learning framework which generalizes rationale extractor optimization as follows: (1) specify architecture for a learned rationale extractor; (2) select explainability objectives (i.e., faithfulness and plausibility criteria); and (3) jointly the train task model and rationale extractor on the task using selected objectives. UNIREX enables replacing prior works’ heuristic design choices with a generic learned rationale extractor in (1) and optimizing it for all three desiderata in (2)-(3). To facilitate comparison between methods w.r.t. multiple desiderata, we introduce the Normalized Relative Gain (NRG) metric. Across five English text classification datasets, our best UNIREX configuration outperforms the strongest baselines by an average of 32.9% NRG. Plus, we find that UNIREX-trained rationale extractors’ faithfulness can even generalize to unseen datasets and tasks.

## introduction large neural language models (lms) have yielded state-of-the-art performance on various natural language processing (nlp) tasks @xcite . however, lms' complex reasoning processes are notoriously opaque @xcite , posing concerns about the societal implications of using lms for high-stakes decision-making * work done while ac was a research intern at meta ai. figure 1 : desiderata of rationale extraction. unlike prior works, unirex enables optimizing for all three desiderata. @xcite . thus, explaining lms' behavior is crucial for promoting trust, ethics, and safety in nlp systems @xcite . given a lm's (i.e., task model's) predicted label on a text classification instance, an extractive rationale is a type of explanation that highlights the tokens that most influenced the model to predict that label @xcite . ideally, rationale extraction should be faithful @xcite and plausible @xcite , without hurting the lm's task performance (deyoung et al., 2019) (fig. 1 ). configuring the rationale extractor and its training can greatly impact these desiderata, yet prior works have commonly adopted two suboptimal heuristics. first, many works rely in some way on attribution algorithms (aas), which extract rationales via handcrafted functions @xcite @xcite . aas cannot be directly trained and tend to be computeintensive @xcite . also, aas can be a bottleneck for plausibility, as producing human-like rationales is a complex objec-tive requiring high capacity rationale extractors @xcite . second, many works use a specialized select-predict pipeline (spp), where a predictor module is trained to solve the task using only tokens chosen by a selector module @xcite @xcite . instead of faithfulness optimization, spps heuristically aim for "faithfulness by construction" by treating the selected tokens as a rationale for the predictor's output (which depends only on those tokens). still, spps typically have worse task performance than vanilla lms since spps hide the full input from the predictor. to tackle this challenge, we propose the unified learning framework for rationale extraction (unirex), which generalizes rationale extractor optimization as follows: (1) specify architecture for a learned rationale extractor; (2) select explainability objectives (i.e., faithfulness and plausibility criteria); and (3) jointly train the task model and rationale extractor on the task using selected objectives (sec. 3). unirex enables replacing prior works' heuristic design choices in (1) with a generic learned rationale extractor and optimizing it for all three desiderata in (2)-(3). unirex provides significant flexibility in performing (1)-(3). for (1), any model architecture is applicable, but we study transformer lm based rationale extractors in this work @xcite . we focus on two architectures: (a) dual lm, where task model and rationale extractor are separate and (b) shared lm, where task model and rationale extractor share parameters. for (2), any faithfulness and plausibility criteria can be used. following @xcite , we focus on comprehensiveness and sufficiency as faithfulness criteria, while using similarity to gold rationales as plausibility criteria. for (3), trade-offs between the three desiderata can be easily managed during rationale extractor optimization by setting arbitrary loss weights for the faithfulness and plausibility objectives. plus, though computing the faithfulness criteria involves discrete (nondifferentiable) token selection, using shared lm can approximate end-to-end training and enable both task model and rationale extractor to be optimized w.r.t. all three desiderata (sec. 3.3). to evaluate all three desiderata in aggregate, we introduce the normalized relative gain (nrg) metric. across five english text classification datasets -sst, movies, cos-e, multirc, and e-snli @xcite our best unirex configuration outperforms the strongest baselines by an average of 32.9% nrg (sec. 4.2), showing that unirex can optimize rationale extractors for all three desiderata. in addition, we verify our unirex design choices via extensive ablation studies (sec. 4.3). furthermore, unirex-trained extractors have high generalization power, yielding high plausiblity with minimal gold rationale supervision (sec. 4.4) and high faithfulness on unseen datasets and tasks (sec. 4.5). finally, our user study shows that humans judge unirex rationales as more plausible than rationales extracted using other methods (sec. 4.6). ## problem formulation rationale extraction let f @xmath0. a.1), where f enc is the text encoder and f task is the task output head. typically, f task has a bert-style architecture @xcite , in which f enc is a transformer @xcite while f task is a linear layer with softmax classifier. let x @xmath1.g., a sentence) for task instance i, and f task (x i ) ∈ r m be the logit vector for the output of the task model. let ŷi = arg max j f task (x i ) j be the class predicted by f task . given f task , x i , and ŷi , the goal of rationale extraction is to output vector s @xmath2, such that each s t @xmath3. let f ext be a rationale extractor, such that s @xmath4, x i , ŷi ). f ext can be a learned or heuristic function. in practice, the final rationale is often obtained by binarizing s i as r @xmath5, 1} n , via the top-k% strategy: r t @xmath6; otherwise, r t @xmath7. for top-k%, let r (k) i be the "important" (i.e., ones) tokens in r i , when using 0 ≤ @xmath8. faithfulness means how well a rationale reflects f task 's true reasoning process for predicting ŷi @xcite . hence, faithfulness metrics measure how much the r (k) i tokens impact p ŷi (x i ), which denotes f task 's confidence probability for ŷi when using x i as input @xcite @xcite . recently, comprehensiveness and sufficiency have emerged as popular faithfulness metrics @xcite . comprehensiveness (comp) measures the change in p ŷi when r ## unirex given task model f task , unirex generalizes rationale extractor optimization as follows: (1) choose architecture for a learned rationale extractor f ext ; (2) select explainability objectives (i.e., faithfulness loss l faith and plausibility loss l plaus ); and (3) jointly train f task and f ext using l task (task loss), l faith , and l plaus . unirex training consists of two backpropagation paths (fig. 2 ). the first path is used to update f task w.r.t. l task and l faith . whereas l task is computed w.r.t. the task target y i , l faith is computed only using the task input x i and the top-k% important tokens r (k) i (obtained via f ext ), based on some combination of comp and suff (sec. 2). the second path is used to update f ext w.r.t. l plaus , which encourages importance scores s i to approximate gold rationale r * i . thus, unirex frames rationale extraction as the following optimization problem: ## experiments we present empirical results demonstrating unirex's effectiveness in managing trade-offs between faithfulness, plausibility, and task performance during rationale extractor optimization. first, our main experiments compare methods w.r.t. faithfulness, plausibility, and task performance (sec. 4.2). second, we perform various ablation studies to verify our design choices for unirex (sec. 4.3). third, we present experiments highlighting unirex's generalization ability, both in terms of limited gold rationale supervision (sec. 4.4) and zero-shot transfer (sec. 4.5). fourth, we conduct a user study to further evaluate unirex rationales' plausibility, relative to those generated by other methods (sec. 4.6). see sec. a.5 for implementation details (lm architecture, aa settings, training). ## related work faithfulness many prior works have tried to improve the faithfulness of extractive rationales through the use of aas @xcite . typically, this involves designing gradientbased @xcite @xcite or perturbation-based @xcite @xcite aas. however, attribution algorithms cannot be optimized and tend to be compute-intensive (often requiring multiple lm forward/backward passes). recently, @xcite addressed the optimization issue by regularizing the task model to yield faithful rationales via the aa, while other works @xcite addressed the compute cost issue by training an lm (requiring only one forward pass) to mimic an aa's behavior. another line of work aims to produce faithful rationales by construction, via spps @xcite @xcite @xcite . still, spps' faithfulness can only guarantee sufficiency -not comprehensiveness @xcite . also, spps generally perform worse than vanilla lms because they hide much of the original text input from the predictor and are hard to train end-to-end. plausibility existing approaches for improving extractive rationale plausibility typically involve supervising lm-based extractors @xcite or spps @xcite @xcite with gold rationales. however, existing lm-based extractors have not been trained for faithfulness, while spps' faithfulness by construction comes at the great cost of task performance. meanwhile, more existing works focus on improving the plausibility of free-text rationales @xcite @xcite , often with task-specific pipelines @xcite . connection to unirex unlike prior works, unirex enables both the task model and rationale extractor to be jointly optimized for faithfulness, plausibility, and task performance. as a result, unirex-trained rationale extractors achieve a better balance of faithfulness and plausibility, without compromising the task model's performance. also, by using a learned rationale extractor, which generally only requires one model forward pass, unirex does not have the computational expenses that limit many aas.