optimize: SPSA and Adam optimizers, generalized call_on_me() in vec.h

optimize.cc

@@ -17,6 +17,7 @@
 */
 
 #include "optimize.h"
+#include "la_random.h"
 
 namespace LA {
 
@@ -203,7 +204,7 @@ return true;
 
 
 template<typename T>
-bool optstate<T>::stop(const T ff, const T gn)
+bool optstate<T>::stop(const T ff, const T gn, const int needed)
 {
 f=ff; gnorm=gn;
 if(hp.my_verbose)
@@ -214,7 +215,7 @@ if(hp.my_verbose)
 	}
 if(hp.converged(*this)) {converged=true; return true;}
 if(iter>=hp.my_maxstep) {if(hp.my_verbose) std::cout <<name<<": maximum number of iterations reached\n"; return true;}
-if(e.exhausted()) {if(hp.my_verbose) std::cout <<name<<": maximum number of function evaluations reached\n"; return true;}
+if(e.remaining()<needed) {if(hp.my_verbose) std::cout <<name<<": maximum number of function evaluations reached\n"; return true;}
 return false;
 }
 
@@ -741,6 +742,209 @@ return st.finish(achieved);
 
 
 
+template<typename T>
+optresult<T> firstorder_minimize(const char *name, firstorder_step<T> &stepper, const OptProblem<T> *o, NRVec<T> &x, const firstorder_hyperparameters<T> &hp, hyperparameters<T> *achieved, gradient_estimator<T> *est)
+{
+gradient_estimator<T> exact;
+gradient_estimator<T> &grad = est ? *est : exact;
+optevaluator<T> e(o,hp.my_maxfunc);
+optstate<T> st(name,hp,e);
+const int n=x.size();
+const T lr0=grad.calibrate(e,x,hp);
+if(hp.my_verbose && lr0!=hp.my_lr) std::cout <<name<<": calibrated learning rate "<<lr0<<std::endl;
+NRVec<T> g(n);
+T f=grad(e,x,g,0);
+
+for(;; ++st.iter)
+	{
+	if(st.stop(f,grad.exact() ? g.norm() : (T)-1,grad.cost())) break;
+	T lr=lr0/std::pow(1+hp.my_lr_decay*st.iter,hp.my_lr_power);
+	NRVec<T> dx = stepper.step(g,st.iter+1);
+	dx *= -lr;
+	if(hp.my_weight_decay!=0) dx.axpy(-lr*hp.my_weight_decay,x);
+	x += dx;
+	T fold=f;
+	NRVec<T> gnew(n); //fresh vector, the stepper may keep the previous gradient
+	f=grad(e,x,gnew,st.iter+1);
+	g=gnew;
+	if(grad.exact()) st.step(f-fold,dx.norm()); //for estimates, small changes are no sign of convergence
+	}
+
+if(hp.my_return_best && e.fbest<f)
+	{
+	if(hp.my_verbose) std::cout <<name<<": returning the best point found, f= "<<e.fbest<<std::endl;
+	x=e.xbest;
+	st.f=e.fbest;
+	}
+optresult<T> r=st.finish(achieved);
+firstorder_hyperparameters<T> *fach=dynamic_cast<firstorder_hyperparameters<T> *>(achieved);
+if(fach) fach->my_lr=lr0;
+return r;
+}
+
+
+//SGD with momentum
+template<typename T>
+class SGD_stepper : public firstorder_step<T> {
+public:
+const SGD_hyperparameters<T> &hp;
+NRVec<T> v; //velocity
+SGD_stepper(const SGD_hyperparameters<T> &h) : hp(h) {};
+NRVec<T> step(const NRVec<T> &g, const int k)
+	{
+	if(k==1) v=g;
+	else {v *= hp.my_momentum; v += g;}
+	if(!hp.my_nesterov) return v;
+	NRVec<T> d(g);
+	d.axpy(hp.my_momentum,v);
+	return d;
+	}
+};
+
+
+template<typename T>
+optresult<T> SGD_Optimizer<T>::minimize(const OptProblem<T> *o, NRVec<T> &x, const hyperparameters<T> *par, hyperparameters<T> *achieved)
+{
+SGD_hyperparameters<T> hp;
+if(par) hp.setfrom(*par);
+if(hp.my_saddle_order) laerror("SGD cannot search for saddle points");
+SGD_stepper<T> stepper(hp);
+return firstorder_minimize(hp.my_nesterov ? "SGD-Nesterov" : "SGD",stepper,o,x,hp,achieved);
+}
+
+
+//elementwise functions for Adam
+template<typename T>
+static T elementmax(const T &a, const T &b) {return a>b ? a : b;}
+
+//m/(sqrt(v)/sqrtc2 + eps) with the bias correction sqrtc2 = sqrt(1-beta2^k)
+template<typename T>
+class adamratio {
+public:
+T sqrtc2, eps;
+adamratio(const T s, const T e) : sqrtc2(s), eps(e) {};
+T operator()(const T &m, const T &v) const {return m/(std::sqrt(v)/sqrtc2+eps);}
+};
+
+
+//Adam, AdamW, AMSGrad
+template<typename T>
+class Adam_stepper : public firstorder_step<T> {
+public:
+const Adam_hyperparameters<T> &hp;
+NRVec<T> m; //first moment estimate
+NRVec<T> v; //second moment estimate
+NRVec<T> vmax; //maximum of the second moments for AMSGrad
+T beta1k, beta2k; //beta1^k, beta2^k
+Adam_stepper(const Adam_hyperparameters<T> &h) : hp(h), beta1k(1), beta2k(1) {};
+NRVec<T> step(const NRVec<T> &g, const int k)
+	{
+	NRVec<T> g2(g);
+	g2 *= g;
+	if(k==1) {m = g*(1-hp.my_beta1); v = g2*(1-hp.my_beta2);}
+	else
+		{
+		m *= hp.my_beta1; m.axpy(1-hp.my_beta1,g);
+		v *= hp.my_beta2; v.axpy(1-hp.my_beta2,g2);
+		}
+	beta1k *= hp.my_beta1;
+	beta2k *= hp.my_beta2;
+	if(hp.my_amsgrad) {if(k==1) vmax=v; else vmax.call_on_me(elementmax<T>,v);}
+	NRVec<T> d = m*(1/(1-beta1k));
+	d.call_on_me(adamratio<T>(std::sqrt(1-beta2k),hp.my_eps), hp.my_amsgrad ? vmax : v);
+	return d;
+	}
+};
+
+
+template<typename T>
+optresult<T> Adam_Optimizer<T>::minimize(const OptProblem<T> *o, NRVec<T> &x, const hyperparameters<T> *par, hyperparameters<T> *achieved)
+{
+Adam_hyperparameters<T> hp;
+if(par) hp.setfrom(*par);
+if(hp.my_saddle_order) laerror("Adam cannot search for saddle points");
+Adam_stepper<T> stepper(hp);
+return firstorder_minimize(hp.my_amsgrad ? "AMSGrad" : (hp.my_weight_decay!=0 ? "AdamW" : "Adam"),stepper,o,x,hp,achieved);
+}
+
+
+
+//SPSA gradient estimate (g_hat)_i = (f(x+c_k Delta) - f(x-c_k Delta))/(2 c_k Delta_i), Delta_i = +-1, averaged over resamplings
+template<typename T>
+class SPSA_estimator : public gradient_estimator<T> {
+public:
+const SPSA_hyperparameters<T> &hp;
+SPSA_estimator(const SPSA_hyperparameters<T> &h) : hp(h) {};
+T operator()(optevaluator<T> &e, const NRVec<T> &x, NRVec<T> &g, const int k);
+int cost() const {return 2*hp.my_resamplings;}
+bool exact() const {return false;}
+T calibrate(optevaluator<T> &e, const NRVec<T> &x, const firstorder_hyperparameters<T> &h);
+};
+
+//returns the average of f(x+c_k Delta) and f(x-c_k Delta) as an estimate of f(x)
+template<typename T>
+T SPSA_estimator<T>::operator()(optevaluator<T> &e, const NRVec<T> &x, NRVec<T> &g, const int k)
+{
+const int n=x.size();
+const T ck=hp.my_c/std::pow((T)(k+1),hp.my_gamma);
+g.clear();
+T fsum=0;
+for(int r=0; r<hp.my_resamplings; ++r)
+	{
+	NRVec<T> delta(n); //fresh vector, element access is safe
+	for(int i=0; i<n; ++i) delta[i]= randdouble()<0.5 ? -1 : 1;
+	NRVec<T> xp(x), xm(x);
+	xp.axpy(ck,delta);
+	xm.axpy(-ck,delta);
+	T fp=e(xp);
+	T fm=e(xm);
+	g.axpy((fp-fm)/(2*ck),delta); //1/Delta_i = Delta_i
+	fsum+=fp+fm;
+	}
+g /= (T)hp.my_resamplings;
+return fsum/(2*hp.my_resamplings);
+}
+
+//learning rate such that the first step lr*|g_hat| has on average the target length
+template<typename T>
+T SPSA_estimator<T>::calibrate(optevaluator<T> &e, const NRVec<T> &x, const firstorder_hyperparameters<T> &h)
+{
+if(hp.my_calibrate<=0) return hp.my_lr;
+T gsum=0;
+int ncal=0;
+for(; ncal<hp.my_calibrate && e.remaining()>=2*cost(); ++ncal) //leave the budget for the first estimate
+	{
+	NRVec<T> g(x.size());
+	(*this)(e,x,g,0);
+	gsum+=g.norm();
+	}
+if(ncal==0 || gsum==0) return hp.my_lr;
+return hp.my_target_step*ncal/gsum;
+}
+
+
+//plain gradient step for SPSA
+template<typename T>
+class plain_stepper : public firstorder_step<T> {
+public:
+NRVec<T> step(const NRVec<T> &g, const int k) {return g;}
+};
+
+
+template<typename T>
+optresult<T> SPSA_Optimizer<T>::minimize(const OptProblem<T> *o, NRVec<T> &x, const hyperparameters<T> *par, hyperparameters<T> *achieved)
+{
+SPSA_hyperparameters<T> hp;
+if(par) hp.setfrom(*par);
+if(hp.my_saddle_order) laerror("SPSA cannot search for saddle points");
+if(hp.my_resamplings<1) laerror("SPSA needs resamplings >= 1");
+plain_stepper<T> stepper;
+SPSA_estimator<T> est(hp);
+return firstorder_minimize("SPSA",stepper,o,x,hp,achieved,&est);
+}
+
+
+
 //force instantiation
 template double linesearch_wolfe(optevaluator<double> &e, NRVec<double> &x, double &f, NRVec<double> &g, const NRVec<double> &p, const double alpha1, const linesearch_hyperparameters<double> &lp);
 template class optstate<double>;
