# When Language Model Meets Private Library

Daoguang Zan1,2, Bei Chen3, Zeqi Lin3, Bei Guan2,4, Yongji Wang2,4,5, Jian-
Guang Lou3
1Cooperative Innovation Center, Institute of Software, Chinese Academy of
Sciences
2University of Chinese Academy of Sciences; 3Microsoft Research Asia
4Integrative Innovation Center, Institute of Software, Chinese Academy of
Sciences
5State Key Laboratory of Computer Science, Institute of Software, Chinese
Academy of Sciences
{daoguang@, guanbei<EMAIL_ADDRESS>
{beichen, zeqi.lin<EMAIL_ADDRESS>
Work done as an intern at Microsoft Research Asia.

###### Abstract

With the rapid development of pre-training techniques, a number of language
models have been pre-trained on large-scale code corpora and perform well in
code generation. In this paper, we investigate how to equip pre-trained
language models with the ability of code generation for private libraries. In
practice, it is common for programmers to write code using private libraries.
However, this is a challenge for language models since they have never seen
private APIs during training. Motivated by the fact that private libraries
usually come with elaborate API documentation, we propose a novel framework
with two modules: the APIRetriever finds useful APIs, and then the APICoder
generates code using these APIs. For APIRetriever, we present a dense
retrieval system and also design a friendly interaction to involve uses. For
APICoder, we can directly use off-the-shelf language models, or continually
pre-train the base model on a code corpus containing API information. Both
modules are trained with data from public libraries and can be generalized to
private ones. Furthermore, we craft three benchmarks for private libraries,
named TorchDataEval, MonkeyEval, and BeatNumEval. Experimental results
demonstrate the impressive performance of our framework.111Our work is
available at https://github.com/microsoft/PyCodeGPT/tree/main/apicoder..

##  1 Introduction

Code generation, automatically generating code snippets based on user
descriptions, is one of the long-standing challenges in the software
engineering and artificial intelligence communities. With the rapid
development of pre-training techniques, a number of language models are pre-
trained on large-scale code corpora and able to generate decent code snippets,
for example, Codex Chen et al. (2021), AlphaCode Li et al. (2022), CodeGen
Nijkamp et al. (2022), and InCoder Fried et al. (2022). They bring fresh
energy to code generation and improve coding efficiency Vaithilingam et al.
(2022); Chen et al. (2022). Although making remarkable progress, these models
may be biased towards generating code that is similar to the training
distribution Chen et al. (2021). What if one wants to generate code beyond the
training distribution? A real-world scenario for programmers is to write code
using a private library, which is very common in practice. For example, for
security and functionality reasons, companies often build private libraries
for internal use only. Private libraries provide a number of private APIs that
have not been seen by the language models and are also not publicly available
on any code hosting platform like GitHub. Therefore, it is worth exploring
whether and how pre-trained language models can generate code using private
libraries.

Figure 1: An example of converting PandasEval (public) to MonkeyEval
(private). The changed parts are highlighted in yellow. The performance of
Codex $12$B and CodeGen-Mono $350$M is shown at the bottom.

It is challenging for existing language models to generate code that uses a
private library directly. A practical evidence is shown in Figure 1. We built
a pseudo private library named Monkey based on a public one named Pandas.
PandasEval Zan et al. (2022a) is a benchmark consisting of $101$ Pandas
programming problems. We convert all Pandas-related keywords in PandasEval
into the new version and construct MonkeyEval (details in Section 4). As seen
in Figure 1, Codex $12$B and CodeGen-Mono $350$M show a significant drop in
performance on the private MonkeyEval compared to their performance on the
public PandasEval. For example, Codex $12$B drops from $18.88$% to $1.47$% on
pass@$1$, showing the inadequacy of the language models in code generation for
private libraries.

To meet the challenge, we propose a framework to equip pre-trained language
models with the ability to generate code that uses private libraries. As is
known, private libraries usually come with elaborate API documentation, which
motivates our main idea to mimic the process of a programmer learning to write
code using private libraries. This process is also known as API practices in
the software engineering field Snodgrass and Winnie (2019): first learning the
private API documentation and then invoking the APIs to implement the needed
functionalities. Analogically, there are two modules in our framework: an
APIRetriever first retrieves the useful APIs based on the programming problem
and the API documentation, and then an APICoder uses these APIs to generate
code. For APIRetriever, we train a dense retriever and also design a friendly
interaction to involve users in the loop optionally. For APICoder, we can
directly use existing language models of code generation, such as CodeGen, to
invoke the private APIs; furthermore, to better teach a language model how to
invoke APIs, we also continually pre-train the base model on a code corpus
containing API information from public libraries, and obtain our reinforced
model called CodeGenAPI. Since we only have access to the data of public
libraries during training, we expect that APIRetriever and APICoder can be
generalized to private libraries via learning.

To evaluate the code generation for private libraries, we craft three
benchmarks, named TorchDataEval, MonkeyEval, and BeatNumEval. TorchDataEval
includes $50$ programming problems using the TorchData library. The last two
are adapted from PandasEval and NumpyEval Zan et al. (2022a), respectively,
each consisting of $101$ programming problems. Extensive experiments on the
three benchmarks have revealed that our framework effectively improves the
performance of pre-trained language models on code generation for private
libraries. We also provide a thorough analysis to facilitate progress in this
direction.

Figure 2: The overview of our proposed framework. Figure 3: The training
process of APIRetriever and CodeGenAPI.

##  2 Framework

First, we would like to define the task of code generation formally. Given
context, the task aims to generate target code. In Figure 1, context and
target code are shown in white and grey backgrounds, respectively. Context
consists of a comment, which is a natural language description of the
programming problem, and a code snippet including import statements, function
header, etc. Target code solves the programming problem in context. We denote
the context by $\mathbf{x}$. Code generation model $\mathcal{M}$ outputs
target code $\mathbf{y}$ based on $\mathbf{x}$. For the task of code
generation for private library, the context $\mathbf{x}$ contains the
instruction for using a private library, such as an import statement. The
target code $\mathbf{y}$ contains the calls of the corresponding private
library APIs.

As mentioned in Section 1, private libraries are usually equipped with
elaborate API documentation. As a technical reference manual outlining how to
use the library, API documentation typically includes a quick start guide,
tutorials, and an instruction for each API (e.g., API name, signature,
description, parameters, and examples). To take advantage of the API
documentation, we propose to mimic the generic process of a programmer coding
with private APIs, and design a framework to generate code that can invoke
private APIs. The framework consists of APIRetriever and APICoder with the
overview shown in Figure 2. Given the context, APIRetriever
$\mathcal{M}_{\rm{R}}$ aims to retrieve possible used APIs from the API
documentation; and APICoder $\mathcal{M}_{\rm{C}}$ is dedicated to generating
code using the retrieved APIs. The process can be formalized as
$\mathcal{A}=\mathcal{M}_{\rm{R}}(\mathbf{x})$ and
$\mathbf{y}=\mathcal{M}_{\rm{C}}(\mathcal{A};\mathbf{x})$, where $\mathcal{A}$
represents the set of information of all proper APIs, and each
$\mathbf{a}\in\mathcal{A}$ is the information of an API. In our
implementation, we design the API information to include the API name,
signature and description. Note that we only use the first sentence of the API
description since it is sufficient to summarize.

##  3 Methodology

We have introduced our framework that provides pre-trained models a fantastic
way to deal with private libraries. In this section, we present the data
collection, followed by the detailed design of our APIRetriever and APICoder.

###  3.1 Data Collection

We collect API information and code files of public libraries due to the fact
that we can only access data from public libraries. Then we train the models
based on the public data with the expectation that the model can be
generalized to private libraries. For API information, we consider the $31$
most popular public libraries in Python (e.g., Pandas, NumPy, and scikit-
learn) according to the popularity ranking on
StackOverFlow222https://stackoverflow.com/tags?tab=popular. For each of the
libraries, we crawled its API documentation and extracted detailed information
about each API, including the API name, signature, description, parameters,
usage examples, and so on. Please refer to Appendix A for the details of the
31 public libraries. For code files, we first collected a $330$GB corpus from
GitHub containing $60.6$M python files and then extracted those files that
involved one or more API calls from the $31$ public libraries. After a bunch
of pre-processing strategies, for example, de-duplicating, cleaning, and
formatting, we finally obtained $4.54$M python files, denoted by
$\mathcal{D}$.

###  3.2 APIRetriever

APIRetriever aims to find the proper APIs based on the description of a
programming problem. We regard it as a dense retrieval task Qu et al. (2020);
Xiong et al. (2020); Santhanam et al. (2021); Zan et al. (2022b); Formal et
al. (2022) and design a simple dual-encoder model Karpukhin et al. (2020) to
retrieve the possible used APIs for each programming problem. To further boost
the retrieval performance, a friendly interaction approach is designed to
involve users.

#### Training.

To train APIRetriever, we need a large amount of pairwise data, natural
language description and API information. We first segment each python file
$\mathbf{d}\in\mathcal{D}$ into $K$ code blocks $(d_{1},d_{2},{\cdots},d_{K})$
using the pip-tools, i.e., redbaron, autopep8, and docformatter, where each
code block is a relatively well-rounded code fragment, such as a function or a
class. For each code block $d_{i}$, we extract all API names and obtain the
corresponding API signatures and descriptions by searching our collected $31$
API documentations333If an API name matches more than one candidate API, we
randomly pick one.. The information of an API includes its name, signature and
description, denoted by $\mathbf{a}\in\mathcal{A}$. Each $\mathbf{a}$ and the
natural language description $\mathbf{p}$ extracted from the same code block
$d_{i}$ are regarded as a positive training sample. For the negative training
sample, we randomly sample an API $\hat{\mathbf{a}}$ that is unrelated to
$d_{i}$ from the same library. In total, we obtained $40.3$M ($\mathbf{p}$,
$\mathbf{a}$, $\hat{\mathbf{a}}_{1}$, $\hat{\mathbf{a}}_{2}$, …) sets as
training samples. As in Figure 3, the left part shows the training process of
APIRetriever. Our APIRetriever is a dual-encoder model. The two dense encoder,
$E_{\mathbf{p}}(.)$ and $E_{\mathbf{a}}(.)$, map $\mathbf{p}$ and $\mathbf{a}$
to $z$-dimensional vectors, respectively. Then, we use the dot product of
their vectors to calculate the similarity score formalized as
$E_{\mathbf{p}}(\mathbf{p})^{\top}E_{\mathbf{a}}(\mathbf{a})$, where
$E_{\mathbf{p}}(.)$ and $E_{\mathbf{a}}(.)$ are implemented by two independent
BERT Devlin et al. (2019) with base-uncased version. We use BERT instead of
CodeBERT Feng et al. (2020) as most tokens in $\mathbf{p}$ and $\mathbf{a}$
are natural language rather than programming language.

#### Inference.

After the training phase, we can use APIRetriever to retrieve private APIs for
each programming problem description. In detail, we apply $E_{\mathbf{a}}$ to
all the APIs and index them by FAISS Johnson et al. (2019) offline. Given a
new programming problem description $\mathbf{p}$ at run-time, we only need to
produce its embedding $v_{\mathbf{p}}=E_{\mathbf{p}}(\mathbf{p})$ and recall
the top-$k$ APIs with the embeddings closest to $v_{\mathbf{p}}$.

Figure 4: Friendly interaction interface for users.

#### Human Interaction with APIRetriever.

In order to further increase the accuracy of API retrieval, we provide a
friendly interaction interface to allow humans in the loop with APIRetriever,
as shown in Figure 4. In the interaction interface, we give the programming
problem and the top-$5$ APIs retrieved by APIRetriever, and let users choose
one or more APIs that may be used in the target code. Note that we only
provide API names and descriptions to users, as we find in our empirical
experiments that providing API signatures has a negative effect on making the
correct choice.

###  3.3 APICoder

APIRetriever finds useful APIs for a programming problem, and then APICoder
aims to generate code that solves the problem with these APIs. We make use of
the most straightforward way for APICoder: prompting API information set
$\mathcal{A}$ in front of the context $\mathbf{x}$. Formally, the APICoder can
be written as
$\mathbf{y}=\mathcal{M}_{\rm{C}}({\rm{Concat}}(\mathcal{A},\mathbf{x}))$,
where ${\rm{Concat}}(\mathcal{A},\mathbf{x})$ means to concatenate the API
information set and the context. Examples can be found in Figure 3. Each API
information is in the form of “name(signature):description”. This is to mimic
programmers learning the APIs properly before writing code using them.

Technically speaking, the off-the-shelf code generation models, such as
CodeT5, CodeGPT, CodeClippy, CodeParrot, CodeGen, and Codex, can be applied
directly to land APICoder. Although these base models can achieve gains in
correctly invoking APIs, they have not learned how to use them as an explicit
training task. To better use the APIs, we devised a fantastic idea of
continually pre-training the base models using code files with API information
inserted.

In practice, we use CodeGen-Mono $350$M Nijkamp et al. (2022) as our base
model, based on which we continually pre-train and obtain our reinforced model
called CodeGenAPI. CodeGen is a GPT-based model skilled at generating code. We
choose it because it is by far the most popular and publicly available model.
As for the training corpus, we use the collected python files $\mathcal{D}$
mentioned in Section 3.1. Firstly, as done for APIRetriever, each file
$\mathbf{d}\in\mathcal{D}$ is split into $K$ code blocks
$(d_{1},d_{2},{\cdots},d_{K})$. For each code block $d_{i}$, we obtain the set
of API information $\mathcal{A}_{i}$. Then, the $K$ code blocks and sets of
API information are cross-merged to output a new file
$\hat{\mathbf{d}}=(\mathcal{A}_{1},d_{1},\mathcal{A}_{2},d_{2},{\cdots},\mathcal{A}_{K},d_{K})$.
This mimics the process of API information as a prompt for each block. Then,
we continually pre-train the base model on the new code files, teaching the
model to write code based on the prompted APIs. In addition, as shown in
Figure 3, to make APICoder more robust, we shuffle the APIs in each set
$\mathcal{A}_{i}$ and also add noise APIs, since APIRetriever does not know
the order of APIs in the target code and often finds incorrect APIs.

During the training phase of CodeGenAPI, unlike the previous settings that
force all files to have the same priority, we design a resampling strategy to
enable high-quality python files to appear more frequently and vice versa. The
strategy considers the star number of the repository, the unit test function
rate of the code file, and the API rate of the code file. More details can be
found in Appendix B.

##  4 Benchmark Construction

Private libraries are commonly used in practice, but few attempts have been
made to evaluate the performance of generating code invoking private APIs. To
fill this gap, we craft three benchmarks, called TorchDataEval, MonkeyEval,
and BeatNumEval. Each programming problem consists of context, target code,
and the corresponding test cases.

To create a realistic benchmark for evaluating code generation for private
library, we use TorchData, a Python library released just recently444Our base
model, CodeGen, is pre-trained with GitHub data before $2021$-$10$. TorchData
was released after this time point and no code files using it are available on
GitHub so far; hence we can consider it as a private library.. We carefully
learnt the official API documentation of TorchData and made sure we were
proficient in all APIs. Then, we manually created $50$ programming problems
based on the API usage examples in the documentation. Two volunteers with
extensive experience in Python were invited to check the correctness of each
problem. We control the difficulty of the programming problems by the number
of APIs in the target code. The percentage of programming problems containing
$1$ API, $2$ APIs, and more APIs is set to $6$:$3$:$1$.

We also construct two benchmarks using pseudo private libraries, named
MonkeyEval and BeatNumEval, each containing $101$ programming problems. They
are modified from PandasEval and NumpyEval, which were proposed for the public
libraries Pandas and Numpy Zan et al. (2022a). In detail, we manually modified
all library-related keywords in PandasEval and NumpyEval. For example, as
shown in Figure 1, pandas is converted to monkey, dataframe is converted to
knowledgeframe, and the API name isin is converted to iscontain. For more
details on keyword conversion, please refer to Appendix C. To craft the API
documentation for Monkey and BeatNum, we manually paraphrased the descriptions
of all the new APIs to ensure that the pre-trained language models have never
seen them.

##  5 Experiments

In this section, we conduct experiments to illustrate the superiority of our
proposed framework.

###  5.1 Experimental Setup

#### API Information.

As shown in the second column (APIs) in Table 1, there are four settings for
prompting API information before the context:

  * •

No API: there is nothing to be prompted;

  * •

Perfect: the information of golden APIs in the target code is prompted;

  * •

Top-$N$: the information of top $N$ APIs retrieved by APIRetriever is
prompted, where $N\in\\{1,2,3,5\\}$;

  * •

Human: the information of the APIs chosen by users is prompted. In our
experiments, we invited three volunteers who are programmers familiar with
Python but without any background in our three benchmarks. As in Figure 4,
they interacted with the APIRetriever and provided their choices for all
programming problems. The final APIs are determined by voting on their
choices.

#### Baselines.

Our contributions can be reviewed in terms of both APIRetriever and APICoder.
For APIRetriever, all models in the No API setting are our baseline, while we
propose the Perfect, Top-$N$, and Human settings. For APICoder, the main
baseline is our base model, CodeGen-Mono $350$M Nijkamp et al. (2022), in the
same API information setting. We use CodeGen for short in the following. In
addition, we include advanced pre-trained code generation models that are
comparable in parameter size: CodeT5 Wang et al. (2021), CodeGPT Lu et al.
(2021), CodeClippy555https://github.com/CodedotAl/gpt-code-clippy and
CodeParrot666https://huggingface.co/lvwerra/codeparrot. Codex $12$B Chen et
al. (2021) is also used to show the performance of giant models.

|  APIs | TorchDataEval | MonkeyEval | BeatNumEval
---|---|---|---|---
3-11 APICoder | pass@$1$ | pass@$10$ | pass@$100$ | pass@$1$ | pass@$10$ | pass@$100$ | pass@$1$ | pass@$10$ | pass@$100$
CodeT5 220M | Top-2 | 0.00 | 0.00 | 0.00 | 0.00 | 0.00 | 0.00 | 0.00 | 0.00 | 0.00
CodeGPT 124M | Top-2 | 0.67 | 2.78 | 7.72 | 0.82 | 0.99 | 1.73 | 0.52 | 1.88 | 4.70
CodeClippy 125M | Top-2 | 0.04 | 0.39 | 2.75 | 0.10 | 0.76 | 1.86 | 0.03 | 0.33 | 2.11
| No API | 4.04 | 7.11 | 13.26 | 0.54 | 2.04 | 7.38 | 2.67 | 7.66 | 18.86
| Perfect | 4.86+0.82 | 8.88+1.77 | 17.25+3.99 | 2.39+1.85 | 3.33+1.29 | 9.99+2.61 | 5.01+2.34 | 11.30+3.64 | 26.36+7.5
| Top-1 | 4.02-0.02 | 8.35+1.24 | 18.17+4.91 | 2.54+2.00 | 3.43+1.39 | 11.39+4.01 | 4.32+1.65 | 9.39+1.73 | 19.91+1.05
| Top-2 | 4.64+0.60 | 8.96+1.85 | 17.48+4.22 | 1.52+0.98 | 2.96+0.92 | 9.32+1.94 | 2.77+0.10 | 8.57+0.91 | 19.74+0.88
| Top-3 | 4.00-0.04 | 7.51+0.40 | 15.13+1.87 | 1.32+0.78 | 3.16+1.12 | 10.59+3.21 | 1.69-0.98 | 9.01+1.35 | 19.90+1.04
| Top-5 | 4.22+0.18 | 7.51+0.40 | 15.43+2.17 | 0.99+0.45 | 2.78+0.74 | 11.76+4.38 | 1.74-0.93 | 8.11+0.45 | 17.54-1.32
CodeParrot 110M | Human | 4.01-0.03 | 7.60+0.49 | 14.47+1.21 | 2.44+1.90 | 3.62+1.58 | 9.83+2.45 | 5.23+2.56 | 11.78+4.12 | 22.81+3.95
| No API | 6.72 | 15.71 | 22.00 | 0.95 | 4.90 | 8.89 | 5.15 | 11.96 | 18.79
| Perfect | 9.84+3.12 | 22.62+6.91 | 34.00+12.00 | 2.14+1.19 | 6.41+1.51 | 11.86+2.97 | 9.47+4.32 | 17.05+5.09 | 28.67+9.88
| Top-1 | 8.72+2.00 | 19.22+3.51 | 27.97+5.97 | 2.22+1.27 | 7.20+2.30 | 12.85+3.96 | 7.52+2.37 | 15.25+3.29 | 24.71+5.92
| Top-2 | 7.52+0.80 | 16.36+0.65 | 26.00+4.00 | 2.46+1.51 | 6.35+1.45 | 9.89+1.00 | 6.65+1.50 | 13.68+1.72 | 22.74+3.95
| Top-3 | 7.92+1.20 | 18.65+2.94 | 28.00+6.00 | 2.02+1.07 | 5.26+0.36 | 8.89+0.00 | 6.26+1.11 | 16.12+4.16 | 24.72+5.93
| Top-5 | 6.08-0.64 | 17.48+1.77 | 25.95+3.95 | 1.58+0.63 | 5.45+0.55 | 9.88+0.99 | 6.34+1.19 | 15.05+3.09 | 21.76+2.97
CodeGen 350M | Human | 8.08+1.36 | 19.85+4.14 | 31.95+9.95 | 2.14+1.19 | 6.14+1.24 | 11.86+2.97 | 9.47+4.32 | 17.12+5.06 | 28.67+9.88
| No API | 7.19 | 16.93 | 23.97 | 1.19 | 4.68 | 7.91 | 4.44 | 8.24 | 13.83
| Perfect | 20.23+13.04 | 33.37+16.44 | 41.97+18.00 | 4.59+3.40 | 9.14+4.46 | 13.85+5.94 | 9.62+5.18 | 16.51+8.27 | 22.75+8.92
| Top-1 | 12.89+5.70 | 24.26+7.33 | 31.97+8.00 | 2.89+1.70 | 8.28+3.60 | 12.86+4.94 | 6.61+2.17 | 12.62+4.38 | 17.80+3.97
| Top-2 | 10.41+3.22 | 23.50+6.57 | 31.98+8.01 | 3.41+2.22 | 8.33+3.65 | 11.87+8.90 | 5.90+1.46 | 11.79+3.55 | 15.83+2.00
| Top-3 | 10.49+3.30 | 25.45+8.52 | 35.98+12.01 | 3.17+1.98 | 7.51+2.83 | 10.88+2.97 | 5.11+0.67 | 11.40+3.16 | 15.82+1.99
| Top-5 | 10.34+3.15 | 23.04+6.11 | 27.99+4.02 | 1.94+0.75 | 4.75+0.07 | 7.91+0.00 | 5.07+0.63 | 9.64+1.40 | 13.84+0.01
CodeGenAPI 350M | Human | 15.57+8.38 | 27.76+10.83 | 33.97+10.00 | 3.76+2.57 | 8.32+3.64 | 12.86+4.95 | 9.39+4.95 | 16.40+8.16 | 23.74+9.91
| No API | 7.16 | 14.46 | 23.75 | 1.47 | 3.53 | 7.31 | 6.95 | 17.54 | 25.57
| Perfect | 25.03+17.87 | 51.26+36.80 | 56.75+33.00 | 3.58+2.11 | 7.48+3.95 | 12.61+5.30 | 8.59+1.64 | 23.75+6.21 | 36.99+11.42
Codex 12B | Top-2 | 17.98+10.82 | 32.75+18.29 | 41.51+17.76 | 1.92+0.45 | 5.91+2.38 | 11.08+3.77 | 9.54+2.59 | 21.77+4.23 | 32.45+6.88

Table 1: Pass$@k$($\%$) results on the three benchmarks. The blue background
means no API as extra prompt; the yellow background means perfect APIs as
extra prompt; the write background means top-$1,2,3,$ or $5$ APIs retrieved by
APIRetriever as extra prompt; and the purple background means the APIs chosen
by human from top-$5$ of APIRetriever as extra prompt. Numbers in red and
green indicate the absolute changes over no API setting.

#### Evaluation Metrics.

Followed by Chen et al. (2021), we regard pass$@k$ as our metric. For each
programming problem, we sample $n\geq k$ code snippets, and then count the
number of correct ones $c$, where passing all test cases is considered as
correct. If $n-c<k$, then pass$@k$ equals $1$; otherwise, equals
$1-\prod\nolimits_{i=n-c+1}^{n}(1-\frac{k}{i})$. In our experiments, $k$ is
set to one of $[1,10,100]$ and $n$ is set to $200$.

#### Implementation Details.

We implement our approach based on PyTorch Paszke et al. (2019) and
Huggingface’s transformers Wolf et al. (2019). We use a dense retrieval
toolkit777https://github.com/luyug/Dense to train APIRetriever by setting the
batch size to $10$ per device, the learning rate to $1$e-$5$, the ratio of
positive vs. negative samples to $1$:$8$, and the vector dimensions $z$ of
$\mathbf{p}$ and $\mathbf{a}$ to $768$. The model uses cross-entropy as the
loss function and Adam Kingma and Ba (2014) as the parameters optimizer. It is
trained for $100$K steps on a cluster of $8$ NVIDIA V$100$ GPUs with $32$GB
memory. The training time is about $3$ days. For pre-training CodeGenAPI, we
set the code block size to $1,024$, the batch size to $4$, the learning rate
to $5$e-$4$, the gradient accumulation steps to $4$, the weight decay to
$0.1$, and the warm up steps to $1,000$. Noise APIs are added at a rate of
$0.05$. It is trained for $100$K steps about $1.6$ days on $8$ $32$GB NVIDIA
V$100$ GPUs. In all of our training phases, we use mixed precision FP$16$ to
speed up. When generating code snippets using pre-trained models, we conduct
various temperatures ranging from $0.1$ to $1.0$ with the interval of $0.1$.
All results are reported with the best values across these hyper-parameters.

###  5.2 Main Results

Table 1 summarizes the performance of our framework and all baselines on
TorchDataEval, MonkeyEval, and BeatNumEval. Based on numerous experimental
results, we derived plausible observations and valuable insights to answer the
following research questions.

_“Is API information useful for private library oriented code generation?”_ As
we can see in Table 1, all models without prompting any APIs (the No API
setting) achieve relatively poor performance on all benchmarks. Especially,
Codex $12$B, a powerful code generation model with large parameters, can only
achieve similar performance to CodeGen and CodeGenAPI $350$M in the No API
setting. This indicates that even with gigantic models, the task of code
generation with private libraries is extremely challenging. Encouragingly,
with prompted API information (the Perfect, Top-$N$, Human settings), both the
off-the-shelf models (e.g., CodeParrot, CodeGen, and Codex) and our
continually pre-trained CodeGenAPI achieve consistent performance gains
compared to those in the No API setting. Moreover, the more powerful the model
itself in code generation (i.e., Codex $12$B > CodeGen $350$M > CodeParrot
$110$M), the more benefits that API information can bring. For example, on
TorchDataEval in the Perfect setting, Codex $12$B brings pass@$10$ an absolute
improvement of $36.89\%$, while CodeParrot $110$M only improves $1.77\%$. This
observation also suggests that prompting API information can unleash the
potential of gigantic models towards invoking private APIs. All the above
results prove the usefulness of API information for code generation for
private libraries.

_“Is the APIRetriever effective in finding useful API information?”_ All
models in the Top-$N$ setting outperform the same models in the No API
setting, suggesting that APIRetriever is able to find useful APIs. For a
certain model, we observe that the Top-$1$/Top-$2$ settings usually perform
better than the Top-$3$/Top-$5$ settings due to the fact that the latter
introduces more noise APIs to the APICoder. In addition, involving humans (the
Human setting) in the selection of APIs can further improve performance,
suggesting the effectiveness of the human interaction we designed. Note that
the Top-$N$ and Human settings are occasionally superior to the Perfect
setting, which is reasonable because the noise APIs exist when training the
model.

Figure 5: The recall rates of retrieved APIs.

_“Is the APICoder effective in invoking private APIs?”_ As shown in Table 1,
off-the-shelf models like CodeGen are capable of handling private library
invocations. To seek more extraordinary performance, we continually pre-train
CodeGen and obtain a new model CodeGenAPI. We can observe that CodeGenAPI
consistently outperforms its base model CodeGen on TorchDataEval and
MonkeyEval, which proves the effectiveness of CodeGenAPI. However, on
BeatNumEval, CodeGenAPI is inferior to CodeGen. After careful troubleshooting,
we reveal that the process of continual pre-training aims to essentially learn
how to invoke the correct APIs with maximum likelihood, while the key obstacle
to using BeatNum modified from Numpy lies in the numerical calculation like
‘a[:,None]+b*2’, instead of invoking the correct APIs. Therefore, CodeGenAPI
fails to yield benefits for BeatNumEval. Overall, APICoder has the capability
to invoke private APIs.

APICoder | TorchDataEval | MonkeyEval | BeatNumEval
---|---|---|---
2-10 | pass@1 | pass@10 | pass@100 | pass@1 | pass@10 | pass@100 | pass@1 | pass@10 | pass@100
CodeGenAPI | 10.41 | 23.50 | 31.98 | 3.41 | 8.33 | 11.87 | 5.90 | 11.79 | 15.83
-w/ noise rate $0$% | 9.41 | 22.88 | 31.08 | 2.69 | 8.03 | 11.18 | 5.77 | 11.01 | 14.52
-w/ noise rate $10$% | 9.19 | 22.87 | 30.98 | 3.04 | 7.67 | 11.10 | 4.99 | 10.80 | 15.18
-w/ noise rate $20$% | 8.92 | 23.04 | 30.57 | 2.00 | 7.39 | 10.64 | 4.48 | 10.97 | 13.41
-w/o resampling | 8.65 | 21.00 | 29.71 | 2.47 | 7.96 | 10.13 | 5.21 | 8.68 | 14.75

Table 2: Ablation studies for CodeGenAPI in the Top-2 setting (top 2 APIs
provided by APIRetriever are prompted). The default setting of CodeGenAPI is
to use the resampling strategy and a noise rate of $5$%.

###  5.3 Closer Analysis

We have demonstrated the effectiveness of our framework. In this subsection,
we provide several closer analyses to inspire future work in this direction.

#### Quality of Retrieved APIs.

Retrieving the correct APIs as prompts can enhance the code generation
performance for private libraries, so we would like to evaluate the
effectiveness of APIRetriever. Figure 5 shows the recall rate of APIRetriever
on five benchmarks. We can see that the recall rates of top-$5$ are already
high, demonstrating that it is reasonable to provide $5$ API candidates for
users to choose from. Furthermore, as shown in Figure 6, we analyze the
accuracy of APIs chosen by users. We observe that it dramatically exceeds the
accuracy of top $1,2$ or $3$ APIs retrieved by APIRetriever. This suggests
that it is feasible to involve humans in the retrieval of APIs.

Figure 6: Accuracy of retrieved APIs. Figure 7: Accuracy of CodeGenAPI and
CodeGen with respect to the number of APIs. The problem is solved if one of
$200$ samples passes all test cases.

#### Different Difficulty.

We would like to explore the performance of CodeGenAPI on varying difficulty
problems. So we calculate its accuracy across various numbers of APIs in
target code $\mathbf{y}$. Each benchmark is divided into $3$ parts, according
to the number of APIs. Figure 7 shows that CodeGenAPI outperforms CodeGen by a
large margin on the problems containing only one API. The trend still holds as
the number of APIs increases. It demonstrates CodeGenAPI can boost the
performance of generating code snippets using private library on varying
difficulty.

#### Noise Rate.

A well-chosen noise rate can improve the robustness of CodeGenAPI against a
variety of APIs. If we set the noise rate too large, it may change the
original distribution of the code corpus, while too small will lose the
capability to deal with noise APIs. The default noise rate is $5$%, and we
also try $0$%, $10$%, and $20$%. As shown in Table 2, both too large and too
small noise rates can degrade the performance.

#### Resampling Strategy.

Making high-quality python files high-priority, and vice versa, is in line
with our intuition. To demonstrate it, we remove the resampling strategy as
mentioned in Section 3.3. As shown in Table 2, we observe a steady decline in
performance on the three benchmarks. Such an observation demonstrates the
effectiveness of the sampling strategy.

####  CodeGenAPI for Public Library.

Technically speaking, CodeGenAPI also can be employed for generating code for
public libraries. So, we do experiments on PandasEval and NumpyEval and show
the results in Table 3. We find that the performance improvement of CodeGenAPI
over the base model on public libraries is not as significant as on private
libraries. One major reason is that the models have seen the public libraries
during pre-training, so prompting API information yields limited benefit. We
can see CodeGenAPI excels over CodeGen when prompting perfect APIs. But when
prompting top-2 APIs, the advantages of CodeGenAPI are not exhibited. This
means that CodeGenAPI can also work on third-party public libraries, but it
depends heavily on the performance of APIRetriever.

APICoder |  APIs | PandasEval
---|---|---
3-5 | pass@1 | pass@10 | pass@100
CodeGen | No API | 14.24 | 30.71 | 46.04
Perfect | 11.21 | 33.59 | 48.47
Top-2 | 9.54 | 29.02 | 40.56
CodeGenAPI | No API | 13.58 | 34.95 | 46.51
Perfect | 19.96 | 42.36 | 53.43
Top-2 | 11.25 | 28.61 | 39.48
|  | NumpyEval
CodeGen | No API | 19.31 | 40.89 | 60.58
Perfect | 21.41 | 41.08 | 56.38
Top-2 | 18.30 | 35.12 | 48.46
CodeGenAPI | No API | 16.55 | 29.48 | 42.52
Perfect | 24.83 | 41.47 | 54.41
Top-2 | 12.67 | 27.32 | 35.62

Table 3: Results of CodeGen and CodeGenAPI on PandasEval and NumpyEval.

##  6 Related Work

###  6.1 Code Generation

Thanks to the recent development of pre-training techniques Devlin et al.
(2019); Raffel et al. (2020); Lewis et al. (2020); Yan et al. (2021); Brown et
al. (2021), a lot of pre-trained language models have been proposed for code-
related tasks. For example, CuBERT Kanade et al. (2020), CodeBERT Feng et al.
(2020), GraphCodeBERT Guo et al. (2020), CodeT5 Wang et al. (2021), CodeGPT Lu
et al. (2021), PLBART Ahmad et al. (2021), PyCodeGPT Zan et al. (2022a),
CodeGen Nijkamp et al. (2022), Codex Chen et al. (2021), AlphaCode Li et al.
(2022), and InCoder Fried et al. (2022). Almost all of them focus on
standalone code, while JigSaw Jain et al. (2021) and CERT Zan et al. (2022a)
are presented for generating code using public libraries. In this paper, we
aim to generate code invoking private APIs, which is a common scenario in
practice. It is more challenging because pre-trained language models have
never seen any information about private libraries. As for benchmarks,
HumanEval Chen et al. (2021), APPs Hendrycks et al. (2021), P3 Schuster et al.
(2021), MBPP Austin et al. (2021), BIG-bench Srivastava et al. (2022), and
CodeContests Li et al. (2022) were proposed to evaluate the performance of
generating standalone code. GSM8K-Python Cobbe et al. (2021) and MathQA-Python
Austin et al. (2021) were engaged in evaluating the capability of solving
mathematical problems. PandasEval and NumpyEval Zan et al. (2022a) were
released to evaluate the code generation for public library. We propose three
benchmarks, called TorchDataEval, MonkeyEval, and BeatNumEval, aiming to
evaluate the performance of code generation for private library.

###  6.2 Retrieval-Based Generation

In the natural language field, retrieval-based generation is a hot topic. A
lot of works Izacard and Grave (2020); Karpukhin et al. (2020); Qu et al.
(2020); Xiong et al. (2020); Santhanam et al. (2021); Zan et al. (2022b);
Formal et al. (2022) have emerged under this topic. Therefore, we refer to the
above methods and design our APIRetriever for private API retrieval. In the
programming language field, there are also several attempts to use retrieval
techniques, such as DeepAPI Gu et al. (2016), REDCODER Parvez et al. (2021),
ReACC Lu et al. (2022), and DocCoder Zhou et al. (2022). Our work is
fundamentally different from them. They all aim to retrieve public code
snippets or other resources on GitHub/StackOverFlow based on the user query,
while our goal is to retrieve APIs from the API documentation of private
library based on code comments. Besides, we design retrieval because we focus
on private APIs, which have not been seen by the pre-trained generative
language models.

##  7 Conclusion

In this paper, we propose a novel framework for code generation for private
library. There are two modules: for a specific programming problem,
APIRetriever first finds out the useful private APIs from API documentation,
and then APICoder leverages these APIs to generate the code. We craft three
benchmarks, including TorchDataEval, MonkeyEval, and BeatNumEval, for better
evaluating private library oriented code generation. The experimental results
and thorough analysis demonstrate the reasonableness and effectiveness of our
framework. In future work, we would like to explore how to make better use of
API documentation for code generation and improve the approach for real use
when programming with private libraries.

## Limitations

While our proposed approach exhibits many advantages, it also has a few
limitations. (1) As stated in Section 5.2, our approach that prompts APIs for
programming problem relies heavily upon the code generation capacity of the
language model itself. The more powerful the model itself, the more benefits
the prompting APIs bring. Likewise, we also find that if a model itself shows
very poor performance, prompting APIs will not bring any benefit to it or even
bring negative effects. (2) As the first navigator to explore code generation
with private library, we have built three private libraries, but they all
include a relatively small number of APIs (<$200$). With these APIs, our
APIRetriever can exhibit decent performance. But we surmise that it may become
more challenging for APIRetriever as the number of APIs increases. (3) It is
extremely challenging to find a real private library and craft a benchmark
like TorchDataEval. To evaluate our idea quickly and cost-effectively, besides
TorchDataEval, we also crafted two pseudo private libraries that are modified
from the existing public ones as mentioned in Section 4. Although we have done
our best to preserve the two pseudo private libraries in line with the real
private library, it may still pose some threats to the fair evaluation of code
generation for private library. (4) We can see from Table 1 that most models
with the Top-N setting fall behind the same model with the Perfect setting.
Such observation demonstrates that APIRetriever we designed has a big room for
improvement. (5) Our experiments show that our framework can enhance the
quality of private library oriented code generation on Python. Limitations may
exist when we generalize it to other programming languages such as Java, C,
and C++ since the characteristics of libraries for different programming
languages are slightly different.

## References

  * Ahmad et al. (2021) Wasi Ahmad, Saikat Chakraborty, Baishakhi Ray, and Kai-Wei Chang. 2021.  Unified pre-training for program understanding and generation.  In _NAACL_ , pages 2655–2668.
  * Austin et al. (2021) Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. 2021.  Program synthesis with large language models.  _arXiv preprint arXiv:2108.07732_.
  * Brown et al. (2021) Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2021.  Language models are few-shot learners.  In _Neural Information Processing Systems_.
  * Chen et al. (2022) Bei Chen, Fengji Zhang, Anh Nguyen, Daoguang Zan, Zeqi Lin, Jian-Guang Lou, and Weizhu Chen. 2022.  CodeT: Code generation with generated tests.  _arXiv preprint arXiv:2207.10397_.
  * Chen et al. (2021) Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. 2021.  Evaluating large language models trained on code.  _arXiv preprint arXiv:2107.03374_.
  * Cobbe et al. (2021) Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. 2021.  Training verifiers to solve math word problems.  _arXiv preprint arXiv:2110.14168_.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019.  BERT: Pre-training of deep bidirectional transformers for language understanding.  In _NAACL_ , pages 4171–4186.
  * Feng et al. (2020) Zhangyin Feng, Daya Guo, Duyu Tang, Nan Duan, Xiaocheng Feng, Ming Gong, Linjun Shou, Bing Qin, Ting Liu, Daxin Jiang, et al. 2020.  CodeBERT: A pre-trained model for programming and natural languages.  In _Findings of EMNLP_ , pages 1536–1547.
  * Formal et al. (2022) Thibault Formal, Carlos Lassance, Benjamin Piwowarski, and Stéphane Clinchant. 2022.  From distillation to hard negative sampling: Making sparse neural ir models more effective.  _arXiv preprint arXiv:2205.04733_.
  * Fried et al. (2022) Daniel Fried, Armen Aghajanyan, Jessy Lin, Sida Wang, Eric Wallace, Freda Shi, Ruiqi Zhong, Wen-tau Yih, Luke Zettlemoyer, and Mike Lewis. 2022.  InCoder: A generative model for code infilling and synthesis.  _arXiv preprint arXiv:2204.05999_.
  * Gu et al. (2016) Xiaodong Gu, Hongyu Zhang, Dongmei Zhang, and Sunghun Kim. 2016.  Deep api learning.  In _ACM SIGSOFT International Symposium on Foundations of Software Engineering_ , pages 631–642.
  * Guo et al. (2020) Daya Guo, Shuo Ren, Shuai Lu, Zhangyin Feng, Duyu Tang, et al. 2020.  GraphCodeBERT: Pre-training code representations with data flow.  In _International Conference on Learning Representations_.
  * Hendrycks et al. (2021) Dan Hendrycks, Steven Basart, Saurav Kadavath, Mantas Mazeika, Akul Arora, Ethan Guo, et al. 2021.  Measuring coding challenge competence with apps.  In _Neural Information Processing Systems Datasets and Benchmarks Track_.
  * Izacard and Grave (2020) Gautier Izacard and Edouard Grave. 2020.  Leveraging passage retrieval with generative models for open domain question answering.  _arXiv preprint arXiv:2007.01282_.
  * Jain et al. (2021) Naman Jain, Skanda Vaidyanath, Arun Iyer, Nagarajan Natarajan, Suresh Parthasarathy, Sriram Rajamani, and Rahul Sharma. 2021.  Jigsaw: Large language models meet program synthesis.  _arXiv preprint arXiv:2112.02969_.
  * Johnson et al. (2019) Jeff Johnson, Matthijs Douze, and Hervé Jégou. 2019.  Billion-scale similarity search with gpus.  _IEEE Transactions on Big Data_ , 7(3):535–547.
  * Kanade et al. (2020) Aditya Kanade, Petros Maniatis, Gogul Balakrishnan, and Kensen Shi. 2020.  Learning and evaluating contextual embedding of source code.  In _International Conference on Machine Learning_ , pages 5110–5121.
  * Karpukhin et al. (2020) Vladimir Karpukhin, Barlas Oğuz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. 2020.  Dense passage retrieval for open-domain question answering.  _arXiv preprint arXiv:2004.04906_.
  * Kingma and Ba (2014) Diederik P Kingma and Jimmy Ba. 2014.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_.
  * Lewis et al. (2020) Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. 2020.  BART: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension.  In _ACL_ , pages 7871–7880.
  * Li et al. (2022) Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, et al. 2022.  Competition-level code generation with alphacode.  _arXiv preprint arXiv:2203.07814_.
  * Lu et al. (2022) Shuai Lu, Nan Duan, Hojae Han, Daya Guo, Seung-won Hwang, and Alexey Svyatkovskiy. 2022.  ReACC: A retrieval-augmented code completion framework.  In _Association for Computational Linguistics_ , pages 6227–6240.
  * Lu et al. (2021) Shuai Lu, Daya Guo, Shuo Ren, Junjie Huang, Alexey Svyatkovskiy, Ambrosio Blanco, Colin Clement, et al. 2021.  CodeXGLUE: A machine learning benchmark dataset for code understanding and generation.  _arXiv preprint arXiv:2102.04664_.
  * Nijkamp et al. (2022) Erik Nijkamp, Bo Pang, Hiroaki Hayashi, Lifu Tu, Huan Wang, Yingbo Zhou, et al. 2022\.  A conversational paradigm for program synthesis.  _arXiv preprint arXiv:2203.13474_.
  * Parvez et al. (2021) Md Rizwan Parvez, Wasi Ahmad, Saikat Chakraborty, Baishakhi Ray, and Kai-Wei Chang. 2021.  Retrieval augmented code generation and summarization.  In _Findings of EMNLP_ , pages 2719–2734.
  * Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, et al. 2019.  PyTorch: An imperative style, high-performance deep learning library.  _Neural Information Processing Systems_.
  * Qu et al. (2020) Yingqi Qu, Yuchen Ding, Jing Liu, Kai Liu, Ruiyang Ren, Wayne Xin Zhao, Daxiang Dong, Hua Wu, and Haifeng Wang. 2020.  RocketQA: An optimized training approach to dense passage retrieval for open-domain question answering.  _arXiv preprint arXiv:2010.08191_.
  * Raffel et al. (2020) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J Liu, et al. 2020.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _J. Mach. Learn. Res._ , 21(140):1–67.
  * Santhanam et al. (2021) Keshav Santhanam, Omar Khattab, Jon Saad-Falcon, Christopher Potts, and Matei Zaharia. 2021.  Colbertv2: Effective and efficient retrieval via lightweight late interaction.  _arXiv preprint arXiv:2112.01488_.
  * Schuster et al. (2021) Tal Schuster, Ashwin Kalyan, Oleksandr Polozov, and Adam Tauman Kalai. 2021.  Programming puzzles.  _arXiv preprint arXiv:2106.05784_.
  * Snodgrass and Winnie (2019) Eric Snodgrass and Soon Winnie. 2019.  Api practices and paradigms: Exploring the protocological parameters of apis as key facilitators of sociotechnical forms of exchange.  _First Monday_ , 24(2).
  * Srivastava et al. (2022) Aarohi Srivastava, Abhinav Rastogi, Abhishek Rao, et al. 2022.  Beyond the imitation game: Quantifying and extrapolating the capabilities of language models.  _arXiv preprint arXiv:2206.04615_.
  * Vaithilingam et al. (2022) Priyan Vaithilingam, Tianyi Zhang, and Elena L Glassman. 2022.  Expectation vs. experience: Evaluating the usability of code generation tools powered by large language models.  In _CHI Conference on Human Factors in Computing Systems Extended Abstracts_ , pages 1–7.
  * Wang et al. (2021) Yue Wang, Weishi Wang, Shafiq Joty, and Steven CH Hoi. 2021.  CodeT5: Identifier-aware unified pre-trained encoder-decoder models for code understanding and generation.  In _Empirical Methods in Natural Language Processing_ , pages 8696–8708.
  * Wolf et al. (2019) Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. 2019.  Huggingface’s transformers: State-of-the-art natural language processing.  _arXiv preprint arXiv:1910.03771_.
  * Xiong et al. (2020) Lee Xiong, Chenyan Xiong, Ye Li, Kwok-Fung Tang, Jialin Liu, Paul Bennett, Junaid Ahmed, and Arnold Overwijk. 2020.  Approximate nearest neighbor negative contrastive learning for dense text retrieval.  _arXiv preprint arXiv:2007.00808_.
  * Yan et al. (2021) Yuanmeng Yan, Rumei Li, Sirui Wang, Hongzhi Zhang, Daoguang Zan, Fuzheng Zhang, Wei Wu, and Weiran Xu. 2021.  Large-scale relation learning for question answering over knowledge bases with pre-trained language models.  In _The 2021 Conference on Empirical Methods in Natural Language Processing_ , pages 3653–3660.
  * Zan et al. (2022a) Daoguang Zan, Bei Chen, Dejian Yang, Zeqi Lin, Minsu Kim, Bei Guan, Yongji Wang, Weizhu Chen, and Jian-Guang Lou. 2022a.  CERT: Continual pre-training on sketches for library-oriented code generation.  In _The 2022 International Joint Conference on Artificial Intelligence_ , pages 3653–3660.
  * Zan et al. (2022b) Daoguang Zan, Wang Sirui, Zhang Hongzhi, Yan Yuanmeng, Wu Wei, Guan Bei, and Wang Yongji. 2022b.  S${}^{\mbox{2}}$QL: Retrieval augmented zero-shot question answering over knowledge graph.  In _PAKDD 2022_ , volume 13282 of _Lecture Notes in Computer Science_ , pages 223–236. Springer.
  * Zhou et al. (2022) Shuyan Zhou, Uri Alon, Frank F Xu, Zhengbao JIang, and Graham Neubig. 2022.  DocCoder: Generating code by retrieving and reading docs.  _arXiv preprint arXiv:2207.05987_.

##  Appendix A Collection of API Documentation

We aim to use data from public libraries for training and generalize the
models to private libraries. Thus, we crawled the API documentation of the
$31$ most popular public libraries in Python. Table 4 summarizes the number of
APIs we extracted for each library.

##  Appendix B Resampling Strategy

The resampling strategy allows high-quality python files to be more frequently
sampled, and vice versa. So the resampling weight (${w}$) of each python file
is defined in the following aspects: the star number of the corresponding
repository ($N_{\rm star}$), the unit test function rate ($R_{{\rm ut}}$) that
is the number of unit test functions divided by the number of all functions,
the number of API name ($N_{\rm api}$) in the file, and the number of APIs
($M_{\rm api}$) considering one API name may match multiple APIs. Formally,
the strategy can be formulated as follows:

$\begin{split}&{w}_{\rm star}=1.0+\log(N_{\rm star}+1).{\rm
clip}(_{0}^{5})\times 0.2,\\\ &{w}_{\rm ut}=(0.5+(1-{R}_{{\rm ut}})).{\rm
clip}(_{0}^{1}),\\\ &{w}_{\rm api}=5.0-\log(\frac{M_{\rm api}}{N_{\rm
api}}).{\rm clip}(_{0}^{5})\times 0.2,\\\ &{w}={w}_{\rm star}\times{w}_{\rm
ut}\times{w}_{\rm api},\end{split}$ (1)

where ${\rm clip}(_{x}^{y})$ limits the value to $[x,y]$.

##  Appendix C Keywords Conversion from Public Library to Private Library.

As mentioned in Section 4, we convert the public library benchmarks
(PandasEval and NumpyEval) to the private library benchmarks (MonkeyEval and
BeatNumEval) by manually modifying all public library-related keywords. In
Table 5, we list all the keywords before and after the conversion.

Pandas | NumPy | sklearn | PyTorch | TensorFlow | Django | selenium | Matplotlib | Flask | SciPy | Seaborn
---|---|---|---|---|---|---|---|---|---|---
7,094 | 12,085 | 53,166 | 124,902 | 32,116 | 24,375 | 4,842 | 439,913 | 31,867 | 153,359 | 161,477
NLTK | BeatifulSoup | pygame | PIL | jieba | Gensim | spaCy | transformers | fairseq | SQLAlchemy | Scrapy
206,816 | 22,519 | 70,396 | 127,212 | 26,620 | 37,331 | 239,945 | 652,913 | 158,721 | 54,765 | 3,537
AllenNLP | datasets | tokenizers | MXNet | imageio | pytest | MetPy | ansible | requests |  |
276,088 | 136,843 | 195 | 142,070 | 175,878 | 1,047 | 27,429 | 40,839 | 39,333 |  |

Table 4: The number of APIs in the 31 public libraries we crawled.

PandasEval-MonkeyEval
---
isnull | mean | pandas | dataframe | df | isin | pd
ifnull | average | monkey | knowledgeframe | kf | incontain | mk
tolist | apply | to_numeric | dropna | append | tail | copy
convert_list | employ | to_num | sipna | adding | last_tail | clone
innull | astype | select_dtypes | iterrows | min | max | map
isnone | totype | choose_dtypes | traversal | get_min | get_max | mapping
last | shift | merge | value_counts | rename_axis | reset_index | sample
final_item | shifting | unioner | counts_value_num | renaming_axis | reseting_index | sample_by_num
concat | to_dict | cumsum | sort_index | to_string | drop_duplicates | duplicated
concating | convert_dict | cumulative_sum | sorting_index | convert_string | remove_duplicates | duplicated_values
round | format | to_pydatetime | div | ceil | assign | intersection
value_round | formating | convert_pydatetime | division | ceiling | allocate | interst
drop | Series | ravel | any | fillna | all | Pandas
sip | Collections | flat_underlying | whatever | fillnone | total_all | Monkey
reindex | get | std | rename | sum | unique | to_datetime
reindexing | getting | standard | renaming | total_sum | distinctive | convert_datetime
applymap | sort_values | DataFrame | groupby | nlargest | replace | len
conduct_map | sort_the_values | KnowledgeFrame | grouper | nbiggest | replacing | length
head | series | isna |  |  |  |
header_num | collections | ifna |  |  |  |
NumpyEval-BeatNumEval
to_numpy | ndarray | array | transpose | numpy | Numpy | np
to_beatnum | ndnumset | numset | switching_places | beatnum | Beatnum | bn
column_stack | concatenate | slice | sum | imag | abs | real
stack_col | connect | piece | total_count | imaginary | absolute | reality
fill_diagonal | all | fromstring | in1d | mean | where | std
pad_diagonal | total | come_from_str | intersection1dim | average | filter_condition | standard_op
add | histogram | fromarrays | reshape | filled | stack | cumsum
add_concat | hist_operation | come_from_arrays | change_shape_to | masked_fill | pile_operation | cumulative_sum
astype | arange | setxor1d | compressed |  | argmin | argmax
convert_type | arr_range | seting_exclusive_or_one_dim | remove_masked_data |  | get_argmin_value | get_argmax
vstack | squeeze | hstack | asarray | repeat | vectorize | split
vertical_stack | sqz | horizontal_stack | asnumset | duplicate | vectorisation | sep_split
diff | unique | unravel_index | flatten | norm | delete | ones
difference | uniq | convert_index_or_arr | convert_into_one_dim | normlizattion | remove_operation | create_ones
append | any | logical_and | bincount | isnan | argpartition | ravel
apd | any_condition | logic_and_element_wise | binoccurrence | ifnan | perform_partition | asview
array_split | inv | insert | searchsorted | min | max | full
split_array | inverse | stick | find_sorted | get_min | get_max | full_value_func

Table 5: The keywords of converting PandasEval to MonkeyEval, and NumpyEval to
BeatNumEval. The grey background means the original keywords, and the white
background means the converted ones.


1University of Chicago, 2Carnegie Mellon University

# EFFUSE: Efficient Self-Supervised Feature Fusion for E2E ASR in Multilingual
and Low Resource Scenarios

###### Abstract

Self-Supervised Learning (SSL) models have demonstrated exceptional
performance in various speech tasks, particularly in low-resource and
multilingual domains. Recent works show that fusing SSL models could achieve
superior performance compared to using one SSL model. However, fusion models
have increased model parameter size, leading to longer inference times. In
this paper, we propose a novel approach of predicting other SSL models’
features from a single SSL model, resulting in a light-weight framework with
competitive performance. Our experiments show that SSL feature prediction
models outperform individual SSL models in multilingual speech recognition
tasks. The leading prediction model achieves an average SUPERB score increase
of 135.4 in ML-SUPERB benchmarks. Moreover, our proposed framework offers an
efficient solution, as it reduces the resulting model parameter size and
inference times compared to previous fusion models.

Index Terms—  Self-supervised learning, feature fusion, multilingual speech
recognition, ML-SUPERB

##  1 Introduction

Conventionally, automatic speech recognition (ASR) has relied on hybrid deep
neural networks - hidden Markov models (DNN-HMM) [1, 2, 3]. However, in the
past decade, the introduction of end-to-end (E2E) models has demonstrated
similar or superior performances compared to the hybrid models [4, 5, 6, 7,
8]. Recent advancements have adapted E2E models by incorporating self-
supervised learning (SSL) models as frontend feature extractors [9, 10, 11,
12, 13]. Despite the encouraging progress in E2E models equipped with an SSL
frontend, there are still many limitations when applying SSL features. One
such restraint is that most SSL models are trained with corpora in a single
English talker scenario. Thus, this setting may not align well with other
languages or recording conditions. Previous works have shown that these models
might not capture enough information for downstream tasks in other languages
or recording environments [14, 15, 16]. There have been methods proposed that
have great performance over different target domains: frontend feature fusion
[17], adapters [18, 19], teacher-student distillation [20, 21, 22, 23, 24],
and continual training [25, 26, 27, 28].

Each of these methods has its own drawbacks. Specifically, fusing SSL
representations with spectral features provides significant benefits in
results, but the spectral information has limited contextual representations
for downstream tasks [11]. Adapters are primarily designed to transfer
knowledge from a pre-trained model to new tasks. However, they might struggle
to adapt to drastically different data distribution or tasks, especially
considering low-resource scenarios. Teacher-student distillation requires
extra effort in training another model from scratch. Moreover, continual
training requires fine-tuning a large model and a reasonable amount of data.

(a) Architecture of the standard SSL fusion model that is used in the fusion
stage. The raw waveform $S$ is passed through the SSL models
$f_{1},f_{2},f_{3}$ to produce representations $f_{1}(S),f_{2}(S),f_{3}(S)$.
These representations are then concatenated as the input to the downstream ASR
model.

(b) Architecture of the proposed prediction stage. The raw waveform $S$ is
passed through the SSL models $f_{1},f_{2},f_{3}$ to produce representations
$f_{1}(S),f_{2}(S),f_{3}(S)$. Linear layers $\mathbf{L}_{2}$ and
$\mathbf{L}_{3}$ are used to predict the representations $f_{2}(S)$ and
$f_{3}(S)$ respectively. An L1-Loss function is used to train the linear
layers. The features
$f_{1}(S),\mathbf{L}_{2}(f_{1}(S)),\mathbf{L}_{3}(f_{1}(S))$ are concatenated
and projected and passed to the ASR model. For inference, we can safely
disregard SSL $f_{2}$ and $f_{3}$.

Fig. 1: The proposed methodology with 3 SSL Models

In this paper, we propose an efficient self-supervised feature fusion (EFFUSE)
methodology to utilize features from multiple frozen SSL models that limits
parameter increase and inference costs. Specifically, EFFUSE consists of a
two-stage training strategy: (1) Fusion stage: train the model with a fusion
of self-supervised features, similar to [17, 29]. (2) Prediction stage:
continue to train the downstream model by using only one SSL model’s features
to predict other SSL model representations by a minimal number of linear
layers.

Our proposed EFFUSE method reduces the number of parameters by using a single
SSL model, thereby scaling down inference costs compared to the original
feature fusion network. It is also lighter than adapters, requiring only a
limited amount of linear layers. Unlike teacher-student distillation and
continual training, our method offers a more lightweight training process for
SSL models. We extensively validate our framework’s effectiveness in the low-
resource and multilingual speech-processing tasks. Our proposed methodology
has an average diminution of 4.5 absolute character error rate (CER) in the
low-resource domain over two benchmarks and 2.8 absolute CER in the
multilingual domain over three benchmarks. Notably, there is an average
relative decrease of 16% in decoding the real-time factor (RTF) of our
proposed methodology, compared to the conventional fusion model.

Our contributions can be summarized as follows:

  * •

We extensively explore the SSL fusion model in low-resource and multilingual
scenarios.

  * •

We propose a novel framework that is able to use one SSL model to predict the
features of other SSL models in the fusion network during inference.

  * •

We demonstrate that our approach can improve the performance of downstream
multilingual tasks while maintaining reasonable inference costs.

##  2 Predictive Ability Between SSL Models

In this section, we discuss the capability of SSL models to predict other SSL
models features. This serves as a basis for the prediction model, which is the
second stage of EFFUSE.

In [29], the authors demonstrate a substantial linear correlation between
features extracted from distinct SSL models in a fusion model. This
observation paves the way for analyzing the predictive ability of SSL
features. Given the complexities of considering different layers collectively,
our attention is primarily drawn toward the weighted sum layer-wise features
of the SSL models, which are widely used in SUPERB [11].

To facilitate this, we initially train an ASR model on a low-resource ASR
corpus, namely, Totonac [17]. The ASR model employs the weighted sum layer-
wise features from frozen SSL models. Then, we extract weighted sum features
from 5,000 utterances and use them to train a linear regression model for
projecting features across HuBERT, WavLM, and Wav2Vec 2.0.

To assess predictive capability, we employ the coefficient of determination
(i.e., R2) as a measure of variance explainability. In order to calculate R2,
we utilize a matrix representation of the calculation, since the target
features more than one dimension. The resultant data is presented in Table 1.
The table exemplifies the SSL models’ capabilities to predict weighted
features with each R2 value being approximately 0.6 or higher. Thus, this
experiment demonstrates that the prediction of features can retain the
performance of fusion models to a great degree.

##  3 EFFUSE Formulation

As discussed in Sec. 1, EFFUSE contains two training stages, namely the fusion
stage and the prediction stage. The first stage follows the fusion methodology
as in [17, 29], while the prediction stage continues the model training with
additional prediction modules. In this section, we introduce both parts of
EFFUSE.

Table 1: The R2 values are interpreted as the layerwise features of the models on the left side predicting weighted-sum features of the models on top, where the weighted-sum features are learned from ASR modeling on Totonac [17]. The table shows that the predictability of weighted-sum features can be achieved reasonably well from other SSLs. | HuBERT | WavLM | Wav2Vec 2.0
---|---|---|---
HuBERT | ✕ | 0.71 | 0.60
WavLM | 0.70 | ✕ | 0.59
Wav2Vec2.0 | 0.63 | 0.65 | ✕

###  3.1 Stage 1: Fusion

We present this first stage in Fig. 1(a), and provide a detailed description
of its inner workings. We begin by letting $S$ denote a sampled and quantized
raw waveform. Let $f_{1},\ldots,f_{n}$ be a set of $n$ SSL models. Suppose
that for an SSL model $f_{i}$, there are $k_{i}$ layers. Then, for each layer,
the SSL model extracts features from $S$ given by

$\\{f_{i}^{k}(S):k\in\\{1,2,\ldots,k_{i}\\}\\}.$ (1)

The significance of each layer’s features is gauged using learnable weights
($w_{i}^{k}$), obtained by fine-tuning for downstream tasks. These weights,
learned during the training process, epitomize the importance of the features
extracted at each layer. Using weights $w_{i}^{k}$, the weighted sum features,
using the layerwise features in Equation (1), are computed as

$f_{i}(S)=\sum_{k=1}^{k_{i}}w_{i}^{k}f_{i}^{k}(S).$ (2)

We derive the set of weighted sum features as in Equation (2) for $n$ SSL
models as

$\mathrm{F}=\\{f_{1}(S),\ldots,f_{n}(S)\\}.$ (3)

Thus, in Fig. 1(a), we note a total of three weighted features:
$f_{1}(S),f_{2}(S)$, and $f_{3}(S)$. Then, let $\mathfrak{L}$ symbolize a
linear layer which projects features into an $m$ dimensional space. Suppose
$\mathrm{\mathbf{T}}$ denotes a transform function. In line with [17], the
fusion feature $f_{\text{FUSE}}$ is represented as follows:

$f_{\text{FUSE}}(S)=\mathfrak{L}(\mathrm{\mathbf{T}}(f_{1}(S),\ldots,f_{n}(S))),$
(4)

where $\mathrm{\mathbf{T}}$ can be any of concatenation, linear,
convolutional, and co-attention modules that fuse the different
representations, as described in [17]. The fusion feature $f_{\text{FUSE}}$ is
fed into the encoder-decoder model as shown in Fig. 1(a). The model that uses
this procedure is deemed as the fusion model.

Table 2: An overview of SSL Models. The table details the variables that were used to select the SSL Models for our experiments. It shows the training objectives, training dataset, parameter size, and if they are trained for specific downstream domains. Model |  Training Objectives |  Training Dataset | Param. Size (M) | Multilingual |  Noise Robust
---|---|---|---|---|---
Wav2Vec 2.0 [30] | Contrastive |  LibriLight [31] | 317 M | ✕ | ✕
XLS-R [32] | Contrastive |  Commonvoice [33] MLS [34] | 317 M | ✓ | ✕
HuBERT [35] |  Masked Prediction |  LibriLight [31] | 317 M | ✕ | ✕
WavLM [36] |  Masked Prediction |  Librispeech [37] LibriLight [31] GigaSpeech [38] | 317 M | ✕ | ✓

###  3.2 Stage 2: Prediction

As discussed in Section 2, SSL models possess the capability to predict the
weighted sum features of other SSL models. In the prediction stage, we
leverage this ability to develop a framework that preserves the exceptional
performance achieved during the fusion stage. Importantly, this framework
significantly reduces the parameter size by utilizing only one SSL model for
inference.

We discuss our methodology, illustrated in Fig. 1(b), in greater depth. In
this stage, the layer-wise weights $w_{i}^{k}$ in each SSL model are frozen to
focus on learning the weighted sum features. An SSL model in (3) is chosen as
the prediction source, denoted as $f_{1}$. We then define $n-1$ linear layers
given by

$\mathbf{L}_{2},\ldots,\mathbf{L}_{n}.$ (5)

Thus, in Fig. 1(b), there are two linear layers $\mathbf{L}_{2}$ and
$\mathbf{L}_{3}$ as there are three SSL models. Next, we utilize these linear
layers to predict the weighted sum representations given in (3). Thus, a
projected representation for the feature $f_{k}(S)$ is obtained as
$\mathrm{\mathbf{L}}_{k}(f_{1}(S))$. Then, the set of predicted features is

$\\{f_{1}(S),\mathrm{\mathbf{L}}_{2}(f_{1}(S)),\ldots,\mathrm{\mathbf{L}}_{n}(f_{1}(S))\\}.$
(6)

Hence, for our example in Fig. 1(b), we observe that the set of predicted
features is
$\\{f_{1}(S),\mathrm{\mathbf{L}}_{2}(f_{1}(S)),\mathrm{\mathbf{L}}_{3}(f_{1}(S))\\}$.
The predicted features in (6) are then fused, using a linear layer
$\mathfrak{L}$ and a transform $\mathbf{T}$ and yield,

$f_{\mathrm{PRED}}=\mathfrak{L}(\mathrm{\mathbf{T}}(f_{1}(S),\mathrm{\mathbf{L}}_{2}(f_{1}(S)),\ldots,\mathrm{\mathbf{L}}_{n}(f_{1}(S)))),$
(7)

where $f_{\mathrm{PRED}}$ denotes the final representation of the waveform
$S$. Note that the difference between (4) and (7) is that we are fusing
weighted features of the SSL models directly in (4). In (7), we specifically
fuse only the predicted features. Finally, to optimize the introduced linear
layers ($\mathbf{L}_{2},...,\mathbf{L}_{n}$), we use the L1-loss function:

$\mathcal{L}_{i}=\mathrm{L1Loss}(\mathrm{\mathbf{L}}_{i}(f_{1}(S)),f_{i}(S)),$
(8)

where the loss is jointly trained with the downstream objective. The models
using this second stage of EFFUSE are called prediction models.

##  4 Experiments

###  4.1 Datasets

To assess the robustness of our method proposed in Sec. 3.1, 3.2, we
investigate its effectiveness in two settings (mentioned in Sec. 1):
multilingual scenarios and low-resource conditions. For the multilingual
setting, we utilize the datasets ML-SUPERB [39] and Fleurs [40]. In the low-
resource setting, we experiment with two endangered languages: the Yoloxochitl
Mixtec (YM) [41] and Totonac [17].

The ML-SUPERB public data includes 123 standard languages and 20 few-shot
languages. We undertake all four categories of experiments: multilingual ASR,
multilingual ASR+LID, LID, and monolingual ASR. Each type of experiment makes
use of a 10-minute training set and a 1-hour training set. We strictly adhere
to the official training, validation, and testing split as provided by the
challenge organizers [39]. We also calculate the SUPERB score (i.e., SUPERBs)
to present the overall performance of EFFUSE in ML-SUPERB. Fleurs is a
multilingual dataset specifically designed as a few-shot learning benchmark.
The dataset comprises a total of 102 languages from 17 different language
families. Notably, this dataset encompasses a total of 1.4k hours of recorded
speech [40].

The endangered language Yoloxochitl Mixtec is spoken in the San Luis Acatlan
region of Guerrero, Mexico. This language has a dataset that includes three
corpora, each serving different applications - ASR, novice correction, and
evaluation. We specifically employ the 100-hour corpus for ASR [42, 41].
Another endangered language in Mexico, Totonac, has a corresponding dataset
consisting of 10 hours of audio recordings [17]. The distinct sizes of these
two datasets (i.e., 100 hours for Yoloxochitl Mixtec and 10 hours for Totonac)
provide two distinct contexts for evaluating our methodology with low-resource
languages.

###  4.2 Experimental Setup

All the experiments are conducted with ESPnet [43] with S3PRL [11] and its
corresponding recipes with the four datasets.111All of our configurations,
sources, and model checkpoints will be publicly released for reproducibility.

####  4.2.1 SSL Models

We select SSL models based on various factors such as their training
objectives, and pretraining domains among others. Our experiments particularly
focus on four SSL models: Wav2Vec 2.0 [30], XLS-R [32], HuBERT [35], and WavLM
[36]. We utilize the large versions of these models. A synopsis of these
models’ key features can be found in Table 2.

Table 3: 10-minute set ML-SUPERB benchmark. Both steps of EFFUSE yield
significant improvements. We see that the HuBERT+XSLR fusion model has
exceptional performance across most of the experiments as given by the SUPERB
score. The Wav2Vec2.0+XLS-R fusion model yields the most optimal performance
for Multilingual ASR experiments.

SSL | Monolingual ASR | Multilingual ASR | LID | Multilingual ASR + LID | SUPERBs
---|---|---|---|---|---
| Normal | Few-shot | Normal | Normal | Few-shot
CER/PER | CER | CER | ACC | ACC | CER | CER
HuBERT | 40.34 | 43.5 | 46.4 | 45.77 | 54.6 | 45.6 | 49.3 | 655.5
HuBERT+XLS-R (H+X) fusion | 37.87 | 31.8 | 42.5 | 59.38 | 72.71 | 31.2 | 41.0 | 933.0
HuBERT$\rightarrow$XLS-R (H+X) predict | 41.46 | 38.4 | 44.8 | 56.49 | 70.12 | 37.8 | 44.8 | 812.9
Wav2Vec2.0 | 44.35 | 41.0 | 45.8 | 34.39 | 55.4 | 45.5 | 50.3 | 590.9
Wav2Vec2.0+XLS-R (WV+X) fusion | 38.52 | 31.0 | 40.9 | 22.99 | 69.74 | 31.4 | 41.8 | 769.2
Wav2Vec2.0$\rightarrow$XLS-R (WV+X) predict | 45.04 | 38.8 | 44.6 | 32.91 | 69.28 | 38.5 | 44.3 | 671.2

Table 4: 1 hour set ML-SUPERB benchmark. The proposed EFFUSE method exhibits
great improvements. We see that the HuBERT+XSLR fusion model has exceptional
performance across all tracks, yielding an exceptional SUPERB score.

SSL | Monolingual ASR | Multilingual ASR | LID | Multilingual ASR + LID | SUPERBs
---|---|---|---|---|---
| Normal | Few-shot | Normal | Normal | Few-shot
CER/PER | CER | CER | ACC | ACC | CER | CER
HuBERT | 32.16 | 37.7 | 43.5 | 63.93 | 77.7 | 35.1 | 42.2 | 774.3
HuBERT+XLS-R (H+X) fusion | 28.94 | 23.5 | 38.3 | 78.95 | 89.8 | 23.3 | 37.6 | 1001.9
HuBERT$\rightarrow$XLS-R (H+X) predict | 33.96 | 28.4 | 41.8 | 78.52 | 87.0 | 29.0 | 41.1 | 887.1
Wav2Vec2.0 | 35.32 | 33.0 | 43.2 | 8.0 | 78.2 | 34.7 | 42.2 | 592.2
Wav2Vec2.0+XLS-R (WV+X) fusion | 29.52 | 24.4 | 40.2 | 13.54 | 87.8 | 24.4 | 39.3 | 759.0
Wav2Vec2.0$\rightarrow$XLS-R (WV+X) predict | 35.57 | 31.0 | 42.6 | 18.88 | 84.3 | 31.4 | 41.9 | 656.5

####  4.2.2 Baseline

For ML-SUPERB, our baseline SSL models include HuBERT and Wav2Vec 2.0. We
follow the same setting as the framework design in [39]. Specifically, we
utilize a transformer encoder with 2 transformer blocks with 8 attention heads
and a 256-dimensional hidden size.

Similar to ML-SUPERB, HuBERT, and Wav2Vec 2.0 are designated as baseline
models for Fleurs. We employ an ESPnet recipe that utilizes a transformer-
based encoder-decoder model, combining Connectionist Temporal Classification
(CTC) and Attentional Encoder-Decoder (AED) training [44]. The transformer
encoder comprises of 12 self-attention blocks with 8 attention heads and a
hidden size of 512. The transformer decoder consists of 6 self-attention
blocks and 8 attention heads.

In our low-resource study, we choose HuBERT, WavLM, and Wav2Vec 2.0 as the
baseline models. These models employ a transformer-based encoder-decoder
framework with CTC-AED training. Both YM and Totonac corpora share a 12-layer
transformer encoder and a 6-layer transformer decoder architecture, with 4
attention heads per layer. The attention dimension is set to 256 and 512 for
YM and Totonac, respectively. Additionally, we utilize data augmentation
techniques like speed perturbation and SpecAugment to enhance the model’s
performance.

####  4.2.3 Fusion Model

For the multilingual setting, we choose the combination of HuBERT and XLS-R
(H+X), along with the combination of Wav2Vec 2.0 and XLS-R (WV+X). In the low-
resource setting, we utilize the combination of HuBERT, WavLM, and Wav2Vec
2.0. The representations of these features are fused as described in Sec. 3.1.

In all our experiments, we adhere to a consistent methodology. As discussed in
Sec. 3.1, we must select $\mathbf{T}$ and the dimension $m$ for
$\mathfrak{L}$. We choose $\mathbf{T}$ as concatenation and apply it over the
features. Finally, we project the concatenated features using $\mathfrak{L}$
into an 80-dimensional space. This procedure is uniformly applied to each
feature across all the experiments we conduct.

####  4.2.4 Prediction Model

As discussed in Sec. 3.2, we must select a prediction source, which is also
denoted as $f_{1}$, from the fusion model. In the multilingual scenario, we
designate HuBERT and Wav2Vec 2.0 as the prediction sources for the first and
second experiments listed in Sec. 4.2.3. For the fusion model in the low
resource setting (see Sec. 4.2.3) there are three corresponding prediction
models: 1) HuBERT is designated as the prediction source [i.e. $f_{1}=$
HuBERT], 2) WavLM is the prediction source, and 3) $f_{1}=$ Wav2Vec 2.0.
Throughout these experiments, we retain the architecture of the fusion models.
However, we incorporate additional predictive linear layers. These added
layers possess an input dimension of 512 and an output dimension of 512,
thereby maintaining consistency with the feature dimensions of the SSL models.

