add temperature

Author: ferrine

pymc3/tests/test_variational_inference.py

@@ -152,7 +152,6 @@ def test_optimizer_with_full_data(self):
                 mu_ = Normal('mu', mu=mu0, sd=sd0, testval=0)
                 Normal('x', mu=mu_, sd=sd, observed=data)
                 inf = self.inference(start={})
-                inf.fit(10)
                 approx = inf.fit(self.NITER,
                                  obj_optimizer=self.optimizer,
                                  callbacks=self.conv_cb,)
@@ -295,11 +294,9 @@ class TestSVGD(TestApproximates.Base):
 
 
 class TestASVGD(TestApproximates.Base):
-    NITER = 15000
-    inference = ASVGD
+    NITER = 5000
+    inference = functools.partial(ASVGD, temperature=1.5)
     test_aevb = _test_aevb
-    optimizer = pm.adagrad_window(learning_rate=0.002)
-    conv_cb = []
 
 
 class TestEmpirical(SeededTest):

pymc3/variational/inference.py

@@ -41,11 +41,15 @@ class Inference(object):
         See (AEVB; Kingma and Welling, 2014) for details
     model : Model
         PyMC3 Model
+    op_kwargs : dict
+        kwargs passed to :class:`Operator`
     kwargs : kwargs
         additional kwargs for :class:`Approximation`
     """
 
-    def __init__(self, op, approx, tf, local_rv=None, model=None, **kwargs):
+    def __init__(self, op, approx, tf, local_rv=None, model=None, op_kwargs=None, **kwargs):
+        if op_kwargs is None:
+            op_kwargs = dict()
         self.hist = np.asarray(())
         if isinstance(approx, type) and issubclass(approx, Approximation):
             approx = approx(
@@ -56,7 +60,7 @@ def __init__(self, op, approx, tf, local_rv=None, model=None, **kwargs):
         else:   # pragma: no cover
             raise TypeError(
                 'approx should be Approximation instance or Approximation subclass')
-        self.objective = op(approx)(tf)
+        self.objective = op(approx, **op_kwargs)(tf)
 
     approx = property(lambda self: self.objective.approx)
 
@@ -535,6 +539,8 @@ class SVGD(Inference):
         PyMC3 model for inference
     kernel : `callable`
         kernel function for KSD :math:`f(histogram) -> (k(x,.), \nabla_x k(x,.))`
+    temperature : float
+        parameter responsible for exploration, higher temperature gives more broad posterior estimate
     scale_cost_to_minibatch : bool, default False
         Scale cost to minibatch instead of full dataset
     start : `dict`
@@ -552,10 +558,14 @@ class SVGD(Inference):
     -   Qiang Liu, Dilin Wang (2016)
         Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm
         arXiv:1608.04471
+
+    -   Yang Liu, Prajit Ramachandran, Qiang Liu, Jian Peng (2017)
+        Stein Variational Policy Gradient
+        arXiv:1704.02399
     """
 
     def __init__(self, n_particles=100, jitter=.01, model=None, kernel=test_functions.rbf,
-                 scale_cost_to_minibatch=False, start=None, histogram=None,
+                 temperature=1, scale_cost_to_minibatch=False, start=None, histogram=None,
                  random_seed=None, local_rv=None):
         if histogram is None:
             histogram = Empirical.from_noise(
@@ -597,6 +607,8 @@ class ASVGD(Inference):
         See (AEVB; Kingma and Welling, 2014) for details
     kernel : `callable`
         kernel function for KSD :math:`f(histogram) -> (k(x,.), \nabla_x k(x,.))`
+    temperature : float
+        parameter responsible for exploration, higher temperature gives more broad posterior estimate
     model : :class:`Model`
     kwargs : kwargs for :class:`Approximation`
 
@@ -608,17 +620,22 @@ class ASVGD(Inference):
 
     -   Dilin Wang, Qiang Liu (2016)
         Learning to Draw Samples: With Application to Amortized MLE for Generative Adversarial Learning
-        https://arxiv.org/abs/1611.01722
+        arXiv:1611.01722
+    
+    -   Yang Liu, Prajit Ramachandran, Qiang Liu, Jian Peng (2017)
+        Stein Variational Policy Gradient
+        arXiv:1704.02399
     """
 
     def __init__(self, approx=FullRank, local_rv=None,
-                 kernel=test_functions.rbf, model=None, **kwargs):
+                 kernel=test_functions.rbf, temperature=1, model=None, **kwargs):
         super(ASVGD, self).__init__(
             op=AKSD,
             approx=approx,
             local_rv=local_rv,
             tf=kernel,
             model=model,
+            op_kwargs=dict(temperature=temperature),
             **kwargs
         )
 

pymc3/variational/operators.py

@@ -1,3 +1,4 @@
+import warnings
 from theano import theano, tensor as tt
 from pymc3.variational.opvi import Operator, ObjectiveFunction, _warn_not_used
 from pymc3.variational.stein import Stein
@@ -95,15 +96,29 @@ class KSD(Operator):
     SUPPORT_AEVB = False
     OBJECTIVE = KSDObjective
 
-    def __init__(self, approx):
+    def __init__(self, approx, temperature=1):
         Operator.__init__(self, approx)
+        self.temperature = temperature
         self.input_matrix = tt.matrix('KSD input matrix')
 
     def apply(self, f):
         # f: kernel function for KSD f(histogram) -> (k(x,.), \nabla_x k(x,.))
-        stein = Stein(self.approx, f, self.input_matrix)
+        stein = Stein(
+            approx=self.approx,
+            kernel=f,
+            input_matrix=self.input_matrix,
+            temperature=self.temperature)
         return pm.floatX(-1) * stein.grad
 
 
 class AKSD(KSD):
+    def __init__(self, approx, temperature=1):
+        warnings.warn('You are using experimental inference Operator. '
+                      'It requires careful choice of temperature, default is 1. '
+                      'Default temperature works well for low dimensional problems and '
+                      'for significant `n_obj_mc`. Temperature > 1 gives more exploration '
+                      'power to algorithm, < 1 leads to undesirable results. Please take '
+                      'it in account when looking at inference result. Posterior variance '
+                      'is often **underestimated** when using temperature = 1.', stacklevel=2)
+        super(AKSD, self).__init__(approx, temperature)
     SUPPORT_AEVB = True

pymc3/variational/stein.py

@@ -1,15 +1,16 @@
 from theano import theano, tensor as tt
 from pymc3.variational.test_functions import rbf
-from pymc3.theanof import memoize
+from pymc3.theanof import memoize, floatX
 
 __all__ = [
     'Stein'
 ]
 
 
 class Stein(object):
-    def __init__(self, approx, kernel=rbf, input_matrix=None):
+    def __init__(self, approx, kernel=rbf, input_matrix=None, temperature=1):
         self.approx = approx
+        self.temperature = floatX(temperature)
         self._kernel_f = kernel
         if input_matrix is None:
             input_matrix = tt.matrix('stein_input_matrix')
@@ -22,8 +23,9 @@ def grad(self):
         t = self.approx.normalizing_constant
         Kxy, dxkxy = self.Kxy, self.dxkxy
         dlogpdx = self.dlogp  # Normalized
-        n = self.input_matrix.shape[0].astype('float32')
-        svgd_grad = (tt.dot(Kxy, dlogpdx) + dxkxy/t) / n
+        n = floatX(self.input_matrix.shape[0])
+        temperature = self.temperature
+        svgd_grad = (tt.dot(Kxy, dlogpdx)/temperature + dxkxy/t) / n
         return svgd_grad
 
     @property

pymc3/variational/test_functions.py

@@ -1,4 +1,4 @@
-from theano import tensor as tt
+from theano import tensor as tt, theano
 from .opvi import TestFunction
 from pymc3.theanof import floatX