@@ -750,6 +954,9 @@ template class LBFGS_Optimizer<double>;
 template class CG_Optimizer<double>;
 template class NelderMead_Optimizer<double>;
 template class Powell_Optimizer<double>;
+template class SGD_Optimizer<double>;
+template class Adam_Optimizer<double>;
+template class SPSA_Optimizer<double>;
 template double linemin_brent(optevaluator<double> &e, NRVec<double> &x, double &f, const NRVec<double> &u, const double t1, const double tol, const double abstol, const int maxiter, const int verbose);
 
 }//namespace

optimize.h

@@ -145,7 +145,7 @@ T dx; //|x change| in the last step
 bool stepdone; //df and dx are valid
 bool converged;
 optstate(const char *nm, const hyperparameters<T> &h, const optevaluator<T> &ee) : name(nm), hp(h), e(ee), iter(0), f(0), gnorm(-1), df(0), dx(0), stepdone(false), converged(false) {};
-bool stop(const T ff, const T gn= -1); //to be called at the beginning of each iteration, returns true if the iterations should finish
+bool stop(const T ff, const T gn= -1, const int needed=1); //to be called at the beginning of each iteration, returns true if the iterations should finish; needed = no. of function evaluations of the next iteration
 void step(const T dff, const T dxx) {df=std::abs(dff); dx=dxx; stepdone=true;} //record a step done
 optresult<T> finish(hyperparameters<T> *achieved) const; //final report, fills achieved if not NULL
 };