##  5 Results

###  5.1 Multilingual ASR Results

####  5.1.1 ML-SUPERB

As introduced in [39], there are four tasks (monolingual ASR, multilingual
ASR, LID, multilingual ASR+LID) with two data configurations (i.e., 10-minute
and 1-hour). We recall that the fusion and prediction models are chosen in
Sec. 4.2.3 and Sec. 4.2.4 respectively. The results for the 10 minute and 1
hour set with four tasks are recorded in Table 3 and Table 4, respectively.

Monolingual ASR Track:

  * •

The fusion models consistently outperform baseline models in the 10-minute and
1-hour sets.

    * –

The H+X fusion model delivers the best results in both sets.

  * •

All the prediction models are substandard compared to the baselines in this
scenario, especially over languages aside from English

We offer a feasible explanation for the prediction model results: the
prediction sources may encounter difficulties in effectively learning the
features of XLS-R due to limited data, given the usage of a smaller dataset in
the monolingual experiments.

Table 5: Fleur Results. The fusion model results are exceptional, whereas the prediction models maintain the improvements of fusion to a certain extent. We see that by a limited increase in parameter size, there are significant gains in the prediction models. Model | WER | CER
---|---|---
Hubert | 42.0 | 12.8
Hubert+XLS-R (H+X) fusion | 35.5 | 10.5
Hubert$\rightarrow$XLS-R (H+X) predict | 40.7 | 12.6
Wav2Vec 2.0 | 43.5 | 14.1
Wav2Vec2.0+XLS-R (WV+X) fusion | 35.6 | 10.5
Wav2Vec2.0$\rightarrow$XLS-R (WV+X) predict | 42.1 | 13.4

