Adding non derivable Wasserstein from scipy to pytorch MSE Works?

username: Aymen_Tlili
cooked: <p>Hi i’m trying to use Optimal transport’s Wasserstein distance in pytorch.Before I had known about Geomloss (<a href="https://www.kernel-operations.io/geomloss/api/install.html" class="inline-onebox" rel="noopener nofollow ugc">Getting started — GeomLoss</a>) and their pytorch implementation of an approximation of it , I tried using the exact wasserstein distance on each predicted vector in a batch . “output[0]” represents the predicted batched and “tg” being the target being approximated .</p>
<p>reading <a href="https://discuss.pytorch.org/t/adding-a-float-variable-as-an-additional-loss-term-works/14363" class="inline-onebox">Adding a float variable as an additional loss term works?</a> I understand that adding constants(in that batch) gives a gradient of 0 so it should not affect back propogation whatsover but for me it does!<br>
A model trained to do regression with MSE alone is being out performed MSE+Scipy’s WS implementation , in fact this combination is even better than the sinkhorn approximation+MSE combo by a notable margin later down the road .</p>
<p>MSE+ws_scipy can be used in loss.backward() and accelerator.backward(loss) but if u use  ws_scipy alone things break because there is no graph behind it telling where to send the gradients to .pure MSE obviously doesn’t have this limitation.</p>
<p>My explination on this is that scipy’s term doesn’t create gradients but it’s presence increases MSE norm and that even in small variation(it’s own gradients alone couple with a big norm) is propagating WS information when it shouldn’t</p>
<p>an analogy is that I imagine the losses being partners .Each has a gradient (money) and a graph of norms(describing where to invest that money at) . Somehow in this combination WS has no money(not derivable by pytorch) and should therefore have no say in things but by injecting it’s norm in MSE’s graph it’s implifying MSE decision by increasing the norm sent through MSE’s graph alone .</p>
<p>if for neurone <span class="hashtag-raw">#250</span> it’s decided MSE small error , WS big error . WS shouldn’t affect the learning of it and should fall flat on it’s own but it’s somehow affecting MSE’s decision and sending big error still</p>
<p>here is code that shouldn’t work but does and outperforms  MSE criterion alone</p>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a> my hero in shining armor show me the way of the graphs (that we shouldn’t detach)</p>
<pre data-code-wrap="from"><code class="lang-from">            vec1_np = output[0].detach().cpu().numpy()
            vec2_np = tg.detach().cpu().numpy()
            # Compute Wasserstein distance for each pair of vectors
            wsd_list = [wasserstein_distance(vec1_np[i], vec2_np[i]) for i in range(vec1_np.shape[0])]
            # If you need an aggregate measure, you can compute the average distance
            average_wsd = sum(wsd_list) / len(wsd_list)
            loss_tr = criterion(output[0],tg)+average_wsd
         accelerator.backward(loss_tr)
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/b/3be8195980c43c171a4a99b6deec88b6fda2429f.png" data-download-href="https://discuss.pytorch.org/uploads/default/3be8195980c43c171a4a99b6deec88b6fda2429f" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/b/3be8195980c43c171a4a99b6deec88b6fda2429f_2_690x454.png" alt="image" data-base62-sha1="8xXrjJwuSRTBxnM4D7YbCz18xUr" width="690" height="454" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/b/3be8195980c43c171a4a99b6deec88b6fda2429f_2_690x454.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/b/3be8195980c43c171a4a99b6deec88b6fda2429f_2_1035x681.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/3/b/3be8195980c43c171a4a99b6deec88b6fda2429f.png 2x" data-dominant-color="252424"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1200×791 86.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div> [figure of MSE alone being the training metric and logging  ws_scipy metric increasing despite sinkhorn’s decreasing ]</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/9/198faec8905204fc92056267b54d605959c7fa42.png" data-download-href="https://discuss.pytorch.org/uploads/default/198faec8905204fc92056267b54d605959c7fa42" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/9/198faec8905204fc92056267b54d605959c7fa42_2_690x455.png" alt="image" data-base62-sha1="3E7M5w3rC5ypamN2rzAs16WWzCO" width="690" height="455" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/9/198faec8905204fc92056267b54d605959c7fa42_2_690x455.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/9/198faec8905204fc92056267b54d605959c7fa42_2_1035x682.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/1/9/198faec8905204fc92056267b54d605959c7fa42.png 2x" data-dominant-color="272424"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1208×797 97.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
[figure of MSE+ws_scipy being significantly different despite running on the same seed]</p>
<pre><code class="lang-auto">seed=74
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
</code></pre>
<p>here’s my github for this <a href="https://github.com/AymenTlili131/Federated-Continual-learning-" class="inline-onebox" rel="noopener nofollow ugc">GitHub - AymenTlili131/Federated-Continual-learning-: PFE</a><br>
Any help is appreciated before I crack open both libraries and start comparing implementations, I’d like to hear the opinions of people who contributed to them</p>
<p>scipy==1.12.0<br>
torch==1.12.1+cu113 (on RTX 3060 ubuntu22)<br>
torch_geometric==2.4.0<br>
torchaudio==0.12.1<br>
torchdata==0.5.1<br>
torcheval==0.0.7<br>
accelerate==0.30.1</p>

