Composed-Program Induction with Latent Program Lattice

Compositional Reasoning

Compositional reasoning is the process of solving new problems by systematically composing primitive operations into complex transformations. This ability hinges on three key components:

1. Learning the compositional structure that governs how primitives interact
2. Inference using this learned structure to generalize from limited examples, and
3. Abstraction discovery that compresses and refines primitives to accelerate inference.

This paper tackles this key components of compositional reasoning by learning and inferring on a latent program lattice.

Program Induction

A program induction task consists of demonstration input output pairs, generated with a consistent hidden rule. The goal is to generate the test output that would be generated if the implied rule is consistently applied to the test input. To achieve this, we should first learn a good representation of the program space.

An example of program induction task on Rubik's cube

1. Learning Consistent Primitive Program Representation

Latent program basis models the effect of applying a program as a uniform field in the latent program space.

First, the encoder that embeds the objects in the demonstration pairs are trained to learn a consistent representation for primitive programs. The difference vector of the input and output embedding represents the latent program effect, and they should be the same if the implied programs are the same.

By enforcing this program consistency loss, the encoder learns to represent each implied programs as uniform vector fields in the latent program space.

2. Learning to Represent Program Compositionality as a Lattice Structure

Latent program representations for distinct programs form the basis of the latent program lattice

Second, the encoder is trained to understand how the programs compose. Just like how word embeddings map semantic relationship as vector translations, our model tries to learn program composition as vector translations.

As a result, any valid discrete composition of the primitive programs are represented as integer linear combinations of the learned latent program basis, spanning a latent program lattice.

3. Reasoning on the Learned Latent Program Lattice

Given unseen demonstration pairs, the average of difference vector in the latent space (latent program effect) is inferred as the latent representation of the implied program.

Then, unseen demonstration pairs are given and the model induces the implied program utilizing the learned latent program lattice. The encoder embeds each demonstration pairs into latent program effects and their average vector represents the implied latent program. Based on the inductive bias that only the lattice points are the valid discrete compositions of primitive programs, the inferred latent program is snapped to its closest lattice point.

The problem of finding the closest lattice point is the Closest Vector Problem, CVP. Hence, the model basically reduces the program induction into CVP.

There are two modes of reasoning to finally generate the inferred output.

3.1. System-1 Reasoning

PLAE encodes observed input-output pairs, takes their difference as a latent program effect, adds this effect to a new input's embedding, and decodes.

First, system-1 reasoning is fast, and intuitive, but approximate approach. The model directly adds the inferred latent program effect to the embedding of the test input and decode to generate the output. This is agnostic to the complexity of the composed program. It is akin to implicitly noticing the mapping from input to output and directly applying to the test input

3.2. System-2 Reasoning

PLAE factorizes the same effect into a sequence of basis and executes them stepwise
—repeatedly decoding/encoding in a guided “lattice walk” that verifies and refines intermediate states

System-2 reasoning is more slow, but deliberate approach. The inferred latent program effect is decomposed into combinations of latent program basis. Instead of directly applying it, the model applies it step by step, walking on the lattice. It repeats encoding, applying latent program basis, and decoding to adjust and refine its intuition on long compositional reasoning chains.

4. Abstraction via Lattice Reduction

Lattice Reduction refines primitive bases to improve efficiency and uncover more fundamental components

At last, abstraction discovery is powered by lattice reduction. Lattice reduction is finding shorter and more orthogonal basis spanning the same lattice points. This significantly improves CVP solver’s efficiency without needing to re-train the encoder. Additionally, the basis on the reduced lattice represents more independent set of primitive programs that still spans the same compositional program space.

This latent program lattice could provide a principled framework for learning and inference in discrete compositional domain, all in a single unified construct.