Multilingual ASR Track:

  * •

The fusion and prediction models outperform the respective baselines

  * •

We specify that the WV+X fusion model is the best-performing result in the 10
minute scenario.

Unlike the monolingual case, the prediction model gives competitive
performance. We attribute this to the increase in training data for this task.
We also highlight that despite Wav2Vec 2.0 sharing the same architecture and
training objectives as XLS-R, it is inferior to HuBERT in predicting XLS-R’s
features. This highlights that feature predictability is not highly influenced
by architecture and objective.

LID Track:

  * •

H+X fusion and prediction models have exceptional performance.

  * •

WV+X fusion and prediction models display poor performance

As mentioned in [39], the Wav2Vec 2.0 Large model does not exhibit good LID
performance. As a result, we expect the fusion and prediction models involving
it to also suffer in performance.

Joint Multilingual+LID Track:

  * •

Fusion models and prediction models outperform all the baselines.

  * •

The H+X models achieve the best performance for both the fusion and prediction
models.

In comparing the two EFFUSE models, we generally see that the prediction
models compete with the performance of the fusion model in the LID task.
However, in the multilingual task, the prediction model does not provide as
competitive results.

We note the trend that the Wav2Vec 2.0 baseline generally outperforms the
HuBERT baseline. This did not translate to the fusion model WV+X achieving the
best results. One possible reason for this is that the introduction of the LID
task caused the model to suffer. Another reason is that in the WV+X fusion
model, the individual SSL models could be learning less predictable features.