username: KFrank
cooked: <p>Hi Aymen!</p>
<aside class="quote no-group" data-username="Aymen_Tlili" data-post="1" data-topic="213191" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aymen_tlili/48/49197_2.png" class="avatar"> Aymen_Tlili:</div>
<blockquote>
<p>I understand that adding constants(in that batch) gives a gradient of 0 so it should not affect back propogation whatsover but for me it does!</p>
</blockquote>
</aside>
<p>This is unlikely to be true (so it is likely that there is something else going on).</p>
<p>It is <em>hypothetically</em> possible that adding a value that is not part of the<br>
computation graph could affect backpropagation and training.  This could<br>
occur due to finite floating-point precision.</p>
<p>Edit:  I’ve changed my mind about this.  Adding a constant (with <code>requires_grad = False</code>)<br>
to a loss value can introduce round-off error to <em>the result</em> of that addition, but I can’t<br>
think of a way that doing so can pollute, or otherwise change, gradients obtained by<br>
backpropagating that result.  Here is a simple illustration:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'2.5.1'
&gt;&gt;&gt; p = torch.tensor ([1.0], requires_grad = True)
&gt;&gt;&gt; (2 * p + 3).backward()
&gt;&gt;&gt; p.grad
tensor([2.])
&gt;&gt;&gt; p.grad = None
&gt;&gt;&gt; (2 * p + float ('inf')).backward()
&gt;&gt;&gt; p.grad
tensor([2.])
&gt;&gt;&gt; p.grad = None
&gt;&gt;&gt; (2 * p + float ('nan')).backward()
&gt;&gt;&gt; p.grad
tensor([2.])
</code></pre>
<p>For example, if your scipy result, <code>average_wsd</code>, were several of orders of<br>
magnitude larger than the differential piece, <code>criterion(output[0],tg)</code>,<br>
the gradient of the differentiable piece could be degraded by round-off error<br>
due to having been added to a much larger value.  If you only see a discrepancy<br>
after many training iterations, it could be that very small differences from a single<br>
pass accumulate and get “amplified” as several rounds of training progress.</p>
<p>Is there any way you could post a super-small, fully-self-contained, runnable<br>
example script that reproduces this behavior?  Ideally, could you capture the<br>
suspicious behavior by seeing differences in the gradients (<code>.grad</code> of various<br>
parameters) after a single backward pass?</p>
<p>Your reasoning that adding a term to the loss function that is not part of the<br>
computation graph will not <em>mathematically</em> affect backpropagation is correct<br>
(although it could affect things numerically).</p>
<p>Best.</p>
<p>K. Frank</p>

