# CODEP: Grammatical Seq2Seq Model for General-Purpose Code Generation

Yihong Dong, Ge Li , Zhi Jin11footnotemark: 1
Key Lab of High Confidence Software Technology (PKU), Ministry of Education
School of Computer Science, Peking University, China
<EMAIL_ADDRESS>{lige<EMAIL_ADDRESS>
Corresponding Author

###### Abstract

General-purpose code generation (GPCG) aims to automatically convert the
natural language description into source code in a general-purpose language
(GPL) like Python. Intrinsically, code generation is a particular type of text
generation that produces grammatically defined text, namely code. However,
existing sequence-to-sequence (Seq2Seq) approaches neglect grammar rules when
generating GPL code. In this paper, we make the first attempt to consider
grammatical Seq2Seq (GSS) models for GPCG and propose CODEP, a GSS code
generation framework equipped with a pushdown automaton (PDA) module. PDA
module (PDAM) contains a PDA and an algorithm to help model generate the
following prediction bounded in a valid set for each generation step, so that
ensuring the grammatical correctness of generated codes. During training,
CODEP additionally incorporates state representation and state prediction
task, which leverages PDA states to assist CODEP in comprehending the parsing
process of PDA. In inference, our method outputs codes satisfying grammatical
constraints with PDAM and the joint prediction of PDA states. Furthermore,
PDAM can be directly applied to Seq2Seq models, i.e., without any need for
training. To evaluate the effectiveness of our proposed method, we construct
the PDA for the most popular GPL Python and conduct extensive experiments on
four benchmark datasets. Experimental results demonstrate the superiority of
CODEP compared to the state-of-the-art approaches without pre-training, and
PDAM also achieves significant improvements over the pre-trained models.

##  1 Introduction

Code generation techniques are critical in artificial intelligence and
software engineering, which are able to help facilitate software development
and revolutionize user programming Ling et al. (2016); Rabinovich et al.
(2017); Yin and Neubig (2018); Wei et al. (2019); Xie et al. (2021); Wang et
al. (2021); Guo et al. (2022). For example, in Fig. 1, a programmer wants to
terminate the program under the specific condition “list $l$ is empty" in
Python, but does not know the relevant prior knowledge to realize their
intention. By using code generation techniques, the target code “if not $l$:
exit(0)” can be obtained with ease. A preeminent code generation technique
requires both semantics and grammar, because semantics ensures that the
functionality of the code is in line with the intent of the developer, while
satisfying the grammar is the basis for the code to be executed. The two
complement each other to produce a satisfactory code.

Figure 1: Illustration of code generation.

In recent years, two kinds of prevailing approaches have been used for code
generation tasks. One is the sequence-to-tree (Seq2Tree) approach Sutskever et
al. (2014), which guarantees the syntactical correctness of generated codes by
deriving the syntactic structure based on the abstract grammar Dong and Lapata
(2016); Yin and Neubig (2017, 2018); Rabinovich et al. (2017); Sun et al.
(2020). Seq2Tree approaches transduce the NL description into a sequence of
abstract syntax tree (AST) nodes, which can be further mapped to the well-
formed code. Generally, the sequence length of AST nodes is much longer than
that of code tokens Sun et al. (2019); Dong et al. (2022), leading to
increased difficulty in the generation. Moreover, collecting training data for
the Seq2Tree approach will be more problematic as it requires the
corresponding AST of the code, and so does its transferability across tasks.
The other approach, called Seq2Seq, inputs the natural language (NL)
description and directly outputs the code token sequence without any
constraints. Due to its accessibility of implementation, transfer, and data
collection, most of the early works Raychev et al. (2014); Jia and Liang
(2016); Ling et al. (2016), models based on specific techniques in natural
language processing (NLP) Wei et al. (2019); Cao et al. (2019), and pre-
trained models Wang et al. (2021); Ahmad et al. (2021); Guo et al. (2022)
adopt Seq2Seq approaches. It is meaningless for models to generate non-
executable code, and grammatically correct code is the basis for code
execution. However, Seq2Seq approaches generate codes with no guarantee of
grammatical correctness (GC), which hinders the usability of generated codes.
Therefore, the GSS method is necessary for the code generation task.

Figure 2: Examples of a parser parsing Python code.

Inspired by the operating principle of programming language parsers in
practice, we argue that code can be efficiently constrained by grammar without
relying on AST. Fig. 2 demonstrates some examples of a parser parsing python
code. (a) and (b) are two typical grammatical errors. (a) indicates that the
end with non-accept state error, while (b) indicates the terminal symbol
mismatch error. Generally, the Seq2Seq models need to recognize grammatical
errors during the code generation process. However, the vast majority of codes
generated during this situation are incomplete, i.e., parsing these codes will
produce grammatical errors like (a). In addition, the parser can only
recognize grammatical errors like (b) if the terminal symbol is input to the
parser. It implies that the Seq2Seq models need to attempt the generated code
token in the candidate set one by one, which is strenuous. 111For example,
there are only 83 syntax-strings in Python, and for a 10k-length vocabulary,
more than 99% of tokens belong to three token-types NAME, STRING, and NUMBER.
If the Seq2Seq models want to generate these types of tokens based on the
semantics but are prevented by the grammar, the cost of attempts could be
overwhelming. Therefore, although the code can be grammatically bounded, it is
still problematic to work with the parser directly on the Seq2Seq model for
code generation. In the practical code generation process, what we really care
about is the terminal symbol mismatch error, while the end with non-accept
state error can be tolerable until the end of the code generation process.
Therefore, we need a valid terminal symbol set for generating the next
prediction and an accept state set for recognizing completed codes.

To this end, we propose a GSS code generation framework equipped with a PDAM
named CODEP. PDAM facilitates CODEP to generate bounded next predictions in a
valid set and end with an accept state. Compared to the traditional code
generation model, which inputs NL and generates code tokens auto-regressively,
CODEP additionally enters the PDA state of the current code token and predicts
that of the following code token to help our model learn better
representations. The PDA state also assists CODEP in jointly generating better
well-formed code. Moreover, with the help of PDAM, CODEP provides a guarantee
for the GC of generated codes. To demonstrate the superiority of our proposed
method, we conduct a series of experiments on several benchmark datasets.
Extensive experimental results and analyses verify the effectiveness and
generality of CODEP.

The main contribution of this paper can be summarized in fourfold:

  * •

we propose CODEP, a GSS code generation framework, which achieves significant
improvements. To the best of our knowledge, we are the first to implement GSS
models for GPCG successfully.

  * •

We propose a PDAM, which can be applied directly to Seq2Seq models without
training to ensure that the generated code satisfies grammatical constraints.

  * •

We propose the state representation, state prediction task, and joint
prediction for the decoder of CODEP. The first one enters additional valid
information into CODEP. The second one enables CODEP to understand the
workflow of PDAM. The last one assists CODEP in avoiding the generation of
code tokens of inappropriate PDA states.

  * •

Extensive experimental results also demonstrate the effectiveness of our
proposed method, which outperforms the previous state-of-the-art (SOTA)
methods on several benchmark datasets.

##  2 Related Work

###  2.1 Seq2Seq-based Methods for GPCG

Ling et al. (2016) considered the code generation task as a conditional text
generation task and adopted the Seq2Seq model to address it. Some later work
followed this generation approach. For example, Wei et al. (2019) proposed a
dual training framework to train both code generation and representation tasks
simultaneously. CodeT5 Wang et al. (2021), UniXcoder Guo et al. (2022),
Codegen Nijkamp et al. (2022), and InCoder Fried et al. (2022) applied the
pre-trained models to the code generation task. In addition, Wang et al.
(2022) considered the compilability of the code. They treated it as one of the
training objectives based on the pre-trained model, reinforcement learning,
and generative adversarial network.

Although COMPCODE proposed in Wang et al. (2022) improved the code compilation
rate, it still fails to guarantee the GC of generated codes.

###  2.2 Seq2Tree-based Methods for GPCG

Dong and Lapata (2016) first considered the Seq2Tree model for code
generation. To exploit the grammatical information of the code, Yin and Neubig
(2017) and Rabinovich et al. (2017) adopted the encoder-decoder architecture
to output the AST node sequence. On this basis, TRANX Yin and Neubig (2018)
was proposed and became an effective and widely used Seq2Tree model. Lots of
work was based on and improved upon TRANX, such as ML-TRANX Xie et al. (2021),
TRANX-RL Jiang et al. (2021), ASED Jiang et al. (2022), and APT Dong et al.
(2022). Moreover, Sun et al. (2019) and Sun et al. (2020) applied the Tree-
based convolutional neural network and Tree-based Transformer for code
generation. Shen et al. (2022) proposed Subtoken-TranX that was adopted by
Alibaba’s BizCook platform, which is the first domain code generation system
adopted in industrial development environments.

Compared to Seq2Tree-based methods, our proposed method not only ensures GC
but also retains the advantages of Seq2Seq-based approaches.

Table 1: Comparisons of Seq2Seq methods and Seq2Tree methods.

| Seq2Seq | Seq2Tree
---|---|---
Data collection | Easy | Hard
(Especially for large-scale data)
Transfer of NLP | Easy | Medium
technologies
Output length | Same as the code | Much longer than the code token
token sequence | sequence (about 1.5-2 times)
Grammatical | No | Yes
correctness

###  2.3 GSS-based Methods for Domain-specific Code Generation

Scholak et al. (2021) proposed PICARD that constrains auto-regressive decoders
of language models through incremental parsing for Structured Query Language
(SQL). PICARD used a parser for filtering in the beam search phase, meaning
its candidate hypotheses might have terminal symbol-matching errors. Poesia et
al. (2022) proposed constrained semantic decoding to address this problem for
domain-specific languages (DSLs). However, it is challenging for these methods
to extend from DSLs like SQL to GPLs like Python, which has more complicated
grammar Poesia et al. (2022).

_For all we know, the GSS model for GPCG has not yet been studied._

##  3 Methodology

Our proposed method can be applied to any Seq2Seq encoder-decoder framework.
In this paper, we employ the Transformer-based encoder-decoder models Vaswani
et al. (2017) to generate GPL codes. Our model mainly modifies the decode side
of Transformer, which adds PDAM, state representation, state prediction task,
and joint prediction to help the model generate better well-formed code. The
details are described in the following.

###  3.1 PDAM

The vast majority of GPLs are context-free languages, such as Python, Java,
C++, etc., which can be accepted by a PDA Chomsky (1962); Schützenberger
(1963); Evey (1963). A PDA ${\displaystyle M}$ can be defined as:

$M=(S\,,\Sigma\,,\Gamma\,,s_{0}\,,g_{0}\,,A\,,\delta\,)$ (1)

where $S$ is a finite set of states, $\Sigma$ is a finite set of input
symbols, $\Gamma$ is a finite set of stack symbols, $s_{0}\in S$ is the start
state, $g_{0}\in\Gamma$ is the starting stack symbol, $A\subseteq S$, where
$A$ is the set of accept states, and $\delta$ is a transition function,
mapping $S\,\times\Gamma\,\times(\Sigma\,\cup\left\\{\varepsilon\,\right\\})$
into the finite subsets of $S\times\Gamma^{*}$, where $\displaystyle*$ is the
Kleene star.

The input symbols in $\Sigma$ are finite, because the grammar of GPLs merges
the same type of tokens. For example, in Python, all the input symbols belong
to one of the 83 syntax-strings and 10 token-types. Fig. 3 illustrates the
diagram of a PDA constructed by Python grammar. As shown in Fig. 3, ‘import’
and ‘as’ are syntax-strings, while ‘numpy’ and ‘np’ belong to the NAME token-
type . Given a valid input symbol, PDA jumps to the valid states and stacks
according to the transition function $\delta$.

Figure 3: Schematic diagram of a PDA constructed by Python grammar.

Given current state $s$ and current stack $g$, PDA ${\displaystyle M}$ is able
to provide the set of valid input symbols according to the transition function
$delta$. If $\varepsilon$ belongs to the set of valid input symbols, the state
and stack transferred after entering $\varepsilon$ should be taken into
account as well. Eventually, we merge all sets of valid input symbols under
each state and stack and remove $\varepsilon$ from them.

Algorithm 1 Pseudocode for PDAM to obtain the valid set.

0: PDA ${\displaystyle M}$, current state $s$ and current stack $g$.

0: The set of valid Input symbols and corresponding states and stacks.

1: Initial empty valid set $V$ and empty queue $Q$.

2: Enqueue($Q$, $(s,g)$).

3: repeat

4: $s,g\leftarrow$ Dequeue($Q$).

5: for $(s,g,\mathcal{I})\in\delta.keys()$ do

6: if $\mathcal{I}=\varepsilon$ then

7: Enqueue($Q$, $\delta(s,g,\varepsilon)$).

8: else

9: $V\leftarrow V\cup(\mathcal{I},\delta(s,g,\mathcal{I}))$.

10: end if

11: end for

12: until $Q$ is empty

13: return $V$

###  3.2 Model Architecture

Fig. 4 shows the model architecture of CODEP. At the encoder side, it maps an
input sequence of NL utterances
$\bm{x}=\left\\{x_{1},x_{2},\cdots,x_{n}\right\\}$ to a sequence of continuous
NL representations $\bm{z}=\left\\{z_{1},z_{2},\cdots,z_{n}\right\\}$. At the
decoder side, given the NL representation sequence $\bm{z}$, it then generates
an output sequence of code tokens one element at a time. At each step, the
model generates the next code token in the auto-regressive manner Graves
(2013), which deploys previous generation results as the additional input.

####  3.2.1 State Representation

The state is a type of valuable information to comprehend the status of the
generated valid prefix in PDAM. In order to exploit states provided by PDAM,
we set both code tokens and corresponding states as the input of the decoder.
Specifically, given the code tokens
$\bm{y}=\left\\{y_{1},y_{2},\cdots,y_{t-1}\right\\}$, we can obtain the
corresponding states $\bm{s}=\left\\{s_{1},s_{2},\cdots,s_{t-1}\right\\}$ from
PDAM. Then, $\bm{y}$ and $\bm{s}$ are converted into $\bm{h}$ as follows:

$\displaystyle\bm{h}_{y}=\bm{e}_{y}+\operatorname{PE}(\bm{y})$ (2)
$\displaystyle\bm{h}_{s}=\bm{e}_{s}+\operatorname{PE}(\bm{s})$ (3)
$\displaystyle\bm{h}_{t}=\operatorname{Concat}(\bm{h}_{y},\bm{h}_{s})$ (4)

where $\bm{e}_{y}$ and $\bm{e}_{s}$ are the embedding of $\bm{y}$ and $\bm{s}$
respectively, $\operatorname{Concat}(\bm{h}_{y},\bm{h}_{s})$ indicates the
concatenation of $\bm{h}_{y}$ and $\bm{h}_{s}$, and $\operatorname{PE}$
indicates the positional encoding Gehring et al. (2017):

$\displaystyle\operatorname{PE}_{(pos,2j)}$
$\displaystyle=\sin\left(pos/10000^{2j/d}\right)$
$\displaystyle\operatorname{PE}_{(pos,2j+1)}$
$\displaystyle=\cos\left(pos/10000^{2j/d}\right)$

where $pos$ is the position, $j$ is the dimension, and $d$ is the number of
dimensions (i.e., embedding size).

Figure 4: Diagram of CODEP.

####  3.2.2 Code Token Prediction Task

To generate code, $\bm{h}_{t}$ and the NL representation sequence $\bm{z}$ are
fed into the decoder:

$\bm{a}_{t}=\operatorname{Decoder}(\bm{h}_{t},\bm{z})$ (5)

The output of the multi-headed self-attention in the decoder layer is computed
via:

$\displaystyle\bm{Q_{i}}=\bm{H}\bm{W_{i}^{Q}},\bm{K}=\bm{H}\bm{W_{i}^{K}},\bm{V}=\bm{H}\bm{W_{i}^{V}},$
(6)
$\displaystyle\bm{head_{i}}=\operatorname{softmax}\left(\frac{\bm{Q_{i}}\bm{K_{i}}^{\top}}{\sqrt{d_{k}}}\right)\bm{V_{i}},$
(7)
$\displaystyle\bm{head}=\operatorname{Concat}(head_{1},head_{2},...,head_{h})\bm{W^{O}}$
(8)

Where $\bm{W_{i}^{Q}}\in{\mathbb{R}}^{d_{x}\times
d_{k}},\bm{W^{K}}\in{\mathbb{R}}^{d_{x}\times
d_{k}},\bm{W^{V}}\in{\mathbb{R}}^{d_{x}\times d_{h}}$ are learnable parameter
matrices. After that, a feed-forward layer and layer normalization are
followed.

According to $\bm{a}_{t}$, we can compute the following probabilities:

$\displaystyle
p(\operatorname{gen}\mid\bm{a}_{t})=\operatorname{softmax}(\textbf{W}_{x}\bm{a}_{t}),$
(9) $\displaystyle
p(\operatorname{copy}\mid\bm{a}_{t})=1-p(\operatorname{gen}\mid\bm{a}_{t}),$
(10) $\displaystyle
p(v_{c}\mid\operatorname{gen},\bm{a}_{t})=\operatorname{softmax}(\bm{e}_{v_{c}}^{T}\textbf{W}_{v}\bm{a}_{t}),$
(11) $\displaystyle
p(x_{i}\mid\operatorname{copy},\bm{a}_{t},\bm{x})=\operatorname{softmax}(\bm{h}_{x_{i}}^{T}\textbf{W}_{x}\bm{a}_{t}).$
(12)

where $\textbf{W}_{g}$, $\textbf{W}_{v}$, and $\textbf{W}_{x}$ are three
different parameter matrices, $v_{c}$ indicates one of code tokens in the
vocabulary, and $\bm{h}_{x_{i}}$ is calculated by the pointer network Vinyals
et al. (2015). Finally, at t step, the probability of the generated code token
$y_{t}$ can be defined as:

$\displaystyle p(y_{t}\mid y_{<t},s_{<t},\bm{x})$
$\displaystyle=p(\operatorname{gen}\mid\bm{a}_{t})p(v_{c}\mid\operatorname{gen},\bm{a}_{t})$
$\displaystyle+p(\operatorname{copy}\mid\bm{a}_{t})p(x_{i}\mid\operatorname{copy},\bm{a}_{t},\bm{x}).$
(13)

Figure 5: An example of the state prediction task.

####  3.2.3 State Prediction Task

To better understand the parsing process of PDAM, we introduce an auxiliary
task of the PDA states prediction. Similar to the code tokens prediction,
given $y_{<t}$, $s_{<t}$, and $\bm{x}$, we can calculate the probability of
the state at $t$ step as follows:

$p(s_{t}\mid
y_{<t},s_{<t},\bm{x})=\operatorname{softmax}(\bm{e}_{v_{s}}^{T}\textbf{W}_{s}\bm{a}_{t}),$
(14)

where $v_{s}$ indicates one of states in the vocabulary, $\textbf{W}_{s}$ is
another parameter matrix, and $\bm{a}_{t}$ is calculated via (5). In this
paper, we adopt the non-terminal symbols of production rules in the GPL
grammar as the states, which can be changed with the PDA you build. Fig. 5
shows an example of the state prediction task. We can see that the state
prediction task also gives guidance for generating code tokens. For example,
‘for’, ‘in’, and ‘:’ are the fixed pairing of ‘for_stmt’ in Python.

###  3.3 Training and Inference

The purpose of the training stage is to minimize the total cross-entropy loss,
which is defined as:

$\displaystyle\mathcal{L}^{\operatorname{ce}}_{total}(\bm{x},\bm{y},\bm{s};\bm{\theta})$
$\displaystyle=\mathcal{L}^{\operatorname{ce}}_{y}(\bm{x},\bm{y},\bm{s};\bm{\theta})$
(15)
$\displaystyle+\alpha\cdot\mathcal{L}^{\operatorname{ce}}_{s}(\bm{x},\bm{y},\bm{s};\bm{\theta}),$
$\displaystyle\mathcal{L}^{\operatorname{ce}}_{y}(\bm{x},\bm{y},\bm{s};\bm{\theta})=-\sum_{t=1}^{T}\log
p\left(y_{t}\mid y_{<t},s_{<t},\bm{x};\bm{\theta}\right),$ (16)
$\displaystyle\mathcal{L}^{\operatorname{ce}}_{s}(\bm{x},\bm{y},\bm{s};\bm{\theta})=-\sum_{t=1}^{T}\log
p\left(s_{t}\mid y_{<t},s_{<t},\bm{x};\bm{\theta}\right),$ (17)

where $\bm{\theta}$ indicates the parameter of our model, and $\alpha$ is a
coefficient used to control the impacts of
$\mathcal{L}^{\operatorname{ce}}_{y}$ and
$\mathcal{L}^{\operatorname{ce}}_{s}$.

In the inference stage, we are able to calculate $p(y_{t}\mid
y_{<t},s_{<t},\bm{x})$ via (13) with additional constraints of PDAM
$y_{t}\in\\{\mathcal{I}|(\mathcal{I},s,g)\in V\\}$, where $V$ is obtain
according to Algorithm 1. Similarly, we can calculate $p(s_{t}\mid
y_{<t},s_{<t},\bm{x})$ via (14). Then, we leverage them to predict the result
jointly as follows:

$\displaystyle p(y^{\prime}_{t}\mid y_{<t},s_{<t},\bm{x})$
$\displaystyle=\frac{1}{1+\alpha}\cdot(p(y_{t}\mid y_{<t},s_{<t},\bm{x})$
$\displaystyle+\frac{\alpha}{1+\alpha}\cdot p(\hat{s_{t}}\mid
y_{<t},s_{<t},\bm{x})),$ (18)

where $\hat{s_{t}}$ indicates the real state of PDAM, which is obtained from
$\delta(s_{t-1},g_{t-1},y_{t})$.

Algorithm 2 The inference procedure of CODEP.

0: The PDA ${\displaystyle M}$, hyperparameter $\alpha$, parameters of CODEP
$\bm{\theta}$, and NL utterances $\bm{x}$.

0: The generated code tokens $\bm{y^{\prime}}$.

1: Initial $t\leftarrow 0$ and obtain $s_{0}$ and $g_{0}$.

2: repeat

3: Obtain $V$ according to Algorithm 1.

4: Calculate $p(y_{t}\mid y_{<t},s_{<t},\bm{x})$ via (13), where
$y_{t}\in\\{\mathcal{I}|(\mathcal{I},s,g)\in V\\}$.

5: Calculate $p(s_{t}\mid y_{<t},s_{<t},\bm{x})$ via (14).

6: Calculate $p(y^{\prime}_{t}\mid y_{<t},s_{<t},\bm{x})$ via (18).

7: $y^{\prime}_{t}\leftarrow\arg\max p(y^{\prime}_{t}\mid
y_{<t},s_{<t},\bm{x})$.

8: $y_{t}\leftarrow y^{\prime}_{t}$.

9: $s_{t+1},g_{t+1}=\delta(s_{t},g_{t},y^{\prime}_{t})$

10: $t\leftarrow t+1$

11: until $s_{t}\in A$

12: return $\bm{y^{\prime}}$

##  4 Experiment Setup

###  4.1 GPL and Datasets

We build the PDA for the most popular GPL Python (both Python2 and Python3)
and conduct experiments on four benchmark datasets as follows:

CONALA Yin et al. (2018) contains 2,879 real-world data of manually annotated
NL questions and their Python3 solutions on STACK OVERFLOW.

DJANGO Oda et al. (2015) contains 18805 NL-annotated python2 code extracted
from the Django Web framework.

JUICE-10K contains 10K training samples randomly selected from the training
set of JUICE Agashe et al. (2019), and the validation and test sets of
JUICE-10K are consistent with those of JUICE. Due to the high demand for
training resources, we use a subset instead of the full JuICe.

MBPP Austin et al. (2021) contains 974 Python programming problems, and each
sample consists of a NL, a code, and 3 test cases.

Table 2: Statistics of datasets.

Dataset | Examples Num | Avg Length
---|---|---
Train | Dev | Test | NL | Code
CONALA | 2,175 | 200 | 500 | 10.2 | 15.1
DJANGO | 16000 | 1000 | 1805 | 14.1 | 10.6
JUICE | 10000 | 1,831 | 2,115 | 40.4 | 43.4
MBPP | - | - | 974 | 48.0 | 61.0

Table 3: Comparison of CODEP and Seq2Tree methods without pre-training.

Model | CONALA | DJANGO | JUICE-10K
---|---|---|---
BLEU | CodeBLEU | ACC. | ACC. | BLEU | CodeBLEU
Seq2Tree methods w/o pre-training |  |  |  |  |  |
TRANX Yin and Neubig (2019) | $24.35\pm 0.4$ | $26.80\pm 0.6$ | $2.5\pm 0.7$ | $77.3\pm 0.4$ | $4.63\pm 0.2$ | $13.09\pm 0.4$
ML-TRANX Xie et al. (2021) | $24.42\pm 0.8$ | $27.59\pm 1.0$ | $2.2\pm 0.4$ | $79.6\pm 0.3$ | $4.75\pm 0.4$ | $13.43\pm 0.5$
TRANX-RL Jiang et al. (2021) | $25.47\pm 0.7$ | $25.14\pm 1.0$ | $2.6\pm 0.4$ | $77.9\pm 0.5$ | $6.08\pm 0.3$ | $13.78\pm 0.6$
APT Dong et al. (2022) | $27.56\pm 0.7$ | $28.54\pm 0.7$ | $2.9\pm 0.7$ | $79.3\pm 0.5$ | $5.78\pm 0.5$ | $14.19\pm 0.5$
Transformer-AST | $26.50\pm 0.4$ | $27.31\pm 0.4$ | $1.8\pm 0.2$ | $77.1\pm 0.5$ | $4.80\pm 0.3$ | $13.57\pm 0.4$
LSTM | $24.23\pm 0.7$ | $25.76\pm 0.8$ | $2.0\pm 0.7$ | $70.5\pm 0.4$ | $4.54\pm 0.5$ | $13.13\pm 0.7$
CODEP (LSTM-based) | $28.22\pm 0.8$ | $30.12\pm 0.9$ | $3.0\pm 0.6$ | $79.6\pm 0.6$ | $6.32\pm 0.6$ | $15.24\pm 0.6$
Transformer | $24.21\pm 0.8$ | $25.40\pm 0.9$ | $2.0\pm 0.6$ | $74.1\pm 0.7$ | $4.67\pm 0.4$ | $14.80\pm 0.5$
CODEP (Transformer-based) | $27.87\pm 0.9$ | $29.49\pm 0.9$ | $2.6\pm 0.8$ | $79.7\pm 0.5$ | $7.26\pm 0.5$ | $16.09\pm 0.6$

###  4.2 Baselines

We select two typical Seq2Seq models (LSTM Hochreiter and Schmidhuber (1997)
and Transformer Vaswani et al. (2017)) as our base model.

####  4.2.1 Non-pre-trained Seq2Tree Methods

TRANX Yin and Neubig (2019) uses a Seq2Tree model to generate the AST as the
intermediate representation of code.
ML-TRANX (Xie et al., 2021) adopts a mutual learning framework to train models
for different traversals-based decodings jointly.
TRANX-RL (Jiang et al., 2021) uses a context-based branch selector to
determine optimal branch expansion orders for multi-branch nodes dynamically.
APT Dong et al. (2022) uses antecedent prioritized loss to help Seq2Tree
models attach importance to antecedent predictions by exploiting the position
information of the generated AST nodes.
Transformer-AST indicates a Seq2Tree method based on Transformer instead of
LSTM adopted in other Seq2Tree methods mentioned above.

####  4.2.2 Pre-trained Seq2Seq Methods

BART Lewis et al. (2020) is a pre-training approach for text generation tasks
that learns to map corrupted documents to the original.
CodeT5 Wang et al. (2021) is a unified pre-trained encoder-decoder Transformer
model that makes better use of the code semantics conveyed from developer-
assigned identifiers.
CodeGen Nijkamp et al. (2022) is a unified generation model that allows left-
to-right code generation and code infilling/editing by the causal mask
language modeling training objective.
InCoder Fried et al. (2022) is a series of large-scale language models trained
on NL and programming data for conversation-based program synthesis.

###  4.3 Evaluation Metrics

To evaluate the effectiveness of different methods, we use three widely-used
evaluation metrics in code generation: the exact matching accuracy (ACC.), the
corpus-level BLEU-4 (BLEU), and CodeBLEU Ren et al. (2020). We also use GC and
grammatical correctness percentage (GCP) to show the percentage of the
grammatical generated codes in all generated codes.

###  4.4 Implementation Details

We train our model with Adam Kingma and Ba (2015) optimizer on a single GPU of
Tesla A100-PCIe-40G. We set the sizes of the word embedding, code embedding,
state embedding, and hidden state as 256, 256, 256, and 512, respectively. The
learning rate is set to $0.0001$ for Transformer-based CODEP and $0.001$ for
LSTM-based CODEP. For additional hyperparameter $\alpha$, we pick
$\alpha\in[0,1]$ on the validation sets. The beam size is set to 15 for all
datasets without pre-training and 2 for the MBPP dataset. For the approaches
mentioned in Section 4.2, we follow the settings in their paper. To mitigate
the instability of the model training, we exhibit the average performance of
the model running five times.

###  4.5 Research Questions

We conduct experiments to answer the following research questions:

  * •

RQ1: How does the proposed CODEP perform compared to the SOTA Seq2Tree
methods, which guarantees the syntactical correctness of generated codes
relying on AST, without pre-training?

  * •

RQ2: How does each part of our proposed method contribute to CODEP?

  * •

RQ3: How does the proposed PDAM improve the effectiveness of the SOTA pre-
trained Seq2Seq methods?

  * •

RQ4: How does CODEP generate grammatical codes with the help of PDAM?

  * •

RQ5: How well do codes generated by CODEP compile, and are there any solutions
to compile errors?

Table 4: Ablation study of CODEP, where SR means state representation, SPT
means state prediction task, and JP means joint prediction.

Model | CONALA | DJANGO | JUICE-10K
---|---|---|---
BLEU | CodeBLEU | GCP | ACC. | GCP | BLEU | CodeBLEU | GCP
CODEP (LSTM-based) | $28.22$ | $30.12$ | $100\%$ | $79.6$ | $100\%$ | $6.32$ | $15.24$ | $100\%$
-JP | $27.81$ | $29.54$ | $100\%$ | $78.7$ | $100\%$ | 5.96 | 15.04 | $100\%$
-SPT | $27.36$ | $28.95$ | $100\%$ | $77.3$ | $100\%$ | 5.38 | 14.26 | $100\%$
-SR | $27.97$ | $29.98$ | $100\%$ | $77.6$ | $100\%$ | 5.54 | 14.47 | $100\%$
-PDA-JP | $26.37$ | $27.49$ | $80.8\%$ | $74.1$ | $91.1\%$ | 5.09 | 13.87 | $54.7\%$
-PDA-SR-SPT-JP | $24.23$ | $25.76$ | $69.4\%$ | $70.5$ | $89.5\%$ | $4.54$ | $13.13$ | $43.9\%$
CODEP (Transformer-based) | $27.87$ | $29.49$ | $100\%$ | $79.7$ | $100\%$ | $7.26$ | $16.09$ | $100\%$
-JP | $27.05$ | $28.85$ | $100\%$ | 78.9 | $100\%$ | $7.12$ | $15.82$ | $100\%$
-SPT | $26.72$ | $27.79$ | $100\%$ | 77.8 | $100\%$ | $6.59$ | $15.52$ | $100\%$
-SR | $27.13$ | $28.83$ | $100\%$ | 78.3 | $100\%$ | $6.97$ | $15.68$ | $100\%$
-PDA-JP | $26.32$ | $27.46$ | $81.4\%$ | 76.7 | $92.3\%$ | $5.45$ | $15.20$ | $57.6\%$
-PDA-SR-SPT-JP | $24.21$ | $25.40$ | $70.6\%$ | $74.1$ | $90.1\%$ | $4.67$ | $14.79$ | $45.3\%$

Table 5: Examples of code generation for each pre-trained model with the zero-
shot setting, where NL is ‘Write a function to find the similar elements from
the given two tuple lists’ and GOLD denotes the reference code.

Model | Code
---|---
BART-base | Write a function to find the similar elements from the given two tuple lists.
CodeT5-small | def find_similar_elements
CodeT5-base | def f_elements( )
CodeT5-large | to find the similar elements ." ," ." ," a function to find the similar elements from the given two tuple lists .
CodeGen-350M | def find_similar ( t1 , t2 ): $\backslash$n """ $\backslash$n Finds the similar elements from the given two tuples. $\backslash$n $\backslash$n
InCoder-1B | def similar_elements ( t1, t2 ): $\backslash$n $\backslash$n $\backslash$n $\backslash$n $\backslash$n $\backslash$n $\backslash$n $\backslash$n $\backslash$n $\backslash$n $\backslash$n $\backslash$n
CodeGen-350M +PDA | def find ( a , b ): $\backslash$n for i in range ( len ( a ) ): $\backslash$n if a [ i ] == b [ i ]: $\backslash$n return a[i]
GOLD | def similar_elements ( test_tup1, test_tup2 ) : $\backslash$n res = tuple ( set ( test_tup1 ) & set ( test_tup2 ) ) $\backslash$n return ( res )

##  5 Experimental Results

###  5.1 RQ1: CODEP vs Seq2Tree methods without pre-training

As shown in Tables 3, the proposed CODEP substantially outperforms the
compared Seq2Tree methods without pre-training on three benchmark datasets in
terms of BLEU, CodeBLEU, and Accuracy. It indicates that although both
Seq2Tree methods and CODEP generate the grammatical code, CODEP takes
advantage in extracting semantics. We also find that Transformer-based CODEP
achieves competitive performance against LSTM-based CODEP on CONALA. On DJANGO
and JUICE-10K, Transformer-based CODEP exhibits its superiority due to the
more extensive training data of DJANGO and JUICE-10K and the longer generated
code token sequence of JUICE-10K. Therefore, for the large dataset and long
code generation, we recommend employing Transformer as the base model for
CODEP. Since the decoder of most Seq2Tree methods without pre-training is
based on LSTM, we implement Transformer-AST to verify that the effect of CODEP
is derived from our proposed method rather than the base model. The
experimental results show that the base model is not the main factor affecting
the performance.

###  5.2 RQ2: Ablation Study

In Table 4, we demonstrate the performance of LSTM-based CODEP and
Transformer-based CODEP with the reduction of some parts of them. The results
indicate that each part of CODEP contributes, and PDAM is the most critical
part of CODEP. When PDAM is removed, GCP drops sharply, especially for the
long code generation, and BLEU and CodeBLEU also have degradation. JP becomes
unavailable due to the dependence of JP on the mapping from the token to the
corresponding state provided by PDAM. In addition, it can be seen that only
reducing JP or SR leads to a relatively small decrease in performance, because
they have some overlap with other parts in helping the model. ‘-PDA-SR-SPT-JP’
means that only the base model is used, which has a significant performance
degradation compared to CODEP (both LSTM-based and Transformer-based). It
indicates the effectiveness of our proposed method.

###  5.3 RQ3: Improvement of pre-trained Seq2Seq methods with PDAM

As can be seen from Table 6, pre-trained models achieve significant
improvements with the help of PDAM. It is apparent from this Table 6 that both
BART and CodeT5 fail to work under the zero-shot setting. The reason may be
that BART is only pre-trained on the data of text generation tasks, while
CodeT5 is pre-trained on the paired NL-code data for six GPLs, but is unable
to distinguish what GPL needs to be generated. In addition, we find that the
size of the pre-trained model with PDAM and its effect is basically positively
correlated, because the large pre-trained model has the advantage of capturing
semantic information and PDAM guarantees GC.

Table 6: Zero-shot Results of each pre-trained Seq2Seq method with PDAM.

Model | MBPP (Zero-shot)
---|---
BLEU | CodeBLEU | GC
Seq2Seq methods w/ pre-training |  |  |
BART-base | 0.02 | 3.74 | False
+PDA | 4.83 | 9.17 | True
CodeT5-small | 0.01 | 2.67 | False
+PDA | 4.61 | 8.62 | True
CodeT5-base | 0.01 | 0.57 | False
+PDA | 4.75 | $9.79$ | True
CodeT5-large | 0.02 | 3.98 | False
+PDA | $6.19$ | 9.31 | True
CodeGen-350M | 1.55 | 3.21 | False
+prompt | 0.33 | 10.13 | False
+PDA | $14.44$ | $21.54$ | True
InCoder-1B | 0.12 | 7.41 | False
+prompt | 0.82 | 7.22 | False
+PDA | 9.84 | 17.38 | True

We demonstrate examples of code generation for each pre-trained model with the
zero-shot setting in Table 5. We can see that the output of BART-base directly
replicates NL, which is related to its pre-training task. CodeT5-small and
CodeT5-base understand some syntactic structures, but the short length of
their generated code is the main factor limiting their performance.
CodeT5-large tries to capture some semantic information at the beginning and
continues to copy later. InCoder only generates the function signature and
fails to generate its main body. CodeGen-350M generates the function signature
and some comments capturing semantic information, which shows the ability to
generate complete code if under the proper guidelines. Under the constraints
of PDAM, CodeGen-350M generates well-formed code. Although the code generated
by CodeGen-350M + PDA still falls short of GOLD, we still think it is an
impressive result considering that it is under the zero-shot setting.

Figure 6: Case study.

###  5.4 RQ4: Case Study

As shown in Fig. 6, PDAM gives the output of the transition function $\delta$
according to the current state and stack, i.e., candidate syntax-strings and
token-types, to which the generated token of CODEP should belong. At step 0,
given the NL, CODEP wants to generate the first token, and PDAM yields the
corresponding candidate tokens. Therefore, at step 1, CODEP generates a token
based on candidate tokens and semantics. Specifically, CODEP picks ‘import’
among candidate tokens according to semantics. At the end of the generation
step, CODEP outputs the token of the ENDMARKER token-type, e.g., </s>.

###  5.5 RQ5: Compilation Error Analysis

Program execution needs to be compiled correctly in addition to being
grammatically correct, and GC is the foundation of compilation. The
compilation rate of CODEP on CONALA dataset is $99.2\%$, and Table 7
demonstrates examples of three types of compilation error. The first type of
compilation error is because the number type cannot be assigned, the second
type of compilation error is the error of the duplicate parameter, and the
third type of compilation error is the mapping error. We found that the above
types of errors also occur in the Seq2Tree methods. For the first two types of
errors, we are able to modify the transfer function, while the last type of
error can be avoided using engineering methods.

Table 7: Three types of compilation error generated by CODEP on CONALA
dataset.

Type | Example
---|---
I | 1 = [ ( i , 2 ) for i in range(1) ]
II | plt.plot(x, var1, color=‘str0’, color=‘str1’)
III | html = response.read(str0, ‘str0’)
slot_map={str0: http://www.example.com/}

##  6 Conclusion and Future work

In this paper, we have proposed a code generation method, namely CODEP, that
first considers GSS models for GPCG. CODEP adds PDAM, state representation,
state prediction task, and joint prediction for the decoder. PDAM guarantees
the GC of generated codes, while the others make use of the information
provided by the PDA to generate codes better. As a result, CODEP outperforms
the SOTA Seq2Tree methods on three benchmark datasets. The ablation study
demonstrates the contribution of each component. With the help of PDAM, pre-
trained models also achieve significant improvements.

The GPLs of generated codes of CODEP are limited by the PDA, which can be
constructed by the grammar of the GPLs. In this paper, we only build the PDA
for the most popular GPL Python, including both Python2 and Python3. In future
work, we envision building PDAs for more GPLs, such as C++, Java, etc.
Furthermore, PDAM shows potential in code generation tasks with the zero-shot
setting. It can help pre-trained models apply to niche GPLs that have little
data.

##  Appendix A Appendix

###  A.1 Syntax-strings and Token-types in Python

In Table 8, we show each of the 83 syntax-strings and 10 token-types in
Python.

Table 8: Terminal symbols of Python grammar. | Python
---|---
Syntax-string | ‘for’, ‘in’, ‘try’, ‘finally’, ‘with’, ‘except’, ‘lambda’, ‘or’, ‘and’, ‘not’, ‘del’, ‘pass’,
‘break’, ‘continue’, ‘return’, ‘raise’, ‘from’, ‘import’, ‘as’, ‘nonlocal’,
‘global’,
‘assert’, ‘if’, ‘else’, ‘elif’, ‘while’, ‘async’, ‘def’, ‘@’, ‘(’, ‘)’, ‘->’,
‘:’, ‘**’, ‘*’, ‘,’,
‘=’, ‘;’, ‘^=’, ‘%=’, ‘//=’, ‘@=’, ‘<<=’, ‘**=’, ‘&=’, ‘*=’, ‘|=’, ‘>>=’,
‘-=’, ‘+=’, ‘/=’,
‘.’, ‘…’, ‘<=’, ‘>=’, ‘is’, ‘==’, ‘<’, ‘>’, ‘<>’, ‘!=’, ‘|’, ‘^’, ‘&’, ‘<<’,
‘>>’, ‘+’, ‘-’, ‘%’,
‘/’, ‘//’, ‘~’, ‘!’, ‘await’, ‘False’, ‘[’, ‘{’, ‘True’, ‘None’, ‘]’, ‘}’,
‘class’, ‘yield’
Token-type | ‘NAME’, ‘STRING’, ‘NUMBER’, ‘INDENT’, ‘DEDENT’, ‘FSTRING_START’,
‘FSTRING_END’, ‘FSTRING_STRING’, ‘NEWLINE’, ‘ENDMARKER’

###  A.2 Deterministic Pushdown Automaton

Due to computer memory constraints, most of GPLs are deterministic context-
free languages, which can be accepted by a deterministic pushdown automaton
(DPDA) Salomaa et al. (2001). A PDA ${\displaystyle M}$ is deterministic only
if both the following two conditions are satisfied:

  * •

$\forall s\in
S,g\in\Gamma,\mathcal{I}\in\Sigma\cup\left\\{\varepsilon\right\\}$,
$|\delta(s,g,\mathcal{I})|=1$ where $|\delta(s,g,\mathcal{I})|$ indicates the
element number of the set $\delta(s,g,\mathcal{I})$.

  * •

$\forall s\in S,g\in\Gamma,\mathcal{I}\in\Sigma$, if
$\delta(s,g,\varepsilon)\not=\emptyset$, then
$\delta\left(s,g,\mathcal{I}\right)=\emptyset$.

For the same deterministic context-free grammar, the general PDA is more
accessible to construct than the DPDA. In addition, deterministic context-free
languages are a proper subset of context-free languages. Therefore, in this
paper, our proposed method is based on the general PDA, which is also
applicable to the DPDA.

Figure 7: Effects of $\alpha$ on the validation sets.

###  A.3 Effect of $\alpha$

The coefficient $\alpha$ is an important hyperparameter that controls the
relative impacts of the state prediction task in the training stage and that
of the joint prediction with the state probability in the inference stage.
Therefore, we investigate the effect of $\alpha$ on our proposed framework in
Fig. 7, which varies $\alpha$ from 0 to 1 with an increment of 0.2 on the
validation set of CONALA, DJANGO, and JUICE-10K datasets. The experimental
results show that as $\alpha$ increases, the performance of CODEP increases
when $\alpha\leq 0.4$, and its tendency is related to the base model and the
dataset when $\alpha>0.4$. What stands out in Fig. 7 is that CODEP performs
relatively well on each of the above datasets when $\alpha$ is set to 1.0,
which indicates the effectiveness of the state prediction task and joint
prediction. For experiments in this paper, we set $\alpha$ as the optimal
$\alpha$ in Fig. 7. In particular, the $\alpha^{\prime}s$ of the LSTM-based
CODEP and the transformer-based CODEP are set to (0.6, 1.0), (0.8, 1.0), and
(1.0, 0.4) on CONALA, DJANGO, and JUICE-10K dataset, respectively.

## References

  * Agashe et al. (2019) Rajas Agashe, Srinivasan Iyer, and Luke Zettlemoyer. 2019.  Juice: A large scale distantly supervised dataset for open domain context-based code generation.  In _EMNLP/IJCNLP (1)_ , pages 5435–5445.
  * Ahmad et al. (2021) Wasi Uddin Ahmad, Saikat Chakraborty, Baishakhi Ray, and Kai-Wei Chang. 2021.  Unified pre-training for program understanding and generation.  In _NAACL-HLT_ , pages 2655–2668. Association for Computational Linguistics.
  * Austin et al. (2021) Jacob Austin, Augustus Odena, Maxwell I. Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie J. Cai, Michael Terry, Quoc V. Le, and Charles Sutton. 2021.  Program synthesis with large language models.  _CoRR_ , abs/2108.07732.
  * Cao et al. (2019) Ruisheng Cao, Su Zhu, Chen Liu, Jieyu Li, and Kai Yu. 2019.  Semantic parsing with dual learning.  In _ACL (1)_ , pages 51–64.
  * Chomsky (1962) Noam Chomsky. 1962.  Context-free grammars and pushdown storage.  _MIT Res. Lab. Electron. Quart. Prog. Report._ , 65:187–194.
  * Dong and Lapata (2016) Li Dong and Mirella Lapata. 2016.  Language to logical form with neural attention.  In _ACL (1)_.
  * Dong et al. (2022) Yihong Dong, Ge Li, and Zhi Jin. 2022.  Antecedent predictions are dominant for tree-based code generation.  _CoRR_ , abs/2208.09998.
  * Evey (1963) R James Evey. 1963.  Application of pushdown-store machines.  In _Proceedings of the November 12-14, 1963, fall joint computer conference_ , pages 215–227.
  * 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.  _CoRR_ , abs/2204.05999.
  * Gehring et al. (2017) Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N. Dauphin. 2017\.  Convolutional sequence to sequence learning.  In _ICML_ , volume 70 of _Proceedings of Machine Learning Research_ , pages 1243–1252. PMLR.
  * Graves (2013) Alex Graves. 2013.  Generating sequences with recurrent neural networks.  _CoRR_ , abs/1308.0850.
  * Guo et al. (2022) Daya Guo, Shuai Lu, Nan Duan, Yanlin Wang, Ming Zhou, and Jian Yin. 2022.  Unixcoder: Unified cross-modal pre-training for code representation.  In _ACL (1)_ , pages 7212–7225. Association for Computational Linguistics.
  * Hochreiter and Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber. 1997.  Long short-term memory.  _Neural Comput._ , 9(8):1735–1780.
  * Jia and Liang (2016) Robin Jia and Percy Liang. 2016.  Data recombination for neural semantic parsing.  In _ACL (1)_.
  * Jiang et al. (2022) Hui Jiang, Linfeng Song, Yubin Ge, Fandong Meng, Junfeng Yao, and Jinsong Su. 2022\.  An AST structure enhanced decoder for code generation.  _IEEE ACM Trans. Audio Speech Lang. Process._ , 30:468–476.
  * Jiang et al. (2021) Hui Jiang, Chulun Zhou, Fandong Meng, Biao Zhang, Jie Zhou, Degen Huang, Qingqiang Wu, and Jinsong Su. 2021.  Exploring dynamic selection of branch expansion orders for code generation.  In _ACL/IJCNLP_ , pages 5076–5085.
  * Kingma and Ba (2015) Diederik P. Kingma and Jimmy Ba. 2015.  Adam: A method for stochastic optimization.  In _ICLR_.
  * 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. Association for Computational Linguistics.
  * Ling et al. (2016) Wang Ling, Phil Blunsom, Edward Grefenstette, Karl Moritz Hermann, Tomás Kociský, Fumin Wang, and Andrew W. Senior. 2016.  Latent predictor networks for code generation.  In _ACL (1)_.
  * Nijkamp et al. (2022) Erik Nijkamp, Bo Pang, Hiroaki Hayashi, Lifu Tu, Huan Wang, Yingbo Zhou, Silvio Savarese, and Caiming Xiong. 2022.  A conversational paradigm for program synthesis.  _CoRR_ , abs/2203.13474.
  * Oda et al. (2015) Yusuke Oda, Hiroyuki Fudaba, Graham Neubig, Hideaki Hata, Sakriani Sakti, Tomoki Toda, and Satoshi Nakamura. 2015.  Learning to generate pseudo-code from source code using statistical machine translation.  In _2015 30th IEEE/ACM International Conference on Automated Software Engineering (ASE)_ , pages 574–584. IEEE.
  * Poesia et al. (2022) Gabriel Poesia, Oleksandr Polozov, Vu Le, Ashish Tiwari, Gustavo Soares, Christopher Meek, and Sumit Gulwani. 2022.  Synchromesh: Reliable code generation from pre-trained language models.  _CoRR_ , abs/2201.11227.
  * Rabinovich et al. (2017) Maxim Rabinovich, Mitchell Stern, and Dan Klein. 2017.  Abstract syntax networks for code generation and semantic parsing.  In _ACL (1)_ , pages 1139–1149.
  * Raychev et al. (2014) Veselin Raychev, Martin T. Vechev, and Eran Yahav. 2014.  Code completion with statistical language models.  In _PLDI_ , pages 419–428.
  * Ren et al. (2020) Shuo Ren, Daya Guo, Shuai Lu, Long Zhou, Shujie Liu, Duyu Tang, Neel Sundaresan, Ming Zhou, Ambrosio Blanco, and Shuai Ma. 2020.  Codebleu: a method for automatic evaluation of code synthesis.  _CoRR_ , abs/2009.10297.
  * Salomaa et al. (2001) Arto Salomaa, Derick Wood, and Sheng Yu. 2001.  _A half-century of automata theory: celebration and inspiration_.  World scientific.
  * Scholak et al. (2021) Torsten Scholak, Nathan Schucher, and Dzmitry Bahdanau. 2021.  PICARD: parsing incrementally for constrained auto-regressive decoding from language models.  In _EMNLP (1)_ , pages 9895–9901. Association for Computational Linguistics.
  * Schützenberger (1963) Marcel Paul Schützenberger. 1963.  On context-free languages and push-down automata.  _Information and control_ , 6(3):246–264.
  * Shen et al. (2022) Sijie Shen, Xiang Zhu, Yihong Dong, Qizhi Guo, Yankun Zhen, and Ge Li. 2022.  Incorporating domain knowledge through task augmentation for front-end javascript code generation.  _CoRR_ , abs/2208.10091.
  * Sun et al. (2019) Zeyu Sun, Qihao Zhu, Lili Mou, Yingfei Xiong, Ge Li, and Lu Zhang. 2019.  A grammar-based structural CNN decoder for code generation.  In _AAAI_ , pages 7055–7062.
  * Sun et al. (2020) Zeyu Sun, Qihao Zhu, Yingfei Xiong, Yican Sun, Lili Mou, and Lu Zhang. 2020.  Treegen: A tree-based transformer architecture for code generation.  In _AAAI_ , pages 8984–8991.
  * Sutskever et al. (2014) Ilya Sutskever, Oriol Vinyals, and Quoc V. Le. 2014.  Sequence to sequence learning with neural networks.  In _NIPS_ , pages 3104–3112.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  In _NIPS_ , pages 5998–6008.
  * Vinyals et al. (2015) Oriol Vinyals, Meire Fortunato, and Navdeep Jaitly. 2015.  Pointer networks.  In _NIPS_ , pages 2692–2700.
  * Wang et al. (2022) Xin Wang, Yasheng Wang, Yao Wan, Fei Mi, Yitong Li, Pingyi Zhou, Jin Liu, Hao Wu, Xin Jiang, and Qun Liu. 2022.  Compilable neural code generation with compiler feedback.  In _ACL (Findings)_ , pages 9–19. Association for Computational Linguistics.
  * Wang et al. (2021) Yue Wang, Weishi Wang, Shafiq R. Joty, and teven C. H. Hoi. 2021.  Codet5: Identifier-aware unified pre-trained encoder-decoder models for code understanding and generation.  In _EMNLP (1)_ , pages 8696–8708.
  * Wei et al. (2019) Bolin Wei, Ge Li, Xin Xia, Zhiyi Fu, and Zhi Jin. 2019.  Code generation as a dual task of code summarization.  In _NeurIPS_ , pages 6559–6569.
  * Xie et al. (2021) Binbin Xie, Jinsong Su, Yubin Ge, Xiang Li, Jianwei Cui, Junfeng Yao, and Bin Wang. 2021.  Improving tree-structured decoder training for code generation via mutual learning.  In _AAAI_ , pages 14121–14128.
  * Yin et al. (2018) Pengcheng Yin, Bowen Deng, Edgar Chen, Bogdan Vasilescu, and Graham Neubig. 2018\.  Learning to mine aligned code and natural language pairs from stack overflow.  In _MSR_ , pages 476–486.
  * Yin and Neubig (2017) Pengcheng Yin and Graham Neubig. 2017.  A syntactic neural model for general-purpose code generation.  In _ACL (1)_ , pages 440–450.
  * Yin and Neubig (2018) Pengcheng Yin and Graham Neubig. 2018.  TRANX: A transition-based neural abstract syntax parser for semantic parsing and code generation.  In _EMNLP (Demonstration)_ , pages 7–12.
  * Yin and Neubig (2019) Pengcheng Yin and Graham Neubig. 2019.  Reranking for neural semantic parsing.  In _ACL_ , pages 4553–4559.


2021

[1] Rasmus Netterstrøm These authors contributed equally to this work.

These authors contributed equally to this work.

[1]Department of computer science, University of Copenhagen,
Universitetsparken, København, 2100, Denmark

2]Department of Neuroscience, University of Copenhagen, Blegdamsvej,
København, 2200, Denmark

3]Department of Clinical Neurophysiology, Rigshospitalet, Blegdamsvej,
København, 2100, Denmark

# Deep Learning-Assisted Localisation of Nanoparticles in synthetically
generated two-photon microscopy images

<EMAIL_ADDRESS>Nikolay Kutuzov<EMAIL_ADDRESS>Sune Darkner
<EMAIL_ADDRESS>Maurits Jørring Pallesen<EMAIL_ADDRESS>Martin
Johannes Lauritzen<EMAIL_ADDRESS>Kenny Erleben<EMAIL_ADDRESS>François
Lauze<EMAIL_ADDRESS>*  [  [

###### Abstract

Tracking single molecules is instrumental for quantifying the transport of
molecules and nanoparticles in biological samples, e.g., in brain drug
delivery studies. Existing intensity-based localisation methods are not
developed for imaging with a scanning microscope, typically used for in vivo
imaging. Low signal-to-noise ratios, movement of molecules out-of-focus, and
high motion blur on images recorded with scanning two-photon microscopy (2PM)
in vivo pose a challenge to the accurate localisation of molecules. Using
data-driven models is challenging due to low data volumes, typical for in vivo
experiments. We developed a 2PM image simulator to supplement scarce training
data. The simulator mimics realistic motion blur, background fluorescence, and
shot noise observed in vivo imaging. Training a data-driven model with
simulated data improves localisation quality in simulated images and shows why
intensity-based methods fail.

###### keywords:

deep learning, single molecule localization, two-photon microscopy.

##  1 Introduction

Tracking single molecules or nanoparticles helps to solve a variety of
biomedical problems, e.g., drug design and delivery Kucharz2021 , studying the
regulation of gene expression Elf2007 , determining mechanisms of viral
infections Liu2012 , and many others. Single-molecule (or particle) tracking
(SMT) is becoming a popular tool to study mechanisms of transport at the
nanoscale in living cells and tissues Godin2017 . In vivo imaging, especially
in scattering tissues like the brain, becomes an ultimate challenge for SMT
due to (i) low photon budgets causing low signal-to-noise ratio on recorded
images, (ii) motion artefacts caused by the animal’s breathing, heartbeat,
movement, and (iii) severe motion blur caused by the slowness of scanning
microscopy typically used in vivo. These factors challenge the accurate
localisation and tracking of molecules on images recorded in vivo, for example
with two-photon microscopy (2PM; Fig. 4a).

Object localisation and detection is a core topic in computer vision. Feature-
based detection, using SIFT SIFT , HoG dalal_triggs:2005 , to cite the most
classically used, were replaced by deep learning (DL) approaches, notably ren-
etal:2015 ; yolo1  and many were built for specialised situations in medical
imaging. On images of single molecules recorded with CCD or EMCCD cameras,
localisation is achievable even for images with low signal-to-noise ratio
or(and) moderate motion blur Mortensen2010 ; Mortensen2021 . The localisation
algorithms, however, require writing down an expression for the distribution
of fluorescence in recorded images, which may not be possible for imaging
moving molecules with 2PM: random motion of molecules overlaid with the laser
beam’s motion (Fig. 4b) can make modelling the distribution of fluorescence on
a recorded image prohibitively complex. Neural networks have the potential to
approximate these complex distributions of fluorescence and return the desired
position of the molecule. Deep learning (DL) has been predominantly used for
the localisation of immobile molecules with wide-field microscopes equipped
with a CCD or EMCCD camera Petrov2017 ; Boyd2018 ; Nehme2018  in various in
vitro experiments but not yet for localisation of moving molecules with 2PM in
living animals.

Experimental data recorded in vivo can be so scarce and expensive to collect
that training of the neural networks is more likely to rely on artificial
data. Simulating large training data sets, informed by well-known theories of
microscope image formation, has been used to train deep networks to localise
immobile objects on images recorded with a CCD or EMCCD camera Petrov2017 ;
Boyd2018 . In vivo imaging in scattering tissues, e.g. in the brain, however,
is typically done with scanning 2PM, which records images pixel-by-pixel,
line-by-line (Fig. 4b), as opposed to collecting all image pixels
simultaneously with a CCD camera. Here we show how to train a neural network
to localise single simulated quantum dots (QDs) and how the commonly used
intensity-based detections fails to localise the QD on pixel-by-pixel and
line-by-line simulated images. We used a combination of realistic simulated
data, based on modelling diffusion of QDs during exposure and realistic motion
of the laser beam. Many DL approaches have been used for the detection and
localisation of objects, both with natural images and medical imaging
modalities too. We demonstrate how a neural network used for object
localisation and detection for natural images, can accommodate the QD
detection and localisation task by training it on simulated data,
incorporating effects of motion blur caused by very different rates between
diffusion motion and camera motion, and then allow localisation of QDs on
images sampled from a simulated test dataset.

##  2 Results

We started by developing a protocol to generate realistic synthetic 2PM images
for training the FRCNN model maskfcnn ; pytorch-keypoint-rcnn ; pytorch ,
later used for the localization of QDs. By combining a simple Brownian-like
motion of QDs with a deterministic motion of the laser beam we reproduced
motion blur characteristics for scanning microscopy (Fig. 2). We
experimentally measured point spread function (PSF) dimensions and estimated
fluorescence background from experimental images ($N_{bkg}$ in Eq. 1) to mimic
real images as close as possible. Finally, with Poisson noise, we simulated
images that looked very similar to images recorded in the brain in vivo (Fig.
2).

In Figure 3, we highlight the primary issue of intensity-based detectors in
four images. The intensity-based detector, TrackPy trackpy , detects high-
intensity pixels even though the high-intensity pixels are not the trajectory
centroid of the QD. Training directly on the centroid coordinate enables FCNN
to detect the trajectory centroid of the QD with higher accuracy. Then we
conducted three experiments to demonstrate further the value of including the
trajectory centroid in the DL training pipeline. Each experiment had a
different diffusion coefficient that provided a different motion blur. The
first experiment had a diffusion coefficient of $0.75\mu m^{2}/s$, the second
of $1\mu m^{2}/s$ and the third of $1.75\mu m^{2}/s$. Besides the different
diffusion coefficients, we tested the robustness of each model by changing the
epsilon-threshold ($\varepsilon$) of the precision-recall curve; The epsilons
were $1$,$2$ and $3$, corresponding to a distance of $70nm$, $140nm$ and
$210nm$ given a pixel size of $0.07\mu m$.

Figure 1: Four artificially simulated images (a-d) with diffusion coefficient going from 0 to $1.75\mu m^{2}/s$. In images b-d, the green circle is the starting point of the trajectory, and the cyan circle is the ending point. (a) show a QD that diffuses with a diffusion coefficient of $0\mu m^{2}/s$, showing the microscope PSF with noise. (b-d) show QD diffusing with increasing diffusion coefficient: $0.75,1\text{ and }1.75\mu m^{2}/s$.  Table 1: This table shows the AUC-score of all three experiments. In the first column, the QD diffuse with $0.75\mu m^{2}/s$, in the second column, the QD diffuse with $1\mu m^{2}/s$ and in the third column, the QD diffuse with $1.75\mu m^{2}/s$. In all three columns, the FCNN model achieves the best AUC-score.  D | $0.75\mu m^{2}/s$ | $1\mu m^{2}/s$ | $1.75\mu m^{2}/s$
---|---|---|---
$nm$ | $70$ | $140$ | $210$ | $70$ | $140$ | $210$ | $70$ | $140$ | $210$
TrackPy | 0.06 | 0.42 | 0.73 | 0.04 | 0.30 | 0.61 | 0.02 | 0.19 | 0.47
FCNN | 0.48 | 0.91 | 0.99 | 0.41 | 0.87 | 0.98 | 0.26 | 0.72 | 0.92
|  |  |  |  |  |  |  |  |

####  2.0.1 Experiment 1 - $0.75\;\mu m^{2}/s$

This experiment concentrates on localising the trajectory centroid of a QD
when point sources of light closely resemble the microscope’s PSF. That
ensures the trajectory centroid is close to the high-intensity pixels, hence
less motion blur. We choose a realistic diffusion coefficient value of
$0.75\mu m^{2}/s$. Figure 1 (b) shows a QD moving with a diffusion coefficient
of $0.75\mu m^{2}/s$, allowing less motion blur. Table 1, column 1
demonstrates that TrackPy and FCNN localise the trajectory centroid for
localisation within $210nm$ with an AUC score of $0.73$ and $0.99$.

####  2.0.2 Experiment 2 - $1\;\mu m^{2}/s$

We increase the diffusion coefficient, thus allowing more motion blur. Figure
1 (c) illustrates how the trajectory steps get larger, and the high-intensity
pixels diverge from the microscope PSF; the distance between the trajectory
centroid and high-intensity pixels becomes unpredictable. Table 1, column 2
tells the immediate effect of more motion blur in the image. The intensity-
based detector TrackPy drops from $0.73$ to $0.61$ AUC score for localisation
within $210nm$, whereas FCNN only drops from an AUC score of $0.99$ to $0.98$
for localisation within $210nm$.

####  2.0.3 Experiment 3 - $1.75\;\mu m^{2}/s$

At last, we make the high-intensity pixels diverge from the microscope PSF as
much as possible for a realistic diffusion coefficient value. We use the
diffusion coefficient of $1.75\mu m^{2}/s$, i.e., ten times larger than the
diffusion coefficient used in experiment one. Figure 1 (d) illustrates how
increasing motion blur scatters the high-intensity pixels across a larger area
and, in some cases, makes the high-intensity pixels interchangeable for noise.
Again we see in table 1, column three, that TrackPy drops in AUC performance,
going from $0.61$ to $0.47$, whereas FCNN keeps an AUC score above $0.90$ for
localisation within $210nm$.

Figure 2: Comparison of three artificially simulated images (a-c) with real
experimental images of QDs in the brain (d-f), recorded with 2PM. Note how the
simulated data reproduces the motion blur: panels a, b, and f, show elongated
fluorescence intensity distributions in the vertical direction due to the
movement of the QDs. Scale bars are 500 nm wide.  Figure 3: Localization of
QDs on simulated images (a-d) Blue and purple crosses show the localisation of
QDs estimated with FCNN and TrackPy, respectively. Green curves show simulated
trajectories; green crosses show trajectory centroids.

##  3 Discussion

To tackle the problem of accurate localization of QD in pixel-by-pixel
recorded images, we developed a simulator that generates a particle undergoing
a Brownian motion recorded by a two-photon laser scanning microscope.
Implementing the motion of the laser beam during pixel-by-pixel image scanning
allows the reproduction of complex motion blur typically observed in
experimental data. Note, that our simulator simplifies local transport in vivo
and 2PM. First, we compute the trajectory of the QD by Brownian motions and
second, we only generate images with one QD showing. To simplify the
simulation, we neglected the excess and readout noises Mortensen2010
introduced by the detector (photomultiplier tube; PMT) and assumed that the
actual number of recorded photons is Poisson-distributed. We show that
realistic simulated data, which incorporates the effect of motion blur, non-
specific background fluorescence, and Poisson shot noise, improves the deep
learning-based localization of QDs, seen in experiments 1 to 3. The primary
result of our experiments demonstrates that the FCNN is robust against motion
blur and keeps an AUC score above $0.90$, whereas TrackPy drops the AUC score
to $0.47$. But we only trained one model. So for future implementations, we
intend to make a five-fold cross-validation. To our knowledge, this is the
first simulator of training data for 2PM. It opens for better and faster
trajectory recovery and analysis, an invaluable tool to understand how
molecules move in complex, turbid, biological tissues, e.g., in the brain.

##  4 Methods

###  4.1 Data description

We simulated a training dataset of 29838 images and three testing datasets of
2570, 2410 and 1870 images. We chose a training dataset size of 29838 to avoid
overfitting, and we chose a test dataset size of 6850 because it is common
practice to choose a test dataset that is $20\%$ of the size of the training
dataset. The following parameters reproduce in vivo experiment parameters: All
images were 24 pixels tall and 25 pixels wide, trajectories of the QDs were
simulated with a set of diffusion coefficients between $0.75\mu m^{2}/s$ and
$1.75\mu m^{2}/s$, the pixel size was $0.07\mu m$, the pixel exposure time was
$2\mu s$ and the retracing time was $1.1ms$. To ensure single-molecule
localisation images contained 1 QD.

###  4.2 Deep Learning Model

We use the open-access, and easy-to-use PyTorch implementation of FRCNN-
keypoint with a resNet backbone resnet  pre-trained on the MS COCO dataset
coco .

###  4.3 Train Configuration

We only trained our model using 20 epochs because the model converged. Of the
20 models, we chose the model with the highest validation accuracy. Due to
limited GPU memory, we chose a batch size of 6 images, a validation batch size
of 4 images, and stochastic gradient descent. We tuned the learning rate to
0.001 pytorch-vision .

Figure 4: Imaging single quantum dots in the brain of living mice with two-
photon microscopy (2PM). a A simplified scheme of the experimental setup: 2PM
images quantum dots (QDs) inside the brain of a living, anaesthetised mouse by
sweeping the laser beam across the brain and collecting fluorescence emitted
by the QDs. b Images are recorded by moving the laser beam (by rotating x- and
y-scan mirrors shown in a) along straight lines (solid red). Between the
lines, the beam is retraced (dashed red): this ensures the same scanning speed
as the laser beam. c Dimensions of the point spread function (PSF; drawn to
scale) of a 2PM inside a brain Kutuzov2018  and a scheme of brain
extracellular spaces — a 3D network of channels between brain cells membranes
— where QDs are moving. In 2PM, the PSF represents the square of the laser
intensity near the focus and it is approximately Gaussian, with standard
deviation along the z-axis being $\approx$6 times larger than along x- and
y-axes.

###  4.4 Materials

We used a commercial two-photon scanning microscope (FVMPE-RS) built by
Olympus (Japan), equipped with a Ti:Sapphire laser and GaAsP detectors
(photomultiplier tubes; PMTs). We used commercial QTracker655 quantum dots
(QDs; ThermoFisher Scientific, USA), which we diluted 15 times in saline
before intravenous injection into anaesthetized mice.

###  4.5 Animals and surgical procedures

All protocols involving mice were approved by the Danish National Committee on
Health Research Ethics, in accord with the guidelines of the European
Council’s Convention for the Protection of Vertebrate Animals used for
Experimental and Other Scientific Purposes, and complied with the ARRIVE
guidelines. We used four male wild type mice (C57bl/6j; Janvier-labs). We
followed standard surgical procedures including tracheotomy, cauterization,
and craniotomy, described in detail elsewhere Kucharz2021 . QDs were injected
intravenously and released in the the brain parenchyma by ablating a small
capillary with high laser power.

###  4.6 Simulating 2PM images of diffusing QDs.

####  4.6.1 Laser beam motion simulation

First, we calculated the trajectory of the laser beam during image acquisition
(Fig. 4b) assuming that the laser beam moves along straight-line segments at
the constant speed of 2 $\mu m$ per pixel (typically used in in vivo
experiments) until it reaches the end of a scanning line. Then the laser beam
is retraced (no photons are recorded during this phase) to the beginning of
the next line, which takes 1.1 ms. The process is repeated until the whole
image is recorded. We estimated the position of the laser beam with the same
time step as for diffusion trajectories ($\delta t$).

####  4.6.2 Random walk

Second, we simulated a random walk trajectory (independently for all three
coordinates, $xyz$) of QDs during image acquisition by cumulatively adding
normally-distributed random numbers $\mathcal{N}(0,\sigma^{2})$, with zero
mean and variance $\sigma^{2}=2D\delta t$, where $D$ is the diffusion
coefficient of the QDs and $\delta$ is a time step. To align the timestep of
the simulator with the timestep $\delta$, we needed to sample from two normal
distributions. The first distribution uses a timestep of $2\mu m$, whereas the
second uses the retracing timestep of $1.1ms$.

####  4.6.3 Point Spread Function

Third, we modelled the expected photon count in individual image pixels.
Assuming that the square of the laser light intensity (two photons are
absorbed simultaneously at 2PM) near the focus is approximately a 3D Gaussian
MPM  with standard deviations $\sigma_{xy}$ in the $xy$-plane and $\sigma_{z}$
(depth) along the optical axis (Fig. 4c), the expected number of photons is

$\displaystyle
N=\frac{N_{qd}}{2\pi\sigma_{xy}^{2}}\exp{\left(-\frac{(X-x)^{2}+(Y-y)^{2}}{2\sigma_{xy}^{2}}\right)}\cdot\frac{1}{\sqrt{2\pi\sigma_{z}^{2}}}\exp{\left(-\frac{(Z-z)^{2}}{2\sigma_{z}^{2}}\right)}+N_{bkg}$
(1)

with $(X,Y,Z)$ the position of the QD, $(x,y,z)$ the position of a laser beam,
$N_{qd}$ is the total possible number of emitted photons by a QD in the focal
plane in the direction of the objective, and $N_{bkg}$ is expected number of
photons of the fluorescent background (autofluorescence). We used the value
$N_{bkg}=8.4$, which was estimated by averaging the background fluorescence on
images recorded in the brain. The value of $N_{qd}$ was adjusted to imitate
the signal-to-noise ratio of the brain images. We used realistic, measured in
vivo, parameters of the 2PM PSF: $\sigma_{xy}=0.21$ $\mu$m and
$\sigma_{z}=1.36$ $\mu$m Kutuzov2018 .

###  4.7 Localization of QDs

Several DL-based detectors have been proposed in recent years redmon-etal:2015
; lin-etal:2017 , summarised in a recent review Youzi-etal:2020 . Many of
these detectors have been tailored to detect complex objects. QDs are simple
practically point sources of light but they do not image as the microscope’s
PSF because of the motion blur. Furthermore, images of QDs can be very noisy
due to the high fluorescence background or movement of the QD out of focus.
For such problems, the Faster R-CNN-Keypoint (FRCNN) maskfcnn ; pytorch-
keypoint-rcnn ; pytorch , which predicts object bounding boxes along with key
points, has proven very reliable and has available implementations and was
adopted in this work.

###  4.8 Metrics

For training the model we have used classical detection metrics for the
bounding boxes and keypoints predicted by FRCNN. Detection speed in frames per
second is of no interest in our problem. Precision and recall
deepLearningSurvey  are based on intersection over union (IoU):

$\displaystyle IoU(\hat{B},B)=\frac{area(\hat{B}\cap B)}{area(\hat{B}\cup B)}$
(2)

with $\hat{B}$ is a predicted bounding box and $B$ the associated ground truth
box. Determining true positives (TP), false positives (FP), true negatives
(TN), false negatives (FN) for both bounding box and keypoints are computed
according to the COCO standard coco ; cocoapi . For evaluation, we needed
another metric since an intensity-based detector such as TrackPy does not
predict bounding boxes trackpy . Therefore, we computed the area under the
curve (AUC) from the precision and the recall. In addition, we modified the
precision and the recall seen in objectDetectionBasedOnDeepLearning  to use an
L2 distance instead of an IoU constraint.

precision $\displaystyle=\frac{tp}{tp+fp}$ (3) recall
$\displaystyle=\frac{tp}{tp+fn}$ (4)

That means we define a prediction to TP if the distance between the predicted
key point and the ground truth key point is below or equal to a threshold
$\varepsilon$. FP is defined as if the distance exceeds the threshold or the
prediction is duplicated. At last, FN is defined as if the detector fails to
predict. We then computed the precision and the recall using equation eqs. 3
and 4. Finally, we computed the area under the curve (AUC) using precision and
recall. Precision tells the probability that a prediction is correct, whereas
recall tells the probability of detecting all possible ground truths. The AUC
combines the metric, e.g. If the AUC score is close to 1, the detector
correctly predicts and detects all possible ground truths. If the AUC score is
close to 0, the detector either poorly predicts, fails to predict or both.

## Declarations

### Funding

Not applicable

### Conflict of interest/Competing interests

Not conflict or competing interests

### Ethics approval

Not applicable

### Consent to participate

Not applicable

### Consent for publication

Not applicable

### Availability of data and materials

We released the data as open-access at the time of publication.

### Code availability

Most of the code in this project is custom-made and will be open-access at the
time of publication. However, we are happy to share the code upon request.

### Authors’ contributions

The first and second co-first authors contributed equally to this work. They
made the primary work of this paper. The third, fifth and sixth authors
contributed with writing, supervision and concept development. The fourth
author helped with the programming

## References

  *   * (1) Kucharz, K., Kristensen, K., Johnsen, K.B., Lund, M.A., Lønstrup, M., Moos, T., Andresen, T.L., Lauritzen, M.J.: Post-capillary venules are the key locus for transcytosis-mediated brain delivery of therapeutic nanoparticles. Nat. Commun. 12(1), 4121 (2021)
  * (2) Elf, J., Li, G.-W., Xie, X.S.: Probing transcription factor dynamics at the single-molecule level in a living cell. Science 316(5828), 1191–1194 (2007)
  * (3) Liu, S.-L., Zhang, Z.-L., Tian, Z.-Q., Zhao, H.-S., Liu, H., Sun, E.-Z., Xiao, G.F., Zhang, W., Wang, H.-Z., Pang, D.-W.: Effectively and efficiently dissecting the infection of influenza virus by quantum-dot-based single-particle tracking. ACS Nano 6(1), 141–150 (2012)
  * (4) Godin, A.G., Varela, J.A., Gao, Z., Danné, N., Dupuis, J.P., Lounis, B., Groc, L., Cognet, L.: Single-nanotube tracking reveals the nanoscale organization of the extracellular space in the live brain. Nat. Nanotechnol. 12(3), 238–243 (2017)
  * (5) Lowe, D.G.: Distinctive image features from scale-invariant keypoints. Int. J. Comput. Vision 60(2), 91–110 (2004). https://doi.org/10.1023/B:VISI.0000029664.99615.94
  * (6) Dalal, N., Triggs, B.: Histograms of oriented gradients for human detection. In: IEEE Computer Society Conference on Computer Vision and Pattern Recognition, vol. 1, pp. 886–893 (2005)
  * (7) Ren, S., He, K., Girshick, R., Sun, J.: Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Methods. IEEE Transaction on Pattern Analysis and Machine Understanding 39(6), 1137–1149 (2017)
  * (8) Redmon, J., Divvala, S.K., Girshick, R.B., Farhadi, A.: You only look once: Unified, real-time object detection. CoRR abs/1506.02640 (2015) 1506.02640
  * (9) Mortensen, K.I., Churchman, L.S., Spudich, J.A., Flyvbjerg, H.: Optimized localization analysis for single-molecule tracking and super-resolution microscopy. Nat. Methods 7(5), 377–381 (2010)
  * (10) Mortensen, K.I., Flyvbjerg, H., Pedersen, J.N.: Confined brownian motion tracked with motion blur: Estimating diffusion coefficient and size of confining space. Frontiers in Physics 8, 601 (2021)
  * (11) Petrov, P.N., Shechtman, Y., Moerner, W.E.: Measurement-based estimation of global pupil functions in 3D localization microscopy. Opt. Express 25(7), 7945–7959 (2017)
  * (12) Boyd, N., Jonas, E., Babcock, H., Recht, B.: DeepLoco: Fast 3D Localization Microscopy Using Neural Networks (2018)
  * (13) Nehme, E., Weiss, L.E., Michaeli, T., Shechtman, Y.: Deep-STORM: super-resolution single-molecule microscopy by deep learning. Optica, OPTICA 5(4), 458–464 (2018)
  * (14) He, K., Gkioxari, G., Dollár, P., Girshick, R.B.: Mask R-CNN. CoRR abs/1703.06870 (2017) 1703.06870
  * (15) PyTorch: Keypoint R-CNN. https://pytorch.org/vision/stable/models/keypoint_rcnn.html Accessed 2023-02-22
  * (16) Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., Chintala, S.: Pytorch: An imperative style, high-performance deep learning library. In: Wallach, H., Larochelle, H., Beygelzimer, A., d Alch’ie Buc, F., Fox, E., Garnett, R. (eds.) Advances in Neural Information Processing Systems 32, pp. 8024–8035. Curran Associates, Inc., ??? (2019). http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf
  * (17) Allan, D.B., Caswell, T., Keim, N.C., van der Wel, C.M., Verweij, R.W.: Soft-matter/trackpy: Trackpy V0.5.0. https://doi.org/10.5281/zenodo.4682814. https://doi.org/10.5281/zenodo.4682814
  * (18) He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: IEEE Conference on Computer Vision and Pattern Recognition, pp. 770–778 (2016)
  * (19) Lin, T., Maire, M., Belongie, S.J., Bourdev, L.D., Girshick, R.B., Hays, J., Perona, P., Ramanan, D., Dollár, P., Zitnick, C.L.: Microsoft COCO: Common Objects in Context, vol. V, pp. 740–755 (2014)
  * (20) Vision, P.: Pytorch Vision. https://github.com/pytorch/vision Accessed 2023-03-17
  * (21) Kutuzov, N., Flyvbjerg, H., Lauritzen, M.: Contributions of the glycocalyx, endothelium, and extravascular compartment to the blood–brain barrier. PNAS 115(40), 9429–9438 (2018)
  * (22) Zipfel, W.R., Williams, R.M.E., Webb, W.W.: Nonlinear magic: multiphoton microscopy in the biosciences. Nature Biotechnology 21, 1369–1377 (2003)
  * (23) Redmon, J., Divvala, S.K., Girshick, R.B., Farhadi, A.: You Only Look Once: Unified, Real-Time Object Detection. In: Proceedings of the IEEE International Conference on Computer Vision and Pattern Recognition, pp. 779–788 (2016)
  * (24) Lin, T.-Y., Goyal, P., Girshick, R.B., He, K., Dollár, P.: Focal Loss for Dense Object Detection. In: International Conference on Computer Vision, pp. 2980–2988 (2017)
  * (25) Youzi, X., Zhiqiang, T., Jiachen, Y., Yinshu, Z., Liu, S., Shaoyi, D., Xuguang, L.: A review of object detection based on deep learning. Multimedia Tools and Applications 79(33-34), 23729–23791 (2020)
  * (26) Liu, L., Ouyang, W., Wang, X., Fieguth, P., Chen, J., Liu, X., Pietikäinen, M.: Deep Learning for Generic Object Detection: A Survey
  * (27) ppwwyyxx, rbg: pycocotools. GitHub (2018)
  * (28) Youzi, X., Zhiqiang, T., Jiachen, Y., Yinshu, Z., Liu, S., Shaoyi, D., Xuguang, L.: A review of object detection based on deep learning. Multimedia Tools and Applications 79(33-34), 23729–23791 (2020). Copyright - © Springer Science+Business Media, LLC, part of Springer Nature 2020; Last updated - 2020-08-22