Overall: In general across the 10 minute scenario, we see that the H+X fusion
model is able to achieve extremely encouraging results, as signified by the
SUPERB score of 933.0. It only fails to outperform the WV+X fusion model for
the Multilingual task. On the other hand, the H+X fusion model yields a SUPERB
score of 1001.9, which is significantly higher than the models trained in
[39]. We highlight that across the two scenarios, the prediction models offer
very competitive performance, except for the Monolingual Task.

Table 6: Results for Low Resource Domain. The fusion models yield optimal results across both datasets, while the prediction models demonstrate high-performance retention with respect to the fusion model with minimal parameter increase. | YM | Totonac
---|---|---
Model | WER | CER | WER | CER
Hubert | 21.8 | 10.5 | 55.0 | 23.2
WavLM | 21.0 | 10.4 | 52.3 | 21.7
Wav2Vec 2.0 | 20.5 | 9.9 | 56.4 | 26.3
Fusion model | 20.3 | 9.6 | 45.2 | 14.5
Hubert Prediction Source | 20.6 | 9.9 | 49.1 | 15.9
WavLM Prediction Source | 20.2 | 9.6 | 46.2 | 14.7
Wav2Vec 2.0 Prediction Source | 20.3 | 9.8 | 46.9 | 15.2

####  5.1.2 Fleurs

