%% *ADAM: A METHOD FOR STOCHASTIC OPTIMIZATION*

%%
%% *Problem formulation*
% The authors propose ADAM, a simple and computationally efficient method for gradient-based optimization of stochastic objective functions. It combines the advantages of two recently popular optimization methods: the ability of AdaGrad to deal with sparse gradients, and the ability of RMSProp to deal with non-stationary objectives. The project evaluated the method on multi-class logistic regression using the MNIST dataset.

%% ADAM ALGORITHM
%%
% |Adam consists of the following steps|
%%
%
% * *Compute gradient at current time.*
% * *Update biased first moment estimate.*
%%
%  
% $$ m_{t} = \beta_1m_{t-1} + (1 - \beta_1) g_t $$
% 
% It is computing a moving average of the derivatives for parameters. The gradient descent with momentum ends up taking steps that produce much smaller oscillations. This allows the algorithm to take a more straightforward path to the minimum. One of the advantages of this exponentially weighted average formula is that it takes very little memory. We need to keep just one row number in computer memory. The formula keeps on overwriting old values with the latest values. It is really efficient compared with the average over all parameters, as it only takes one line of code , and the storage for the memory of a single raw number.
%
% 
% * *Update biased second raw moment estimate.*
% 
% $$v_{t} = \beta_2v_{t-1} + (1 - \beta_2) g_t^2$$
% 
% We want the learning rate to be fast to converge, and we want to slow down the oscillations. So, with the term $v_t$, what we're hoping is that it will be relatively small. Parameters that gets few or small updates receive higher learning rates. One effect is that we can use a larger learning rate that decreases over time, and get faster learning without diverging in the vertical direction.
% 
% * *Compute bias-corrected first moment estimate.*
%
% $$\hat{m}_{t} = \frac{m_t}{1 - \beta_1^t}$$
% 
% * Compute bias-corrected second raw moment estimate.
%
% $$\hat{v}_{t} = \frac{v_t}{1 - \beta_2^t}$$
% 
% When we are implementing a moving average, the initialization value of $v$ is zero. Then the first several values of $v$ are lower. It is not a very good estimate of the gradients. Bias-corrected first momentum is a way to modify this estimate to make it more accurate, especially during this initial phase of the estimate. $1-\beta^t$ will remove the bias. As iteration $t$ is low, the denominator is low and $v$ becomes larger.
% 
% * *Update parameters.*
%
% $$\theta_t = \theta_{t-1} - \alpha \frac{\hat{m}_{t}}{\sqrt{\hat{v}_{t}} + \epsilon}$$
% 
% 
%% Logistic Regression
% % 
% minimize $l(\theta)$
% 
% where $\theta \in R^{(n+1)}$ is the variable, and $l$ is the logistic loss
% 
% $$l(\theta) = logL(\theta) = \sum_{i=1}^m y^{(i)} log h(x^{(i)}) + (1-y^{(i)})log(1-h(x^{(i)}))$$
% 
% where $m$ is the number of training data $(x^{(i)}, y^{(i)})$, and $h(x) = g(\theta^Tx) = \frac{1}{1+e^{-\theta^Tx}}$
% 
% The multiclass classification is to choose the maximum value among binary classifiers for single input data corresponds to the output class label and can be expressed as follows:
% 
% $h(x) = max_{i=1,2,...,k}(g(\theta_i^Tx)) = max_{i=1,2,...,k} \frac{1}{1+e^{-\theta_i^Tx}}$
% 
% where $k$ is the number of the binary classifiers corresponding to the number of classes
% 
% Gradients:
% For i = 1 to m, j = 1 to n
% $$\frac{\partial L}{\partial h^{(i)}} = -\frac{y^{(i)}}{h^{(i)}} + \frac{1-y^{(i)}}{1-h^{(i)}}$$
% 
% $$\frac{\partial L}{\partial \theta_j^{(i)}} = \frac{\partial l}{\partial
% h^{(i)}} \frac{\partial h^{(i)}}{\partial \theta_j^{(i)}}=
% x_j^{(i)}(h^{(i)} - y^{(i)})$$
% 
% $$\frac{\partial L}{\partial \theta_0^{(i)}} = h^{(i)} - y^{(i)}$$
% 
% where m is the number of samples and n is the number of dimensions.
% 
% Update parameters $\theta$ using gradient with Adam
%%

