Let the Rule Speak: Enhancing In-context Learning Debiasing with Interpretability

The classification outputs by in-context learning (ICL) are described as biased when they exhibit imbalanced per-class prediction accuracy. Addressing such imbalances while improving overall accuracy is seen as a category of debiasing. Concretely, the skewness in the output space can be alleviated by targeted corrections on output logits or probabilities, with or without explicitly modeling the per-class accuracy differences, i.e., COBias [1]. However, while effective, prior methods could lack straightforward explanations on why and which certain classes need corrections. What’s more challenging is to have a tailored per-sample, per-class correction.

A direct cause of COBias is that ICL tends to assign specific ranges of output probabilities to each class. When some classes always receive high probabilities for any input example, others may have lower or mixed probability ranges. The consequence is that latter classes are less frequently predicted than the former, resulting in consistently lower accuracies and calling for probability corrections. In addition, among all examples of a class A, the subset of examples whose in-context learned probability of answer A is relatively low often receive a lower test accuracy, compared to the subset whose class A probability is higher, suggesting that different probability ranges within a class need different corrections.

Taking these overlooked aspects into account, a correction should be tailored for each class and for each sample. To achieve this, a helpful correction should be able to asymmetrically amplify or reduce different ranges of a class’s probabilities. In this paper, we address the pressing need for enhanced understandings in how biased ICL predictions happen, and propose two research questions about a main concern, yet a potential direction, in interpretable ICL output corrections.

The key intuition here is that a membership function can asymmetrically amplify or reduce different ranges of inputs. Therefore, a fuzzy rule based debiaser for class probability $p_{mi}$ can be written as $f_{A_{i}}(p_{mi})$, where $A_{i}$ is a fuzzy set for class $i$, and its membership function $f_{A_{i}}$ maps the probability to a corrected $p^{\prime}_{mi}:=f_{A_{i}}(p_{mi})$. Then $\boldsymbol{p}^{\prime}_{m}$ consists of corrected per-class probabilities.

Alternatively, the debiaser can be viewed as a single rule:

$$\text{If}\underbrace{\text{ class $1$ is $A_{1}$ and ... and class $N$ is $A_{N}$}}_{\text{Antecedent}}\text{ then }\underbrace{\text{predict $\textstyle\operatorname*{arg\,max}_{j}f_{A_{j}}(p_{mj})$}}_{\text{Consequent}}$$

(1)

Our goal is to optimize the rule, i.e., select fuzzy sets/membership functions for every class in the antecedent, towards mitigating COBias and improving overall prediction accuracy. Specially, we include a Don’t Change membership function that will keep a class unchanged, suggesting that the LLM in-context learns an accurate probability for the class. When a correction is needed, the membership function detects the probability range that a class’s probability belongs to, and updates it with the returned function value. The problem becomes jointly selecting a set of membership functions for each class towards improving multi-objectives based on COBias and accuracy.

To this end, we propose a Fuzzy Rule Optimization based Debiasing method, FuRud, which demonstrate via extensive experiments (Section 4) and discussions (Section 5) that it achieves good improvements over accuracy and COBias while providing sample-level interpretability.

In a nutshell, FuRud uses an optimization set of samples for membership function selection. The optimization set’s questions are prompted in 1-shot manner, and probabilities are measured across answer classes for each question. These probabilities and ground-truth answers across all questions are aggregated in the multi-objective model, to jointly learn an optimal membership function for each class. At inference, a test example’s class probabilities are obtained similarly. Then we apply the learned membership functions to perform tailored corrections at each class’s probability in the given test sample. An overview of FuRud is shown in Figure 1, illustrating desired corrections and performance improvements.

To highlight, the membership functions learned by FuRud enable sample-level interpretability. FuRud enables us to know whether the LLM in-context learns an accurate probability for a class within a given sample. This is achieved by learning a correction function (membership function) for each class, towards the multi-objectives of reducing COBias and enhancing accuracy. If the Don’t Change function is learned for a class, it means the LLM in-context learns an accurate probability for the class; otherwise, a tailored correction is performed by the membership function. The source code will be released upon paper publication. In summary, our messages are:

