Condition IAF flows on context input h?

megaloman · August 25, 2019, 9:56pm

I was looking at the normalizing flows module as I am hoping to use them on a project I’ve been looking at. Is it possible to pass additional inputs to the IAF layers before creating a TransformedDistribution? I was re-reading the IAF paper and they mention passing as input a context vector h to each flow besides the previous samples z_{t-1}.

From looking at the Pyro code for it, it doesn’t seem they support giving a context h to each flow. Is this actually the case, or did I just miss something? If it ISN"T supported, is there any suggested practices I should be aware of to make an IAF flow that does allow this?

So maybe something like:

#necessary imports (Gaussian, torch.nn , etc. etc.)
iaf = IAFCondition(...)
mu_0 eps_0, h = nn(X)
iaf.condition(h)
dist = Gaussian(mu_0, sigma_0)
dist = TransformedDistribution(dist, [iaf])
#the rest of my pyro goodness

stefanwebb · September 26, 2019, 10:18pm

Hi @megaloman, thanks for your interest! I’m glad to hear that there are users who would like to see conditional flows, which is something I have been working on recently.

We’ve just added the capability to represent conditional distributions and transforms (i.e. that input an additional context variable) as well as conditional MADE. See:

github.com/pyro-ppl/pyro

Transform and TransformedDistribution that can forward optional keyword args (#1937)

committed 03:45PM - 26 Sep 19 UTC

stefanwebb

+234 -43

* Transform and TransformedDistribution that can take option keyword arguments a…nd forward to methods * PEP8 and nit * Fixed docstring * Fritz's suggested conditional type design * Tests and conditional Planar flow * Bug fix to conditional transformed distribution * Fix to docstring of conditional Planar flow * Refactored conditional planar flow according to Fritz' comments * Added entries to docs * Fix for docs * Trying again to fix docstring * ABC for abstract classes and link to class in docstring * Fix to docstrings

Currently, there is only an implementation for conditional PlanarFlow… But a conditional IAF/NAF etc. will be added to Pyro in the very near future (and I’ll do this next for you)! It is quite easy now that conditional transforms and a conditional MADE have been added to Pyro.

megaloman · September 27, 2019, 7:39am

That’s exciting news, and thank you for the update! To get using it right away, I supposed I should install from source correct? Looking forward to hearing about the further updates for other flows

stefanwebb · September 27, 2019, 11:22am

Yes that’s right! Grab the dev branch and look at

github.com

pyro-ppl/pyro/blob/dev/pyro/distributions/transforms/planar.py#L145


      
          bijective = True
          
          
def __init__(self, input_dim):
              super().__init__(self._params)
          
          
    self.bias = nn.Parameter(
                  torch.Tensor(
                      1,
                  )
              )
              self.u = nn.Parameter(
                  torch.Tensor(
                      input_dim,
                  )
              )
              self.w = nn.Parameter(
                  torch.Tensor(
                      input_dim,
                  )
              )

for an example of how to do conditional planar flow

stefanwebb · September 27, 2019, 11:23am

And you can follow work on the conditional flows by checking out my PRs: Pull requests · pyro-ppl/pyro · GitHub