What are good solution recipes to overcome the problem that in neural network learning, when the solution space has symmetries, learning may eventually stall due to the sum of the gradients over the network inputs converging to a tiny sum even though the individual gradients for each input to the network separately stay relatively big?
As a simple example, consider a neural network with a two-dimensional input and a singly dimensional output with a fully connected first layer, and a MaxPool layer as the second layer. We want the network to learn the function $\max(x+2y,3x+y)$ for network input tuples $(x,y)$. For this function, two different suitable weight matrices for the linear layer exist as the order of the input values to the MaxPool layer does not matter.
With a standard Adam optimizer, learning sometimes stalls before the MSE is very close to $0$. Small variations in the setting can cause learning to succeed, such as using a linear layer without a bias function or different random seeds. However, perhaps there are some commonly used approaches to solve such symmetry problems in a more systematic way?
The following example code in PyTorch exemplifies the problem:
import torch
import numpy
from livelossplot import PlotLosses
# Generate data to learn from
numpy.random.seed(42)
torch.manual_seed(123)
num_samples = 1000
X = torch.rand(num_samples, 2)
Y = torch.max(torch.cat((X[:,0:1]+2*X[:,1:2],3*X[:,0:1]+X[:,1:2]),dim=1),dim=1,keepdim=True)[0]
network = torch.nn.Sequential(
torch.nn.Linear(X.shape[1],2,bias=True),
torch.nn.MaxPool1d(2,1)
)
liveloss = PlotLosses()
opt = torch.optim.Adam(network.parameters(), lr=0.02)
lossfn = lambda x,y: ((x-y)**2).mean()
for i in range(4000):
opt.zero_grad()
out = network(X)
l = lossfn(Y,out)
current_loss = l.item()
(l).backward()
opt.step()
logs = {}
logs["err"] = torch.mean(torch.abs(Y - out)).item()
liveloss.update(logs)
liveloss.draw()
The values of the loss function over time (MSE) look as follows:
The final MSE is about 0.106. It can be verified that indeed the function to learn has a perfect representation in this network shape by initializing the weight matrix by hand:
with torch.no_grad():
network[0].weight[0, 0] = 1
network[0].weight[0, 1] = 2
network[0].weight[1, 0] = 3
network[0].weight[1, 1] = 1
network[0].bias[0] = 0
network[0].bias[1] = 0
When doing so before the learning process, the MSE starts with 0 and stays 0.
To provide some context for why this problem is interesting, the above example is just exemplifying a problem that occurs in a different context, namely encoding domain knowledge in a neural network with a structure adapted to this knowledge shape and then using neural network learning to fine-tune the model. The domain knowledge encoding may induce many such symmetries in the solution space, and hence just wiggling some aspects of the network or the learning process in an unsystematic way is unlikely to enforce learning across all such symmetries at the same time.
Edit with additional experiments based on comments.
- With a learning rate of 0.02 and without dropout and L2-regularization, the MSE plateaus at a MSE of 0.106 after ~10 epochs
- With a learning rate of 0.02 with dropout of 0.05 before the first (linear) layer and without L2-regularization, the MSE jumps between ~0.2 and ~0.25 after ~10 epochs
- With a learning rate of 0.02 with dropout of 0.05 after the first (linear) layer and without L2-regularization, the MSE jumps between ~0.17 and ~0.21 after ~10 epochs
In all cases, the learned model doesn't look anywhere close to the perfect model with the linear layer weight matrix $\left(\begin{matrix}1 & 2 \\ 3 & 1 \end{matrix}\right)$ and the perfect bias vector $(0,0)^T$ (for the data set).