username: Aymen_Tlili
cooked: <p>Hey Frank ,<br>
thanks for the input and for reminding me of my netiquette to give a small reproducible example like this : <a href="https://colab.research.google.com/drive/1OaGwXFt8hbcSa1foXwe_4hh9q3B4BgQ3?usp=sharing" class="inline-onebox" rel="noopener nofollow ugc">Google Colab</a> if sharing works if not here is class declarations :</p>
<pre><code class="lang-auto">import torch
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
import torch.optim as optim
import numpy as np
import random
import copy

seed=74
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)


# Create a PyTorch-ready dataset
class ToyRegressionDataset(Dataset):
    def __init__(self, n_samples=10000, noise=0.1):
        # Generate random input data (features)
        self.x = torch.linspace(-10, 10, n_samples).unsqueeze(1)  # Shape: (n_samples, 1)
        # Generate labels (targets) based on the quadratic equation with noise
        self.y = 3 * self.x**2 + 2 * self.x + 1 + noise * torch.randn_like(self.x)
    
    def __len__(self):
        return len(self.x)
    
    def __getitem__(self, index):
        return self.x[index], self.y[index]

# Instantiate the dataset and dataloader
dataset = ToyRegressionDataset()
dataloader = DataLoader(dataset, batch_size=25, shuffle=True)

# Define a simple regression model
class RegressionModel(nn.Module):
    def __init__(self):
        super(RegressionModel, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(1, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 1)
        )
    
    def forward(self, x):
        return self.model(x)
</code></pre>
<p>and the cell that mixes it all (MSE + regularizing ws to scame scale as MSE X ws)</p>
<pre><code class="lang-auto"># Initialize the model, loss function, and optimizer
model_ws_reg=copy.deepcopy(model1)
criterion = nn.MSELoss()
optimizer = optim.Adam(model_ws_reg.parameters(), lr=0.01)

# Training loop
n_epochs = 1001
for epoch in range(n_epochs):
    b_i=0
    for x_batch, y_batch in dataloader:
        # Forward pass
        predictions = model_ws_reg(x_batch)
        loss = criterion(predictions, y_batch)
        vec1_np = predictions.detach().cpu().numpy()
        vec2_np = y_batch.detach().cpu().numpy()

        wsd_list = [wasserstein_distance(vec1_np[i], vec2_np[i]) for i in range(vec1_np.shape[0])]
        # If you need an aggregate measure, you can compute the average distance
        average_wsd = sum(wsd_list) / len(wsd_list)
        loss = criterion(predictions,y_batch)+average_wsd*0.0001
        # Backward pass and optimization
        if b_i==0 and epoch==0 :
          for name, param in model_ws_reg.named_parameters():
            if param.grad is not None:
                print(f"Parameter: {name} {param}, Gradient: {param.grad[:10]}")
                break
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if b_i==0 and epoch==0 :
          for name, param in model_ws_reg.named_parameters():
            if param.grad is not None:
                print(f"Parameter: {name} {param}, Gradient: {param.grad[:10]}")
                break
        b_i=b_i+1
    if (epoch) % 50 == 0:
        print(f"Epoch [{epoch+1}/{n_epochs}], Loss: {loss.item():.4f}")
for name, param in model.named_parameters():
  if param.grad is not None:
    print(f"Parameter: {name} {param}, Gradient: {param.grad[:10]}")
    break
# Test the model on new data
x_test = torch.linspace(-10, 10, 10000).unsqueeze(1)
y_test = 3 * x_test**2 + 2 * x_test + 1
y_pred = model_ws_reg(x_test).detach()

# Plot the results
import matplotlib.pyplot as plt