• •

  We propose an interpretable fuzzy rule optimization based debiasing method (FuRud), to account for both inter-class surface biases and intra-class range-wise influences.

• •

  We formulate a multi-objective programming model to jointly optimize a set of triangular membership functions for each class. The functions are human-readable, which can asymmetrically correct probabilities of different ranges that are misrepresented.

• •

  Across seven benchmarks, FuRud demonstrates its effectiveness for improved overall accuracy, reduced per-class accuracy imbalance, and enhanced interpretability. For example, it improves ICL accuracy by a relative increase of 21% and reduces COBias by a relative decrease of 56%; it achieves higher accuracy (avg. accuracy reaching 72.0%) and competitive COBias (avg. COBias dropping to 17.8%) over state-of-the-art debiasing methods.

The core idea is to handle the imbalanced per-class accuracy issue with fuzzy membership functions. In the fuzzy rule setting, for $N$ classes, each class selects a fuzzy set $A_{i}$, or equivalently, a membership function $f_{A_{i}}$, from a family of $K$ fixed fuzzy sets. We let $F=\{f_{1},...,f_{k},...,f_{K}\}$ denote the family of membership functions. The membership function selection problem can be solved using combinatorial optimization. To this end, we introduce FuRud, a Fuzzy Rule Optimization Based debiasing method. The FuRud optimization is performed on a set of labeled examples, and the selected membership functions are directly applied to transform test-time class probabilities.

Figure 2 shows 19 triangular membership functions of four fuzzy partitions, including the Don’t Change membership function - the identity function (slope=1). Other than Don’t Change, each membership function represents a sharp or smooth transformation of the input variable. Details of the functions are discussed in Appendix A. The general form of a triangular membership function $f_{k}(\cdot)$ can be written as:

