Implementing Tangent Prop Paper

username: Zhihan_Yang
cooked: <p>I’m trying to implement an old paper: <em>Tangent Prop - A formalism for specifying selected invariances in an adaptive network</em>. .</p>
<p>For example, let’s consider the following regression problem. Suppose we have 1-dimensional input (x) and 1-dimensional output (y). In this case, the Jacobian matrix is just 1-by-1 (a scalar) that describes the derivative of y w.r.t. x. So far so good.</p>
<p>So we iteratively perform gradient descent for the regression sum-of-squares loss + adjustable parameter * the squared Jacobian. Ideally, when this loss is minimized, we get the learned function as the solid line (see diagram below) given that our neural net has enough capacity.</p>
<p>Without the squared Jacobian term, we would get the learned function as the dashed line (see diagram below).</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/9/d98f813e7b98da8fec49c5fe9387c8e34379ddac.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/d98f813e7b98da8fec49c5fe9387c8e34379ddac" title="Screen Shot 2020-05-31 at 11.54.24 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/9/d98f813e7b98da8fec49c5fe9387c8e34379ddac_2_317x375.jpeg" alt="Screen Shot 2020-05-31 at 11.54.24 AM" data-base62-sha1="v2D4gAqgMpVI0Y6VzZ6pySQ4cS8" width="317" height="375" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/9/d98f813e7b98da8fec49c5fe9387c8e34379ddac_2_317x375.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/d/9/d98f813e7b98da8fec49c5fe9387c8e34379ddac_2_475x562.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/9/d98f813e7b98da8fec49c5fe9387c8e34379ddac_2_634x750.jpeg 2x" data-dominant-color="F1F1F1"><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">Screen Shot 2020-05-31 at 11.54.24 AM</span><span class="informations">922×1088 117 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>However, what I described above fails when implemented in PyTorch as follows:</p>
<pre><code class="lang-auto">class Net(nn.Module):
    
    def __init__(self):
        super().__init__()
        
        self.main = nn.Sequential(
            nn.Linear(1, 100),
            nn.ReLU(),
            nn.Linear(100, 100),
            nn.ReLU(),
            nn.Linear(100, 1)
        )
        
    def forward(self, x):
        return self.main(x)

def get_model():
    model = Net()
    return model, optim.Adam(model.parameters(), lr=1e-3)

def sos(ypred, y):
    "Sum of squares loss"
    return torch.sum((ypred - y) ** 2)

def get_jacobian(net, x, noutputs):
    x.requires_grad = True
    y = net(x)
    grad_params = torch.autograd.grad(y, x, create_graph=True)
    return grad_params[0]

# just a straight line with slope 1 and intercept 0
data = [0., 0.25, 0.5, 0.75, 1.0] 
labels = [0., 0.25, 0.5, 0.75, 1.0] 

model, opt = get_model()
jacobian_losses = []
reg_losses = []
for i in range(100):

    ypred = model(torch.tensor(data).view(-1, 1))
    reg_loss = sos(ypred, torch.tensor(labels).view(-1, 1))
    reg_losses.append(float(loss))
    
    jacobian_loss = 0
    for x in [0., 0.25, 0.5, 0.75, 1.0]:
        jacobian = get_jacobian(model, torch.tensor([[x]]), 1)
        temp = torch.norm(jacobian) ** 2
        jacobian_loss += temp
    jacobian_losses.append(float(jacobian_loss))
    
    total_loss = reg_loss + 1e-1 * jacobian_loss
    
    total_loss.backward()
    opt.step()
    opt.zero_grad()
</code></pre>
<p>Here’s the losses plotted together:</p>
<pre><code class="lang-auto">plt.plot(jacobian_losses, label='Jacobian')
plt.plot(reg_losses, label='Regression')
plt.xlabel('Epoch'); plt.ylabel('Loss')
plt.legend()
plt.show()
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/3/b3be4a5a6afb002336934d41a23e5dc63eb84509.png" alt="image" data-base62-sha1="pE5a52IubmkeY8zDJv6U0kJMoVb" width="392" height="262"></p>
<p>Here’s the resulting regression line:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/6/76dee639a987adc4395fea43b3acb356d5eac2ca.png" alt="image" data-base62-sha1="gXzTRORXizYpmthVyQe6Klsibr4" width="372" height="252"></p>
<p>And we don’t see the polynomial-like regression line like we saw above in the figure given by the paper. Can anyone help point out my mistake(s)? I’ve tried multiplying Jacobian loss by a smaller weight but I still don’t see the curvatures.</p>

