SlowGuess/NeurIPS_2023_OCR huggingface dataset:
[56] Michihiro Yasunaga, Jure Leskovec, and Percy Liang. Linkbert: Pretraining language models with document links. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio, editors, Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022, pages 8003-8016. Association for Computational Linguistics, 2022. URL https://doi.org/10.18653/v1/2022.acl-long.551.
[57] Chulhee Yun, Suvrit Sra, and Ali Jadbabaie. Small relu networks are powerful memorizers: a tight analysis of memorization capacity. Advances in Neural Information Processing Systems, 32, 2019.
[58] Jianyi Zhang, Aashiq Muhamed, Aditya Anantharaman, Guoyin Wang, Changyou Chen, Kai Zhong, Qingjun Cui, Yi Xu, Belinda Zeng, Trishul Chilimbi, and Yiran Chen. Reaugkd: Retrieval-augmented knowledge distillation for pre-trained language models. In Anna Rogers, Jordan L. Boyd-Graber, and Naoaki Okazaki, editors, Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), ACL 2023, Toronto, Canada, July 9-14, 2023, pages 1128-1136. Association for Computational Linguistics, 2023. URL https://doi.org/10.18653/v1/2023.acl-short.97.
[59] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona T. Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. OPT: open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022. URL https://doi.org/10.48550/arXiv.2205.01068.
A Motivation: Effect of Knowledge-Augmentation on Memorization
A.1 Additional Details and Discussion
We adopted the abstracted language problem, i.e., the next-symbol prediction problem with reference strings  $c_{j} \sim \mathrm{Uniform}\left(\{0,1\}^{d}\right)$ , considered in the main text of [4] with no symbol corruption. In this problem, the random process  $P \sim q$  is defined by drawing  $c_{j} \in \{0,1\}^{d}$  uniformly at random as  $c_{j} \sim \mathrm{Uniform}\left(\{0,1\}^{d}\right)$  for all  $j \in [N]$ : a set  $\{c_{j}\}_{j=1}^{N}$  corresponds to a  $P$ . Then, the random process  $(Z,Y) \sim P$  is defined by sampling  $j \sim \mathrm{Uniform}\left([N]\right)$  and  $\ell \sim \mathrm{Uniform}\left(\{0,1,2,\dots,d-1\}\right)$  and by setting  $Z = (j,c_{j}(1:\ell))$  and  $Y = c_{j}(\ell+1)$ , where  $c_{j}(1:l)$  denotes the first  $l$  symbols of  $c_{j}$  and  $c_{j}(l+1)$  is  $(l+1)$ -th symbol of  $c_{j}$ . The training data  $X \sim P^{\otimes n}$  contains  $((Z_{i},Y_{i}))_{i=1}^{n}$  generated by the same process  $n$  times independently.
Recall that the inference algorithm  $\varphi$  uses a KB denoted by  $S$  such that  $|S| = N + R$  and  $\{c_j\}_{j=1}^N \subseteq S$ . Here,  $S$  is a set and not ordered; therefore we do not know the identifier  $j$  of each  $s \in S$  and which  $s \in S$  is useful given each  $(Z,Y)$ . Thus, this still requires learning and memorizing some information from the training data  $X$ .
In Theorem 2,  $N$  (or  $R$ ) is the number of useful (or unuseful) documents in KB to be extracted by a retriever. Thus,  $N + R$  is the total number of possible documents to be retrieved. As  $N + R$  increases, we need to memorize more information of the training data to get the best match at test time. Theorem 2 shows that this increase happens only in the log rate,  $\log_2(N + R)$ . Thus, we can get improvements when the size of KB  $(d)$  is larger than the log of the number of possible choices to be retrieved by the retriever  $(\log_2(N + R))$ .
Moreover, unlike Theorem 1, we can remove the dependence on  $n$  and have  $I(X; \mathcal{A}(X)|P) = O(N \log_2(N + R))$  in Theorem 2 when the training size  $n$  grows at a faster rate than the number of the useful documents  $N$ . This is because our proof does not rely on the high probability of having a singleton training sample per each reference  $c_j$  unlike the proof of [4].
Proof. Let  $\epsilon > 0$  and  $m = \log_2\left((1 - \left(\frac{N - 1}{N}\right)^n)\frac{(N + R)^2 - (N + R)}{2\epsilon}\right)$ . Since the distribution over subpopulation is uniform,
$$
\operatorname {e r r} _ {q, n} ^ {\varphi} (\mathcal {A}) = \sum_ {j = 1} ^ {N} \Pr (Q _ {j}) \Pr (E _ {0} | Q _ {j}) = \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr (E _ {0} | Q _ {j}),
$$
where  $E_0$  is the event of  $\varphi(Z, \mathcal{A}(X), S) \neq Y$  and  $Q_j$  is the event that the subpopulation of the test sample  $(Z, Y) \sim P$  is  $j$ . Let  $E_1^j$  be the event of having at least one training data sample  $((Z_i, Y_i) \in X)$  in the subpopulation  $j$  of the test sample  $(Z, Y) \sim P$ . Then,
$$
\operatorname * {P r} (E _ {0} | Q _ {j}) = \operatorname * {P r} (E _ {1} ^ {j} | Q _ {j}) \operatorname * {P r} (E _ {0} | Q _ {j}, E _ {1} ^ {j}) + \operatorname * {P r} (\bar {E} _ {1} ^ {j} | Q _ {j}) \operatorname * {P r} (E _ {0} | Q _ {j}, \bar {E} _ {1} ^ {j})
$$
where  $\bar{E}_1^j$  is the complement of  $E_1^j$ . Denote by  $\ell_x^j$  the length of training sample, and by  $\ell_t^j$  the length of the test sample, in the subpopulation  $j$ . Let  $E_2^j$  be the event of  $\ell_x^j \geq m$  for at least one training sample in the subpopulation  $j$ , and  $E_3^j$  be the event of  $\ell_t^j < \ell_x^j$  for at least one training sample in the subpopulation  $j$ . Then,
$$
\begin{array}{l} \operatorname {e r r} _ {q, n} ^ {\varphi} (\mathcal {A}) = \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr (E _ {1} ^ {j} | Q _ {j}) \Pr (E _ {2} ^ {j} | Q _ {j}, E _ {1} ^ {j}) \Pr (E _ {0} | Q _ {j}, E _ {1} ^ {j}, E _ {2} ^ {j}) \\ + \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr (E _ {1} ^ {j} | Q _ {j}) \Pr (\bar {E} _ {2} ^ {j} | Q _ {j}, E _ {1} ^ {j}) \Pr (E _ {3} ^ {j} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}) \Pr (E _ {0} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}, E _ {3} ^ {j}) \\ + \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr \left(E _ {1} ^ {j} | Q _ {j}\right) \Pr \left(\bar {E} _ {2} ^ {j} | Q _ {j}, E _ {1} ^ {j}\right) \Pr \left(\bar {E} _ {3} ^ {j} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}\right) \Pr \left(E _ {0} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}, \bar {E} _ {3} ^ {j}\right) \\ \end{array}
$$
$$
+ \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr \left(\bar {E} _ {1} ^ {j} | Q _ {j}\right) \Pr \left(E _ {0} | Q _ {j}, \bar {E} _ {1} ^ {j}\right).
$$
Define  $E_0^*$  to be the event of  $M_{\mathrm{OPT}}(Z) \neq Y$  where  $M_{\mathrm{OPT}} = \mathcal{A}_{\mathrm{OPT}}(X)$ . Then, the same decomposition holds for  $\operatorname{err}_{q,n}(\mathcal{A}_{\mathrm{OPT}})$  yielding that
$$
\begin{array}{l} \mathrm {e r r} _ {q, n} (\mathcal {A} _ {\mathrm {O P T}}) = \frac {1}{N} \sum_ {j = 1} ^ {N} \mathrm {P r} (E _ {1} ^ {j} | Q _ {j}) \mathrm {P r} (E _ {2} ^ {j} | Q _ {j}, E _ {1} ^ {j}) \mathrm {P r} (E _ {0} ^ {*} | Q _ {j}, E _ {1} ^ {j}, E _ {2} ^ {j}) \\ + \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr (E _ {1} ^ {j} | Q _ {j}) \Pr (\bar {E} _ {2} ^ {j} | Q _ {j}, E _ {1} ^ {j}) \Pr (E _ {3} ^ {j} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}) \Pr (E _ {0} ^ {*} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}, E _ {3} ^ {j}) \\ + \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr (E _ {1} ^ {j} | Q _ {j}) \Pr (\bar {E} _ {2} ^ {j} | Q _ {j}, E _ {1} ^ {j}) \Pr (\bar {E} _ {3} ^ {j} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}) \Pr (E _ {0} ^ {*} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}, \bar {E} _ {3} ^ {j}) \\ + \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr (\bar {E} _ {1} ^ {j} | Q _ {j}) \Pr (E _ {0} ^ {*} | Q _ {j}, \bar {E} _ {1} ^ {j}). \\ \end{array}
$$
Since the probabilities are non-negative, by ignoring some terms, we have that
$$
\begin{array}{l} \operatorname {e r r} _ {q, n} (\mathcal {A} _ {\mathrm {O P T}}) \geq \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr (\bar {E} _ {1} ^ {j} | Q _ {j}) \Pr (E _ {0} ^ {*} | Q _ {j}, \bar {E} _ {1} ^ {j}) \\ + \frac {1}{N} \sum_ {j = 1} ^ {N} \operatorname * {P r} (E _ {1} ^ {j} | Q _ {j}) \operatorname * {P r} (\bar {E} _ {2} ^ {j} | Q _ {j}, E _ {1} ^ {j}) \operatorname * {P r} (\bar {E} _ {3} ^ {j} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}) \operatorname * {P r} (E _ {0} ^ {*} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}, \bar {E} _ {3} ^ {j}). \\ \end{array}
$$
Since these two terms  $\operatorname{Pr}(E_0^| Q_j, \bar{E}_1^j)$  and  $\operatorname{Pr}(E_0^| Q_j, E_1^j, \bar{E}_2^j, \bar{E}_3^j)$  correspond to the errors when the test data comes into the place with no information from the training data, the memorizing data does not help and the random guess achieves the best performance. Moreover, these two cases can be detected at test time by  $(\varphi, \mathcal{A})$  by memorizing the first  $\min(m, \ell_x^j)$  bits of an input sequence in the training data  $X$  and  $j$  (i.e.,  $\min(m, \ell_x^j) + 1$  bits in total) for at most one training sample per  $j \in [N]$ . Thus, we choose  $\mathcal{A}$  to memorize the first  $\min(m, \ell_x^j)$  bits and  $j$  for at most one training sample per  $j \in [N]$ ; i.e., we memorize  $\min(m, \ell_x^j) + 1 \leq m + 1$  bits for at most  $\min(N, n)$  training samples. In total, we memorise at most  $\min(N, n)(m + 1)$  bits. Then, these two terms are the same for  $\mathrm{err}_{q,n}(\mathcal{A}_{\mathrm{OPT}})$  and  $\mathrm{err}_{q,n}^{\varphi}(\mathcal{A})$  by detecting these cases (with memorization) and generating a random guess for these cases. Thus,
$$
\begin{array}{l} \operatorname {e r r} _ {q, n} \left(\mathcal {A} _ {\mathrm {O P T}}\right) \geq \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr \left(\bar {E} _ {1} ^ {j} | Q _ {j}\right) \Pr \left(E _ {0} | Q _ {j}, \bar {E} _ {1} ^ {j}\right) \\ + \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr (E _ {1} ^ {j} | Q _ {j}) \Pr (\bar {E} _ {2} ^ {j} | Q _ {j}, E _ {1} ^ {j}) \Pr (\bar {E} _ {3} ^ {j} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}) \Pr (E _ {0} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}, \bar {E} _ {3} ^ {j}). \\ \end{array}
$$
$$
\begin{array}{l} \operatorname {e r r} _ {q, n} ^ {\varphi} (\mathcal {A}) - \operatorname {e r r} _ {q, n} (\mathcal {A} _ {\mathrm {O P T}}) \\ \leq \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr (E _ {1} ^ {j} | Q _ {j}) \Pr (E _ {2} ^ {j} | Q _ {j}, E _ {1} ^ {j}) \Pr (E _ {0} | Q _ {j}, E _ {1} ^ {j}, E _ {2} ^ {j}) \\ + \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr (E _ {1} ^ {j} | Q _ {j}) \Pr (\bar {E} _ {2} ^ {j} | Q _ {j}, E _ {1} ^ {j}) \Pr (E _ {3} ^ {j} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}) \Pr (E _ {0} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}, E _ {3} ^ {j}) \\ + \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr (E _ {1} ^ {j} | Q _ {j}) \Pr (\bar {E} _ {2} ^ {j} | Q _ {j}, E _ {1} ^ {j}) \Pr (\bar {E} _ {3} ^ {j} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}) \Pr (E _ {0} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}, \bar {E} _ {3} ^ {j}) \\ \end{array}
$$
$$
\begin{array}{l} + \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr (\bar {E} _ {1} ^ {j} | Q _ {j}) \Pr (E _ {0} | Q _ {j}, \bar {E} _ {1} ^ {j}) - \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr (\bar {E} _ {1} ^ {j} | Q _ {j}) \Pr (E _ {0} | Q _ {j}, \bar {E} _ {1} ^ {j}) \\ - \frac {1}{N} \sum_ {j = 1} ^ {N} \operatorname * {P r} (E _ {1} ^ {j} | Q _ {j}) \operatorname * {P r} (\bar {E} _ {2} ^ {j} | Q _ {j}, E _ {1} ^ {j}) \operatorname * {P r} (\bar {E} _ {3} ^ {j} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}) \operatorname * {P r} (E _ {0} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}, \bar {E} _ {3} ^ {j}). \\ = \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr \left(E _ {1} ^ {j} \mid Q _ {j}\right) \Pr \left(E _ {2} ^ {j} \mid Q _ {j}, E _ {1} ^ {j}\right) \Pr \left(E _ {0} \mid Q _ {j}, E _ {1} ^ {j}, E _ {2} ^ {j}\right) \\ + \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr (E _ {1} ^ {j} | Q _ {j}) \Pr (\bar {E} _ {2} ^ {j} | Q _ {j}, E _ {1} ^ {j}) \Pr (E _ {3} ^ {j} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}) \Pr (E _ {0} | Q _ {j}, E _ {1} ^ {j}, \bar {E} _ {2} ^ {j}, E _ {3} ^ {j}) \\ \end{array}
$$
Here, we notice that  $\operatorname*{Pr}(E_0|Q_j,E_1^j,\bar{E}_2^j,E_3^j) = 0$  because  $\mathcal{A}$  memorizes the first  $\min (m,\ell_x^j)$  bits, where  $\min (m,\ell_x^j) = \ell_x^j$  conditioned on  $\bar{E}_2^j$ , and because this is conditioned on the events  $E_1^j$  and  $E_3^j$ . That is,  $\operatorname*{Pr}(E_0|Q_j,E_1^j,\bar{E}_2^j,E_3^j)$  is the error probability when we have a training sample with the length  $\ell_x^j$  larger than the test sample length  $\ell_t^j$  on the same subpopulation and  $\ell_x^j = \min (m,\ell_x^j)$  bits of the training samples are memorized. In this case, we can simply output the next symbol for the test sample from the memorized training sample, yielding  $\operatorname*{Pr}(E_0|Q_j,E_1^j,\bar{E}_2^j,E_3^j) = 0$ . Therefore,
$$
\operatorname {e r r} _ {q, n} ^ {\varphi} (\mathcal {A}) - \operatorname {e r r} _ {q, n} \left(\mathcal {A} _ {\mathrm {O P T}}\right) \leq \frac {1}{N} \sum_ {j = 1} ^ {N} \Pr \left(E _ {1} ^ {j} \mid Q _ {j}\right) \Pr \left(E _ {2} ^ {j} \mid Q _ {j}, E _ {1} ^ {j}\right) \Pr \left(E _ {0} \mid Q _ {j}, E _ {1} ^ {j}, E _ {2} ^ {j}\right). \tag {4}
$$
In the case of the events  $Q_{j}, E_{1}^{j}$  and  $E_{2}^{j}$ ,  $\mathcal{A}(X)$  memorized  $\min(m, \ell_x^j) = m$  bits of the input sequence in the training sample for at least one training sample in the subpopulation  $j$  of the test sample. In this case, at test time,  $\varphi$  first chooses one training sample in the same subpopulation  $j$  of the test sample (we can do this because the non-label part  $Z$  contains the subpopulation information  $j$  in both test sample  $(Z, Y) \sim P$  and in the training sample; see Brown et al. 4). Then,  $\varphi$  uses the memorized  $m$  bits of this training sample to search  $S$ ; i.e., it picks  $s \in S$  such that the first  $m$  bits of  $s$  matches with the memorized  $m$  bits of the training sample. If there are more than one  $s \in S$  that satisfies this condition, then it randomly picks one of them. Note that in this case of the events  $Q_{j}, E_{1}$  and  $E_{2}$ , there is at least one match; there exists  $s \in S$  the first  $m$  bits of  $s$  matches with the memorized  $m$  bits of the training sample.
Thus, if there is no more than one  $s \in S$  that satisfies this condition, then the inference algorithm  $\varphi$  uniquely determines the reference string  $c_{j}$  of the test sample  $(Z,Y) \sim P$ . With the reference string  $c_{j}$  of the test sample,  $\varphi$  then predict the next symbol of the test sample by outputting  $(\ell_t^j + 1)$ -th symbol of the reference string  $c_{j}$ , where  $\ell_t^j$  is computable by computing the length of the test input  $Z$  without any information of training samples. Therefore, if there is no more than one  $s \in S$  that satisfies the above condition, then  $\varphi$  makes a correct prediction. This implies that if  $s$  and  $s'$  do not have the same first  $m$  bits for all  $s, s' \in S$  with  $s \neq s'$ , then  $\varphi$  makes a correct prediction. By taking complement, this means that the error occurs only if there exist  $s, s' \in S$  with  $s \neq s'$  such that  $s$  and  $s'$  have the same first  $m$  bits. Therefore,
$\operatorname*{Pr}(E_0|Q_j,E_1^j,E_2^j)\leq \operatorname*{Pr}(\exists s,s'\in S$  with  $s\neq s^{\prime}$  such that  $s$  and  $s^\prime$  have the same first  $m$  bits)
$$
\begin{array}{l} = \Pr \left(\cup_ {\{s, s ^ {\prime} \} \in \mathcal {T}} \left\{s \text {a n d} s ^ {\prime} \text {h a v e t h e s a m e f i r s t} m \text {b i t s} \right\}\right) \\ \leq \sum_ {\{s, s ^ {\prime} \} \in \mathcal {T}} \Pr (\{s \text {a n d} s ^ {\prime} \text {h a v e t h e s a m e f i r s t} m \text {b i t s} \}) \\ \leq \binom {N + R} {2} \frac {1}{2 ^ {m}} = \frac {(N + R) ^ {2} - (N + R)}{2} \frac {1}{2 ^ {m}} \\ \end{array}
$$
where  $\mathcal{T} = \{\{s,s^{\prime}\} |s,s^{\prime}\in S,s\neq s^{\prime}\}$ . Since  $m = \log_2((1 - (\frac{N - 1}{N})^n)\frac{(N + R)^2 - (N + R)}{2\epsilon})$
$$
\begin{array}{l} \Pr (E _ {0} | Q _ {j}, E _ {1} ^ {j}, E _ {2} ^ {j}) \leq \frac {(N + R) ^ {2} - (N + R)}{2} \frac {1}{(1 - (\frac {N - 1}{N}) ^ {n}) \frac {(N + R) ^ {2} - (N + R)}{2 \epsilon}} \\ = \frac {\epsilon}{1 - (\frac {N - 1}{N}) ^ {n}}. \\ \end{array}
$$
Plugging this into equation 4,
$$
\mathrm {e r r} _ {q, n} ^ {\varphi} (\mathcal {A}) - \mathrm {e r r} _ {q, n} (\mathcal {A} _ {\mathrm {O P T}}) \leq \frac {1}{N} \sum_ {j = 1} ^ {N} \mathrm {P r} (E _ {1} ^ {j} | Q _ {j}) \mathrm {P r} (E _ {2} ^ {j} | Q _ {j}, E _ {1} ^ {j}) \frac {\epsilon}{1 - (\frac {N - 1}{N}) ^ {n}}.
$$
Since  $\operatorname*{Pr}(E_2^j |Q_j,E_1^j)\leq 1$  and  $\operatorname*{Pr}(E_1^j |Q_j) = 1 - \operatorname*{Pr}(\bar{E}_1^j |Q_j^j) = 1 - (\frac{N - 1}{N})^n,$
$$
\operatorname {e r r} _ {q, n} ^ {\varphi} (\mathcal {A}) - \operatorname {e r r} _ {q, n} (\mathcal {A} _ {\mathrm {O P T}}) \leq \frac {1}{N} \sum_ {j = 1} ^ {N} \epsilon = \epsilon .
$$
Since we only require  $m + 1$  bits for at most  $\min(N, n)$  training samples from the above construction of  $\mathcal{A}$  and  $\varphi$ , we have that  $I(X; \mathcal{A}(X)|P) = O(\min(N, n)m)$  while achieving  $\operatorname{err}_{q,n}^{\varphi}(\mathcal{A}) - \operatorname{err}_{q,n}(\mathcal{A}_{\mathrm{OPT}}) \leq \epsilon$  for any  $\epsilon > 0$ .
Rationale Generation As for the teacher LLM, we employ GPT-3.5-turbo (ChatGPT) [42] through the public API. We demonstrate the prompt we used for rationale generation in Table 11, 12, and 13. Specifically, we utilize the instruction and 5-shot examples from Singhal et al. [48] for MedQA-USMLE, and the chain-of-thought prompt Kojima et al. [28] for StrategyQA and Open-bookQA. We generate multiple  $l$  rationales for each training sample with the LLM. This allows training the small LM with a diverse set of rationales. Furthermore, we utilize the filtering method [17] to remove incorrect rationales from the training set but we use a small Flan-T5 base to filter the rationales which lead to incorrect prediction.
Training For all our experiments, we fine-tune the small language model for 3 epochs with a batch size of 32 using the AdamW optimizer [36] and a learning rate of  $10^{-4}$ . Each model utilizes a maximum of 96GB GPU memory with 4 NVIDIA Titan RTX GPUs for fine-tuning. In the StrategyQA and OpenbookQA experiments, we use the T5 model instead of Flan-T5 to prevent any potential data contamination with the corresponding test set, as Flan-T5 is fine-tuned on both datasets during instruction tuning. For the number of documents used for knowledge augmentation during KARD training, we set  $k = 1$  for MedQA-USMLE and StrategyQA and  $k = 3$  for OpenbookQA; specifically, we append documents retrieved from the retriever  $\rho$  along with each training sample to construct the input for training. See Tables 14 and 15 for examples of the input and output used in KARD training. For the train-test split of dataset, we use the official split for MedQA-USMLE [23] and OpenbookQA [39]. For strategyQA, we split the training set into  $7:3$  ratio to build the in-house test set following Ho et al. [17].
Inference When it comes to methods that require rationale generation, such as Chain-of-Thought and reasoning distillation, we employ a technique called self-consistency [52] during inference. Specifically, for each question, a model generates multiple rationales and corresponding predictions during the inference, followed by a majority voting to choose the final answer among the predictions.
Retriever We use Wikipedia as the external knowledge base for both all of datasets. For the retriever  $\rho$ , we use BM25 [46] which is the term-frequency-based sparse retrieval method. To implement BM25, we use the pyserini library<sup>2</sup> which provides a reproducible information retrieval framework.
Reranker To implement a neural reranker, we adopt the scoring method used in ColBERT [26]. We use BioLinkBERT-base and LinkBERT-base [56] as the backbone language model for MedQA-USMLE and StrategyQA, respectively. For reranker training, we utilize LoRA [19] for efficient training. For hyperparameters, we set  $\tau_{1} = 1$  and  $\tau_{2} = 100$ . For all datasets, we train the reranker for 3 epochs with a batch size of 16 using AdamW optimizer and a learning rate of  $10^{-4}$ . We set  $\kappa_{1} = 8, \kappa_{2} = 0$  for MedQA-USMLE and  $\kappa_{1} = 4, \kappa_{2} = 4$  for StrategyQA and OpenbookQA.
$^{2}$ https://github.com/castorini/pyserini
Table 8: Analysis on retriever.
Flan-T5 Base
Retriever	Wikipedia	Pubmed
BM25	33.14	31.58
DPR	29.77	-
Reranker	38.15	36.84
Silver (oracle)	40.30	45.48
Table 9: Analysis on  ${\kappa }_{1}$  and  ${\kappa }_{2}$  .
κ1	κ2	Flan-T5 Base
κ1 = 4	κ2 = 4	36.84
κ1 = 0	κ2 = 8	34.96
κ1 = 8	κ2 = 0	37.71
Table 10: Analysis on queries.
Query	Flan-T5 Base
BM25	Reranker
xi	31.81	35.51
rij	32.91	37.71
DAPT In Section 5.3 Experiments with DAPT, we conduct domain-adaptive pre-training [15] to assess the impact of further pre-training prior to the reasoning distillation. We further pre-train FlanT5 base model on the LM objective as discussed in the original T5 and prompt-tuning paper [45, 31]. We use two different corpus for this experiment. One is the PubMed abstracts<sup>3</sup> which contains the 1.8 abstracts extracted from the biomedical papers uploaded in PubMed. Another is MedWiki which is the subset of the Wikipedia passages containing the biomedical knowledge [35]. We further train the model for 3 epochs with a batch size of 128 using the AdamW optimizer and a learning rate of  $10^{-4}$ .
C More Analysis
C.1 Retriever and Knowledge Base
In our main experiments, we utilize BM25 [46] as the retriever and Wikipedia as the knowledge base. Regarding these choices, several questions arise: (1) Does using a dense retriever provide more advantages than a sparse retriever? (2) Does utilizing a domain-specific knowledge base offer more benefits than a general knowledge base?
To address these questions, we perform an analysis on the choice of retriever and knowledge base. First, we substitute the external knowledge base with Pubmed abstracts corpus<sup>4</sup>, which consists of abstracts from publicly available medical papers containing specialized knowledge. On the other hand, we maintain the knowledge base as Wikipedia but replace the BM25 retriever with DPR [25], one of the popular dense retrievers in open-domain question answering tasks.
In Table 8, we present the results of our analysis. We observe that the use of PubMed does not result in performance improvement compared to Wikipedia. We speculate that incorporating a reranker with PubMed corpus could potentially yield further improvements considering that using silver knowledge from PubMed yields superior performance against using silver knowledge from Wikipedia; however, we leave it as a future work. Furthermore, using DPR rather leads to degraded performance. We hypothesize that DPR may struggle to generalize to our specific task since it is primarily designed for open-domain QA tasks. It is plausible that employing a more adaptable retriever such as contriever [21] might offer greater benefits, but we also leave this as an avenue for future research.
C.2  $\kappa_{1}$  and  $\kappa_{2}$  for Reranker Training
In the training of the reranker (as described in Section 4.2 of the main paper), the reranker learns to prioritize passages that are relevant to the rationale. We normalize the score  $\rho (\pmb {d}|\pmb{r}_{ij};\mathcal{D})$  of the passages in the candidate set through softmax. Therefore, the reranker  $f_{\phi}$  learns a relative score  $f_{\phi}(\pmb {d},\pmb {x}_i)$  between passages in the candidate set during training. Due to computational constraints, the number of passages in the candidate set is limited. As a result, it is necessary to determine which passages should be included in each candidate set  $\tilde{\mathcal{D}}_{ij}$ .