Term Autoencoder

Overview

A stochastic autoencoder on simply typed lambda terms, built from the four structural declarations: a signature fixing the multi-sorted term algebra, an Encoder folding each term bottom-up into one fixed-width code vector, a Decoder corecursively expanding a code back into a term, and a loss head scoring the round trip by reconstruction negative log-likelihood.

This example wires all four declarations together in one module. It shows the keyword-led encoder body (an explicit op rule alongside compiler-scaffolded defaults), the per-encoder / per-decoder dim override, a binder that threads a de Bruijn type annotation through the recursion, and a loss attached to a named site via [on=...]. The #! doc comments attach to the declaration that follows and are carried on its AST node's docs field.

QVR Source

# Term Autoencoder
#
# A stochastic autoencoder on simply typed lambda terms: a
# multi-sorted signature fixes the term algebra, an encoder folds
# each term bottom-up into one fixed-width code vector, a decoder
# corecursively expands a code back into a term, and a loss head
# scores the round trip by reconstruction negative log-likelihood.

#! Simply typed lambda terms: Term and Type are the recursively
#! decoded object sorts; Name is a data sort with a closed
#! vocabulary that the decoder samples and scores leaves from.
signature STLC
    sorts
        Term : object
        Type : object [dim=24]
        Name : data   [dim=24, vocab=["f", "x", "plus", "base"]]
    constructors
        Const : Name -> Term
        App   : Term, Term -> Term
        Base  : Name -> Type
    binders
        Lam : binds (x : Term : ty : Type) in (body : Term) -> Term

#! Fold a term bottom-up into one code vector. App carries an
#! explicit op rule; every other operator keeps its learned
#! scaffold. Bound variables enter the context as their type's
#! annotation embedding.
encoder Enc : STLC
    dim Term = 24
    op App(fun, arg) |-> relu(fun + arg)
    var_init Term from Type as ty |-> ty

#! Expand a code vector back into a term, scoring constructor,
#! binder, and bound-variable candidates at every node down to a
#! bounded depth.
decoder Dec : STLC [depth=6]
    dim Term = 24
    body |-> recursive

#! Reconstruction negative log-likelihood of the round trip.
loss reconstruct [weight=1.0, on=decoder(Dec)]
    -Dec.log_prob(term, Enc(term))

Walkthrough

The signature

signature STLC
    sorts
        Term : object
        Type : object [dim=24]
        Name : data   [dim=24, vocab=["f", "x", "plus", "base"]]
    constructors
        Const : Name -> Term
        App   : Term, Term -> Term
        Base  : Name -> Type
    binders
        Lam : binds (x : Term : ty : Type) in (body : Term) -> Term

Term and Type are object sorts, decoded recursively; Name is a data sort whose closed vocab supplies the tokens the decoder samples and scores leaves from. Term deliberately omits its dim: each encoder or decoder over the signature must then supply one via a dim Term = 24 body entry, and the compiler raises when neither the signature nor the block does. The Lam binder introduces a Term-sorted variable annotated by a Type; the framework threads that annotation through a typed de Bruijn context during both encoding and decoding.

The encoder

encoder Enc : STLC
    dim Term = 24
    op App(fun, arg) |-> relu(fun + arg)
    var_init Term from Type as ty |-> ty

Every encoder body entry is keyword-led. The op keyword introduces a per-operation rewrite rule: App's two child embeddings are summed and passed through relu. Operators without a rule (Const, Base, Lam) receive a compiler-scaffolded 2-layer MLP with the correct per-argument dimensions, so a body may override as few or as many operations as the model demands. The var_init entry sets the embedding a Lam-bound variable enters the context with; here the type annotation's embedding is passed through unchanged.

The decoder

decoder Dec : STLC [depth=6]
    dim Term = 24
    body |-> recursive

The decoder header names its signature after :, exactly like the encoder header. body |-> recursive accepts the scaffolded corecursion: at every node the runtime scores constructor, binder, and bound-variable candidates for the current sort, splits the parent vector into per-child vectors, and recurses down to the depth bound. Each slot (structure, primitive, factor, binder selection) can instead be overridden with an explicit rule, as described in the decoders and losses guide.

The loss

loss reconstruct [weight=1.0, on=decoder(Dec)]
    -Dec.log_prob(term, Enc(term))

A loss declaration registers a weighted scalar head in the module's LossRegistry. The body is compiled as a let-expression closure over an environment: Enc and Dec resolve to the compiled module's own artifacts, while term is supplied by the caller at evaluation time. The on=decoder(Dec) option attaches the entry to a named site, so a training driver may evaluate it selectively via evaluate_on; weight=1.0 sets the multiplier applied by evaluate.

Try it

Compile, encode, decode

Compile the module with quivers.dsl.load, build a term with make_term and bound_var, and run it through the round trip.

import torch
from quivers.dsl import load
from quivers.structural import bound_var, make_term

torch.manual_seed(0)
prog = load("docs/examples/source/term_autoencoder.qvr")
enc = prog.encoders["Enc"]
dec = prog.decoders["Dec"]

# \x : base. plus x
term = make_term(
    "Lam",
    make_term("Base", "base"),
    make_term("App", make_term("Const", "plus"), bound_var(0)),
)
code = enc(term)
print("code shape:", tuple(code.shape))     # (24,)

recon = dec(code)
print("reconstruction root op:", recon.op)

Score and train the round trip

Decoder.log_prob scores an observed term against a code vector; the registry's evaluate sums every registered loss under a caller-supplied environment. A few optimizer steps on the encoder and decoder parameters drive the single-term reconstruction loss toward zero.

import torch
from quivers.dsl import load
from quivers.structural import bound_var, make_term

torch.manual_seed(0)
prog = load("docs/examples/source/term_autoencoder.qvr")
enc = prog.encoders["Enc"]
dec = prog.decoders["Dec"]

term = make_term(
    "Lam",
    make_term("Base", "base"),
    make_term("App", make_term("Const", "plus"), bound_var(0)),
)

loss0 = prog.losses.evaluate({"term": term})
params = list(enc.parameters()) + list(dec.parameters())
opt = torch.optim.Adam(params, lr=1e-2)
for _ in range(20):
    opt.zero_grad()
    loss = prog.losses.evaluate({"term": term})
    loss.backward()
    opt.step()
print(f"reconstruction loss: {float(loss0.detach()):.2f} -> {float(loss.detach()):.2f}")

Note that the encoder's data-leaf embedding tables allocate parameters lazily on first lookup, so the initial evaluate call precedes the optimizer's parameter collection.

Categorical Perspective

The encoder is a Σ-algebra homomorphism \(T_\Sigma \to \mathrm{Vec}_D\): the framework supplies the recursion over the term tree, and the analyst supplies (or accepts scaffolds for) one parametric function per operation. The decoder is a Kleisli coalgebra \(\mathrm{Vec}_D \to \mathrm{Kern}(T_\Sigma)\), corecursively generating a distribution over terms with a per-term log_prob. Their composite is a stochastic endomorphism on \(T_\Sigma\), and the loss head scores how far that endomorphism sits from the identity on the observed term.

See Also