p(t) in Algorithm 2

[boxaio]

In Algorithm 2 (in practice), you sample the time points t \sim p(t), where p(t) can be viewed as a proposal importance sampling distribution. One can take p(t) to be estimated using Eq.(85), as you mentioned in Appendix.C. But in this code repository (see losses.py)
you have
p_t = time_sampler.invdensity(t)
which is according to an uniform distribution (see dynamics.utils).
So I wonder how you actually implemented your claim in Algorithm 2 in Appendix C.

[necludov]

Hi @boxaio!
I haven't implemented this in the JAX version. I have a pytorch implementation of the same algorithm, which is not public.
Here is the code for the class implementing the loss function

class AdaptiveLoss:
    def __init__(self, net, config, n=100, beta=0.99):
        self.t0, self.t1 = config.model.t0, config.model.t1
        self.alpha, self.beta = config.train.alpha, beta
        self.timesteps = np.linspace(self.t0, self.t1, n)
        self.dt = (self.t1-self.t0)/(n-1)
        self.rank = get_rank()
        self.ws = get_world_size()
        
        self.q_t, self.w, self.dwdt = get_q(config)
        self.boundary_conditions = (self.w(torch.tensor(self.t0)).item() != 0.0,
                                    self.w(torch.tensor(self.t1)).item() != 0.0)
        print('boundary conditions are: ', self.boundary_conditions)
        config.train.boundary_conditions = self.boundary_conditions
        
        self.s = get_s(net, config)
        
        self.buffer = {'values': [],
                       'times': [],
                       'size': 100,
                       'mean': np.zeros_like(self.timesteps),
                       'var': np.ones_like(self.timesteps),
                       'p': np.ones_like(self.timesteps),
                       'u0': 0.5}
        self.construct_dist()
        
        meters = [DDPAverageMeter('train_loss'),
                  DDPAverageMeter('dsdx_std'),
                  DDPAverageMeter('dsdt_std'),
                  DDPAverageMeter('s_1_std'),
                  DDPAverageMeter('s_0_std'),
                  DDPAverageMeter('s_std')]
        self.meters = dict((m.name,m) for m in meters)
        
    def load_state_dict(self, buffer_dict):
        self.buffer = buffer_dict
        self.construct_dist()
        
    def state_dict(self):
        return self.buffer
        
    def construct_dist(self):
        dt, t = self.dt, self.timesteps
        p = self.buffer['p']
        self.fp = scipy.interpolate.interp1d(t, p, kind='linear')
        self.dpdt = scipy.interpolate.interp1d(t, np.concatenate([p[1:]-p[:-1], p[-1:]-p[-2:-1]])/dt, kind='zero')
        intercept = lambda t: self.fp(t)-self.dpdt(t)*t
        t0_interval = scipy.interpolate.interp1d(t, t, kind='zero')
        mass = np.concatenate([np.zeros([1]), ((p[1:]+p[:-1])*dt/2).cumsum()[:-1], np.ones([1])])
        F0_interval = scipy.interpolate.interp1d(t, mass, kind='zero')
        F0_inv = scipy.interpolate.interp1d(mass, t, kind='zero')
        def F(t):
            t0_ = t0_interval(t)
            F0_ = F0_interval(t)
            k, b = self.dpdt(t), intercept(t)
            output = 0.5*k*(t**2-t0_**2) + b*(t-t0_)
            return F0_ + output 

        def F_inv(y):
            t0_ = F0_inv(y)
            F0_ = F0_interval(t0_)
            k, b = self.dpdt(t0_), intercept(t0_)
            c = y - F0_
            c = c + 0.5*k*t0_**2 + b*t0_
            D = np.sqrt(b**2 + 2*k*c)
            output = (-b + D) * (np.abs(k) > 0)  + c/b * (np.abs(k) == 0.0)
            output[np.abs(k) > 0] /= k[np.abs(k) > 0]
            return output
        
        self.F_inv = F_inv
        
    def sample_t(self, n, device):
        u = (self.buffer['u0'] + np.sqrt(2)*np.arange(n*self.ws)) % 1
        self.buffer['u0'] = (self.buffer['u0'] + np.sqrt(2)*n*self.ws) % 1
        u = u[self.rank*n:(self.rank+1)*n]
        t = self.F_inv(u)
        assert ((t < 0.0).sum() == 0) and ((t > 1.0).sum() == 0)
        p_t, dpdt = self.fp(t), self.dpdt(t)
        p_0, p_1 = self.fp(self.t0*np.ones_like(t)), self.fp(self.t1*np.ones_like(t))
        t = torch.from_numpy(t).to(device).float()
        p_t, dpdt = torch.from_numpy(p_t).to(device).float(), torch.from_numpy(dpdt).to(device).float()
        p_0, p_1 = torch.from_numpy(p_0).to(device).float(), torch.from_numpy(p_1).to(device).float()
        return t, p_t, dpdt

    def update_history(self, new_p, t, p_t):
        new_p, t, p_t = new_p.cpu().numpy().flatten(), t.cpu().numpy().flatten(), p_t.cpu().numpy().flatten()
        weights = np.exp(-np.abs(self.timesteps.reshape(-1, 1) - t.reshape(1,-1))*1e2)
        weights = weights/weights.sum(1,keepdims=True)
        self.buffer['mean'] = self.beta*self.buffer['mean'] + (1-self.beta)*(weights@new_p)
        mean_func = scipy.interpolate.interp1d(self.timesteps, self.buffer['mean'], kind='linear')
        self.buffer['var'] = self.beta*self.buffer['var'] + (1-self.beta)*(weights@((mean_func(t) - new_p)**2))
        p = np.sqrt(self.buffer['var'])
        p = (1.0-self.alpha)*p/((p[1:]+p[:-1])*self.dt/2).sum() + self.alpha/(self.t1-self.t0)
        self.buffer['p'] = p
        self.construct_dist()
    
    def eval_loss(self, x):
        q_t, w, dwdt, s = self.q_t, self.w, self.dwdt, self.s
        assert (2 == x.dim())
        t_0, t_1 = self.t0, self.t1
        device = x.device
        bs = x.shape[0]
        t, p_t, dpdt = self.sample_t(bs, device)
        while (x.dim() > t.dim()): t = t.unsqueeze(-1)
        x_t, _ = q_t(x, t)
        x_t.requires_grad, t.requires_grad = True, True
        s_t = s(t, x_t)
        assert (2 == s_t.dim())
        dsdt, dsdx = torch.autograd.grad(s_t.sum(), [t, x_t], create_graph=True, retain_graph=True)
        x_t.requires_grad, t.requires_grad = False, False
        
        loss = (0.5*(dsdx**2).sum(1, keepdim=True) + dsdt.sum(1, keepdim=True))*w(t)
        self.meters['dsdx_std'].update((0.5*(dsdx**2).sum(1)*w(t).squeeze()).detach().cpu().std())
        self.meters['dsdt_std'].update((dsdt.sum(1)*w(t).squeeze()).detach().cpu().std())
        loss = loss + s_t*dwdt(t)
        self.meters['s_std'].update((s_t*dwdt(t)).squeeze().detach().cpu().std())
        loss = loss.squeeze()/p_t
        time_loss = loss.detach()*p_t
            
        s_1_std, s_0_std = 0.0, 0.0
        if self.boundary_conditions[0]:
            t_0 = t_0*torch.ones([bs, 1], device=device)
            x_0, _ = q_t(x, t_0)
            left_bound = (s(t_0,x_0)*w(t_0)).squeeze()
            loss = loss + left_bound
            self.meters['s_0_std'].update(left_bound.detach().cpu().std())
        if self.boundary_conditions[1]:
            t_1 = t_1*torch.ones([bs, 1], device=device)
            x_1, _ = q_t(x, t_1)
            right_bound = (-s(t_1,x_1)*w(t_1)).squeeze()
            loss = loss + right_bound
            self.meters['s_1_std'].update(right_bound.detach().cpu().std())
            
        self.meters['train_loss'].update(loss.detach().mean().cpu())
        self.update_history(gather(time_loss), gather(t), gather(p_t))
        return loss.mean(), self.meters
    
    def get_dxdt(self):
        def dxdt(t, x):
            return torch.autograd.grad(self.s(t, x).sum(), x, create_graph=True, retain_graph=True)[0]
        return dxdt