%%
%% Code
% 
clear; clc;
images = loadMNISTImages('train-images-idx3-ubyte');
labels = loadMNISTLabels('train-labels-idx1-ubyte');
% Create a shuffled version of the training set
p = randperm(size(images,1));
I_shuffle = zeros(60000,784);
labels_shuffle= categorical(zeros(60000,1));
%for i = 1:size(I_shuffle,1)
 %    I_shuffle(i,:) = images(p(i),:);
%end
% Create mini-batches
for j = 1:468
    I(j).image = I_shuffle(128*j-127:128*j,:);
    I(j).labels = labels(128*j-127:128*j,:);
end
I(469).image = I_shuffle(59905:60000,:);
I(469).labels = labels_shuffle(59905:60000,:);

function [x, fval, exitflag, output] = adam(fun, x0, stepSize, beta1, beta2, epsilon, nEpochSize, options)

%Usage: [x, fval, exitflag, output] = fmin_adam(fun, x0 <, stepSize, beta1, beta2, epsilon, nEpochSize, options>)
%
% 'output' will be a structure containing information about the
% optimisation process:
%     .stepsize      - Norm of current parameter step
%     .gradient      - Vector of current gradients
%     .funccount     - Number of function calls to 'fun' made so far
%     .iteration     - Current iteration of algorithm
%     .fval          - Current value returned by 'fun'
%     .exitflag      - Flag indicating termination reason
%     .improvement   - Current estimated improvement in 'fun'
%
% The optional parameters 'stepSize', 'beta1', 'beta2' and 'epsilon' are
% parameters of the Adam optimisation algorithm (see [1]). Default values
% of {1e-3, 0.9, 0.999, sqrt(eps)} are reasonable for most problems.
%
% The optional argument 'nEpochSize' specifies how many iterations comprise
% an epoch. This is used in the convergence detection code.
%
%% - Default parameters

DEF_stepSize = 0.001;
DEF_beta1 = 0.9;
DEF_beta2 = 0.999;
DEF_epsilon = sqrt(eps);

% - Default options
if (isequal(fun, 'defaults'))
   x = struct('Display', 'final', ...
      'GradObj', 'on', ...
      'DerivativeCheck', 'off', ...
      'MaxFunEvals', 1e4, ...
      'MaxIter', 1e6, ...
      'TolFun', 1e-6, ...
      'TolX', 1e-5, ...
      'UseParallel', false);
   return;
end


%% - Check arguments and assign defaults

if (nargin < 2)
   help fmin_adam;
   error('*** fmin_adam: Incorrect usage.');
end


if (~exist('stepSize', 'var') || isempty(stepSize))
   stepSize = DEF_stepSize;
end

if (~exist('beta1', 'var') || isempty(beta1))
   beta1 = DEF_beta1;
end

if (~exist('beta2', 'var') || isempty(beta2))
   beta2 = DEF_beta2;
end

if (~exist('epsilon', 'var') || isempty(epsilon))
   epsilon = DEF_epsilon;
end

if (~exist('options', 'var') || isempty(options))
   options = optimset(@adam);
end


%% - Parse options structure

numberofvariables = numel(x0);

% - Are analytical gradients provided?
if (isequal(options.GradObj, 'on'))
   % - Check supplied cost function
   if (nargout(fun) < 2) && (nargout(fun) ~= -1)
      error('*** fmin_adam: The supplied cost function must return analytical gradients, if options.GradObj = ''on''.');
   end
   
   bUseAnalyticalGradients = true;
   nEvalsPerIter = 2;
else
   bUseAnalyticalGradients = false;
   
   % - Wrap cost function for complex step gradients
   fun = @(x, nIter)FA_FunComplexStepGrad(fun, x, nIter);
   nEvalsPerIter = numberofvariables + 2;
end

% - Should we check analytical gradients?
bCheckAnalyticalGradients = isequal(options.DerivativeCheck, 'on');

% - Get iteration and termination options
MaxIter = FA_eval(options.MaxIter);
options.MaxIter = MaxIter;
options.MaxFunEvals = FA_eval(options.MaxFunEvals);

% - Parallel operation is not yet implements
if (options.UseParallel)
   warning('--- fmin_adam: Warning: ''UseParallel'' is not yet implemented.');
end


%% - Check supplied function