@@ -197,6 +197,44 @@ template<typename T>
 optresult<T> linesearch_minimize(const char *name, linesearch_direction<T> &dir, const OptProblem<T> *o, NRVec<T> &x, const linesearch_hyperparameters<T> &hp, hyperparameters<T> *achieved);
 
 
+//parameters of first-order methods with a learning rate instead of line search, suitable also for noisy (stochastic) functions and gradients
+template<typename T>
+class firstorder_hyperparameters : public gradient_hyperparameters<T> {
+public:
+T my_lr; //learning rate
+T my_lr_decay; //learning rate in iteration k is lr/(1+lr_decay*k)^lr_power
+T my_lr_power;
+T my_weight_decay; //decoupled weight decay x -= lr*weight_decay*x (Loshchilov & Hutter, ICLR 2019)
+bool my_return_best; //return the point with the lowest function value found instead of the last one; only for (nearly) deterministic f, the lowest of noisy values is biased by the noise
+firstorder_hyperparameters() : my_lr(1e-3), my_lr_decay(0), my_lr_power(1), my_weight_decay(0), my_return_best(false) {this->my_maxstep=1000;};
+virtual void setfrom(const hyperparameters<T> &p) {const firstorder_hyperparameters *q=dynamic_cast<const firstorder_hyperparameters *>(&p); if(q) *this= *q; else gradient_hyperparameters<T>::setfrom(p);}
+};
+
+//update rule of a first-order method, used by the common driver firstorder_minimize()
+template<typename T>
+class firstorder_step {
+public:
+virtual NRVec<T> step(const NRVec<T> &g, const int k) =0; //direction d of the step x -= lr*d in iteration k=1,2,... with gradient g
+virtual ~firstorder_step() {};
+};
+
+//function value and gradient for firstorder_minimize(), the default is the exact evaluation, derived classes estimate them (SPSA)
+template<typename T>
+class gradient_estimator {
+public:
+virtual T operator()(optevaluator<T> &e, const NRVec<T> &x, NRVec<T> &g, const int k) {return e(x,&g);} //returns f at x and the gradient g, k=0,1,... is the iteration
+virtual int cost() const {return 1;} //no. of function evaluations of one call
+virtual bool exact() const {return true;} //with estimates no convergence test is done, small gradients or changes can occur by chance
+virtual T calibrate(optevaluator<T> &e, const NRVec<T> &x, const firstorder_hyperparameters<T> &hp) {return hp.my_lr;} //learning rate to be used
+virtual ~gradient_estimator() {};
+};
+
+//common driver of first-order methods; est==NULL means exact gradients
+//the learning rate actually used is returned in achieved, if it is of firstorder_hyperparameters type
+template<typename T>
+optresult<T> firstorder_minimize(const char *name, firstorder_step<T> &stepper, const OptProblem<T> *o, NRVec<T> &x, const firstorder_hyperparameters<T> &hp, hyperparameters<T> *achieved, gradient_estimator<T> *est=NULL);
+
+
 //specific methods
 
 //parameters common to quasi-Newton methods
