Backpropagate model based on performance of another model

manifold: <p>Hi,</p>
<p>I currently have a problem where I have two models, <code>model_1</code> and <code>model_2</code>, and I would like to optimise one based on the performance of the other.</p>
<ul>
<li>Both models are of different types.</li>
<li>I would like to optimise <code>model_2</code> based on the performance of <code>model_1</code>. The loss function for <code>model_1</code> is well defined (CrossEntropyLoss), but for <code>model_2</code> it is not.</li>
<li><code>model_2</code> produces an output, which implicitly effects the performance of <code>model_1</code>. The output of <code>model_2</code> is <strong>not</strong> directly given to <code>model_1</code>. Overall, it’s objective is to minimise the loss of <code>model_1</code>.</li>
</ul>
<p>The problem I have is that I cannot back-propagate <code>model_2</code> based on the loss of <code>model_1</code> (being a scalar value). Is there a way to perform this operation?</p>

KFrank: <p>Hi Manifold!</p>
<aside class="quote no-group quote-modified" data-username="manifold" data-post="1" data-topic="196075" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/9de0a6/48.png" class="avatar"> manifold:</div>
<blockquote>
<p>I would like to optimise one based on the performance of the other.<br>
…</p>
<ul>
<li><code>model_2</code> produces an output, which implicitly effects the performance of <code>model_1</code>. The output of <code>model_2</code> is <strong>not</strong> directly given to <code>model_1</code>. Overall, it’s objective is to minimise the loss of <code>model_1</code>.<br>
…<br>
Is there a way to perform this operation?</li>
</ul>
</blockquote>
</aside>
<p>Yes, you should be able to do something like this.</p>
<p>However, you say that “the output of <code>model_2</code> is <strong>not</strong> directly given to<br>
<code>model_1</code>,” but that it “implicitly effects the performance of <code>model_1</code>.”</p>
<p>Unless you can tell us in concrete terms (perhaps with a fully-self-contained,<br>
runnable example script) how <code>model_2</code> is connected to <code>model_1</code>, we can’t<br>
really offer advice as to how (or whether) you might backpropagate through<br>
that connection.</p>
<p>Best.</p>
<p>K. Frank</p>

manifold: <p>Hi K. Frank!</p>
<p>I really appreciate the response. Apologies for leaving out some details, I thought it would simplify the question overall. I’m working in a Federated Learning setting, so in fact there are many models in play. I have <code>model_1</code>, <code>model_2</code> and k client models {<code>model_client_1</code>, <code>model_client_2</code>, … <code>model_client_k</code>}.</p>
<ul>
<li>The output of <code>model_2</code> is used in a regularisation term for each client model’s loss function.</li>
<li>Each client model trains on it’s own local private dataset.</li>
<li><code>model_1</code> is the average (all weights averaged via parameter matching) of all the client models.</li>
<li><code>model_2</code> then needs to be optimised to ensure the performance of <code>model_1</code> is minimised. However, as the output of <code>model_2</code> is used as part of a regularisation term for the client models, backtracking isn’t possible and there is no direct computation map between <code>model_1</code> and <code> model_2</code>.</li>
</ul>
<p>Hopefully this makes sense. I can provide more details if still unclear <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

KFrank: <p>Hi Manifold!</p>
<aside class="quote no-group" data-username="manifold" data-post="3" data-topic="196075" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/9de0a6/48.png" class="avatar"> manifold:</div>
<blockquote>
<ul>
<li>The output of <code>model_2</code> is used in a regularisation term for each client model’s loss function.</li>
<li>Each client model trains on it’s own local private dataset.</li>
</ul>
</blockquote>
</aside>
<p>As I understand it, you train your client models – that is, you update the<br>
parameters of your client models by taking optimization steps.  And the<br>
loss you use to train your client models consists of some normal sort of<br>
loss term plus a regulator that depends on the output of <code>model_2</code>.</p>
<aside class="quote no-group">
<blockquote>
<ul>
<li><code>model_1</code> is the average (all weights averaged via parameter matching) of all the client models.</li>
</ul>
</blockquote>
</aside>
<p><code>model_1</code> depends on (because its the average of) the trained client<br>
models.  It therefore depends on <code>model_2</code> because the optimization<br>
steps used to train the client models depend on <code>model_2</code> (because<br>
<code>model_2</code> was used for the regulator term in the client models’ training<br>
losses).</p>
<aside class="quote no-group">
<blockquote>
<ul>
<li><code>model_2</code> then needs to be optimised to ensure the performance of <code>model_1</code> is minimised. However, as the output of <code>model_2</code> is used as part of a regularisation term for the client models, backtracking isn’t possible and there is no direct computation map between <code>model_1</code> and <code> model_2</code>.</li>
</ul>
</blockquote>
</aside>
<p>The issue is that pytorch doesn’t offer functionality that lets you<br>
backpropagate <code>model_1</code>’s performance back through the optimization<br>
step that depended on <code>model_2</code>.</p>
<p>If I understand your use case properly, you would like to compute the<br>
gradient of <code>model_1</code>’s “performance” with respect to <code>model_2</code> (and<br>
then use that gradient to optimize <code>model_2</code> so that its output has the<br>
desired influence on the training of the client models via the regulator<br>
term).</p>
<p>I believe that you can do this under certain conditions (that hopefully<br>
will be consistent with your use case).</p>
<p>First, we will compute the gradient with respect to <code>model_2</code> for only<br>
a single optimization step of the client models at a time (rather than<br>
for a chain of optimization steps all at once).</p>
<p>Second, we will use plain-vanilla <code>SGD</code> to take the client-model optimization<br>
steps.  This is because the <code>SGD</code> optimization step:</p>
<pre data-code-wrap="python"><code class="lang-python">param = param -= lr * param.grad
</code></pre>
<p>is easy to differentiate:  <code>d_param / d_param.grad = -lr</code>.  This manual<br>
differentiation of the optimization step is the “glue” that will let us<br>
backpropagate <code>model_1</code>’s performance back through the optimization<br>
steps of the client models back to <code>model_2</code>’s parameters.</p>
<p>I’ve written a toy version of what I understand to be your use case and<br>
use this differentiation of the <code>SGD</code> optimization step, together with<br>
autograd (including computing the gradient of a gradient), to compute<br>
the gradient of <code>model_1</code> (after the client models have each undergone<br>
one regulated optimization step) with respect to <code>model_2</code>.</p>
<p>Start by looking at the function <code>fDGrad()</code> in the example script.  This<br>
finite-difference estimate of <code>model_1</code>’s gradient with respect to <code>model_2</code><br>
can be understood as a concrete operational definition of what we mean<br>
by that gradient.  The function <code>autogradGrad()</code> then shows how to<br>
compute that gradient using autograd (and our “manual” differentiation<br>
of the <code>SGD</code> optimization step).</p>
<p>Here is the full script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
torch.__version__