The results for Fleurs are shown in Table 5.

  * •

The fusion models consistently outperform the baselines. Both the H+X and WV+X
models achieve a CER of 10.5.

As far as we know, our fusion model yields the best result for Fleurs with a
transformer model [44]. The performance improvement is attributed to using
XLS-R, a multilingual SSL model that performs well over Fleurs. Moreover, the
combination of English and multilingual SSL models help capture different
information pieces because of their distinct pretrained settings.

  * •

The prediction model demonstrates increased performance with an average
decrease of 0.45 CER from the baseline

In general, the prediction models often exhibit lower performance, compared to
the fusion model. We suggest that this is due to the difficulty of accurately
predicting the representations of XLS-R using English-based SSL models.

Table 7: A summary of parameter efficiency. The table presents the parameter size and SUPERB score with respect to the baselines with a single SSL frontend. The fusion models show a prodigious increase in the SUPERB score, but the change in parameter size is immense. The parameter size increase for the prediction models is minimal and gives exceptional boosts to the SUPERB score. Model | Num. Param. (M) |  $\Delta$ Num. Param. (M) |  $\Delta$ SUPERBs
---|---|---|---
HuBERT | 323 | / | /
HuBERT + XLS-R (10min fusion) | 642 | +319 | +277.5
HuBERT $\rightarrow$ XLS-R (10min predict) | 325 | +2 | +157.4
Wav2Vec 2.0 | 324 | / | /
Wav2Vec 2.0 + XLS-R (10min fusion) | 643 | +319 | +178.3
Wav2Vec 2.0 $\rightarrow$ XLS-R (10min predict) | 326 | +2 | +80.3

