SGHMC optimizer and closure()

bstaber: <p>Hello,</p>
<p>I’m trying to implement a generic optimizer that performs SGHMC (see algorithm  2 in <a href="http://proceedings.mlr.press/v32/cheni14.pdf" rel="noopener nofollow ugc">Chen et al.</a>). In this algorithm, we can perform <strong>m</strong> updates of the parameters and momentum while keeping the same mini-batch. I would like to keep my training and validation loops independant of the optimizer I’m using, and thus I’m trying to add this <strong>m</strong> loop within the SGHMC optimizer.</p>
<p>Within the <code>step</code> method of my optimizer, I added a loop <code>for t in range(self.trajectory_param)</code> where <code>trajectory_param</code> corresponds to the parameter <strong>m</strong> (see below).</p>
<p>At the end of each iteration of this loop, I would need to update the gradients, so I decided to call the <code>closure</code> after each update (see the code below the comment <code>#recomputes the gradients ...</code>.</p>
<p>Is this the right way to go ?</p>
<p>Thanks for your help.</p>
<p>Edit: this implementation is probably wrong because I’m calling the closure inside the loop over the model parameters</p>
<pre><code class="lang-auto">    @torch.no_grad()
    def step(self, closure=None):
        
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:

            # get parameters
            lr = group["lr"]
            alpha = group["alpha"]
            weight_decay = group["weight_decay"]
            cv_flag = group['cv_flag']
            
            for parameter in group["params"]:

                if parameter.grad is None:
                    continue

                state = self.state[parameter]
                if len(state) == 0:
                    state["iteration"] = 0
                    state["momentum"] = torch.zeros_like(parameter)
                    if cv_flag:
                        p.total_grad = torch.zeros_like(p)
                        p.map_grad = torch.zeros_like(p)
                state["iteration"] += 1

                # get v_k 
                momentum = state["momentum"]

                # grad
                gradient = parameter.grad
                if weight_decay!=0.0:
                    gradient.add_(parameter, alpha=weight_decay)

                if state["iteration"] &gt; self.num_warmup_steps:
                    if self.resample_momentum:
                        state["momentum"] = gradient.new(gradient.size()).normal_(mean=0,std=1)
                        #state["momentum"] = torch.normal(mean=torch.zeros_like(gradient), std=1.0)

                    sigma = np.sqrt(alpha)
                    for t in range(self.trajectory_param):
                        #w_scaled ~ N(0,2*alpha) / sqrt(lr), so that lr * w_scaled = N(0,alpha*lr)
                        sample_t = parameter.new(parameter.size()).normal_(mean=0.0, std=sigma) / np.sqrt(lr)

                        #v_k+1 = (1-alpha)*v_k - lr*grad + sqrt(alpha*lr)*w
                        if cv_flag:
                            d_p.add_(p.map_grad, alpha=-1.0)
                            d_p.add_(p.total_grad, alpha=1.0)
                        momentum_t = momentum.mul_(1.0-alpha).add_(0.5*gradient - sample_t, alpha=-lr)

                        #theta_k+1 = theta_k + v_k
                        parameter.add_(momentum_t)
                        
                        if t&gt;0:
                            # recomputes the gradients -&gt; need to test this
                            with torch.enable_grad():
                                loss = closure()
                            gradient = parameter.grad
                            if weight_decay!=0.0:
                                gradient.add_(parameter, alpha=weight_decay)
                else:
                    momentum_t = momentum.mul_(1.0-alpha).add_(0.5*gradient, alpha=-lr)
                    parameter.add_(momentum_t)

        return loss
</code></pre>

googlebot: <p>Aren’t your loops in the wrong order? I think you want them like:</p>
<p>for t in range(self.trajectory_param):<br>
loss=closure()<br>
for p in parameters:<br>
…</p>
<p>I checked <a href="https://github.com/pyro-ppl/pyro/blob/dev/pyro/ops/integrator.py" rel="noopener nofollow ugc">pyro implementation</a> (used from infer.mcmc.hmc), it seems to agree, doing one closure evalutation per step.</p>
<p>Otherwise, this closure approach should work, in principle. But you should zero gradients between calls.</p>

bstaber: <p>Hey, thanks for your reply. Indeed, the <strong>m</strong> loop was in the wrong place. I’m going to try the implementation below. I only need to figure out how to do a single SGD momentum step when the number of iterations is still below the number of warmup steps.</p>
<p>Thanks!</p>
<pre><code class="lang-auto">    @torch.no_grad()
    def step(self, closure=None):

        for t in range(self.trajectory_param):

            self.zero_grad()

            loss = None
            if closure is not None:
                with torch.enable_grad():
                    loss = closure()

            for group in self.param_groups:

                # get parameters
                lr = group["lr"]
                alpha = group["alpha"]
                weight_decay = group["weight_decay"]
                cv_flag = group['cv_flag']

                for parameter in group["params"]:

                    if parameter.grad is None:
                        continue

                    state = self.state[parameter]
                    if len(state) == 0:
                        state["iteration"] = 0
                        state["momentum"] = torch.zeros_like(parameter)
                        if cv_flag:
                            parameter.total_grad = torch.zeros_like(parameter)
                            parameter.map_grad = torch.zeros_like(parameter)
                    state["iteration"] += 1

                    # get v_k 
                    momentum = state["momentum"]

                    # grad
                    gradient = parameter.grad
                    if weight_decay!=0.0:
                        gradient.add_(parameter, alpha=weight_decay)

                    if state["iteration"] &gt; self.num_warmup_steps:

                        if self.resample_momentum:
                            state["momentum"] = gradient.new(gradient.size()).normal_(mean=0,std=1)
                            #state["momentum"] = torch.normal(mean=torch.zeros_like(gradient), std=1.0)

                        sigma = np.sqrt(alpha)
                        #w_scaled ~ N(0,2*alpha) / sqrt(lr), so that lr * w_scaled = N(0,alpha*lr)
                        sample_t = parameter.new(parameter.size()).normal_(mean=0.0, std=sigma) / np.sqrt(lr)
                        #v_k+1 = (1-alpha)*v_k - lr*grad + sqrt(alpha*lr)*w
                        if cv_flag:
                            gradient.add_(parameter.map_grad, alpha=-1.0)
                            gradient.add_(parameter.total_grad, alpha=1.0)
                        momentum_t = momentum.mul_(1.0-alpha).add_(0.5*gradient - sample_t, alpha=-lr)
                        #theta_k+1 = theta_k + v_k
                        parameter.add_(momentum_t)
                    else:
                        # single SGD momentum step if still in burnin stage
                        momentum_t = momentum.mul_(1.0-alpha).add_(0.5*gradient, alpha=-lr)
                        parameter.add_(momentum_t)
                        #break ?

        return loss
</code></pre>