torch.random.manual_seed (2024)

nDim = 4                       # length of parameters, inputs, etc.
nClient = 3                    # number of client models
doDouble = True                # use double precision to improve finite-difference stability

print ('nDim:     ', nDim)
print ('nClient:  ', nClient)
print ('doDouble: ', doDouble)

# some helper functions

def applyParams (p, x):        # apply (warped) model parameters to input x -- returns scalar "loss"
    return  torch.nn.functional.mse_loss (x, torch.nn.functional.tanhshrink (p))

def getParams (m):
    return  m.p.detach().clone()

def setParams (m, params, gradToNone = True):
    with torch.no_grad():
        m.p.copy_ (params)
    if  gradToNone:
        m.p.grad = None
    return

def applyAverage (mList, x):   # apply average of models to x
    pList = [m.p for m in mList]
    pAvg = torch.stack (pList).mean (dim = 0)
    return  applyParams (pAvg, x)

class Mod (torch.nn.Module):   # simple "model"
    def __init__ (self):
        super().__init__()
        self.p = torch.nn.Parameter (torch.randn (nDim))
    def forward (self, x):
        return  applyParams (self.p, x)

# instantiate models

# model_1 -- the average of the client models -- is never explicitly instantiated
modCList = [Mod() for _ in range (nClient)]
mod2 = Mod()

x = torch.randn (nDim)         # use the same fixed input for all models

# lossReg will be the loss of a single client model together with its mod2 regulator
# lossOpt will be the loss of the average of the client models after they have each taken an optimization step

if  doDouble:
    mod2.double()
    for  m in modCList:
        m.double()
    x = x.double()
    # randDir = randDir.double()

# save model parameters so that models can be reset to original state
modCParams = [getParams (m) for m in modCList]
mod2Params = getParams (mod2)

# instantiate client-model optimizers
lr = 0.1
modCOpts = [torch.optim.SGD (m.parameters(), lr = lr) for m in modCList]

# note that the following functions depend on global variables such as modCList, mod2, etc.

def lossOptParams (params):    # helper function that computes lossOpt for given values of mod2 parameters
    setParams (mod2, params)
    for  m, p in zip (modCList, modCParams):
        setParams (m, p)
    m2Reg = mod2 (x).detach()
    for  m in modCList:
        loss = m (x)
        lossReg = loss + m2Reg * (m.p**2).sum()
        m.zero_grad()
        lossReg.backward()
    for  o in modCOpts:
        o.step()
    lossOpt = applyAverage (modCList, x)
    return lossOpt

# helper function that computes the two-sided finite-difference estimate of the
# directional derivative of lossOpt wrt mod2's parameters in the direction dDir
def fDDir (dDir):
    eps = torch.finfo (dDir.dtype).eps   # use different step sizes for different precisions
    h = eps**(1.0 / 3.0)                 # finite-difference step for single precision
    lossPlus  = lossOptParams (mod2Params + h * dDir)
    lossMinus = lossOptParams (mod2Params - h * dDir)
    return  (lossPlus - lossMinus) / (2.0 * h)