if (nargin(fun) < 2)
   % - Function does not make use of the 'nIter' argument, so make a wrapper
   fun = @(x, nIter)fun(x);
end

% - Check that gradients are identical for a given nIter
if (~bUseAnalyticalGradients)
   [~, vfGrad0] = fun(x0, 1);
   [~, vfGrad1] = fun(x0, 1);
   
   if (max(abs(vfGrad0 - vfGrad1)) > eps(max(max(abs(vfGrad0), abs(vfGrad1)))))
      error('*** fmin_adam: Cost function must return identical stochastic gradients for a given ''nIter'', when analytical gradients are not provided.');
   end
end

% - Check analytical gradients
if (bUseAnalyticalGradients && bCheckAnalyticalGradients)
   FA_CheckGradients(fun, x0);
end

% - Check user function for errors
try
   [fval0, vfCdX0] = fun(x0, 1);
   
catch mErr
   error('*** fmin_adam: Error when evaluating function to minimise.');
end

% - Check that initial point is reasonable
if (isinf(fval0) || isnan(fval0) || any(isinf(vfCdX0) | isnan(vfCdX0)))
   error('*** fmin_adam: Invalid starting point for user function. NaN or Inf returned.');
end


%% - Initialise algorithm

% - Preallocate cost and parameters
xHist = zeros(numberofvariables, MaxIter+1);%MappedTensor(numberofvariables, MaxIter+1);
xHist(:, 1) = x0;
x = x0;
vfCost = zeros(1, MaxIter);

if (~exist('nEpochSize', 'var') || isempty(nEpochSize))
   nEpochSize = numberofvariables;
end

vfCost(1) = fval0;
fLastCost = fval0;
fval = fval0;

% - Initialise moment estimates
m = zeros(numberofvariables, 1);
v = zeros(numberofvariables, 1);

% - Initialise optimization values
optimValues = struct('fval', vfCost(1), ...
   'funccount', nEvalsPerIter, ...
   'gradient', vfCdX0, ...
   'iteration', 0, ...
   'improvement', inf, ...
   'stepsize', 0);

% - Initial display
FA_Display(options, x, optimValues, 'init', nEpochSize);
FA_Display(options, x, optimValues, 'iter', nEpochSize);

% - Initialise plot and output functions
FA_CallOutputFunctions(options, x0, optimValues, 'init');
FA_CallOutputFunctions(options, x0, optimValues, 'iter');


%% - Optimisation loop
while true
   %% - Next iteration
   optimValues.iteration = optimValues.iteration + 1;
   nIter = optimValues.iteration;
   
   % - Update biased 1st moment estimate
   m = beta1.*m + (1 - beta1) .* optimValues.gradient(:);
   % - Update biased 2nd raw moment estimate
   v = beta2.*v + (1 - beta2) .* (optimValues.gradient(:).^2);
   
   % - Compute bias-corrected 1st moment estimate
   mHat = m./(1 - beta1^nIter);
   % - Compute bias-corrected 2nd raw moment estimate
   vHat = v./(1 - beta2^nIter);
   
   % - Determine step to take at this iteration
   vfStep = stepSize.*mHat./(sqrt(vHat) + epsilon);
   
   % - Test step for true improvement, reject bad steps
   if (fun(x(:) - vfStep(:), nIter) <= fval)
      x = x(:) - vfStep(:);
      optimValues.stepsize = max(abs(vfStep));
   end
   
   % - Get next cost and gradient
   [fval, optimValues.gradient] = fun(x, nIter+1);
   vfCost(nIter + 1) = fval;
   optimValues.funccount = optimValues.funccount + nEvalsPerIter;
   
   % - Call display, output and plot functions
   bStop = FA_Display(options, x, optimValues, 'iter', nEpochSize);
   bStop = bStop | FA_CallOutputFunctions(options, x, optimValues, 'iter');
   
   % - Store historical x
   xHist(:, nIter + 1) = x;
   
   % - Update covergence variables
   nFirstCost = max(1, nIter + 1-nEpochSize);
   fEstCost = mean(vfCost(nFirstCost:nIter+1));
   fImprEst = abs(fEstCost - fLastCost);
   optimValues.improvement = fImprEst;
   fLastCost = fEstCost;
   optimValues.fval = fEstCost;
   
   %% - Check termination criteria
   if (bStop)
      optimValues.exitflag = 0;
      break;
   end
   
   if (nIter > nEpochSize)
      if (fImprEst < options.TolFun / nEpochSize)
         optimValues.exitflag = 1;
         break;
      end
      
      if (optimValues.stepsize < options.TolX)
         optimValues.exitflag = 2;
         break;
      end
      
      if (nIter >= options.MaxIter-1)
         optimValues.exitflag = 3;
         break;
      end
      
      if (optimValues.funccount > options.MaxFunEvals)
         optimValues.exitflag = 4;
         break;
      end
   end