plt.figure(figsize=(8, 6))
plt.scatter(dataset.x, dataset.y, label="Original Data", color="blue", alpha=0.5)
plt.plot(x_test, y_test, label="True Function", color="green")
plt.plot(x_test, y_pred, label="Model Prediction", color="red")
plt.legend()
plt.title("Regression Model Fit")
plt.show()
</code></pre>
<p>The print outputs are as follows (from left to right : MSE , MSE+ws_scipy , MSE+0.0001ws_scipy )<br>
for 1000 data points and a batch size of 2 intended to create many training steps the results are <strong>identical and 100 epochs .</strong><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/7/57f0096447a983a92b321d60838de28822d860a0.png" data-download-href="https://discuss.pytorch.org/uploads/default/57f0096447a983a92b321d60838de28822d860a0" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/7/57f0096447a983a92b321d60838de28822d860a0_2_690x283.png" alt="image" data-base62-sha1="cxVOBjgP3ZYb6nWEsiOITBSEUz6" width="690" height="283" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/7/57f0096447a983a92b321d60838de28822d860a0_2_690x283.png, https://discuss.pytorch.org/uploads/default/optimized/3X/5/7/57f0096447a983a92b321d60838de28822d860a0_2_1035x424.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/5/7/57f0096447a983a92b321d60838de28822d860a0.png 2x" data-dominant-color="373737"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1267×521 45.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
in the example i’m running now i’ve tried to increase those by 10s or 100 folds but even with that and for context key dissimalirities between this and my first post . my first post is working on batchs of 25 1d vectors of lengths 2464 . there’s approximately 13500 examples of these 1D vectors in the trainset and the model is a double input transformer subject to <strong>gradient clipping</strong> and differing learning rates for each part of it .</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/7/77510ca7d66a4e3d861ba630f4d7867206a037b0.png" data-download-href="https://discuss.pytorch.org/uploads/default/77510ca7d66a4e3d861ba630f4d7867206a037b0" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/7/77510ca7d66a4e3d861ba630f4d7867206a037b0.png" alt="image" data-base62-sha1="h1wsVBEGp17iOk819C6rNBTkne0" width="599" height="129"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">599×129 6.59 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
and learning rates<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/e/2ef789468d576c6ac1575bda6dcf322163402f05.png" data-download-href="https://discuss.pytorch.org/uploads/default/2ef789468d576c6ac1575bda6dcf322163402f05" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/e/2ef789468d576c6ac1575bda6dcf322163402f05_2_689x133.png" alt="image" data-base62-sha1="6Huj4jndfKE1E2lWuULetGBmfxH" width="689" height="133" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/e/2ef789468d576c6ac1575bda6dcf322163402f05_2_689x133.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/e/2ef789468d576c6ac1575bda6dcf322163402f05_2_1033x199.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/e/2ef789468d576c6ac1575bda6dcf322163402f05.png 2x" data-dominant-color="F7F7F8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1293×250 55.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>the toy example is trying to fit this curve but I don’t know if the scale of values is comparable but i understand your motivation to showcase the change happening after one step , which there is none of here<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/a/8ad3e13ed16f5ee345e3a329a33837308a673c24.png" data-download-href="https://discuss.pytorch.org/uploads/default/8ad3e13ed16f5ee345e3a329a33837308a673c24" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/8/a/8ad3e13ed16f5ee345e3a329a33837308a673c24.png" alt="image" data-base62-sha1="jO7QauolV1ezVxKNSVfP4LClCpC" width="608" height="470"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">608×470 38.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Hope the info helps and gives more context</p>
<p><strong>update after the increases 100folds :</strong> the code I’m posting here is about 1000 epochs instead and here is the same comparison left to right . I’m assuming tracking the first named neurone is enough as params are put in an ordered dict data structure to my knowledge,correct me if not.<br>
(note also i didn’t enforce my package versions to deny that this problem stems from a certain old interaction that might’ve already been debugged)</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/d/fd28176f6d9a952b3f2ccc2c31de521f4bd2cde3.png" data-download-href="https://discuss.pytorch.org/uploads/default/fd28176f6d9a952b3f2ccc2c31de521f4bd2cde3" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/d/fd28176f6d9a952b3f2ccc2c31de521f4bd2cde3_2_690x328.png" alt="image" data-base62-sha1="A7wH8vcWMuUBmrghg6IZeFyLMeT" width="690" height="328" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/d/fd28176f6d9a952b3f2ccc2c31de521f4bd2cde3_2_690x328.png, https://discuss.pytorch.org/uploads/default/optimized/3X/f/d/fd28176f6d9a952b3f2ccc2c31de521f4bd2cde3_2_1035x492.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/f/d/fd28176f6d9a952b3f2ccc2c31de521f4bd2cde3.png 2x" data-dominant-color="3A3A3A"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1281×610 70.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>so the end model is the same but u can notice at the start of training that <strong>the gradients were different</strong> . I think the fact that the toy example is simple led to this unique solution .Arguably for a much complex task that gradient should cause differing end results if caught before convergence (i.e 900 epochs for a giga transformer)</p>
<p>now the tragic part about this is that the lack of regularization and the WS oversaturating the numerical result is causing better results for my original experiment <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=12" title=":sweat_smile:" class="emoji" alt=":sweat_smile:" loading="lazy" width="20" height="20"></p>