username: albanD
cooked: <p>Hi,</p>
<p>Given your loss, I guess it is expected that the curve does not look good right? It would. only look like your first image if the loss goes to 0 right?</p>
<p>Maybe your model is not expressive enough? Have you checked?</p>

username: Zhihan_Yang
cooked: <p>Hi albanD,</p>
<p>Thanks for the comment. You’re right: the loss plot reflects the fact that the curve does not look good.</p>
<p>So I tried a model with a higher capacity as follows and the curve still looks as bad.</p>
<p>Do you find any mistake in my <code>get_jacobian</code> function?</p>
<pre><code class="lang-auto">class Net(nn.Module):
    
    def __init__(self):
        super().__init__()
        
        self.main = nn.Sequential(
            nn.Linear(1, 100),
            nn.ReLU(),
            nn.Linear(100, 100),
            nn.ReLU(),
            nn.Linear(100, 100),
            nn.ReLU(),
            nn.Linear(100, 100),
            nn.ReLU(),
            nn.Linear(100, 100),
            nn.ReLU(),
            nn.Linear(100, 1)
        )
        
    def forward(self, x):
        return self.main(x)
</code></pre>

username: albanD
cooked: <p>Hi.</p>
<p>Changing the weight for the jacobian part to 100 makes it go to 0 nicely.<br>
Similarly, putting it to a very low value makes the regular loss go to 0.</p>
<p>So it is an expression power of your model that fails to optimize both at the same time. Maybe try different activation? I am not sure that a relu network can do this task.</p>

username: Zhihan_Yang
cooked: <p>I tried the same model on the following data without increasing the expressive power of the model, yet the model was able to learn curvature with no problem (which is strange). I needed to increase the number of epochs for the model below to learn curvature, but adding the number of epochs did not work for minimizing regression loss + Jacobian loss.</p>
<p>Maybe there are some hidden mistakes in how I computed the Jacobian and the gradients?</p>
<p>Additional things I tried:</p>
<ul>
<li>I tried adding 3 more layers to my model and increasing the number of neurons per layer, but nothing worked.</li>
<li>I tried changing the ReLU activations to Sigmoid activations. But for Sigmoids to perform well, BatchNorm1d was needed (probably due to vanishing gradients). However, if you look at my <code>get_jacobian</code> function, I’m passing through the network only ONE example, so BatchNorm1d would raise an error since batch-wise standard deviation can’t be calculated from a batch of one example.</li>
</ul>
<pre><code class="lang-auto">data = [0, 0.20, 0.25, 0.30, 0.5, 0.75, 1.0]
labels = [0, 0.25, 0.25, 0.25, 0.5, 0.75, 1.0]

model, opt = get_model()
jacobian_losses = []
reg_losses = []
for i in range(300):
    
    ypred = model(torch.tensor(data).view(-1, 1).float())
    loss = sos(ypred, torch.tensor(labels).view(-1, 1))
    reg_losses.append(float(loss))
    
    loss.backward()
    opt.step()
    opt.zero_grad()

preds = model(torch.arange(-0.1, 1.1, 0.01).view(-1, 1)).view(-1).detach().numpy()
plt.scatter(data, labels)
plt.plot(np.arange(-0.1, 1.1, 0.01), preds)
plt.ylim(0, 1)
plt.show()
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/f/2/f24c22d541372ef1b10b9a0a44649aaf37126753.png" alt="image" data-base62-sha1="yzsFsaAXsCi6D1qvImzq0mnRW9R" width="372" height="252"></p>

username: Zhihan_Yang
cooked: <p>More than a year later, I decided to give this another go. I found that changing ReLU to Sigmoid / Softplus solved the issue. I was able to get the following plot:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/8/987eeeab8a068f889c1b7370df05f85b4445e35e.png" alt="Screen Shot 2021-12-29 at 4.11.35 PM" data-base62-sha1="lL2wEHw15m6BjGvqQYOkFJ2YHbE" width="330" height="324"></p>