@@ -311,6 +349,69 @@ virtual ~Powell_Optimizer() {};
 };
 
 
+//stochastic gradient descent with momentum v = momentum*v + g, step d = v or d = g + momentum*v (Nesterov) (Sutskever et al., ICML 2013)
+template<typename T>
+class SGD_hyperparameters : public firstorder_hyperparameters<T> {
+public:
+T my_momentum;
+bool my_nesterov;
+SGD_hyperparameters() : my_momentum(0.9), my_nesterov(false) {this->my_lr=1e-2;};
+virtual void setfrom(const hyperparameters<T> &p) {const SGD_hyperparameters *q=dynamic_cast<const SGD_hyperparameters *>(&p); if(q) *this= *q; else firstorder_hyperparameters<T>::setfrom(p);}
+};
+
+template<typename T>
+class SGD_Optimizer : public Optimizer<T> {
+public:
+virtual optresult<T> minimize(const OptProblem<T> *o, NRVec<T> &x, const hyperparameters<T> *par=NULL, hyperparameters<T> *achieved=NULL);
+virtual ~SGD_Optimizer() {};
+};
+
+
+//Adam (Kingma & Ba, ICLR 2015), with weight_decay>0 AdamW, optionally AMSGrad (Reddi et al., ICLR 2018); beta1=0 gives RMSprop
+template<typename T>
+class Adam_hyperparameters : public firstorder_hyperparameters<T> {
+public:
+T my_beta1; //decay rate of the first moment estimate
+T my_beta2; //decay rate of the second moment estimate
+T my_eps; //regularization of the division by the second moment
+bool my_amsgrad; //use the maximum of the past second moments
+Adam_hyperparameters() : my_beta1(0.9), my_beta2(0.999), my_eps(1e-8), my_amsgrad(false) {};
+virtual void setfrom(const hyperparameters<T> &p) {const Adam_hyperparameters *q=dynamic_cast<const Adam_hyperparameters *>(&p); if(q) *this= *q; else firstorder_hyperparameters<T>::setfrom(p);}
+};
+
+template<typename T>
+class Adam_Optimizer : public Optimizer<T> {
+public:
+virtual optresult<T> minimize(const OptProblem<T> *o, NRVec<T> &x, const hyperparameters<T> *par=NULL, hyperparameters<T> *achieved=NULL);
+virtual ~Adam_Optimizer() {};
+};
+
+
+//simultaneous perturbation stochastic approximation (Spall, IEEE Trans. Autom. Control 37, 332 (1992); IEEE Trans. Aerosp. Electron. Syst. 34, 817 (1998))
+//derivative-free, 2 function evaluations per gradient estimate independently of the dimension, robust to noise (VQE)
+//gain sequences: a_k = lr/(1+lr_decay*k)^lr_power, i.e. Spall's a/(k+1+A)^alpha with lr=a/(1+A)^alpha, lr_decay=1/(1+A), lr_power=alpha; c_k = c/(k+1)^gamma
+//the gradient and f are estimates, therefore no convergence test is done and the iterations run until maxstep or maxfunc
+//the first estimate at the starting point is always done, even if maxfunc is smaller than its 2*resamplings evaluations
+template<typename T>
+class SPSA_hyperparameters : public firstorder_hyperparameters<T> {
+public:
+T my_c; //perturbation size, of the order of the noise of f
+T my_gamma; //decay exponent of the perturbation size
+int my_resamplings; //no. of gradient estimates averaged in one iteration
+int my_calibrate; //no. of gradient estimates at the starting point for the calibration of lr, 0 = no calibration
+T my_target_step; //calibration: lr is set so that the first step has this length on average
+SPSA_hyperparameters() : my_c(0.1), my_gamma(0.101), my_resamplings(1), my_calibrate(20), my_target_step(0.1) {this->my_lr_decay=0.01; this->my_lr_power=0.602; this->my_g_thr=0;};
+virtual void setfrom(const hyperparameters<T> &p) {const SPSA_hyperparameters *q=dynamic_cast<const SPSA_hyperparameters *>(&p); if(q) *this= *q; else firstorder_hyperparameters<T>::setfrom(p);}
+};
+
+template<typename T>
+class SPSA_Optimizer : public Optimizer<T> {
+public:
+virtual optresult<T> minimize(const OptProblem<T> *o, NRVec<T> &x, const hyperparameters<T> *par=NULL, hyperparameters<T> *achieved=NULL);
+virtual ~SPSA_Optimizer() {};
+};
+
+
 
 }//namespace
 #endif /* _OPTIMIZE_H_ */

