Implement Theano graph optimization to remove normalizing constants from logp

[twiecki]

Our logp's include normalization terms which are unnecessary for sampling as that is fine with an unnormalized logp. As such, we can optimize the model logp evaluation if we take the Theano computation graph and apply a custom graph optimization that identifies and removes these terms. This should lead to quite a nice speed-up as it saves unnecessary computations.

I think the way to do this would be to traverse the logp graph, break the large sum into individual terms, and throw out any terms that do not depend on value nodes.

For more information on how to optimize the theano graph, see https://theano-pymc.readthedocs.io/en/latest/extending/optimization.html and https://theano-pymc.readthedocs.io/en/latest/optimizations.html.

[twiecki]

I wrote a small NB to set up the environment to start experimenting with this: https://gist.github.com/twiecki/e758db2c3d2df5f3368fc49e6087e58f

[brandonwillard]

Here's an implementation that's nearly 75% of the way there:

import numpy as np
import theano.tensor as tt
from theano import config
from theano.compile import optdb
from theano.gof.fg import FunctionGraph
from theano.gof.graph import inputs as tt_inputs
from theano.gof.opt import EquilibriumOptimizer, PatternSub
from theano.gof.optdb import Query
from theano.printing import debugprint as tt_dprint
from theano.tensor.opt import get_clients

# We don't need to waste time compiling graphs to C
config.cxx = ""


# a / b -> a * 1/b, for a != 1 and b != 1
div_to_mul_pattern = PatternSub(
    (tt.true_div, "a", "b"),
    (tt.mul, "a", (tt.inv, "b")),
    allow_multiple_clients=True,
    name="div_to_mul",
    tracks=[tt.true_div],
    get_nodes=get_clients,
)

# a - b -> a + (-b)
sub_to_add_pattern = PatternSub(
    (tt.sub, "a", "b"),
    (tt.add, "a", (tt.neg, "b")),
    allow_multiple_clients=True,
    name="sub_to_add",
    tracks=[tt.sub],
    get_nodes=get_clients,
)

# a * (x + y) -> a * x + a * y
distribute_mul_pattern = PatternSub(
    (tt.mul, "a", (tt.add, "x", "y")),
    (tt.add, (tt.mul, "a", "x"), (tt.mul, "a", "y")),
    allow_multiple_clients=True,
    name="distribute_mul",
    tracks=[tt.mul],
    get_nodes=get_clients,
)

expand_opt = EquilibriumOptimizer(
    [div_to_mul_pattern, distribute_mul_pattern, sub_to_add_pattern],
    ignore_newtrees=False,
    tracks_on_change_inputs=True,
    max_use_ratio=config.optdb__max_use_ratio,
)


def optimize_graph(fgraph, include=["canonicalize"], custom_opt=None, **kwargs):
    if not isinstance(fgraph, FunctionGraph):
        inputs = tt_inputs([fgraph])
        fgraph = FunctionGraph(inputs, [fgraph], clone=False)

    canonicalize_opt = optdb.query(Query(include=include, **kwargs))
    _ = canonicalize_opt.optimize(fgraph)

    if custom_opt:
        custom_opt.optimize(fgraph)

    return fgraph


tau = tt.dscalar("tau")
value = tt.dscalar("value")
mu = tt.dscalar("mu")

logp = (-tau * (value - mu) ** 2 + tt.log(tau / np.pi / 2.0)) / 2.0

logp_fg = optimize_graph(logp, custom_opt=expand_opt)

tt_dprint(logp_fg)

# TODO: Remove additive terms that do not contain the desired terms (e.g. `mu`
# and `tau` when the likelihood is only a function of `mu`, `tau`)

# This is what we want from the optimization
# logp_goal = -tau * (value - mu) ** 2

[twiecki]

Updated the gist: https://gist.github.com/e758db2c3d2df5f3368fc49e6087e58f

[chandan5362]

Hey @twiecki, I am currently working on this issue. I went through the documentation and I could not find any inbuilt method like replace_validate to remove a node from the FunctionGraph. shall I have to write the usual graph node removal code or is there already an inbuilt method ? If there is, I might have missed that. So, do let me know.
Thanks in advance :)

[twiecki]

That's great @chandan5362 - this is a super high-impact issue. Good questions, I'm not sure. Probably @brandonwillard knows.

Did you make any progress on the algorithm itself?

[chandan5362]

Did you make any progress on the algorithm itself?

As per my understanding, most of the optimizations have already been done. only things that are left is to remove the node form the FunctionGraph and addressing the same in optdb every time RemoveNormalizingConstants runs until no more node is left to optimize with respect to expand_opt. . I am just looking for a concrete method to delete the unnecessary node.
I might need a little help here @brandonwillard .

[twiecki]

I think there is still some complexity to figure out. What's not implemented yet is identifying sub-terms that contain value nodes. And you can't just break everything according to additions or subtractions as for example (value - x)**2 still needs to stay together. Does that make sense?

[chandan5362]

Now I guess, I will need a little more information about the nodes that can be removed or can't be. where do we actually need to remove the normalisation term ? Just to have a clear insight, could you please provide a reference to such sampling method ?

[brandonwillard]

The outline in that Gist already covers most of it.

There's more than one way to accomplish this, but, given what's already in the Gist, one can use the MergeOptimizer after expanding all the terms—per the example—and guarantee that all the arguments to Add and Mul are de-duplicated. From there, terms that are present in the arguments of every Mul in the Add and do not contain the arguments of interest can be factored out. Determining whether or not a term contains an argument of interest is simple (albeit rather inefficient): arg in theano.graph.basic.ancestors([term]).

Here's a simple illustration: graph = (sqrt(b) * log(x) + 2**3 * sqrt(b) * log(x - 1)) / a and our argument of interest is x (i.e. graph is proportional to some f(x)).

1.) Apply the existing transforms: inv(a) * sqrt(b) * log(x) + inv(a) * 2**3 * sqrt(b) * log(x - 1). We now have an Add with the arguments a1 = inv(a) * sqrt(b) * log(x) and a2 = inv(a) * 2**3 * sqrt(b) * log(x - 1), i.e. Add(a1, a2). Each argument is a Mul with its own arguments, i.e. a1 = Mul(inv(a), sqrt(b), log(x)) and a2 = Mul(inv(a), 2**3, sqrt(b), log(x - 1)).

2.) Now, we know that we can "gather"/"collect" the inv(a) and sqrt(b) terms, because they appear in both Muls; however, Theano doesn't know that the inv(a) in a1 is the same inv(a) in a2—at least not with ==—and that's why we need to apply the MergeOptimizer. After applying it, a1.inputs[0] is a2.inputs[0] and so on, so we can programmatically determine which terms are present in all the Muls.

3.) Of the terms that are present in all Muls, we can only remove the ones that do not contain x, and that's when we can use an expression like x not in theano.graph.basic.ancestors([term]) for each shared term. Once we've obtained the set of terms that are shared amongst all the Muls and not x, we can remove them or create a new Add/Mul graph that simply omits them.

All of the graph manipulation should be done within the context of an Optimizer class, and any replacements to an existing graph are performed using the methods in FunctionGraph. There are many ways one could approach this using existing Optimizer classes, but creating a GlobalOptimizer from scratch is probably the most conceptually straightforward approach.

[chandan5362]

Thanks for the help @brandonwillard .
So, according to you, what should be the final output for the above expression?
And also the output for the one in the Gist.
According to me, for the example in the Gist, the output should be (value - mu) ** 2 and tau should be omitted as well. What do you say?

[brandonwillard]

Yes, tau doesn't need to be in the final output, just the term (value - mu)**2.