def fDGrad():         # function that computes finite-difference estimate gradient of lossOpt wrt mod2's parameters
    dDirs = torch.eye (len (mod2Params), dtype = mod2Params.dtype)
    result = torch.empty_like (mod2Params)
    for  i, dDir in enumerate (dDirs):
        result[i] = fDDir (dDir)
    return  result

def autogradGrad():   # function that uses autograd to compute the gradient of lossOpt wrt mod2's parameters
    # reset models
    setParams (mod2, mod2Params)
    for  m, p in zip (modCList, modCParams):
        setParams (m, p)
    
    # take a single optimization step for each client model with m2Reg as requlator
    m2Reg = mod2 (x)
    lossReg = 0.0
    for  m in modCList:
        lossReg += m (x) + m2Reg * (m.p**2).sum()
        m.zero_grad()
    # gradients are linear so perform a single backward for the sum of losses
    lossReg.backward (retain_graph = True)   # retain graph from mod2Reg to mod2 parameters
    for  o in modCOpts:
        o.step()
    
    # compute lossOpt and its gradient wrt client models
    lossOpt = applyAverage (modCList, x)
    modCPList = [p for m in modCList for p in m.parameters()]
    lossOptGrad = torch.autograd.grad (lossOpt, modCPList)
    
    # reset client models to their pre-optimization-step state
    for  m, p in zip (modCList, modCParams):
        setParams (m, p)
    
    # compute lossReg and its gradient wrt client models
    lossReg = 0.0
    for  m in modCList:
        lossReg += m (x) + m2Reg * (m.p**2).sum()   # lossReg depends on both modCList and mod2
    lossRegGrad = torch.autograd.grad (lossReg, modCPList, create_graph = True)
    
    # compute the vector-jacobian-product of lossOptGrad with the derivative of lossRegGrad wrt mod2
    # this -- with a factor of the optimizer lr -- gives the gradient lossOpt wrt mod2
    # chaining these gradients together effects backpropagation through the SGD optimization step
    lossOptGradJvp = -lr * torch.autograd.grad (lossRegGrad, mod2.parameters(), lossOptGrad)[0]
    
    return  lossOptGradJvp

mod2GradFD = fDGrad()
mod2GradAuto = autogradGrad()

# compare autograd result with  finite-difference estimate
torch.set_printoptions (precision = 15)
print ('mod2GradFD = ...')
print (mod2GradFD)
print ('mod2GradAuto = ...')
print (mod2GradAuto)
print ('mod2GradAuto / mod2GradFD = ...')
print (mod2GradAuto / mod2GradFD)
print ('torch.allclose (mod2GradAuto, mod2GradFD, rtol = 1.e-7, atol = 1.e-12) =',
        torch.allclose (mod2GradAuto, mod2GradFD, rtol = 1.e-7, atol = 1.e-12))
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">nDim:      4
nClient:   3
doDouble:  True
mod2GradFD = ...
tensor([ 0.001068887811096, -0.001403252429497, -0.003088872513755,
        -0.000195400080379], dtype=torch.float64, grad_fn=&lt;CopySlices&gt;)
mod2GradAuto = ...
tensor([ 0.001068887812660, -0.001403252420098, -0.003088872494453,
        -0.000195400066177], dtype=torch.float64)
mod2GradAuto / mod2GradFD = ...
tensor([1.000000001463032, 0.999999993302146, 0.999999993750954,
        0.999999927319424], dtype=torch.float64, grad_fn=&lt;DivBackward0&gt;)
torch.allclose (mod2GradAuto, mod2GradFD, rtol = 1.e-7, atol = 1.e-12) = True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

KFrank: <p>Hi Manifold!</p>
<p>Let me add a potentially useful bit of information:</p>
<aside class="quote no-group" data-username="KFrank" data-post="4" data-topic="196075" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>The issue is that pytorch doesn’t offer functionality that lets you<br>
backpropagate <code>model_1</code>’s performance back through the optimization<br>
step that depended on <code>model_2</code>.</p>
</blockquote>
</aside>
<p>It appears that (at least some of) pytorch’s optimizers offer a feature that<br>
may (or may not) do this for you.</p>
<p>As of pytorch version 1.13, <code>SGD</code> (and at least also <code>Adam</code>) offers a<br>
<code>differentiable</code> <a href="https://pytorch.org/docs/2.2/generated/torch.optim.SGD.html" rel="noopener nofollow ugc">constructor argument:</a></p>
<blockquote>
<p><strong>differentiable</strong> (<a href="https://docs.python.org/3/library/functions.html#bool" rel="noopener nofollow ugc"><em>bool</em></a><em>,</em> <em>optional</em>) – whether autograd should occur through the optimizer step in training. Otherwise, the step() function runs in a torch.no_grad() context. Setting to True can impair performance, so leave it False if you don’t intend to run autograd through this instance (default: False)</p>
</blockquote>
<p>I’ve never tried using it and I don’t know how (or whether) it works, but it<br>
might be worth looking into for your use case.</p>
<p>Best.</p>
<p>K. Frank</p>

manifold: <p>Hi KFrank!</p>
<p>Thanks a lot for this, it really helped me a lot.</p>
<p>I’ve managed to get it working now, from your suggestions!</p>