neuroprob/utils/pytorch.py
def compute_nat_grads(param_pairs):
    """
    Compute the natural gradients and replace the .grad of the relevant parameters. [1]
    Note compared to the Newton method, the gradient here is computed in parallel for each output dimension
    while the Newton method uses the exact Hessian involving cross-terms.

    References:

    [1] `Natural Gradients in Practice: Non-Conjugate Variational Inference in Gaussian Process Models`, 
        Hugh Salimbeni, Stefanos Eleftheriadis, James Hensman (2018)

    """
    for k in range(len(param_pairs) // 2):
        m = param_pairs[2*k]
        L = param_pairs[2*k+1]
        grad_m = m.grad # N x M
        grad_L = torch.tril(L.grad) # N x M x M
        
        M = L.shape[-1]
        N = L.shape[0]
        
        # convert to grad_S
        A = (L.permute(0, 2, 1) @ grad_L)
        A.view(N, -1)[:, ::M+1] -= 0.5*A[:, range(M), range(M)] # apply phi()
        b = torch.eye(M, device=L.device)
        L_inv = torch.triangular_solve(b, L, upper=False)[0]
        S = L_inv.permute(0, 2, 1) @ A @ L_inv
        grad_S = S + S.permute(0, 2, 1)
        grad_S.view(N, -1)[:, ::M+1] -= 0.5*grad_S[:, range(M), range(M)] # apply phi()
        
        dL_dxi = torch.cat((grad_m.flatten(), grad_L.flatten())).double()

        # compute expectation parameters to xi
        eta_m = m.data
        eta_S = torch.matmul(L.data, L.data.permute(0, 2, 1)) + m.data[..., None] * m.data[:, None, :]
        eta_params = torch.cat((eta_m.flatten(), eta_S.flatten())).double()
        eta_params.requires_grad = True

        eta_m = eta_params[:N*M].view(N, M)
        eta_S = eta_params[N*M:].view(N, M, M)
        chsk_term = eta_S - eta_m[..., None]*eta_m[:, None, :]
        chsk_term.data.view(N, -1)[:, ::M+1] += 1e-1
        L_back = torch.cholesky(chsk_term, upper=False)
        #print(((L.data-L_back)**2).sum())
        xi_back = torch.cat((eta_m.flatten(), L_back.flatten()))
        dL_deta, = torch.autograd.grad(xi_back, eta_params, dL_dxi, retain_graph=False, create_graph=False)

        # compute theta (natural) to xi transform
        b = torch.eye(M, device=L.device).double()
        S_inv = torch.cholesky_solve(b, L.data.double(), upper=False)
        nat_m = (S_inv @ m.data[..., None].double())[..., 0]
        nat_S = -0.5*S_inv.contiguous()
        nat_params = torch.cat((nat_m.flatten(), nat_S.flatten()))
        nat_params.requires_grad = True

        nat_m = nat_params[:N*M].view(N, M)
        nat_S = nat_params[N*M:].view(N, M, M)
        chsk_term = -2.*nat_S
        #chsk_term.data.view(N, -1)[:, ::M+1] += 1e-3
        L_inv = torch.cholesky(chsk_term, upper=False)
        L__ = torch.triangular_solve(b, L_inv, upper=False)[0]
        S = L__.permute(0, 2, 1) @ L__
        #S = torch.inverse(-2.*nat_S)
        m_ = (S @ nat_m[..., None])[..., 0]
        L_ = torch.cholesky(S, upper=False)
        xi = torch.cat((m_.flatten(), L_.flatten()))

        # Jacobian-vector product trick
        v = torch.ones(len(xi), device=L.device, requires_grad=True)
        dvxi_dnat, = torch.autograd.grad(xi, nat_params, v, retain_graph=True, create_graph=True)
        dJ_dv, = torch.autograd.grad(dvxi_dnat, v, dL_deta, retain_graph=False, create_graph=False)

        # assign gradients
        m.grad = dJ_dv[:N*M].view(N, M).clone().type(self.rate_model[0].tensor_type)
        L.grad = torch.tril(dJ_dv[N*M:].view(N, M, M).clone()).type(self.rate_model[0].tensor_type)

    Compute the natural gradients and replace the .grad of the relevant parameters. [1]
    Note compared to the Newton method, the gradient here is computed in parallel for each output dimension
    while the Newton method uses the exact Hessian involving cross-terms.

    References:

    [1] `Natural Gradients in Practice: Non-Conjugate Variational Inference in Gaussian Process Models`, 
        Hugh Salimbeni, Stefanos Eleftheriadis, James Hensman (2018)

    
Compute the natural gradients and replace the .grad of the relevant parameters. [1]
Note compared to the Newton method, the gradient here is computed in parallel for each output dimension
while the Newton method uses the exact Hessian involving cross-terms.



[1] `Natural Gradients in Practice: Non-Conjugate Variational Inference in Gaussian Process Models`,
Hugh Salimbeni, Stefanos Eleftheriadis, James Hensman (2018)
Compute the natural gradients and replace the . grad of the relevant parameters . [ 1 ] Note compared to the Newton method the gradient here is computed in parallel for each output dimension while the Newton method uses the exact Hessian involving cross - terms . [ 1 ] ` Natural Gradients in Practice : Non - Conjugate Variational Inference in Gaussian Process Models ` Hugh Salimbeni Stefanos Eleftheriadis James Hensman ( 2018 )
def compute_nat_grads(param_pairs):
    for k in range(len(param_pairs) // 2):
        m = param_pairs[2*k]
        L = param_pairs[2*k+1]
        grad_m = m.grad 
        grad_L = torch.tril(L.grad)  x M
        M = L.shape[-1]
        N = L.shape[0]
        A = (L.permute(0, 2, 1) @ grad_L)
        A.view(N, -1)[:, ::M+1] -= 0.5*A[:, range(M), range(M)] 
        b = torch.eye(M, device=L.device)
        L_inv = torch.triangular_solve(b, L, upper=False)[0]
        S = L_inv.permute(0, 2, 1) @ A @ L_inv
        grad_S = S + S.permute(0, 2, 1)
        grad_S.view(N, -1)[:, ::M+1] -= 0.5*grad_S[:, range(M), range(M)] 
        dL_dxi = torch.cat((grad_m.flatten(), grad_L.flatten())).double()
        eta_m = m.data
        eta_S = torch.matmul(L.data, L.data.permute(0, 2, 1)) + m.data[..., None] * m.data[:, None, :]
        eta_params = torch.cat((eta_m.flatten(), eta_S.flatten())).double()
        eta_params.requires_grad = True
        eta_m = eta_params[:N*M].view(N, M)
        eta_S = eta_params[N*M:].view(N, M, M)
        chsk_term = eta_S - eta_m[..., None]*eta_m[:, None, :]
        chsk_term.data.view(N, -1)[:, ::M+1] += 1e-1
        L_back = torch.cholesky(chsk_term, upper=False)
        xi_back = torch.cat((eta_m.flatten(), L_back.flatten()))
        dL_deta, = torch.autograd.grad(xi_back, eta_params, dL_dxi, retain_graph=False, create_graph=False)
        b = torch.eye(M, device=L.device).double()
        S_inv = torch.cholesky_solve(b, L.data.double(), upper=False)
        nat_m = (S_inv @ m.data[..., None].double())[..., 0]
        nat_S = -0.5*S_inv.contiguous()
        nat_params = torch.cat((nat_m.flatten(), nat_S.flatten()))
        nat_params.requires_grad = True
        nat_m = nat_params[:N*M].view(N, M)
        nat_S = nat_params[N*M:].view(N, M, M)
        chsk_term = -2.*nat_S
        L_inv = torch.cholesky(chsk_term, upper=False)
        L__ = torch.triangular_solve(b, L_inv, upper=False)[0]
        S = L__.permute(0, 2, 1) @ L__
        m_ = (S @ nat_m[..., None])[..., 0]
        L_ = torch.cholesky(S, upper=False)
        xi = torch.cat((m_.flatten(), L_.flatten()))
        v = torch.ones(len(xi), device=L.device, requires_grad=True)
        dvxi_dnat, = torch.autograd.grad(xi, nat_params, v, retain_graph=True, create_graph=True)
        dJ_dv, = torch.autograd.grad(dvxi_dnat, v, dL_deta, retain_graph=False, create_graph=False)
        m.grad = dJ_dv[:N*M].view(N, M).clone().type(self.rate_model[0].tensor_type)
        L.grad = torch.tril(dJ_dv[N*M:].view(N, M, M).clone()).type(self.rate_model[0].tensor_type)
def compute_nat_grads ( param_pairs ) : for k in range ( len ( param_pairs ) // 2 ) : m = param_pairs [ 2 * k ] L = param_pairs [ 2 * k + 1 ] grad_m = m . grad grad_L = torch . tril ( L . grad ) M = L . shape [ - 1 ] N = L . shape [ 0 ] A = ( L . permute ( 0 , 2 , 1 ) @ grad_L ) A . view ( N , - 1 ) [ : , : : M + 1 ] -= 0.5 * A [ : , range ( M ) , range ( M ) ] b = torch . eye ( M , device = L . device ) L_inv = torch . triangular_solve ( b , L , upper = False ) [ 0 ] S = L_inv . permute ( 0 , 2 , 1 ) @ A @ L_inv grad_S = S + S . permute ( 0 , 2 , 1 ) grad_S . view ( N , - 1 ) [ : , : : M + 1 ] -= 0.5 * grad_S [ : , range ( M ) , range ( M ) ] dL_dxi = torch . cat ( ( grad_m . flatten ( ) , grad_L . flatten ( ) ) ) . double ( ) eta_m = m . data eta_S = torch . matmul ( L . data , L . data . permute ( 0 , 2 , 1 ) ) + m . data [ ... , None ] * m . data [ : , None , : ] eta_params = torch . cat ( ( eta_m . flatten ( ) , eta_S . flatten ( ) ) ) . double ( ) eta_params . requires_grad = True eta_m = eta_params [ : N * M ] . view ( N , M ) eta_S = eta_params [ N * M : ] . view ( N , M , M ) chsk_term = eta_S - eta_m [ ... , None ] * eta_m [ : , None , : ] chsk_term . data . view ( N , - 1 ) [ : , : : M + 1 ] += 1e-1 L_back = torch . cholesky ( chsk_term , upper = False ) xi_back = torch . cat ( ( eta_m . flatten ( ) , L_back . flatten ( ) ) ) dL_deta , = torch . autograd . grad ( xi_back , eta_params , dL_dxi , retain_graph = False , create_graph = False ) b = torch . eye ( M , device = L . device ) . double ( ) S_inv = torch . cholesky_solve ( b , L . data . double ( ) , upper = False ) nat_m = ( S_inv @ m . data [ ... , None ] . double ( ) ) [ ... , 0 ] nat_S = - 0.5 * S_inv . contiguous ( ) nat_params = torch . cat ( ( nat_m . flatten ( ) , nat_S . flatten ( ) ) ) nat_params . requires_grad = True nat_m = nat_params [ : N * M ] . view ( N , M ) nat_S = nat_params [ N * M : ] . view ( N , M , M ) chsk_term = - 2. * nat_S L_inv = torch . cholesky ( chsk_term , upper = False ) L__ = torch . triangular_solve ( b , L_inv , upper = False ) [ 0 ] S = L__ . permute ( 0 , 2 , 1 ) @ L__ m_ = ( S @ nat_m [ ... , None ] ) [ ... , 0 ] L_ = torch . cholesky ( S , upper = False ) xi = torch . cat ( ( m_ . flatten ( ) , L_ . flatten ( ) ) ) v = torch . ones ( len ( xi ) , device = L . device , requires_grad = True ) dvxi_dnat , = torch . autograd . grad ( xi , nat_params , v , retain_graph = True , create_graph = True ) dJ_dv , = torch . autograd . grad ( dvxi_dnat , v , dL_deta , retain_graph = False , create_graph = False ) m . grad = dJ_dv [ : N * M ] . view ( N , M ) . clone ( ) . type ( self . rate_model [ 0 ] . tensor_type ) L . grad = torch . tril ( dJ_dv [ N * M : ] . view ( N , M , M ) . clone ( ) ) . type ( self . rate_model [ 0 ] . tensor_type )
Compute the natural gradients and replace the .grad of the relevant parameters.
Compute the natural gradients and replace the . grad of the relevant parameters .
"""
    Compute the natural gradients and replace the .grad of the relevant parameters. [1]
    Note compared to the Newton method, the gradient here is computed in parallel for each output dimension
    while the Newton method uses the exact Hessian involving cross-terms.

    References:

    [1] `Natural Gradients in Practice: Non-Conjugate Variational Inference in Gaussian Process Models`, 
        Hugh Salimbeni, Stefanos Eleftheriadis, James Hensman (2018)

    """ # N x M # N x M x M # convert to grad_S # apply phi() # apply phi() # compute expectation parameters to xi #print(((L.data-L_back)**2).sum()) # compute theta (natural) to xi transform #chsk_term.data.view(N, -1)[:, ::M+1] += 1e-3 #S = torch.inverse(-2.*nat_S) # Jacobian-vector product trick # assign gradients
param_pairs