t.cc

@@ -4847,11 +4847,16 @@ testproblem<double> problem;
 //LBFGS_hyperparameters<double> param; //tested
 //CG_hyperparameters<double> param; //tested
 //NelderMead_hyperparameters<double> param; //tested
-Powell_hyperparameters<double> param;
+//Powell_hyperparameters<double> param; //tested
+//SGD_hyperparameters<double> param; //tested
+//Adam_hyperparameters<double> param; //tested
+SPSA_hyperparameters<double> param;
 param.my_maxstep=100;
 param.my_f_thr=1e-6;
 param.my_x_thr=1e-3;
-//param.my_g_thr=1e-3; //gradient methods
+param.my_g_thr=1e-3; //gradient methods
+param.my_lr=1e-2; //first-order methods
+param.my_maxstep=2000; //first-order methods
 //param.my_update=QN_BFGS; //tested
 //param.my_update=QN_DFP; //tested
 //param.my_update=QN_SR1; //tested
@@ -4860,13 +4865,16 @@ cout << "initial " <<x<<endl;
 //LBFGS_Optimizer<double> o; //tested
 //CG_Optimizer<double> o; //tested
 //NelderMead_Optimizer<double> o; //tested
-Powell_Optimizer<double> o;
-//gradient_hyperparameters<double> achieved; //gradient methods
-hyperparameters<double> achieved;
+//Powell_Optimizer<double> o; //tested
+//SGD_Optimizer<double> o; //tested
+//Adam_Optimizer<double> o; //tested
+SPSA_Optimizer<double> o;
+gradient_hyperparameters<double> achieved; //gradient methods
+//hyperparameters<double> achieved; //derivative-free methods
 optresult<double> res=o.minimize(&problem,x,&param,&achieved);
 cout << "optimized " <<x<<endl;