username: KFrank
cooked: <p>Hi Aymen!</p>
<aside class="quote no-group" data-username="Aymen_Tlili" data-post="3" data-topic="213191" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aymen_tlili/48/49197_2.png" class="avatar"> Aymen_Tlili:</div>
<blockquote>
<p>a small reproducible example like this :</p>
</blockquote>
</aside>
<p>That’s too much for me to wade through.</p>
<p>Your claim is that adding a scipy result – so not part of the computation graph – to<br>
your loss function changes the result of backpropagation.  If this claim is true, you<br>
should be able to demonstrate it with a single forward and backward pass.  So you<br>
shouldn’t need many (or any) epochs of training to see it.</p>
<p>Can you reproduce this with something truly minimal?  (No <code>DataLoader</code>, no <code>optim</code>,<br>
etc., just a single forward and backward pass with a super-simple toy model and<br>
hard-coded example data.)</p>
<p>Use a divide-and-conquer strategy to “capture” where the effect first shows up and<br>
the data that causes it.  Then post just the relevant code in a trimmed-down script<br>
together with the results you get when you run it.</p>
<p>Also, please don’t post screenshots of textual data.  Doing so breaks accessibility,<br>
searchability, and copy-paste.</p>
<p>Note, I no longer think that round-off error could be a possible explanation for what<br>
you’re seeing.  See the edit to my previous post.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Aymen_Tlili
cooked: <p>Hey again Frank Thanks for the feedback and keeping the best for everyone in mind .</p>
<p>I got rid of optimizers and dataloaders and just created fixed tensors to work x and y<br>
note I tried to use a batch size of 25 for both experiments but for one I tried 100 elements in the second dimension and in the second I increased that number of columns to 2464 .</p>
<p>The first step seems to be in all ways i can measure identical between the 2 distances in large and small column numberscenarios so I increased the number of iteration after confirming on 1 step</p>
<p>What i noticed when training this singular neurone is that for 100 columns torch.rand(25, 100) the results were identical even after 10 million iterations when comparing mse to mse+ws_scipy</p>
<p>And for the second experiment running this code on torch rands torch.rand(25, 2464) the results are identical.</p>
<p>So the effect does not appear at this scale for a single neurone or even after a single pass .<br>
I’ll try to up things a notch and use a dense layer and bit by bit reproduce the toy example proved the phenomenon in my last response .</p>
<pre><code class="lang-auto"># Set seed for reproducibility
torch.manual_seed(42)
# Sample toy dataset
x = torch.rand(25, 2464)# Input features
y = torch.rand(25, 2464)# Ground truth (targets)
# Initialize parameters
w = torch.tensor(0.5, requires_grad=True)  # Weight
b = torch.tensor(0.5, requires_grad=True)  # Bias
learning_rate = 0.0001
for iteration in range(1000):
  y_pred = w * x + b
  vec1_np = y_pred.detach().cpu().numpy()
  vec2_np = y.detach().cpu().numpy()
  wsd_list = [wasserstein_distance(vec1_np[i], vec2_np[i]) for i in range(vec1_np.shape[0])]
  
  average_wsd = sum(wsd_list) / len(wsd_list)

  loss = torch.mean((y_pred - y) ** 2)#+average_wsd
  if iteration%100==0:
    print(iteration , "Before loss backward \t",w.grad , b.grad)
  loss.backward()
  if iteration%100==0:
    print(iteration , "After loss backward \t",w.grad , b.grad ,"\t Loss \t", average_wsd , loss)
  with torch.no_grad():
      w -= learning_rate * w.grad
      b -= learning_rate * b.grad
  w.grad.zero_()
  b.grad.zero_()
  if iteration%100==0:
    print(iteration,"After with torch no_grad \t",w.grad , b.grad)