end
% - Determine best solution
vfCost = vfCost(1:nIter+1);
[~, nBestParams] = nanmin(vfCost);
x = xHist(:, nBestParams);
fval = vfCost(nBestParams);
exitflag = optimValues.exitflag;
output = optimValues;

% - Final display
FA_Display(options, x, optimValues, 'done', nEpochSize);
FA_CallOutputFunctions(options, x, optimValues, 'done');

end

%% Utility functions

% FA_FunComplexStepGrad - FUNCTION Compute complex step finite difference
% gradient estimates for an analytial function
function [fVal, vfCdX] = FA_FunComplexStepGrad(fun, x, nIter)
   % - Step size
   fStep = sqrt(eps);
   
   % - Get nominal value of function
   fVal = fun(x, nIter);
   
   % - Estimate gradients with complex step
   for (nParam = numel(x):-1:1)
      xStep = x;
      xStep(nParam) = xStep(nParam) + fStep * 1i;
      vfCdX(nParam, 1) = imag(fun(xStep, nIter)) ./ fStep; 
   end
end

% FA_CheckGradients - FUNCTION Check that analytical gradients match finite
% difference estimates
function FA_CheckGradients(fun, x0)
   % - Get analytical gradients
   [~, vfCdXAnalytical] = fun(x0, 1);
   
   % - Get complex-step finite-difference gradient estimates
   [~, vfCdXFDE] = FA_FunComplexStepGrad(fun, x0, 1);
   
   disp('--- fmin_adam: Checking analytical gradients...');
   
   % - Compare gradients
   vfGradDiff = abs(vfCdXAnalytical - vfCdXFDE);
   [fMaxDiff, nDiffIndex] = max(vfGradDiff);
   fTolGrad = eps(max(max(abs(vfCdXFDE), abs(vfCdXAnalytical))));
   if (fMaxDiff > fTolGrad)
      fprintf('   Gradient check failed.\n');
      fprintf('   Maximum difference between analytical and finite-step estimate: %.2g\n', fMaxDiff);
      fprintf('   Analytical: %.2g; Finite-step: %.2g\n', vfCdXAnalytical(nDiffIndex), vfCdXFDE(nDiffIndex));
      fprintf('   Tolerance: %.2g\n', fTolGrad);
      fprintf('   Gradient indicies violating tolerance: [');
      fprintf('%d, ', find(vfGradDiff > fTolGrad));
      fprintf(']\n');
      
      error('*** fmin_adam: Gradient check failed.');
   end
   
   disp('   Gradient check passed. Well done!');
end

% FA_Display - FUNCTION Display the current state of the optimisation
% algorithm
function bStop = FA_Display(options, x, optimValues, state, nEpochSize) %#ok<INUSL>
   bStop = false;

   % - Should we display anything?
   if (isequal(options.Display, 'none'))
      return;
   end

   % - Determine what to display
   switch (state)
      case 'init'
         if (isequal(options.Display, 'iter'))
            fprintf('\n\n%10s   %10s   %10s   %10s\n', ...
               'Iteration', 'Func-count', 'f(x)', 'Improvement', 'Step-size');
            fprintf('%10s   %10s   %10s   %10s   %10s\n', ...
               '----------', '----------', '----------', '----------', '----------');
         end

      case 'iter'
         if (isequal(options.Display, 'iter') && (mod(optimValues.iteration, nEpochSize) == 0))
            fprintf('%10d   %10d   %10.2g   %10.2g   %10.2g\n', ...
               optimValues.iteration, optimValues.funccount, ...
               optimValues.fval, optimValues.improvement, optimValues.stepsize);
         end

      case 'done'
         bPrintSummary = isequal(options.Display, 'final') | ...
            isequal(options.Display, 'iter') | ...
            (isequal(options.Display, 'notify') & (optimValues.exitflag ~= 1) & (optimValues.exitflag ~= 2));

         if (bPrintSummary)
            fprintf('\n\n%10s   %10s   %10s   %10s   %10s\n', ...
               'Iteration', 'Func-count', 'f(x)', 'Improvement', 'Step-size');
            fprintf('%10s   %10s   %10s   %10s   %10s\n', ...
               '----------', '----------', '----------', '----------', '----------');
            fprintf('%10d   %10d   %10.2g   %10.2g   %10.2g\n', ...
               optimValues.iteration, optimValues.funccount, ...
               optimValues.fval, optimValues.improvement, optimValues.stepsize);
            fprintf('%10s   %10s   %10s   %10s   %10s\n', ...
               '----------', '----------', '----------', '----------', '----------');

            strExitMessage = FA_GetExitMessage(optimValues, options);
            fprintf('\nFinished optimization.\n   Reason: %s\n\n', strExitMessage);
         end
   end