-//cout << "converged "<<res.converged<<" f= "<<res.f<<" iterations "<<achieved.my_maxstep<<" function evaluations "<<achieved.my_maxfunc<<" last df "<<achieved.my_f_thr<<" last dx "<<achieved.my_x_thr<<" |g| "<<achieved.my_g_thr<<endl; //gradient methods
-cout << "converged "<<res.converged<<" f= "<<res.f<<" iterations "<<achieved.my_maxstep<<" function evaluations "<<achieved.my_maxfunc<<" last df "<<achieved.my_f_thr<<" last dx "<<achieved.my_x_thr<<endl;
+cout << "converged "<<res.converged<<" f= "<<res.f<<" iterations "<<achieved.my_maxstep<<" function evaluations "<<achieved.my_maxfunc<<" last df "<<achieved.my_f_thr<<" last dx "<<achieved.my_x_thr<<" |g| "<<achieved.my_g_thr<<endl; //gradient methods
+//cout << "converged "<<res.converged<<" f= "<<res.f<<" iterations "<<achieved.my_maxstep<<" function evaluations "<<achieved.my_maxfunc<<" last df "<<achieved.my_f_thr<<" last dx "<<achieved.my_x_thr<<endl; //derivative-free methods
 double r=problem.funcgrad(x,NULL);
 cout <<r<<endl;
 }

vec.h

@@ -506,6 +506,33 @@ public:
 		return *this;
 	};
 
+	//! apply given function object (functor, lambda) to each element
+	//NOTE: the non-template version above is kept on purpose: for the name of an overloaded function the template argument
+	//cannot be deduced, while the fixed parameter type T (*)(const T &) selects the right overload; for a single function
+	//both match and the non-template one is preferred, so existing code behaves as before
+	template<typename F>
+	NRVec& call_on_me(F _F){
+		NOT_GPU(*this);
+
+		copyonwrite();
+		for(int i=0; i<nn; ++i) v[i] = _F(v[i]);
+		return *this;
+	};
+
+	//! apply given function (or function object) to each element and the corresponding element of rhs: v[i] = F(v[i],rhs[i])
+	template<typename F>
+	NRVec& call_on_me(F _F, const NRVec &rhs){
+#ifdef DEBUG
+		if(nn!=rhs.nn) laerror("incompatible dimensions in call_on_me");
+#endif
+		NOT_GPU(*this);
+		NOT_GPU(rhs);
+
+		copyonwrite();
+		for(int i=0; i<nn; ++i) v[i] = _F(v[i],rhs.v[i]);
+		return *this;
+	};
+
 	void swap(NRVec &rhs) //more efficient swap than via tmp and constructors and operator=
 		{
 		int tmpnn=nn; nn=rhs.nn; rhs.nn=tmpnn;