# Gradient Origin Networks with Zygote/Flux?

**URL:** <https://discourse.julialang.org/t/gradient-origin-networks-with-zygote-flux/73421>\
**Category:** Machine Learning\
**Tags:** flux, zygote\
**Created:** [December 21, 2021, 8:07am UTC](https://discourse.julialang.org/t/gradient-origin-networks-with-zygote-flux/73421 "2021-12-21T08:07:39Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![afishy](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/afishy/32/18870_2.png) [@afishy](https://discourse.julialang.org/u/afishy)\
**Post date:** [December 21, 2021, 8:07am UTC](https://discourse.julialang.org/t/gradient-origin-networks-with-zygote-flux/73421/1 "2021-12-21T08:07:39Z")

</div>

Hi! I’d like to train an autoencoder according to [this paper](https://arxiv.org/abs/2007.02798). Briefly, it requires first taking a gradient step to get an estimate of the latent vector `z`, and then evaluating the output of the decoder using `z`. The difference between this and predictive/sparse coding is that you should keep the gradient of the loss wrt the model params from the `z` evaluation step when evaluating the loss after `Decoder(z)`. In the Pytorch implementation, this looks like [this](https://github.com/cwkx/GON/blob/master/GON.py)

```julia

   # compute the gradients of the inner loss with respect to zeros (gradient origin)
    z = torch.zeros(batch_size, nz, 1, 1).to(device).requires_grad_()
    g = F(z)
    inner_loss = ((g - x)**2).sum(1).mean()
    grad = torch.autograd.grad(inner_loss, [z], create_graph=True, retain_graph=True)[0]
    z = (-grad)

    # now with z as our new latent points, optimise the data fitting loss
    g = F(z)
    outer_loss = ((g - x)**2).sum(1).mean()
    optim.zero_grad()
    outer_loss.backward()
    optim.step()

```

Is it possible to do this with Zygote?

Thanks!

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [December 21, 2021, 10:35pm UTC](https://discourse.julialang.org/t/gradient-origin-networks-with-zygote-flux/73421/2 "2021-12-21T22:35:01Z")

</div>

For better or worse, the only way to know is to try. In theory it should be possible to do this kind of thing (GANs with gradient penalties are another example that has come up before) by nesting `gradient`/`pullback` with itself or some other autodiff, but success is highly dependent on what operations are being used (many are not second derivative friendly for Zygote + Zygote).