As you mentioned, the reason for the error is the mismatch between expected shape of output data. This hapens because the default loss function that you are using expects similar output (scalar for regression tasks) from the Neural Network as of your target data YTrain.
To train a Mixture Density Network (MDN) using trainNetwork in MATLAB, you need to implement a custom loss function to compute the negative log likelihood of the Gaussian mixture model.
A MDN typically outputs parameters: the means, variances, and weights for each component of the mixture. Assuming your MDN has K mixture components and each component is a Gaussian with D dimensions, the network's output layer should have K * (2 * D + 1) units.
Refer to the below code snippet which shows a way to write the custom loss function:
function loss = mdnLoss(Y, T, K)
% Y: Network output (Nx(K*3) matrix)
% T: Target data (Nx1 vector)
% K: Number of mixture components
% Extract means, variances, and weights from Y
N = size(T, 1);
D = 1; % Assuming 1D output for simplicity
% Reshape Y into means, variances, and weights
mu = reshape(Y(:, 1:K*D), N, K);
sigma = reshape(Y(:, K*D+1:2*K*D), N, K);
pi = reshape(Y(:, 2*K*D+1:end), N, K);
% Ensure variances are positive
sigma = exp(sigma);
% Apply softmax to weights to ensure they sum to 1
pi = softmax(pi, 2);
% Compute the Gaussian probability for each component
gaussians = exp(-0.5 * ((T - mu).^2) ./ (sigma.^2)) ./ (sqrt(2 * pi) * sigma);
% Compute the weighted sum of Gaussian probabilities
mixture_prob = sum(pi .* gaussians, 2);
% Compute the negative log-likelihood loss
loss = -sum(log(mixture_prob)) / N;
end
The code has some essential cheks for the oputput parameters:
- Use of the softmax function to ensure that the weights (pi_k) sum to 1.
- Use of an exponential function to ensure variances are positive.
Refer to the following documentation link to read more about defining custom loss functions:
