[Paper Review] Memory Augmented Policy Optimization for Program Synthesis and Semantic Parsing

Liang Chen, Mohammad Norouzi+3 more|arXiv (Cornell University)|Jul 6, 2018

Parallel Computing and Optimization TechniquesComputer Science99 citations

TL;DR

MAPO introduces a memory-augmented policy gradient method that uses a memory buffer of high-reward trajectories to reduce gradient variance, achieving strong results on weakly supervised semantic parsing tasks.

ABSTRACT

We present Memory Augmented Policy Optimization (MAPO), a simple and novel way to leverage a memory buffer of promising trajectories to reduce the variance of policy gradient estimate. MAPO is applicable to deterministic environments with discrete actions, such as structured prediction and combinatorial optimization tasks. We express the expected return objective as a weighted sum of two terms: an expectation over the high-reward trajectories inside the memory buffer, and a separate expectation over trajectories outside the buffer. To make an efficient algorithm of MAPO, we propose: (1) memory weight clipping to accelerate and stabilize training; (2) systematic exploration to discover high-reward trajectories; (3) distributed sampling from inside and outside of the memory buffer to scale up training. MAPO improves the sample efficiency and robustness of policy gradient, especially on tasks with sparse rewards. We evaluate MAPO on weakly supervised program synthesis from natural language (semantic parsing). On the WikiTableQuestions benchmark, we improve the state-of-the-art by 2.6%, achieving an accuracy of 46.3%. On the WikiSQL benchmark, MAPO achieves an accuracy of 74.9% with only weak supervision, outperforming several strong baselines with full supervision. Our source code is available at https://github.com/crazydonkey200/neural-symbolic-machines

Motivation & Objective

• Motivate and address high-variance policy gradient in deterministic, discrete-action settings like program synthesis.
• Leverage a memory buffer of promising trajectories to decompose the objective into inside-buffer and outside-buffer expectations.
• Propose mechanisms (memory weight clipping, systematic exploration, distributed sampling) to stabilize and scale training.
• Evaluate MAPO on weakly supervised semantic parsing benchmarks to assess improvements in sample efficiency and robustness.

Proposed method

• Express the expected return as a weighted sum of two terms: inside-memory-buffer expectation and outside-memory-buffer expectation.
• Define memory buffer B with total probability pi_B and residual outside-buffer probability 1 - pi_B.
• Use pi_B-containing gradient plus a second gradient from outside-buffer samples, with pi_B acting as a weight (Equation 7).
• Introduce memory weight clipping pi_B^c = max(pi_B, alpha) to stabilize cold-start training (Equation 8).
• Implement systematic exploration to discover high-reward trajectories using a bloom-filter-based set of fully explored prefixes (Algorithm 1).
• Employ distributed actor-learner sampling to parallelize data collection and gradient computation (Algorithm 2).
• Provide an exact or stratified approximation of the inside-buffer expectation depending on buffer size (exact enumeration or sampling).
• Use rejection sampling to sample outside-buffer trajectories from the current policy when computing the outside expectation.

Experimental results

Research questions

• RQ1 How can memory-based replay be integrated into policy gradient methods for deterministic, discrete-action domains to reduce gradient variance?
• RQ2 Can a principled decomposition into inside-buffer and outside-buffer expectations, coupled with memory-weighted gradients, improve sample efficiency in weakly supervised program synthesis?
• RQ3 Do mechanisms like memory weight clipping, systematic exploration, and distributed sampling yield scalable and robust training for MAPO across semantic parsing benchmarks?

Key findings

• MAPO achieved dev/test accuracies of 42.7/43.8 on WikiTableQuestions with a single run and 46.3 when ensembling (as reported).
• MAPO achieved 72.6 test accuracy on WikiSQL under weak supervision, and 74.9 with ensemble on the same benchmark.
• Ablation shows removing systematic exploration (SE) or memory weight clipping (MWC) significantly degrades performance.
• MAPO outperformed several baselines in both WikiTableQuestions and WikiSQL, including methods trained with full supervision.
• Distributed sampling with 30 actors yielded about a 20x speedup in sampling, demonstrating scalable training.
• MAPO demonstrated strong robustness and sample efficiency over traditional REINFORCE and other memory-based approaches.