KNIFE: D ISTILLING M ETA -R EASONING K NOWLEDGE WITH F REE -T EXT R ATIONALES
A BSTRACT
Recent works have explored using free-text rationales (FTRs)— i.e. , natural language explanations of a task output—to teach language models (LMs) how to solve NLP tasks.
In these works, the LM is often finetuned or prompted to jointly generate the FTR and task output.
However, this approach either involves finetuning LMs on possibly conflicting objectives or prompting prohibitively large
LMs.
To address this, we propose KNIFE, which guides LM reasoning via FTR knowledge distillation, instead of via FTR generation.
KNIFE first finetunes an
FTR-augmented teacher LM to predict the task output, then finetunes a student LM so that its hidden states are aligned with the teacher’s.
As a result, the student LM learns general reasoning knowledge from the FTRs and can be used for inference, without FTR generation or large LMs.
On two question answering datasets, we show that KNIFE outperforms various baselines in both fully-supervised and low-resource settings.
Also, using two more datasets, we analyze KNIFE’s failure modes and identify FTR quality as critical to KNIFE performance.
I NTRODUCTION
A free-text rationale (FTR) explains how to solve a given task instance.
Collectively, FTR-augmented task instances ( i.e., I nput, O utput, FT R ) can provide meta-reasoning knowledge for solving unseen instances.
Still, it remains unclear how to effectively inject this knowledge into LMs.
Whereas conventional supervised learning only gives feedback on a language model’s (LM’s) task output correctness, explanation-based learning (EBL) aims to improve LM generalization by additionally explaining the correct reasoning process behind a given correct output (Hase & Bansal, 2021;
Narang et al., 2020; Joshi et al., 2022).
In particular, there is growing interest in learning from free-text rationales (FTRs), which use natural language to explain the reasoning process for solving a given task instance (Narang et al., 2020; Rajani et al., 2019; Camburu et al., 2018; Wei et al., 2022).
Prior works have considered three paradigms for FTR learning (Fig. 2).
In the input augmentation paradigm, the LM is finetuned to generate the task output given both the task input and an FTR (Sun et al., 2022;
Wang et al., 2022; Wiegreffe et al., 2021;
Hase et al., 2020).
In the self-rationalization paradigm, the LM is finetuned or prompted to generate both the task output and an FTR (Narang et al., 2020; Brahman et al., 2021;
Li et al., 2022; Wei et al., 2022; Marasovi´c et al., 2022; Zelikman et al., 2022; Liu et al., 2018).
In the pipeline rationalization paradigm, a finetuned rationalizing LM first generates the
FTR, which is then used as input for a finetuned reasoning LM to predict the task output (Kumar &
Talukdar, 2020; Rajani et al., 2019; Wiegreffe et al., 2021; Hase et al., 2020).
However, these FTR learning paradigms either struggle to improve the LM’s task performance (input augmentation, finetuned self-rationalization, pipeline rationalization) or require prohibitively large ∗ Equal contribution.
LMs ( i.e., >100B) to work well (prompted self-rationalization).
While an individual FTR provides instance-level reasoning knowledge, a set of FTRs can collectively convey task-level meta-reasoning knowledge that generalizes to many unseen task instances (Fig. 1).
With this in mind, our goal is to extract meta-reasoning knowledge from a set of FTRs and effectively inject this knowledge into a small-scale ( i.e., <1B)
LM, in order to improve its generalization ability.
We propose KN owledge D
I stillation with F ree-Text Rational E s (
KNIFE ).
Instead of simply inserting FTRs into the LM’s input or target output, KNIFE guides the LM’s reasoning by distilling meta-reasoning knowledge from the FTRs to the LM’s hidden states.
First, given both the task input and an FTR, a teacher LM is finetuned to predict the task output.
Second, given only the task input, a student LM is finetuned so that its task input/output (hidden) states align with the teacher’s.
Since the student does not take FTR inputs, the teacher’s forward computation includes a bottleneck stage where its FTR (hidden) states are masked out.
Beyond this bottleneck, the teacher must use only its task input/output states to predict the task output, thus routing knowledge from the FTR to the teacher’s task input/output states.
Unlike input augmentation, KNIFE does not need FTRs for inference or cause an input distribution shift, since the student is given only the task input during both training and inference.
Unlike finetuned self-rationalization, KNIFE does not involve jointly optimizing task and FTR generation losses (which may conflict), since the student is only trained with KNIFE’s distillation losses.
Unlike prompted self-rationalization, KNIFE does not require large LMs, since the teacher and student are finetuned instead of relying only on their pretrained knowledge.
Unlike pipeline rationalization,
KNIFE does not require multiple inference LMs or create a non-differentiable path between them, since only the student is used for inference.
Plus, KNIFE generally has lower inference-time costs, since the student does not process additional FTR inputs or generate additional FTR tokens.
On two question answering (QA) datasets (OBQA, StrategyQA), we show that KNIFE outperforms various baselines on both fully-supervised and low-resource settings, using either gold FTRs or generated FTRs (§4.5, §A.4).
Furthermore, we validate our KNIFE design choices via extensive ablation studies (§A.5).
Finally, using two additional datasets (ECQA, QuaRTz), we analyze KNIFE’s failure modes and identify FTR quality as a critical factor in KNIFE performance (§A.6).
B ACKGROUND
Problem Definition Let F denote an LM for text classification.
Given task input x and class labels Y =
{ y i } , F ’s goal is to predict a confidence score ρ ( x , y i ) for each ( x , y i ) pair, so the predicted label ˆ y = arg max
y i ∈ Y ρ ( x , y i ) matches the gold ( i.e., correct) label y ∗ ∈ Y .
We consider multi-choice ( Y can vary across instances) and closed-set ( Y is fixed) text classification.
Free-Text Rationales Given a ( x , y i ) pair, a free-text rationale (FTR) r is a natural language text sequence that explains the reasoning process for predicting label
y i (Camburu et al., 2018; Rajani et al., 2019).
Compared to extractive rationales (Sundararajan
2022b) , FTRs may be more intuitive to humans, can reference things beyond the task input, and support high flexibility in content, style, and length (Wiegreffe et al., 2021; Chan et al., 2022a).
As a result, recent works have explored generating FTRs to explain LM behavior (Camburu et al., 2018;
Narang et al., 2020; Wei et al., 2022) and utilizing FTRs to improve LM decision-making (Sun et al., 2022; Wang et al., 2022; Li et al., 2022).
KNIFE
Existing works aim to learn from FTRs by using approaches like input augmentation, selfrationalization, or pipeline rationalization.
However, such approaches may hurt task performance or require prohibitively large LMs to work well.
To address this, we propose KNIFE, a knowledge distillation (KD) approach for injecting FTR-based meta-reasoning knowledge into LMs (Fig. 3).
3. LM D ESIGNS
KNIFE consists of a teacher LM and a student LM.
Below, we discuss the design of each LM type.
Basic LM Design Following prior works, we use text-to-text ( i.e., encoder-decoder) Transformer (Vaswani et al., 2017) architectures for both teacher and student LMs (Raffel et al., 2020;
Narang et al., 2020).
Building upon §2, let the task input be denoted as n x -token sequence x =
[ x (1) , x (2) , ..., x ( n x ) ] , while each label is denoted as n y i
-token
sequence y i
=
[ y (1) i , y (2) i , ..., y ( n yi ) i ] ∈ Y .
Given an LM F , we can decouple F as encoder F enc and decoder
F dec .
By default, F enc outputs hidden states F enc ( x ) =
[ e (1) x , e (2) x , ...
e ( n x ) x ] .
Typically, decoding is done by taking only a special start token as F dec ’s input, then using greedy search to sequentially generate the remaining tokens until a special end token is generated (Raffel et al., 2020; Vaswani et al., 2017).
However, when using this decoding strategy for classification, it is possible to generate an output sequence that does not match any of the labels in Y .
To avoid this issue, we instead produce a dedicated decoder output for each label sequence y i , by separately teacher-forcing each y
i to be F dec ’s input (Wang et al., 2022).
This gives us a conditional probability
P ( y ( j ) i | y (1) i , ..., y ( j − 1) i , x ) for each token y ( j ) i in y i .
Then, following Wang et al.
(2022) and Shwartz et al.
, we compute ρ i
=
ρ ( x , y i ) , the confidence score for y i , by aggregating these token probabilities as: ρ i = 1 n yi (cid:80) n yi j =1 log P ( y ( j ) i | y (1) i , ..., y ( j − 1) i , x ) .
Finally, we use the softmax function to normalize ρ i as label confidence probability P ( y i | x ) =
e ρ i / (cid:80) | Y | j =1 e ρ j .
Given gold label y ∗ , the goal of the downstream classification task is to train F such that P ( y ∗ | x ) is maximized.
Teacher LM Design
KNIFE first finetunes an FTR-augmented teacher LM to solve the task, using both the task input x and an FTR r as input.
We assume that the FTR expresses a sufficiently correct reasoning process corresponding to the gold label y ∗ .
By default, r is a gold FTR, which is humanannotated to support y ∗ .
Let T denote the teacher LM, which has encoder T enc and decoder T dec .
Let the FTR be denoted as n r -token sequence r =
[ r (1) , r (2) , ..., r ( n r ) ] .
Hence, T ’s input is denoted as [ x , r ] , where | [ x , r ] | = n x + n r .
As a result, T enc outputs hidden states e T
=
T enc ([ x , r ]) =
[ e T x , e T r ] .
In particular, we refer to e T x as the task input states and e T r as the FTR states .
Unlike the basic LM design described earlier, T has a bottleneck in the cross-attention between T enc and T dec .
Here, e T r is masked out so that T dec only has access to e T x during decoding.
This means all
FTR knowledge must be funneled to T dec through e T x .
Thus, for each label y i , T dec produces hidden states d T y i
= T dec ( y i , e T x ) =
[ d (1) T y i , d (2) T
y i , ..., d ( n yi )
T y i ] , which we call the task output states .
The bottleneck is designed to move knowledge from the FTR to e
T x and d T y i , through which knowledge is distilled to the student (which has no e T r ).
Note that the same r (which supports gold label y ∗ ) is used for decoding each label y i .
We can interpret T dec ’s decoding process as predicting the compatibility between y i and the correct reasoning process expressed by r .
We assume that, among these labels, only y ∗ is compatible with r .
Student LM Design After finetuning T , KNIFE finetunes a student LM based on KD objectives, given only x .
Let S be the student LM, with encoder S enc and decoder S dec .
Since S does not have access to r , S enc only outputs task input states e S
= S enc ( x ) =
[ e (1) S x , e (2) S x , ..., e ( n x ) S x ] and can be used for inference.
Hence, S does not have a bottleneck.
For each label y i , S dec produces task output states d S
y i
= [Equation] dec
[Equation]
[Equation] =
[Equation]
[Equation]
3. K NOWLEDGE D ISTILLATION
Recall that T is trained to predict the compatibility of each ( y i , r ) pair.
T ’s bottleneck design causes the FTR knowledge to be routed to e T x and d T y i .
First, since T dec can only use e T x as context,
T enc must be trained to distill sufficient knowledge from r into the e T x .
Second, since T dec is the component that actually performs the reasoning process to compute the final compatibility scores, d T y i must also contain useful knowledge.
Thus, in KNIFE, KD is done by training the student LM
S so that its task input states and/or task output states are aligned with the teacher LM’s.
Besides the
KD losses, we can also train S with the same task loss used to train T .
However, we show in our experiments that S actually performs better when trained with only KD losses (§A.5).
Now, we formally define these learning objectives.
Given predicted label distribution P and target label distribution Q , we first define the task loss L task as cross-entropy loss L task =
− (cid:80) y
i ∈ Y Q ( y i | x ) log P ( y i | x ) , where Q ( y i | x ) is 1 if y i =
y ∗ and 0 otherwise.
Next, let dist denote an arbitrary distance function, e.g., mean squared error (MSE).
Let L KD-in denote KNIFE’s task input states based KD loss, which pushes the student LM’s task input states ( e ( j ) S x ) to be closer to the teacher LM’s
( e ( j ) T x ): L KD-in = 1 n x (cid:80) n x j =1 dist ( e ( j ) S x , e ( j ) T x ) .
Similarly, let L KD-out denote KNIFE’s task output states based KD loss, which pushes the student LM’s task output states ( e ( j ) S y i ) to be closer to the teacher LM’s
( e ( j ) T y i ): L KD-out = 1 | Y | n yi (cid:80) y i ∈ Y (cid:80) n yi j =1 dist ( e ( j ) S
y i , e ( j ) T y i ) .
Finally, let L denote the total loss, defined as L = λ task L task + λ KD-in L KD-in + λ KD-out
L KD-out , with loss weights λ task , λ KD-in , and λ
KD-out .
The teacher is trained only with L task , whereas the student may or may not use L task .
E XPERIMENTS
This section (along with the appendix) presents experiments showing KNIFE’s effectiveness in FTR learning.
First, in both fully-supervised and low-resource settings, we show that KNIFE outperforms various baselines, using either gold FTRs or generated FTRs (§4.5, §A.4).
Second, we validate our
KNIFE design choices via extensive ablation studies (§A.5).
Third, we analyze KNIFE’s failure modes and identify FTR quality as critical to KNIFE performance (§A.6).
4. E VALUATION P ROTOCOL
Datasets Since FTRs are commonly annotated for QA datasets (Wiegreffe & Marasovic, 2021), we focus on two QA datasets: OBQA and StrategyQA.
OBQA (Mihaylov et al., 2018) is a multi-choice ( i.e., four-choice) QA dataset that simulates science exams.
StrategyQA
(Geva et al., 2021) is a closed-set ( i.e., binary yes/no) QA dataset that tests open-domain knowledge.
Since StrategyQA does not provide public test set labels, we use the data split from Wang et al.
Presentation of Results For all results, we report mean and standard deviation accuracy over three random seeds ( mean ± std ).
For each table, we use horizontal lines to partition the table into sub-tables.
Each sub-table contains results for methods that have comparable settings, so Result should only be compared to Result2 if there is no horizontal line between them.
In each sub-table, we highlight the best performing method in red and the second-best performing method in blue .
4. B ASELINES
Standard Finetuning does not use FTRs or KD.
FT (I → O) finetunes an LM to generate the task output, given the task input.
This is equivalent to training KNIFE’s student LM without KD.
Input Augmentation appends an FTR to a finetuned LM’s input.
FT (IR → O) finetunes an LM to generate the task output, given the task input and an FTR.
However, for fair comparison, the
FTR input is omitted during inference.
This is equivalent to training KNIFE’s teacher LM without the bottleneck.
FT Dropout (IR → O) mitigates FT (IR → O)’s input distribution shift by randomly dropping out the FTR input during training.
Finetuned Self-Rationalization appends an FTR to a finetuned LM’s target output.
FT (I → OR) finetunes an LM to generate the task output followed by the FTR, given the task input. FT (I → RO) finetunes an LM to generate the FTR followed by the task output, given the task input.
Prompted Self-Rationalization uses chain-of-thought (CoT) prompting (Wei et al., 2022).
CoT (I → RO) prompts an LM to generate the FTR followed by the task output, given the task input.
Pipeline Rationalization finetunes two LMs as a pipeline.
For FT
(I → R → O) , the first LM is finetuned to generate the FTR given the task input, while the second LMis finetuned to generate the task output given the first LM’s generated FTR.
Since FT (I → R → O) is known to not perform well (Wiegreffe et al., 2021; Wang et al., 2022), we only consider FT (I → R → O) in a limited set of settings.
In these settings, we simply present the results reported in Wang et al.
FT Teacher Init.
(I → O) modifies FT (I → O) by initializing the LM with the KNIFE teacher’s parameters, which is equivalent to training KNIFE’s student with task loss only.
This is used to verify that the KNIFE student’s gains come from KD, not from approximating the teacher’s parameters.
4. KNIFE V ARIANTS
We consider three KNIFE variants, each with a different combination of the L KD-in and L KD-out losses (§3.2).
KNIFE (In) trains the student LM only using L KD-in .
KNIFE (Out) trains the student
LM only using L KD-out .
KNIFE (In+Out) trains the student LM using both L KD-in and L KD-out .
By default, we use KNIFE (In+Out) and do not train the student with task loss L task (§A.5).
4. I MPLEMENTATION D ETAILS
Following prior works (Narang et al., 2020;
Sun et al., 2022; Wiegreffe et al., 2021), we use T5-Base and T5-Large (Raffel et al., 2020).
For KD methods, T5-Base → T5-Base means teacher and student use T5-Base, T5-Large → T5-Large means teacher and student use T5-Large, and T5-Large → T5Base means T5-Large teacher and T5-Base student.
For non-KD methods, T5-Base → T5-Base,
T5-Large → T5-Large, and T5-Large → T5-Base mean the LM uses T5-Base, T5-Large, and T5-Base, respectively.
For KNIFE, the student uses the teacher’s language modeling head.
Also, for KNIFE, if the student and teacher have the same architecture, the student’s parameters are initialized as the teacher’s.
See §A.2-§A.3 for more about hyperparameters and implementation details, respectively.
4. M AIN R ESULTS
Table 1 presents our main results.
Here, LMs are finetuned on the entire training set, using gold FTRs if applicable.
We observe that KNIFE (In+Out) consistently outperforms all baselines.
Unlike FT (I → O), KNIFE (In+Out) benefits from training with FTR knowledge.
Unlike FT (I → OR) and
FT (I → RO), KNIFE (In+Out) does not have an FTR generation loss to compete with the task loss.
Unlike FT (I → R → O), KNIFE (In+Out)’s LM has direct access to the task input when predicting the task output.
Unlike FT (IR → O) and FT Dropout (I → RO), KNIFE (In+Out) avoids input distribution shift since its student only takes task input in training and inference.
Unlike FT Teacher Init.
(I → O), KNIFE (In+Out) explicitly trains the student to follow the teacher’s reasoning process
[Table]
Also, we report CoT (I → RO) results for GPT-NeoX (20B) (Black et al., 2022) and GPT- (text-davinci-003, 175B) (Brown et al., 2020)
Since GPT-NeoX and GPT-3 are much larger than
T5-Base (220M) and T5-Large (770M), it is unfair to expect other methods to perform as well as CoT (I → RO).
Even so, we find that KNIFE (In+Out) greatly outperforms GPT-NeoX on all settings, while KNIFE (In+Out) with T5-Large achieves similar performance to GPT-3 on
StrategyQA.
See §A.4 for results with GPT-NeoX generated FTRs and low-resource learning.
R ELATED W ORK
Learning from Free-Text Rationales There are three main FTR learning paradigms: input augmentation, self-rationalization, and pipeline rationalization.
In input augmentation, the LM is finetuned to generate the task output given both the task input and an FTR (Sun et al., 2022; Wang et al., 2022;
Wiegreffe et al., 2021; Hase et al., 2020).
Still, this either assumes access to gold (or large-LMgenerated) FTRs during inference, or introduces an input distribution shift between training and inference when FTRs are unavailable during inference.
In self-rationalization, the LM is finetuned or prompted to generate both the task output and an FTR (Narang et al., 2020; Brahman et al., 2021;
Li et al., 2022; Wei et al., 2022; Marasovi´c et al., 2022; Zelikman et al., 2022; Liu et al., 2018; Majumder et al., 2022).
However, for self-rationalization, finetuning may create conflict between the task and
FTR objectives, while prompting requires very large LMs to work well.
In pipeline rationalization, a finetuned rationalizing LM first generates the FTR, which is then used as input for a finetuned reasoning LM to predict the task output (Kumar & Talukdar, 2020; Rajani et al., 2019;
Wiegreffe et al., 2021; Hase et al., 2020).
Here, the generated FTR forms a non-differentiable path between the two LMs, which complicates end-to-end training and can hurt task performance.
Unlike prior works,
KNIFE distills FTR knowledge from an FTR-augmented teacher LM to a student LM that does not have FTR inputs, in order to improve the student’s generalization.
Knowledge Distillation Knowledge distillation has been widely used to transfer knowledge from a larger teacher to a smaller student model.
(Hinton et al.
; Sanh et al., 2019; Jiao et al., 2020; Mirzadeh et al., 2020, etc).
Instead of aiming for this typical goal, KNIFE distills the FTR knowledge learned by a teacher model to a student model, which has no direct access to FTRs.
Similar to the line of work that incorporates knowledge distillation with privileged information (Lopez-Paz et al., 2015;
Vapnik et al., 2015; Fukuda et al., 2017; Wang et al., 2018), where student models benefit from privilege information, the student model in KNIFE essentially gains additional knowledge from FTRs rather than relying on the larger teacher model capacity.
Snell et al.
(2022) propose to internalize the in-context learning ability such that the performance gains can keep without context tokens.
It does not directly distill the knowledge from FTRs and requires dedicated prompt designs.
Shridhar et al.
; Magister et al.
; Ho et al.
(2022) propose to distill reasoning abilities from larger language models to smaller models.
They require large-scale language models with such abilities, while KNIFE can work well with small models
A A PPENDIX A. B ACKGROUND (E XTENDED )
Problem Definition Let F denote an LM for text classification.
Given a task input x and a set of class labels Y =
{ y i } , F ’s goal is to predict a confidence score ρ ( x , y i ) for each ( x , y i ) pair, so that the predicted label ˆ y = arg max
y i ∈ Y ρ ( x , y i ) matches the gold ( i.e., correct) label y ∗ ∈ Y .
In this work, we consider both multi-choice ( Y can vary across task instances) and closed-set
( Y is the same for all task instances) text classification.
Free-Text Rationales Given some ( x , y i ) pair, a free-text rationale
(FTR) r is a natural language text sequence that explains the reasoning process for predicting label
y i (Camburu et al., 2018; Rajani et al., 2019)
Most explainability works have focused on extractive rationales (Sundararajan et al., 2017;
Li et al., 2016; Chan et al., 2022b), which highlight important features ( e.g., tokens) in x , but there is growing interest in FTRs (Narang et al., 2020; Kumar & Talukdar, 2020;
Wei et al., 2022).
First, compared to extractive rationales, FTRs may be more intuitive and understandable to humans, since they follow how humans communicate in natural language (Camburu et al., 2018; Wiegreffe et al., 2021).
Second, in contrast to extractive rationales, FTRs can reference things beyond the task input (Rajani et al., 2019; Wiegreffe et al., 2021).
Third, FTRs support high flexibility in content, style, and length (Chan et al., 2022a).
Hence, FTRs can potentially provide much richer knowledge than extractive rationales.
As a result, a number of recent works have explored generating FTRs to explain LM behavior (Narang et al., 2020; Rajani et al., 2019; Camburu et al., 2018) and utilizing
FTRs to improve LM decision-making (Sun et al., 2022; Wang et al., 2022; Li et al., 2022).
Learning from Free-Text Rationales Since extractive rationales assign an importance score to each input feature, there is a one-to-one correspondence between the features and scores (Sundararajan et al., 2017;
Li et al., 2016; Chan et al., 2022b;
Joshi et al., 2022).
Given this structure, prior works regularize LM behavior by aligning the LM’s extractive rationales with human-annotated extractive rationales (Joshi et al., 2022; Ross et al., 2017; Rieger et al., 2020; Liu & Avci, 2019).
However, unlike extractive rationales, FTRs do not have such a feature-score correspondence to leverage.
Without this kind of structure, it becomes less straightforward to effectively learn from FTRs.
The FTR learning literature revolves around three main paradigms: input augmentation (Sun et al., 2022; Wang et al., 2022; Wiegreffe et al., 2021), self-rationalization (Narang et al., 2020; Wei et al., 2022;
Marasovi ´ c et al., 2022), and pipeline rationalization (Kumar & Talukdar, 2020; Rajani et al., 2019; Hase et al., 2020).
These paradigms explicitly add the FTR to the LM’s input or target output, but this general approach has key limitations.
First, adding FTRs to the LM’s input assumes access to sufficiently “correct” ( i.e., explains correct reasoning process) FTRs during inference.
Yet, this is unrealistic because having correct FTRs would already be tantamount to solving the task.
Meanwhile, simply omitting the FTR only during inference introduces an input distribution shift between training and inference, thus hurting inference performance.
Second, adding FTRs to the LM’s target output assumes that task prediction and FTR generation are complementary objectives.
In practice, this is often not the case, especially if the generated FTR does not support the LM’s predicted label.
Plus, if the LM is trained to generate the predicted label and FTR as a single output sequence (Narang et al., 2020; Wiegreffe et al., 2021), then it is not uncommon for the LM to generate a predicted label that does not match any valid label candidates or even omit the predicted label entirely.
A. H YPERPARAMETERS
For KD, we used sweeps of λ task =
[0 , 1] , λ KD-in = [0 , 1] , and λ KD-out = [0 , 1] .
We always use
AdamW (Loshchilov & Hutter, 2017) as the optimizer.
For early stopping, we stop training when the model performance on the development set has not improved for five epochs.
The maximum number of epochs is 10 .
For OBQA with T5-Base, we train the teacher model with learning rate of 1e − 4 and batch size of 64 .
We train the student model (by KD) with batch size of 64 .
For OBQA with T5-Large, we train the teacher model with learning rate of 5e − 5 and batch size of 64 .
We train the student model with batch size of 48 .
For the KD training student model, we always search the learning rate in { 1e − 4 , 2e − 4 , 3e − 4 , 4e − 4 , 5e − 4 }.
[Table]
For StrategyQA, we always set the warmup rate as 0 . 06 .
For T5-Base, we train the teacher model with learning rate of 3e − 4 and batch size of 16 .
For T5-Large, we train the teacher model with learning rate of 5e − 5 and batch size of 16 .
For the KD training student model, the batch size is always 16 , and we always search the learning rate in { 1e − 4 , 2e − 4 , 3e − 4 , 4e − 4 , 5e − 4 }.
FT (I → O), FT (I → OR), FT (I → RO), and FT (IR → O) use the same hyperparameters as the teacher models in the same settings do, except that we always search the learning rate in { 1e − 5 , 2e − 5 , 5e − 5 1e − 4 , 2e − 4 , 3e − 4 , 4e − 4 , 5e − 4 }.
FT Dropout (IR → O) uses the same hyperparameters as
FT (IR → O) does, and we search the dropout rate in { 0 . 1 , 0 . 2 , 0 . 3 , 0 . 4 , 0 . 5 , 0 . 6 , 0 .
FT
Teacher Init.
(I → O) uses the same hyperparameters as FT (I → O) does.
A. I MPLEMENTATION D ETAILS
For T5-Base, the parameter number of a single backbone model is around 220M. As we have two backbone models for the teacher and student model, the total number is around 440M. For T5Large, the parameter number of a single backbone model is around 770M, and the total number is around 1.5B. GPT-NeoX has 20B parameters.
We use NVIDIA Quadro RTX8000 GPUs for all experiments, which take around 700 GPU hours in total.
We implement the models using Hugging
Face Transformers, PyTorch, and Lightning.
FT (I → OR) and FT (I → RO) use greedy decoding (§3.1) because, unlike task labels, there do not exist multiple FTR choices (besides gold FTR y ∗ ) to separately teacher-force into the decoder.
When using greedy decoding for classification, it is possible for the LM to generate an output sequence that does not match any of the labels in Y .
Thus, the LM may perform much worse than random chance.
For T5-Large → T5-Base, the teacher and student have different embedding dimensionalities.
Thus, for the task input states and task output states, we train a linear projection layer transforming the student LM’s states to have the same dimensionality as the teacher LM’s.
Note that only the mean performance is available for FT (I → R → O), since these results were obtained from Wang et al.
All CoT (I → RO) results were obtained from Wang et al.
(2022), except GPT-3 on OBQA, which was obtained from Huang et al.
A. M AIN R ESULTS (E XTENDED )
In Table 2, we repeat these experiments for GPT-NeoX generated FTRs (instead of gold FTRs) and obtain the same conclusions.
This makes sense since GPT-NeoX’s FTRs are known to fluently convey useful knowledge (Black et al., 2022).
Interestingly, KNIFE with GPT-NeoX FTRs still considerably
[Table] [Table]
outperforms CoT (I → RO) with GPT-NeoX, despite KNIFE using much smaller LMs.
Furthermore, in Table
[Table]
we consider a low-resource setting where LMs are finetuned on only 10% of the training data, using T5-Base → T5-Base on OBQA.
We find that KNIFE beats all baselines in the low-resource setting, showing that KNIFE more efficiently leverages the FTR learning signal.
A. A BLATION S TUDIES
To justify design choices made for KNIFE and understand why it works
[Table]
we present five KNIFE ablation studies, analyzing the impacts of KD objective, FTR usage, FTR quality, teacher bottleneck, and student task loss.
KD Objectives Table 4 compares the performance of KNIFE (In), KNIFE (Out)
[Table]
and KNIFE (In+Out).
For both gold FTRs and GPT-NeoX FTRs, we find that KNIFE (In+Out) generally achieves the highest performance.
This suggests that useful FTR knowledge can be distilled via both task input states and task output states, so it is best to use both.
Although KNIFE (In) and KNIFE (Out) perform similarly for T5-Base → T5-Base and T5-Large → T5-Large, KNIFE (In) performs much worse for T5-Large → T5-Base.
Since KNIFE (In) only performs KD via the task input states, the student’s decoder cannot be trained without task loss.
Yet, for T5-Large → T5-Base, the student cannot be initialized with the teacher’s parameters, leaving the decoder with T5-Base’s pretrained parameters.
While this could be addressed by training the student with task loss, it shows a major disadvantage of KNIFE (In) compared to other KNIFE variants.
FTR Usage This experiment verifies the importance of FTR usage in KD by comparing KNIFE to non-FTR KD baselines, where the teacher is FT (I → O).
Non-FTR KD (Logit) finetunes the student so its logit distribution aligns with the teacher’s.
Non-FTR KD (In+Out) finetunes the student so its task input states and task output states aligns with the teacher
’s.
Since non-FTR KD generally requires the teacher to be larger than the student, we only use T5-Large → T5-Base here.
In Fig.
4a, both KNIFE (In+Out) + Gold and KNIFE (In+Out) + GPT-NeoX outperform the non-FTR KD baselines, showing that FTR-based KD is helpful.
Plus, Non-FTR KD (Logit) performs much worse than Non-FTR KD (In+Out), which further validates KNIFE’s use of representation-based KD.
FTR Quality By default, KNIFE uses gold FTRs to train the teacher, and we have also explored using GPT-NeoX FTRs.
Using T5-Base → T5-Base, we investigate the relationship between KNIFE performance and FTR quality by considering KNIFE variants that train the teacher on noisy FTRs:
Replace creates noisy FTRs by replacing each gold FTR token with a random token from the token vocabulary.
Shuffle creates noisy FTRs by shuffling gold FTRs across all training instances in the dataset.
In Fig.
4b, KNIFE (In+Out)’s performance with Gold and GPT-NeoX is much higher than with Replace and Shuffle.
Plus, Replace performs worse than Shuffle since Replace’s tokenlevel randomization fully corrupts the FTR.
This suggests that KNIFE’s performance is positively correlated with FTR quality.
[Figure]
Teacher Bottleneck We validate KNIFE’s teacher bottleneck by considering a KNIFE variant with no bottleneck.
In Fig.
4c, for T5-Base → T5-Base, KNIFE (In+Out) w/ Bottleneck greatly outperforms KNIFE (In+Out) w/o Bottleneck.
This demonstrates the bottleneck’s effectiveness in routing FTR knowledge to the teacher’s task input/output states, through which FTR knowledge can be distilled to the student.
In Table 8, we show that the bottleneck is also very important for KNIFE (In), but not as critical for KNIFE (Out).
Student Task Loss By default, KNIFE trains the student with only KD losses.
We justify this design choice by comparing to KNIFE variants where the student is also trained on the task loss.
In
Fig.
4d, we see that omitting the task loss consistently yields higher performance, which indicates that KD losses are not always compatible with task loss.
If the KD and task losses produce conflicting gradients during optimization, the student may get confused and learn a suboptimal “in-between” reasoning process leading to even worse generalization.
A. F AILURE A NALYSIS
Although KNIFE performs well on OBQA and StrategyQA, it yield negative results on other QA datasets like ECQA (Aggarwal et al., 2021) and QuaRTz (Tafjord et al., 2019).
Using T5-Base → T5Base, we compare the performance of KNIFE and most of the baselines considered in the main results.
In Table 5, we see that KNIFE generally outperforms all FTR-based baselines, sometimes by a very large margin.
Still, none of the FTR-based methods (including all KNIFE variants) are able to significantly outperform FT (I → O).
Since KNIFE distills FTR knowledge to the student LM, the student’s performance is expected to depend on the amount and quality of meta-reasoning knowledge stored in the FTRs.
Thus, to investigate these negative results, we conducted a case study to qualitatively analyze the gold FTRs in
[Table] [Table]
OBQA, StrategyQA, ECQA, and QuaRTz.
Overall, we found that FTRs in OBQA and StrategyQA are more informative than those in ECQA and QuaRTz.
For OBQA and StrategyQA, we find that gold FTRs tend to have the following properties.
First, they describe a logically sufficient reasoning process for getting from the question (input) to the answer (output).
Second, they provide general and self-contained knowledge that goes beyond the information given in the question and answer.
This means they do not simply rephrase the question and/or answer.
Meanwhile, FTRs from ECQA and QuaRTz tend to exhibit the opposite properties.
To illustrate, we give a representative example of a good OBQA FTR: “ Question : There is most likely going to be fog around: Answer Choices : (A) a marsh, (B) a tundra, (C) the plains (D) a desert.
Gold FTR : fog is formed by water vapor condensing in the air. ”
We also give a representative example of a bad ECQA FTR: “ Question : What might a person see at the scene of a brutal killing?
Answer Choices : (A) bloody mess, (B) pleasure, (C) being imprisoned, (D) feeling of guilt, (E) cake.
Gold FTR : Bloody mess is covered or stained with blood.
A person might see a bloody mess at the scene of a brutal killing. ”
A. A BLATION S TUDIES (E XTENDED )
We present the full results of ablation studies in the Appendix.
Table 6 shows the full results of ablation studies on FTR Usage.
Table 7 shows the full results of ablation studies on FTR Quality.
Table 8 shows the full results of ablation studies on teacher bottleneck.
Table 9 shows the full results of ablation studies on student task loss.
The details of ablation studies are in A.5.
In the last three tables, we always use T5-Base → T5-Base.
[Table]
[Table]
[Table]