$$f_{k}(p_{mi};a_{k},b_{k},c_{k})=\left\{\begin{aligned} &0,&&\text{if}\ p_{mi}\leq a_{k}\\ &\frac{p_{mi}-a_{k}}{b_{k}-a_{k}},&&a_{k}\leq p_{mi}\leq b_{k}\\ &\frac{c_{k}-p_{mi}}{c_{k}-b_{k}},&&b_{k}\leq p_{mi}\leq c_{k}\\ &0,&&\text{otherwise}\end{aligned}\right.$$

(2)

where $a_{k},b_{k},c_{k}$ are the left endpoint, the input value where the peak is reached, and the right endpoint of $f_{k}$. For example, for $f_{11}$, the $a_{k},b_{k},c_{k}$ values are 0.125, 0.25, 0.375 respectively.

Then, we compute the updated probability $p^{\prime}_{mi}$ by:

$$p^{\prime}_{mi}=\left\{\begin{aligned} &p_{mi},&&\text{if $\textstyle\sum_{i=1}^{N}p^{\prime}_{mi}=0$}\\ \sum\limits_{k}&f_{k}(p_{mi})\mathbbm{1}(\kappa_{i}=k),&&\text{otherwise}\end{aligned}\right.$$

(3)

where $\kappa_{i}$ is the integer selection variable for class $i$. $\mathbbm{1}(\cdot)$ evaluates to 1 if the condition inside is satisfied, otherwise 0. Furthermore, in case $p^{\prime}_{mi}=0$ for all classes, we reset each to be its original probability in $\boldsymbol{p}_{m}$. Therefore, $\hat{y}_{m}=\operatorname*{arg\,max}_{i}p^{\prime}_{mi}$.

The first objective is to improve overall accuracy:

$$\max Z\textsuperscript{Acc}=\frac{1}{|S\textsuperscript{Opt}|}\sum\nolimits_{m\in S\textsuperscript{Opt}}\mathbbm{1}\{\hat{y}_{m}=y_{m}\}$$

(4)

where S^Opt is the indices of examples used for optimization.

Furthermore, we balance the class accuracy difference by explicitly modeling COBias, which accounts for an overall difference between pairwise per-class accuracies. Minimizing COBias helps address low-accuracy classes from ICL outputs. Therefore, the second objective is:

$$\min Z\textsuperscript{COBias}=\frac{1}{\textsubscript{N}C_{2}}\sum\nolimits_{i=1}^{N-1}\sum_{j=i+1}^{N}\bigl{|}\text{Acc}_{i}-\text{Acc}_{j}\bigr{|}$$

(5)

where $\textsubscript{N}C_{2}=N(N-1)/2$, $\text{Acc}_{i}$ is the accuracy score for optimization examples in class $i$.

To further handle extreme cases of low class accuracies, we penalize classes that fail to reach an accuracy threshold, and minimize the loss between the threshold and per-class accuracy (cut off at 0). The third objective is:

$$\min Z\textsuperscript{Extreme}=\sum\nolimits_{i=1}^{N}\max\{0,\lambda-\text{Acc}_{i}\}$$

(6)

where $\lambda$ is a fixed threshold value.

The above objective functions are a mix of minimization and maximization, and the resulted multi-objective programming model requires integer variables. Each of them alone corresponds to an integer programming problem, which is NP-complete [21]. Classic solutions for integer programming use operational research techniques, such as Branch-and-Bound, often used for linear integer programming problems. It could be difficult for such methods to handle nonlinear integer programming models which contain non-differentiable functions. Consequently, a series of metaheuristic algorithms have emerged, such as Simulated Annealing (SA), and each metaheuristic has their own strengths and limitations. We use one of the metaheuristics, SA, to tackle the proposed mathematical model. The SA implementation follows [1]. Since it is difficult to solve each one as an individual optimization problem and force an optimal solution, our strategy is instead to compute a weighted sum of $1-Z\textsuperscript{Acc},Z\textsuperscript{COBias},Z\textsuperscript{Extreme}$ as a single energy function $E$ to be optimized using SA. Hence, the multi-objectives are combined into a total minimization objective:

$$\min_{\kappa}E(\kappa;\lambda,\boldsymbol{p}^{\prime})$$

(7)

where $E(\kappa;\lambda,\boldsymbol{p}^{\prime})=\omega+\textstyle\sum_{h\in S\textsuperscript{Obj}}\gamma^{h}Z^{h}$, $S\textsuperscript{Obj}$ is the names of the penalty functions corresponding to the individual objectives, and $\omega,\gamma^{h}$s are penalty parameters. Therefore, the SA algorithm optimizes on $E$ to obtain an optimal set of membership functions.

In summary, the class corrections aim at reducing COBias and improving accuracy. Each equation from 4 to 6 exactly targets one of these two goals. In detail, Eq. 4 targets maximizing overall accuracy, Eq. 6 targets minimizing COBias, and Eq. 6 targets maximizing per-class accuracy, which enforces it to meet a threshold; Eq. 7 combines the three objectives as a multi-objective function. Details on how Eq. 7 is optimized are described in experimental setups (Section 4.1).

In this work, we present a fuzzy rule optimization based debiasing method to enhance ICL output class representations with interpretability. FuRud learns a per-class correction function, i.e., a membership function, which decides if and how a class’s probability needs correction for each sample. If correction is needed, the corrected class probability will be tailored by the membership function, which is a main innovation of this paper. On a diverse set of text classification benchmarks, FuRud greatly improves the average test accuracy and test COBias over ICL, by a relative increase of 21% and a relative reduction of 56%, outperforming state-of-the-art methods. Moreover, FuRud can work for prompt formats that may lead to highly skewed predictions, e.g., letter options. Furthermore, FuRud can optimize a downstream task with as few as 10 optimization examples.

In the future, more versatile rules can be explored, and we may also examine the tradeoff between the accuracy and rule complexity. Simpler rules are easier to understand, but the transformations may fail to catch the intricate interactions between class predictions. More complex rules may have better modeling capabilities, but they are harder to read. In addition, this work focuses on evaluating text classification, and we will extend interpretable ICL debiasing to more language tasks, modalities, and model architectures.