[Paper Review] Learning to learn generative programs with Memoised Wake-Sleep

Luke Hewitt, Tuan Le+1 more|arXiv (Cornell University)|Jul 6, 2020

Artificial Intelligence in Games25 references4 citations

TL;DR

The paper proposes Memoised Wake-Sleep (MWS), a novel algorithm that enhances program induction in neuro-symbolic generative models by caching and reusing the best programs discovered during training. By maintaining a finite memory of high-quality latent programs, MWS accelerates inference and improves learning efficiency and accuracy in structured generation tasks such as stroke-based character modeling, cellular automata, and few-shot string concept learning.

ABSTRACT

We study a class of neuro-symbolic generative models in which neural networks are used both for inference and as priors over symbolic, data-generating programs. As generative models, these programs capture compositional structures in a naturally explainable form. To tackle the challenge of performing program induction as an 'inner-loop' to learning, we propose the Memoised Wake-Sleep (MWS) algorithm, which extends Wake Sleep by explicitly storing and reusing the best programs discovered by the inference network throughout training. We use MWS to learn accurate, explainable models in three challenging domains: stroke-based character modelling, cellular automata, and few-shot learning in a novel dataset of real-world string concepts.

Motivation & Objective

• Address the challenge of performing computationally expensive program induction repeatedly during training in neuro-symbolic generative models.
• Improve generalization and interpretability by learning symbolic, compositional programs as latent variables.
• Develop a scalable training algorithm that reduces reliance on repeated, costly inference by reusing previously discovered high-quality programs.
• Enable joint learning of program priors and model parameters in structured domains with sparse, discrete latent spaces.
• Demonstrate effectiveness across diverse domains, including character generation, cellular automata, and real-world string concept induction.

Proposed method

• Introduce Memoised Wake-Sleep (MWS), an extension of the Wake-Sleep algorithm that maintains a finite set $\mathcal{Z}_i$ of the top $K$ programs found for each input $x_i$.
• Use a recognition network $r_\psi(z|x)$ to propose candidate programs and update the memory set $\mathcal{Z}_i$ at each training iteration.
• Optimize a variational lower bound on the data likelihood where the variational distribution $Q_i(z_i)$ is supported on $\mathcal{Z}_i$.
• Train the generative model $p_\theta(z)p_\phi(x|z)$ and the recognition network $r_\psi(z|x)$ in an alternating fashion, using cached programs as effective inference targets.
• Allow for end-to-end learning of both program priors and continuous parameters (e.g., noise levels in cellular automata).
• Apply MWS to three domains: stroke-based character generation, cellular automata with latent rules, and a novel few-shot string concept dataset.

Experimental results

Research questions

• RQ1 Can caching high-quality programs during training significantly reduce the computational burden of repeated program induction in neuro-symbolic models?
• RQ2 Does MWS improve the accuracy and convergence speed of learning symbolic generative programs compared to standard Wake-Sleep and related methods?
• RQ3 Can MWS effectively learn meaningful, interpretable program priors and model parameters (e.g., noise levels) in structured domains like cellular automata?
• RQ4 How well does MWS generalize to few-shot learning of real-world string concepts from minimal examples?
• RQ5 To what extent does MWS outperform advanced baselines like Reweighted Wake-Sleep (RWS) and VIMCO in terms of both speed and accuracy?

Key findings

• MWS significantly outperforms RWS and VIMCO in learning accuracy on synthetic program induction tasks, reducing the distance to the true model by up to 75% on hard tasks.
• On the hard task, MWS achieved a distance of 0.75 to the true model with $K=10$, compared to 1.26 for VIMCO and 2.43 for RWS.
• In the cellular automata domain, MWS accurately estimated the true noise parameter $\epsilon$ (2%) with minimal error, outperforming baselines in both easy and hard rule variants.
• On the novel String-Concepts dataset of 1500 few-shot problems, MWS demonstrated strong few-shot generalization, learning real-world string patterns like dates and emails from minimal examples.
• MWS achieved a 21.8% accuracy on character generation tasks, matching or exceeding state-of-the-art performance despite the complexity of stroke-based program induction.
• The algorithm provided substantial speedups over RWS and VIMCO by reducing the need for repeated, expensive inference across training steps.