[Paper Review] Learning Graph Structure With A Finite-State Automaton Layer

Daniel D. Johnson, Hugo Larochelle+1 more|arXiv (Cornell University)|Jul 9, 2020

Machine Learning and Algorithms40 references4 citations

TL;DR

This paper proposes a differentiable Graph Finite-State Automaton (GFSA) layer that learns to derive abstract relational edges in graph-structured data by modeling paths accepted by finite-state automata. Using implicit differentiation on a graph-based POMDP, the GFSA layer end-to-end learns to add semantically meaningful, task-optimized edges—demonstrating superior performance on program analysis and grid-world reasoning tasks compared to handcrafted or learned baselines.

ABSTRACT

Graph-based neural network models are producing strong results in a number of domains, in part because graphs provide flexibility to encode domain knowledge in the form of relational structure (edges) between nodes in the graph. In practice, edges are used both to represent intrinsic structure (e.g., abstract syntax trees of programs) and more abstract relations that aid reasoning for a downstream task (e.g., results of relevant program analyses). In this work, we study the problem of learning to derive abstract relations from the intrinsic graph structure. Motivated by their power in program analyses, we consider relations defined by paths on the base graph accepted by a finite-state automaton. We show how to learn these relations end-to-end by relaxing the problem into learning finite-state automata policies on a graph-based POMDP and then training these policies using implicit differentiation. The result is a differentiable Graph Finite-State Automaton (GFSA) layer that adds a new edge type (expressed as a weighted adjacency matrix) to a base graph. We demonstrate that this layer can find shortcuts in grid-world graphs and reproduce simple static analyses on Python programs. Additionally, we combine the GFSA layer with a larger graph-based model trained end-to-end on the variable misuse program understanding task, and find that using the GFSA layer leads to better performance than using hand-engineered semantic edges or other baseline methods for adding learned edge types.

Motivation & Objective

• To address the challenge of automatically learning high-level, abstract relations in graph-structured data that improve downstream model performance.
• To model derived relations as path-based reachability constraints using finite-state automata, inspired by program analysis techniques.
• To enable end-to-end training of these relations via differentiable policy learning on a partially observable Markov decision process (POMDP) over graphs.
• To demonstrate that learned relations via GFSA outperform hand-engineered edges and other baselines in program understanding and navigation tasks.
• To provide a differentiable, interpretable mechanism for injecting domain-specific relational inductive biases into graph neural networks.

Proposed method

• The GFSA layer defines a finite-state automaton (FSA) policy that traverses a base graph, where transitions are conditioned on node and edge labels.
• The FSA policy is trained via implicit differentiation on a graph-based POMDP, enabling backpropagation through the automaton's path selection.
• The layer outputs a weighted adjacency matrix representing new edge types corresponding to accepted paths in the FSA.
• The method relaxes discrete path selection into a continuous, differentiable optimization problem using stochastic relaxation and gradient estimation.
• The FSA policy is trained to maximize downstream task performance, with supervision only on start and end nodes, not intermediate steps.
• The approach integrates seamlessly into existing graph neural networks (e.g., GGNN, RAT, GREAT), enabling end-to-end training with the GFSA layer as a learnable component.

Experimental results

Research questions

• RQ1 Can a differentiable layer learn abstract, task-relevant relations in graphs without explicit supervision on intermediate path steps?
• RQ2 Does learning relations via finite-state automata provide better inductive biases than random walks or reinforcement learning baselines in graph reasoning tasks?
• RQ3 Can the GFSA layer reproduce known program analysis patterns (e.g., data dependence) in Python code graphs?
• RQ4 How does the GFSA layer compare to hand-engineered semantic edges in terms of performance on program understanding tasks?
• RQ5 Can the GFSA layer discover useful shortcuts in structured environments like grid-worlds, even without explicit supervision on path structure?

Key findings

• The GFSA layer achieved 94.058 ± 0.064 accuracy on the localization task in the RAT model, outperforming hand-engineered edges (93.100 ± 0.069) and all baselines.
• On the GREAT model, the GFSA layer reached 93.266 ± 0.068 accuracy on localization, surpassing the next best method (NRI head at 92.446 ± 0.072).
• In the grid-world task, the GFSA layer achieved 93.612 ± 0.066 accuracy when trained from middle positions, outperforming RL ablation (90.643 ± 0.079) and random walk (92.644 ± 0.071).
• The GFSA layer successfully reproduced LastRead program analysis patterns in Python ASTs, with visualized automaton behavior matching expected reverse-execution semantics.
• The GFSA layer improved performance on variable misuse detection by up to 1.5% over hand-engineered edges and 3% over base AST-only models in some settings.
• The GFSA layer demonstrated strong generalization, learning effective path-finding policies in unseen grid-world layouts and correctly identifying relevant program dependencies in code.