Why Is Not Your Neural Network Learning? Insights and Solutions

I recently stumbled upon a perplexing issue where my neural network, tasked with a regression problem, stubbornly learned to predict the average value of the dataset. After hours of meticulous investigation, I managed to close this issue.

In this blog, I’ll delve into the crux of the problem and share some insights I obtained for training neural networks till now.

Problem reformulation

This section chronicles my debugging journey. Readers eager to skip the gritty details can head straight to the next section.

Mind-map when debugging

Initially, my goal was straightforward: to perform a regression task, mapping 11-dimensional input data to a single-dimensional label. Upon initiating the training process, the absence of errors or warnings was reassuring. However, the loss curves quickly converged to a disappointingly high value. To my astonishment, the neural network’s predictions were consistently the average label value from the training set, regardless of the input.

My first hypothesis was that the network had settled into a local minimum, which the average value could represent in this scenario. I promptly experimented with various optimizer settings, including adjusting the learning rate and testing the second-order optimizer L-BFGS, as well as reinitializing the network multiple times. To my dismay, these efforts yielded no improvement, with the network stubbornly outputting the average label value.

Then I hypothesized that perhaps the model lacked the complexity necessary to capture the intricacies of the mapping. To test this, I expanded the network’s architecture, increasing both its depth and width, and even incorporated residual connections to enhance learning capabilities. Despite these ambitious modifications, the network persisted in its undesired behavior, outputting the average label value without fail.

The repeated failures after hours of meticulous tweaking were incredibly frustrating. It was time to take a deeper dive into the network’s inner workings. I began to scrutinize the intermediate outputs and gradients, searching for clues that could illuminate why the network was failing to learn beyond the simplest form of regression.

Caught the culprit

The Eureka Moment is that I discovered the root cause of the issue: the auto-broadcasting mechanism was at play!

Specifically, as an illustration example, if we run this code

import torch
x = torch.tensor([1.,2.,3.],requires_grad=True)
y = torch.tensor([2.,2.,2.],requires_grad=True)
res = torch.nn.MSELoss()(x,y.unsqueeze(-1))
res.backward()
print(f"res:{res}")
print(f"x.grad:{x.grad}")
print(f"y.grad:{y.grad}")

What printed on the screen is

res:0.6666666865348816
x.grad:tensor([-0.6667,  0.0000,  0.6667])
y.grad:tensor([0., 0., 0.])

See the problem?

If we take a closer look at the gradient with respect to y, it is weird that all entries are 0 unexpectedly. The correct answer would be the opposite of x.grad.

Here is the bug-free version:

import torch
x = torch.tensor([1.,2.,3.],requires_grad=True)
y = torch.tensor([2.,2.,2.],requires_grad=True)
res = torch.nn.MSELoss()(x,y)
res.backward()
print(f"res:{res}")
print(f"x.grad:{x.grad}")
print(f"y.grad:{y.grad}")

with the desired output

res:0.6666666865348816
x.grad:tensor([-0.6667,  0.0000,  0.6667])
y.grad:tensor([ 0.6667,  0.0000, -0.6667])

What is the subtlety here?

Both codes have the same value of “res”, which is the key element that misled me.

The pivotal difference here is the data shape, the first version computes the MSE loss of data of shape [3] and [3,1] while the right code fix this and compute MSE loss of data of shape [3] and [3]. And for different shapes, the pytorch would execute the so-called “Broadcasting” mechanism to resolve this.

Till now, I still can’t get the detailed mechanism of the implement of backward() in pytorch. However, it taught me a tough lesson that even with flexible language like python, we should always check the data type first before everything to get things done as expected. The broadcasting mechnism may not always be trusted.

Suggestions for training neural networks

In the process of resolving the aforementioned issue, I delved into numerous online resources and would like to compile these insights along with my former experience here.

The following are the general guidelines for troubleshooting learning issues in neural networks.

1. Verify the code is bug-free

It is not easy as the name suggested. A code running without errors does not guarantee it’s bug-free. Also, absence of warnings is not a green light. Sometimes, the code could run and give a result to cheat on you. But the process inside may not be the one you intended. Like the same loss but with wrong gradient issue I encountered.

To test the code, prioritize unit testing to ensure every module works as intended.

Common issues check list:

• Tensor Operations: Understand the difference between clone(), detach(), and .clone().detach().
• Optimizer Usage: Be aware of the correct sequence of zero_grad(), backward(), and step() when using optimizers.
• Loss Functions: Choose the appropriate loss function; for instance, avoid Mean Absolute Error (MAE) in networks that require second derivatives.
• Activation Functions: Select activation functions that are suitable for your network’s architecture and task.
• Broadcasting Issues: Be cautious about broadcasting issues, as they can lead to incorrect computations, as I experienced.
• Regularization Techniques: Properly apply regularizations like batch normalization, layer normalization, and dropout.
• Data Integrity: Check for NA, NaN, or Inf values in your dataset, both during preprocessing and training.
• Data Scaling: Scale your training data appropriately but ensure the testing data remains in its original form for accurate evaluation.

2. Start from small

The configuration of a neural network defines the solution space within which the optimization process operates.

A technique I often employ is to begin with a simple, small neural network rather than jumping straight into complex architectures. This initial step is not about achieving a minimal loss but about establishing a robust pipeline that runs error-free, thus reducing the training time and overhead.

Benefits of starting small:

• Pipeline Validation: A small network helps ensure that the entire pipeline, from data preprocessing to training, operates smoothly.
• Loss Benchmarking: Even if the loss from a small network is not satisfactory, it provides a baseline to gauge the complexity of the problem.
• Dataset Familiarization: Training a small network allows you to understand the characteristics of your dataset and the behavior of the optimizer, which is invaluable when scaling up.
• Learning rate calibration: It’s an opportunity to determine a suitable learning rate for your specific settings, which will be beneficial when training larger models.

From small to big, be aware of the following:

• Neuron Count: The number of neurons per layer significantly impacts performance. While theory suggests that three layers with an infinite number of neurons can approximate any function, in practice, finding the right balance is key. Too few neurons might lead to underfitting, while an excess can complicate training and increase the risk of overfitting. However, some over parameterization is generally beneficial for neural networks.
• Layer Depth: The number of layers is also crucial. Deep neural networks have achieved remarkable results, but going too deep without skip connections can exacerbate the gradient vanishing problem.
• Network Connectivity: Beyond fully connected networks, consider architectures that employ weight-sharing mechanisms to reduce parameter count and ease training.

3. Give your optimizer some patience

The essence of deep learning lies in parameterizing the solution space into a finite set, making it amenable to optimization. You might have heard that convex optimization boasts desirable properties, such as the equivalence of local and global minimum, the uniqueness of the global minimum, and ease of optimization.

However, neural networks are typically non-convex, with the exception of the linear regression case, which is not the norm. In the face of non-convex optimization landscapes, there’s no guaranteed way to find the global optimum, and the vast number of parameters exacerbates the training challenge. So, please be patient with your optimizer, as it’s tasked with a formidable job.

Sharing experiences:

• Be prepared to experiment with different settings and configurations. What works for one problem may not work for another.

• Try different optimizers: SGD, Adam, L-BFGS, Adam+L-BFGS

• Schedule the learning rate

  Instead of fix the learning rate, there are other ways more efficient to train your network

  • Assign different learning rates for different layers

    optimizer = optim.SGD([
            {"params":model.net2.parameters(),"lr":1e-5},],
            lr=1e-2,) #Default
    

  • Decrease the learning rate every step-size

    torch.optim.lr_scheduler.StepLR(optimizer, step_size, gamma=0.1, last_epoch=-1)
        
    

  • Cosine Learning rate

    torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max, eta_min=0, last_epoch=-1, verbose=False)
    

  • Adaptive leanring rate according to the loss

    torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min',patience=5)
    

Keep your intermediate results

• Checkpoints: Always maintain several checkpoints of your model at different stages of training. Ensure the filenames are descriptive and understandable even after a considerable time has passed. This practice is invaluable, especially for large projects, as it prevents the need to restart training from scratch if an issue arises.
• Loss Logs: Keep a detailed log of your training loss. When debugging, the loss curve can be a treasure trove of information, revealing insights into how your model is learning or where it might be encountering problems.

Final Remarks

• Sanity Checks: Conduct sanity checks to quickly verify the correct configuration of your network.

  Begin with a single data point, then progress to two distinct data points, and gradually increase. If your network cannot achieve near-zero loss even with overfitting on a few points, it’s likely an indicator of an underlying bug. Double-check your setup before scaling up to the entire dataset.

• Problem Learnability: Before committing to training a neural network, use traditional tools to analyze your dataset.
  • Perform a linear regression to assess linear relationships or employ other classical methods like Decision Trees or SVMs to gauge the feasibility of learning from your data.
  • If these techniques fail to provide solutions, it may indicate issues with the dataset itself. If the labels are entirely random and unrelated to the input features, a model will struggle to learn any meaningful patterns.
• Further Resources: For additional guidance on training neural networks, consider exploring reputable sources and tutorials, such as the insightful blog post by Andrej Karpathy, which offers a comprehensive recipe for neural network training.