end

% FA_CallOutputFunctions - FUNCTION Call output and plot functions during
% optimisation
function bStop = FA_CallOutputFunctions(options, x, optimValues, state)
   bStop = false;

   if (~isempty(options.OutputFcn))
      bStop = bStop | options.OutputFcn(x, optimValues, state);
      drawnow;
   end

   if (~isempty(options.PlotFcns))
      if (iscell(options.PlotFcns))
         bStop = bStop | any(cellfun(@(fh)fh(x, optimValues, state), options.PlotFcns));
      else
         bStop = bStop | options.PlotFcns(x, optimValues, state);
      end
      drawnow;
   end
end

% FA_eval - FUNCTION Evaluate a string or return a value
function oResult = FA_eval(oInput)
   if (ischar(oInput))
      oResult = evalin('caller', oInput);
   else
      oResult = oInput;
   end
end

% FA_GetExitMessage - FUNCTION Return the message describing why the
% algorithm terminated
function strMessage = FA_GetExitMessage(optimValues, options)
   switch (optimValues.exitflag)
      case 0
         strMessage = 'Terminated due to output or plot function.';

      case 1
         strMessage = sprintf('Function improvement [%.2g] less than TolFun [%.2g].', optimValues.improvement, options.TolFun);

      case 2
         strMessage = sprintf('Step size [%.2g] less than TolX [%.2g].', optimValues.stepsize, options.TolX);
         
      case 3
         strMessage = sprintf('Number of iterations reached MaxIter [%d].', options.MaxIter);

      case 4
         strMessage = sprintf('Number of function evaluations reached MaxFunEvals [%d].', options.MaxFunEvals);

      otherwise
         strMessage = 'Unknown termination reason.';
   end
end

% --- END ---
%%
%
%% Analysis and conclusions
% Adam is an optimization algorithm that can used instead of the classical stochastic gradient descent procedure to update network weights iteratively based on training data. There are attractive benefits of using Adam as follows:
%%
% * Straightforward to implement.
% * Computationally efficient.
% * Little memory requirements.
% * Well suited for problems that are large in terms of data.
% * Appropriate for non-stationary objectives.
% * Appropriate for problems with very sparse gradients.
% * Hyper-parameters have intuitive interpretation and typically require little tuning.
%%
% |In reproducing the results of the paper "ADAM: A Method for Stochastic Optimization", we found our algorithm for ADAM had very similar results. Our implementation converges similarly to the SGD Nesterov method for the MNIST dataset. As well we also observed that both of those methods converge faster than ADAGrad for the MNIST database.|
%%
%% References
% 
% # Kingma, D.P. and Ba, J., 2014. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980.
% # Duchi, J., Hazan, E. and Singer, Y., 2011. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(Jul), pp.2121-2159.
% # Y. LeCun, C. Cortes and C. Burges, "MNIST handwritten digit database", Yann.lecun.com, 2010. [Online]. Available: http://yann.lecun.com/exdb/mnist/. [Accessed: 14- Apr- 2018].
%%
%% Submitted by
%
% Gunjan Shah
%
% Neha Bhatia
%
% *ECE - 602: Introduction to Optimization*
%
% Prof. Oleg Michailovich
%
% |Research Paper - ADAM: A Method for Stochastic Optimization by Diederik P. Kingma and Jimmy Lei Ba [1]| <https://arxiv.org/pdf/1412.6980.pdf>