###  5.2 Low Resource ASR

Our results for our low resource study are summarized in Table 6. We highlight
that the fusion stage delivers improvements over all SSL model baselines. The
results indicate that a combination of SSL models is able to learn broader
pieces of information holistically in comparison to singular SSL models. In
other words, the SSL models are learning distinctive, relevant features in the
fusion model.

  * •

Prediction models show good performance over the baselines and fusion model

    * –

YM: The WavLM prediction model demonstrates a decrease of 0.1 in WER while
maintaining the CER from the fusion model

      * *

This result is significant since it surpasses the performance of the fusion
model and indicates that WavLM is able to learn the representations extremely
well

    * –

Totonac: WavLM prediction source exhibits the best performance, with a 0.2
increase in CER from the fusion model

Fig. 2: CER vs. Real-Time Factor (RTF) for Totonac. The individual SSL models
demonstrate excellent performance in terms of RTF, but there is a trade-off
with CER. The prediction models strike a balance by maintaining high
performance in both RTF and
CER.$14$$16$$18$$20$$22$$24$$26$$0.2$$0.25$$0.3$CERRTFHuBERTWavLMWav2Vec
2.0FusionHuBERT PredWavLM PredWav2Vec 2.0 Pred

###  5.3 Further Discussion

As shown in Sec. 5.1 and 5.2, both stages of EFFUSE exhibit substantial boosts
over baselines for numerous benchmarks. Table 7 provides a summary of
parameter sizes for the fusion and prediction models in ML-SUPERB (see Sec.
5.1.1) and their impact on the SUPERB score. Notably, fusion models increase
the SUPERB score by an average of 227.9, while prediction models achieve an
average increase of 118.9. However, the prediction models exhibit exceptional
parameter efficiency, with only a minimal increase of 1.97 M, compared to the
fusion models’ enormous parameter growth of 318.97 M. With minimal parameter
growth, the prediction models deliver competitive results that significantly
outperform the baselines.

Additionally, we analyze the efficiency of inference in Figure 2 over the
Totonac benchmark. We observe that the baseline models achieve the fastest
inference speeds but have substandard performance. The fusion model, on the
other hand, significantly suffers in terms of inference but yields the optimal
CER. The prediction model functions as the most economical framework,
retaining the performance of the fusion model. Importantly, these models
experience a noteworthy speedup that is only slightly slower than the singular
SSL models.

##  6 Conclusion

We propose a novel methodology, EFFUSE, to leverage the predictive
capabilities of SSL models. EFFUSE uses two stages: 1) the SSL fusion model is
used as an initial training phase 2) a prediction source is employed to
successfully predict SSL-weighted features. We painstakingly check the
performance of our method over three multilingual benchmarks and two low-
resource domain benchmarks. Across all benchmarks, we found the EFFUSE
methodology provides significant benefits in performance and limits parameter
increase. Future work would involve leveraging the predictive ability of
layer-wise features that are dependent on downstream tasks.

##  7 References

## References

  * [1] Ossama Abdel-Hamid et al. “Applying Convolutional Neural Networks concepts to hybrid NN-HMM model for speech recognition” In _ICASSP_ , 2012, pp. 4277–4280 DOI: 10.1109/ICASSP.2012.6288864
  * [2] Geoffrey Hinton et al. “Deep Neural Networks for Acoustic Modeling in Speech Recognition: The Shared Views of Four Research Groups” In _IEEE Signal Processing Magazine_ 29.6, 2012, pp. 82–97 DOI: 10.1109/MSP.2012.2205597
  * [3] George Dahl et al. “Context-Dependent Pre-Trained Deep Neural Networks for Large-Vocabulary Speech Recognition” In _Audio, Speech, and Language Processing, IEEE Transactions on_ 20, 2012, pp. 30–42 DOI: 10.1109/TASL.2011.2134090
  * [4] Chung-Cheng Chiu et al. “State-of-the-Art Speech Recognition with Sequence-to-Sequence Models” In _ICASSP_ , 2018, pp. 4774–4778 DOI: 10.1109/ICASSP.2018.8462105
  * [5] Alex Graves et al. “Connectionist temporal classification: Labelling unsegmented sequence data with recurrent neural ’networks” In _ICML_ 2006, 2006, pp. 369–376 DOI: 10.1145/1143844.1143891
  * [6] Alex Graves, Abdel-rahman Mohamed and Geoffrey Hinton “Speech Recognition with Deep Recurrent Neural Networks” In _ICASSP_ , 2013, pp. 6645–6649
  * [7] William Chan et al. “Listen, attend and spell: A neural network for large vocabulary conversational speech recognition” In _ICASSP_ , 2016, pp. 4960–4964
  * [8] Rohit Prabhavalkar et al. “End-to-End Speech Recognition: A Survey” In _ArXiv_ abs/2303.03329, 2023
  * [9] Xuankai Chang et al. “An Exploration of Self-Supervised Pretrained Representations for End-to-End Speech Recognition” In _Proc. ASRU_ , 2021, pp. 228–235
  * [10] Abdelrahman Mohamed et al. “Self-Supervised Speech Representation Learning: A Review” In _IEEE Journal of Selected Topics in Signal Processing_ 16.6, 2022, pp. 1179–1210 DOI: 10.1109/JSTSP.2022.3207050
  * [11] Shu-wen Yang et al. “SUPERB: Speech processing Universal PERformance Benchmark” In _Interspeech_ , 2021
  * [12] N KrishnaD., Pinyi Wang and Bruno Bozza “Using Large Self-Supervised Models for Low-Resource Speech Recognition” In _Interspeech_ , 2021
  * [13] Jing Zhao and Wei-Qiang Zhang “Improving Automatic Speech Recognition Performance for Low-Resource Languages With Self-Supervised Models” In _IEEE Journal of Selected Topics in Signal Processing_ 16.6, 2022, pp. 1227–1241 DOI: 10.1109/JSTSP.2022.3184480
  * [14] Ramon Sanabria et al. “Measuring the Impact of Individual Domain Factors in Self-Supervised Pre-Training” In _ArXiv_ abs/2203.00648, 2022
  * [15] Yen Meng et al. “Don’t Speak Too Fast: The Impact of Data Bias on Self-Supervised Speech Models” In _ICAASP_ , 2022, pp. 3258–3262
  * [16] Juan Zuluaga-Gómez et al. “How Does Pre-Trained Wav2Vec 2.0 Perform on Domain-Shifted Asr? an Extensive Benchmark on Air Traffic Control Communications” In _SLT_ , 2022, pp. 205–212
  * [17] Dan Berrebbi et al. “Combining Spectral and Self-Supervised Features for Low Resource Speech Recognition and Translation” In _Interspeech_ , 2022, pp. 3533–3537 DOI: 10.21437/Interspeech.2022-10796
  * [18] Xie Chen et al. “Front-End Adapter: Adapting Front-End Input of Speech Based Self-Supervised Learning for Speech Recognition” In _ICASSP_ , 2023, pp. 1–5 DOI: 10.1109/ICASSP49357.2023.10095130
  * [19] Bethan Thomas, Samuel Kessler and Salah Karout “Efficient Adapter Transfer of Self-Supervised Speech Models for Automatic Speech Recognition” In _ICASSP_ , 2022, pp. 7102–7106 DOI: 10.1109/ICASSP43922.2022.9746223
  * [20] Zilun Peng et al. “Shrinking Bigfoot: Reducing wav2vec 2.0 footprint” In _SUSTAINLP_ , 2021
  * [21] Heng-Jui Chang, Shu-wen Yang and Hung-yi Lee “Distilhubert: Speech Representation Learning by Layer-Wise Distillation of Hidden-Unit Bert” In _ICASSP_ , 2022, pp. 7087–7091 DOI: 10.1109/ICASSP43922.2022.9747490
  * [22] Yifan Peng et al. “DPHuBERT: Joint Distillation and Pruning of Self-Supervised Speech Models” In _ArXiv_ abs/2305.17651, 2023
  * [23] Yeonghyeon Lee et al. “FitHuBERT: Going Thinner and Deeper for Knowledge Distillation of Speech Self-Supervised Models” In _Proc. Interspeech 2022_ , 2022, pp. 3588–3592 DOI: 10.21437/Interspeech.2022-11112
  * [24] Xiaoyu Yang, Qiujia Li and Philip C. Woodland “Knowledge Distillation for Neural Transducers from Large Self-Supervised Pre-Trained Models” In _ICASSP_ , 2021, pp. 8527–8531
  * [25] Kuan-Po Huang et al. “Improving Distortion Robustness of Self-supervised Speech Processing Tasks with Domain Adaptation” In _Interspeech_ ISCA, 2022, pp. 2193–2197 DOI: 10.21437/Interspeech.2022-519
  * [26] Steven Vander Eeckt and Hugo Van Hamme “Continual Learning for Monolingual End-to-End Automatic Speech Recognition” In _EUSIPCO_ , 2022, pp. 459–463 DOI: 10.23919/EUSIPCO55093.2022.9909589
  * [27] Wei-Ning Hsu et al. “Robust wav2vec 2.0: Analyzing Domain Shift in Self-Supervised Pre-Training” In _Interspeech_ , 2021
  * [28] Bo Li et al. “Massively Multilingual ASR: A Lifelong Learning Solution” In _ICASSP_ , 2022, pp. 6397–6401 DOI: 10.1109/ICASSP43922.2022.9746594
  * [29] Szu-Jui Chen, Jiamin Xie and John Hansen “FeaRLESS: Feature Refinement Loss for Ensembling Self-Supervised Learning Features in Robust End-to-end Speech Recognition” In _Interspeech_ , 2022, pp. 3058–3062 DOI: 10.21437/Interspeech.2022-10917
  * [30] Alexei Baevski et al. “wav2vec 2.0: A Framework for Self-Supervised Learning of Speech Representations” In _NeurIPS_ 33, 2020, pp. 12449–12460 URL: https://proceedings.neurips.cc/paper_files/paper/2020/file/92d1e1eb1cd6f9fba3227870bb6d7f07-Paper.pdf
  * [31] J. Kahn et al. “Libri-Light: A Benchmark for ASR with Limited or No Supervision” In _ICASSP_ , 2020, pp. 7669–7673 DOI: 10.1109/ICASSP40776.2020.9052942
  * [32] Alexis Conneau et al. “Unsupervised Cross-lingual Representation Learning for Speech Recognition” In _Interspeech_ , 2020
  * [33] Rosana Ardila et al. “Common Voice: A Massively-Multilingual Speech Corpus” In _Proceedings of the Twelfth Language Resources and Evaluation Conference_ Marseille, France: European Language Resources Association, 2020, pp. 4218–4222 URL: https://aclanthology.org/2020.lrec-1.520
  * [34] Vineel Pratap et al. “MLS: A Large-Scale Multilingual Dataset for Speech Research” In _Interspeech_ abs/2012.03411, 2020
  * [35] Wei-Ning Hsu et al. “HuBERT: Self-Supervised Speech Representation Learning by Masked Prediction of Hidden Units” In _IEEE/ACM Transactions on Audio, Speech, and Language Processing_ 29, 2021, pp. 3451–3460 DOI: 10.1109/TASLP.2021.3122291
  * [36] Sanyuan Chen et al. “WavLM: Large-Scale Self-Supervised Pre-Training for Full Stack Speech Processing” In _IEEE Journal of Selected Topics in Signal Processing_ 16, 2021, pp. 1505–1518
  * [37] Vassil Panayotov et al. “Librispeech: An ASR corpus based on public domain audio books” In _2015 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , 2015, pp. 5206–5210 DOI: 10.1109/ICASSP.2015.7178964
  * [38] Guoguo Chen et al. “GigaSpeech: An Evolving, Multi-Domain ASR Corpus with 10,000 Hours of Transcribed Audio” In _Interspeech_ , 2021, pp. 3670–3674
  * [39] Jiatong Shi et al. “ML-SUPERB: Multilingual Speech Universal PERformance Benchmark”, 2023 arXiv:2305.10615 [cs.SD]
  * [40] Alexis Conneau et al. “FLEURS: FEW-Shot Learning Evaluation of Universal Representations of Speech” In _2022 IEEE Spoken Language Technology Workshop (SLT)_ , 2022, pp. 798–805
  * [41] Jiatong Shi et al. “Leveraging End-to-End ASR for Endangered Language Documentation: An Empirical Study on Yolóxochitl Mixtec” In _Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume_ , 2021, pp. 1134–1145
  * [42] Jiatong Shi Johnathan D. and Rey Castillo García “End-to-end automatic speech recognition: Its impact on the workflow in documenting Yoloxóchitl Mixtec” In _Proceedings of the First Workshop on Natural Language Processing for Indigenous Languages of the Americas_ , 2021, pp. 64–80 DOI: 10.18653/v1/2021.americasnlp-1.8
  * [43] Shinji Watanabe et al. “ESPnet: End-to-End Speech Processing Toolkit” In _Proceedings of Interspeech_ , 2018, pp. 2207–2211 DOI: 10.21437/Interspeech.2018-1456
  * [44] William Chen et al. “Improving massively multilingual asr with auxiliary ctc objectives” In _ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , 2023, pp. 1–5 IEEE