ICSE-2022 AST-trans：code summarization with efficient tree-structured attention

2022-07-05 · ICSE· Software · code-summarization

Paper: AST-trans：code summarization with efficient tree-structured attention

Code: https://github.com/zetang94/ICSE2022_AST_Trans

AST-trans: Tree attention on parent–child and sibling nodes

Abstract

Abstract syntax trees (ASTs) carry structural information but are often much longer than the corresponding source code. Existing approaches ignore length limits and feed entire ASTs into encoders, which makes it hard to exploit relational dependencies effectively and leads to large computational cost. This work proposes AST-trans, which uses two node relations—ancestor–descendant and sibling—and tree-structured attention to weight relevant and irrelevant nodes. It further gives an efficient implementation that supports fast parallel computation of tree-structured attention.

Introduction

Linearized ASTs are much longer than the corresponding source code; some linearization schemes increase length further—for example, linearized SBT can expand the sequence by several times—so models struggle to recover useful dependencies. This also imposes heavy computation: in transformer-based state-of-the-art methods, attention cost grows with sequence length.

The paper assumes each AST node’s state is shaped mainly by two factors: (1) ancestor–descendant links, which capture cross-block relations at different hierarchy levels; (2) sibling links, which capture sequential relations within a block. Ancestor–descendant structure supports high-level semantics; sibling structure supports low-level semantics. Together they are sufficient for comment generation and for the attention each node needs over the tree.

Main contributions:

AST-trans has linear complexity and encodes long ASTs more efficiently than a standard transformer’s quadratic cost
A full analysis of computational complexity for alternative implementations, supported by theory and experiments
State-of-the-art results on Java and Python summarization datasets
A comparison of representative AST encoding methods and discussion of their trade-offs

AST-TRANS

AST Linearization

Pre-order traversal (POT), structure-based traversal (SBT), and path decomposition (PD).

Relationship Matrices

Ancestor–descendant relations use matrix A; sibling relations use matrix S. Entries in A are shortest-path distances; entries in S are horizontal sibling distances in the AST. P is the relative-distance threshold.

Tree-Structured Attention

Self-Attention

Relative position embedding

Equation (2) has no explicit position terms. The original transformer used absolute position embeddings; here the model uses relative position encoding (more effective for code summarization). Relative position $\delta(i,j)$ is the pairwise distance between nodes $n_i$ and $n_j$.

Disentangled Attention

Attention with Tree-Structured Relationships

The method replaces the relative distance defined on a linear sequence with $\delta_R(i,j)$, where $R$ is either ancestor–descendant relation A or sibling relation S in the tree.

There are two relation types; each maps to one attention head, so the top-level transformer adds no extra parameters.

$V^P$ is the value projection matrix for relative distances; $V^P_{R_{ij}}$ is row $R_{ij}$ of $V^P$. Restricting computation to $\delta_R(i,j) > 0$ reduces the time and space cost of self-attention.

Efficient Implementation

A standard transformer’s cost grows with sequence length. AST-Trans can lower cost by computing attention only where $\delta_R(i,j) > 0$, in the spirit of a sliding window that limits attention to a fixed distance.

With a sliding window, node pairs in a sequence can be laid out linearly (dropping pairs with $\delta(i,j) = 0$ or $2P - 1$) and computed in parallel via matrix tiling. That trick does not transfer directly here, because the positions of related nodes change from tree to tree.

The paper describes five implementation options for AST-Trans and their trade-offs:

Mask

Compute full attention over all node pairs, then zero out scores where $\delta_R(i,j) = 0$. Time and space remain quadratic, as in a standard transformer.

Loop

Iterate over pairs with $\delta_R(i,j) > 0$ and compute attention scores. This is efficient in terms of work but does not parallelize well.

Sparse

Store $\delta_R$ as a sparse matrix $ST(\delta_R)$; frameworks such as PyTorch skip zero entries automatically. This works for content-to-position and position-to-content terms, but not for content-to-content, which still requires mask or loop.

Gather with COO (GC)

Building on the sparse view, content-to-content attention can be optimized with extra gather operations. The idea is to list each query–key pair that must be scored in a one-to-one layout and materialize them as dense matrices.

The total number of gathers is four times the number of nonzeros in $\delta_R$: two for content and two for position.

Gather with decomposed COO (GDC)

To cut gathers in GC, the matrix for $\delta_R$ is decomposed into three submatrices; content embeddings are gathered as follows.

Finally, all $\alpha_{\mathrm{coo}}$ terms are summed.

This yields three benefits:

$K^P$ and $Q^P$ can be reused: each $Q_{\mathrm{row}s}$ and $K{\mathrm{row}_s}$ shares the same relative distance $s$, so the position embedding for $s$ can be fused into content without a gather
Only one quarter as many gather operations
Only one dot product is needed; the second can be reused