def backward(ctx, *grad_outputs):
ts, flat_params, *ans = ctx.saved_tensors
sde, dt, bm, adjoint_method, adaptive, rtol, atol, dt_min, adjoint_options = (
ctx.sde, ctx.dt, ctx.bm, ctx.adjoint_method, ctx.adaptive,
ctx.rtol, ctx.atol, ctx.dt_min, ctx.adjoint_options)
params = misc.make_seq_requires_grad(sde.parameters())
n_tensors, n_params = len(ans), len(params)
# TODO: Make use of adjoint_method.
aug_bm = lambda t: tuple(-bmi for bmi in bm(-t))
adjoint_sde, adjoint_method, adjoint_adaptive = _get_adjoint_params(
sde=sde, params=params, adaptive=adaptive, logqp=True)
T = ans[0].size(0)
with torch.no_grad():
adj_y = tuple(grad_outputs_[-1]
for grad_outputs_ in grad_outputs[:n_tensors])
adj_l = tuple(grad_outputs_[-1]
for grad_outputs_ in grad_outputs[n_tensors:])
adj_params = torch.zeros_like(flat_params)
for i in range(T - 1, 0, -1):
ans_i = tuple(ans_[i] for ans_ in ans)
aug_y0 = (*ans_i, *adj_y, *adj_l, adj_params)
aug_ans = sdeint.integrate(sde=adjoint_sde,
y0=aug_y0,
ts=torch.tensor(
[-ts[i], -ts[i - 1]]).to(ts),
bm=aug_bm,
method=adjoint_method,
dt=dt,
adaptive=adjoint_adaptive,
rtol=rtol,
atol=atol,
dt_min=dt_min,
options=adjoint_options)
adj_y = aug_ans[n_tensors:2 * n_tensors]
adj_params = aug_ans[-1]
adj_y = tuple(adj_y_[1] for adj_y_ in adj_y)
adj_params = adj_params[1]
adj_y = misc.seq_add(
adj_y,
tuple(grad_outputs_[i - 1]
for grad_outputs_ in grad_outputs[:n_tensors]))
adj_l = tuple(grad_outputs_[i - 1]
for grad_outputs_ in grad_outputs[n_tensors:])
del aug_y0, aug_ans
return (*adj_y, None, None, adj_params, None, None, None, None, None,
None, None, None, None, None)
def f(self, t, y_aug):
sde, params, n_tensors = self._base_sde, self.params, len(y_aug) // 3
y, adj_y, adj_l = y_aug[:n_tensors], y_aug[
n_tensors:2 * n_tensors], y_aug[2 * n_tensors:3 * n_tensors]
vjp_l = tuple(torch.zeros_like(adj_l_) for adj_l_ in adj_l)
with torch.enable_grad():
y = tuple(y_.detach().requires_grad_(True) for y_ in y)
adj_y = tuple(adj_y_.detach() for adj_y_ in adj_y)
f_eval = sde.f(-t, y)
f_eval = tuple(-f_eval_ for f_eval_ in f_eval)
f_eval = misc.make_seq_requires_grad_y(f_eval, y)
vjp_y_and_params = torch.autograd.grad(
outputs=f_eval,
inputs=y + params,
grad_outputs=tuple(-adj_y_ for adj_y_ in adj_y),
allow_unused=True,
create_graph=True)
vjp_y = vjp_y_and_params[:n_tensors]
vjp_y = misc.convert_none_to_zeros(vjp_y, y)
vjp_params = vjp_y_and_params[n_tensors:]
vjp_params = misc.flatten_convert_none_to_zeros(vjp_params, params)
# Vector field change due to log-ratio term, i.e. ||u||^2 / 2.
g_eval = sde.g(-t, y)
h_eval = sde.h(-t, y)
ginv_eval = tuple(torch.pinverse(g_eval_) for g_eval_ in g_eval)
u_eval = misc.seq_sub(f_eval, h_eval)
u_eval = tuple(
torch.bmm(ginv_eval_, u_eval_)
for ginv_eval_, u_eval_ in zip(ginv_eval, u_eval))
log_ratio_correction = tuple(.5 * torch.sum(u_eval_**2., dim=1)
for u_eval_ in u_eval)
log_ratio_correction = misc.make_seq_requires_grad_y(
log_ratio_correction, y)
corr_vjp_y_and_params = torch.autograd.grad(
outputs=log_ratio_correction,
inputs=y + params,
grad_outputs=adj_l,
allow_unused=True,
)
corr_vjp_y = corr_vjp_y_and_params[:n_tensors]
corr_vjp_y = misc.convert_none_to_zeros(corr_vjp_y, y)
corr_vjp_params = corr_vjp_y_and_params[n_tensors:]
corr_vjp_params = misc.flatten_convert_none_to_zeros(
corr_vjp_params, params)
vjp_y = misc.seq_add(vjp_y, corr_vjp_y)
vjp_params = vjp_params + corr_vjp_params
return (*f_eval, *vjp_y, *vjp_l, vjp_params)
def f(self, t, y_aug):
sde, params, n_tensors = self._base_sde, self.params, len(y_aug) // 2
y, adj_y = y_aug[:n_tensors], y_aug[n_tensors:2 * n_tensors]
with torch.enable_grad():
y = tuple(y_.detach().requires_grad_(True) for y_ in y)
adj_y = tuple(adj_y_.detach() for adj_y_ in adj_y)
g_eval = sde.g(-t, y)
g_eval = misc.make_seq_requires_grad_y(g_eval, y)
gdg = torch.autograd.grad(outputs=g_eval,
inputs=y,
grad_outputs=g_eval,
allow_unused=True,
create_graph=True)
gdg = misc.convert_none_to_zeros(gdg, y)
f_eval = sde.f(-t, y)
f_eval_corrected = misc.seq_sub(
gdg, f_eval) # Stratonovich correction for reverse-time.
f_eval_corrected = misc.make_seq_requires_grad_y(
f_eval_corrected, y)
vjp_y_and_params = torch.autograd.grad(
outputs=f_eval_corrected,
inputs=y + params,
grad_outputs=tuple(-adj_y_ for adj_y_ in adj_y),
allow_unused=True,
create_graph=True)
vjp_y = vjp_y_and_params[:n_tensors]
vjp_y = misc.convert_none_to_zeros(vjp_y, y)
vjp_params = vjp_y_and_params[n_tensors:]
vjp_params = misc.flatten_convert_none_to_zeros(vjp_params, params)
adj_times_dgdx = torch.autograd.grad(outputs=g_eval,
inputs=y,
grad_outputs=adj_y,
allow_unused=True,
create_graph=True)
adj_times_dgdx = misc.convert_none_to_zeros(adj_times_dgdx, y)
# This extra term is due to converting the *adjoint* Stratonovich backward SDE to Itô.
extra_vjp_y_and_params = torch.autograd.grad(
outputs=g_eval,
inputs=y + params,
grad_outputs=adj_times_dgdx,
allow_unused=True,
)
extra_vjp_y = extra_vjp_y_and_params[:n_tensors]
extra_vjp_y = misc.convert_none_to_zeros(extra_vjp_y, y)
extra_vjp_params = extra_vjp_y_and_params[n_tensors:]
extra_vjp_params = misc.flatten_convert_none_to_zeros(
extra_vjp_params, params)
vjp_y = misc.seq_add(vjp_y, extra_vjp_y)
vjp_params = vjp_params + extra_vjp_params
return (*f_eval_corrected, *vjp_y, vjp_params)