</code></pre>
<p>the results for a single forward backward pass’s and many more after are the gradients (first 2 numbers one for weight and one for bias)</p>
<p><strong>with ws_scipy</strong></p>
<pre><code class="lang-auto">0 Before loss backward 	 None None
0 After loss backward 	 tensor(0.3343) tensor(0.4997) 	 Loss 	 0.2498602782064058 tensor(0.4166, grad_fn=&lt;AddBackward0&gt;)
0 After with torch no_grad 	 tensor(0.) tensor(0.)
1000 Before loss backward 	 tensor(0.) tensor(0.)
1000 After loss backward 	 tensor(0.2703) tensor(0.3819) 	 Loss 	 0.20119250731857183 tensor(0.3395, grad_fn=&lt;AddBackward0&gt;)
1000 After with torch no_grad 	 tensor(0.) tensor(0.)
.
.
.
9000 Before loss backward 	 tensor(0.) tensor(0.)
9000 After loss backward 	 tensor(0.0746) tensor(0.0292) 	 Loss 	 0.1610148009207922 tensor(0.2552, grad_fn=&lt;AddBackward0&gt;)
9000 After with torch no_grad 	 tensor(0.) tensor(0.)
</code></pre>
<p><strong>without</strong></p>
<pre><code class="lang-auto">0 Before loss backward 	 None None
0 After loss backward 	 tensor(0.3343) tensor(0.4997) 	 Loss 	 0.2498602782064058 tensor(0.1667, grad_fn=&lt;MeanBackward0&gt;)
0 After with torch no_grad 	 tensor(0.) tensor(0.)
1000 Before loss backward 	 tensor(0.) tensor(0.)
1000 After loss backward 	 tensor(0.2703) tensor(0.3819) 	 Loss 	 0.20119250731857183 tensor(0.1383, grad_fn=&lt;MeanBackward0&gt;)
1000 After with torch no_grad 	 tensor(0.) tensor(0.)
.
.
9000 Before loss backward 	 tensor(0.) tensor(0.)
9000 After loss backward 	 tensor(0.0746) tensor(0.0292) 	 Loss 	 0.1610148009207922 tensor(0.0942, grad_fn=&lt;MeanBackward0&gt;)
9000 After with torch no_grad 	 tensor(0.) tensor(0.)
</code></pre>
<p>the only difference is that the norm of the loss is bigger but from the gradients and identical results we can infer that only the graph having MSE is back propagated .</p>
<p>I’ll upscale things with an optimizer and later a dense layer and follow the divide and conquer method . Thanks , any hypothesis or experiment ideas are welcome as we unravel this<br>
Kind regards</p>