Maximum likelihood estimation

We already described how to compute a MLE in the previous tutorial, which is the set of parameters \(\hat{\theta}\) that maximizes the probability of the observed data \(P(D|m,\theta)\). We slightly augment our notation from before to clarify that the MLE is with respect to a specific model \(m\), since we will be considering multiple models (e.g., \(m_1, m_2, \ldots, m_k\))

Let’s start by loading in some simulated from the Tutorial 2 Q-learning model and treat this as if it were participant data for now.

simDF <- readRDS('data/simChoicesQlearning.Rds') #load data that was generated in Tutorial 2
simDF <- subset(simDF, chosen==1) #only include chosen actions
simDF$action <- factor(simDF$action) #factorize choices into discrete options
roundDF <- subset(simDF, round == 1) #plot one illustrative round

#Plot choices and rewards of agent 1 based on social info
ggplot(subset(roundDF, agent==1), aes(x = trial, y = reward, color = action))+
  geom_point(data = subset(roundDF, agent == 1), aes(x = trial, y = -11, color = action))+ #agent choices
  geom_text(x=3, y = -9, label='Agent Choices', color = 'black')+
  geom_point(data = subset(roundDF, agent != 1), aes(x = trial, y = agent+10, color = action))+ #socially observable choices
  geom_text(x=5, y = 16, label='Socially Observable Choices', color = 'black')+ #reward
  stat_summary(fun = mean, geom='line', color = 'black', size = 1)+
  coord_cartesian(ylim=c(-12,16))+
  theme_classic()+
  xlab('Trial') +
  ylab('Reward')+
  scale_color_viridis_d(name = 'option', drop=F)+
  ggtitle('Q-Learning agent')+
  theme(legend.position='right')

## Warning: Using `size` aesthetic for lines was deprecated in ggplot2 3.4.0.
## ℹ Please use `linewidth` instead.
## This warning is displayed once every 8 hours.
## Call `lifecycle::last_lifecycle_warnings()` to see where this warning was
## generated.

Simulated data from a Q-learning agent. The black line shows rewards earned over trials. The colored dots on the bottom indicate the pattern of choices by the agent, while the dots at the top indicate the pattern of socially observable choices made by the other agents.

The plot above shows the data we want to provide to our model, where specifically, we want to describe the sequence of choices made by the agent (dots at the bottom).

In order to describe how the participant/agent makes these choices and learns from the environment, we want to provide the model with the same information that your participants had access to in the task. To describe the mechanisms of individual learning, the model has access to the reward observations for each choice (black line). And since we also want to describe social learning mechanisms, we may additionally provide the model with the socially observable choices made by other agents (dots at the top). Here, this data is generated by a Q-learning model, which simply ignores social information. Yet, other models describe explicit mechanisms where social information influences individual choices (e.g., decision-biasing or value shaping).

We can now compute the MLE for the correct underlying Q-learning model, an incorrect value shaping model, or any model we might want to consider.

MLE using Q-learning model:

#Modify the simulated data

#Q learning likelihood function
Qlikelihood <- function(params, data, Q0=0, k=10){ #We assume that prior value estimates Q0 are fixed to 0 and not estimated as part of the model
  names(params) <- c('alpha', 'beta') #name parameter vec
  nLL <- 0 #Initialize negative log likelihood
  rounds <- unique(data$round)
  trials <- max(data$trial)
  for (r in rounds){ #loop through rounds
    Qvec <- rep(Q0,k) #reset Q-values each new round
    for (t in 1:trials){ #loop through trials
      p <- softmax(params['beta'], Qvec) #compute softmax policy
      trueAction <- subset(data,trial==t & round == r)$action
      negativeloglikelihood <- -log(p[trueAction]) #compute negative log likelihood
      nLL <- nLL + negativeloglikelihood #update running count
      Qvec[trueAction] <- Qvec[trueAction] + params['alpha']*(subset(data, trial==t & round == r)$reward - Qvec[trueAction]) #update q-values
    }
  }
  return(nLL)
}

#let's optimize the parameters for the Q-learning model
init <- c(1,1) #initial guesses
lower <- c(0,0) #lower and upper limits
upper <- c(1,Inf)

MLE_q <- optim(par=init, fn = Qlikelihood,lower = lower, upper=upper, method = 'L-BFGS-B', data = subset(simDF,agent == 1 ))

MLE using VS model:

#VS likelihood function
VSlikelihood <- function(params, data, targetAgent=1, Q0=0, k=10, nAgents = 4){ #We assume that prior value estimates Q0 are fixed to 0 and not estimated as part of the model
  names(params) <- c('alpha', 'beta', 'eta') #name parameter vec
  nLL <- 0 #Initialize negative log likelihood
  rounds <- unique(data$round)
  trials <- max(data$trial)
  for (r in rounds){ #loop through rounds
    Qvec <- rep(Q0,k) #reset Q-values each new round
    socialObs <- rep(NA, nAgents) #social observations
    for (t in 1:trials){ #loop through trials
      socialFreq <- table(factor(socialObs[-targetAgent], levels = 1:k)) + .0001 #Frequency of each socially observed action + a very small number
      Qvec <- Qvec+ (params['eta']*socialFreq) #value bonus
      p <- softmax(params['beta'], Qvec) #compute softmax policy
      trueAction <- subset(data,agent == targetAgent & trial==t & round == r)$action
      negativeloglikelihood <- -log(p[trueAction]) #compute negative log likelihood
      nLL <- nLL + negativeloglikelihood #update running count
      Qvec[trueAction] <- Qvec[trueAction] + params['alpha']*(subset(data, agent == targetAgent & trial==t & round == r)$reward - Qvec[trueAction]) #update q-values
    }
  }
  return(nLL)
}

#let's optimize the parameters for the VS model
init <- c(1,1,1) #initial guesses
lower <- c(0,0, 0) #lower and upper limits
upper <- c(1,Inf, Inf)

MLE_VS <- optim(par=init, fn = VSlikelihood,lower = lower, upper=upper, method = 'L-BFGS-B', data = simDF)

Let’s quickly examine these results. The original parameters of the agent were \(\theta = \{\alpha = .9, \beta = 1\}\). The MLE of the (correct) Q-learning model returned \(\hat{\theta}_{\text{Q-learning}} = \{\hat{\alpha} =0.947, \hat{\beta} = 1.096\}\), while the MLE of the (mismatched) value shaping model returned \(\hat{\theta}_{\text{VS}} = \{\hat{\alpha} =0.947, \hat{\beta} = 1.096, \hat{\eta} = 1.000\}\).

We can describe the MLE fits of both models in terms of their negative log likelihood \(nLL = - \log P(D|\hat{\theta})\), where a lower value (i.e., lower log loss) suggests a better fit. \(nLL_\text{Q-learning} = 109.565902\) vs. \(nLL_\text{VS} = 109.568116\). We see only slight differences here. However, comparing only the nLL fails to account for overfitting. Note that a VS model can also mimic a Q-learning model when \(\eta \rightarrow 0\). Thus, we can compute the AIC or BIC by penalizing for the number of parameters.

Akaike Information Criterion (AIC)

AIC provides an estimate of the predictive accuracy of a model, by adding a penalty for the number of parameters \(k\): \[ \begin{equation} \text{AIC} = -2 \log P(D|\hat{\theta} ) + 2k \end{equation} \] The first term is simply twice the nLL (sometimes known as the deviance), while the second term adds an additional loss proportional to the number of parameters. As with nLL, lower values indicate better results, since it is a measure of the error or loss.

The definition of AIC comes from information theory, and describes the amount of information that is shared between some fixed reality and an imperfect model. A more down-to-earth interpretation is that AIC is an approximation of leave-one-out cross validation error, with an exact equivalence for linear regression and mixed-effects regression models in the limit of infinite data (Stone 1977). However, AIC is also commonly used for models that are not linear regression models and certainly always in settings with finite data. Thus, AIC can more practically be regarded as one of the more lax goodness of fit measures we describe, which is more prone to prefer overly complex models.

Bayesian Information Criterion (BIC)

BIC is an alternative method for penalizing for complexity. While it has a Bayesian interpretation, in practice, computing BIC looks very similar to AIC:

\[ \begin{equation} \text{BIC} = -2 \log P(D|\hat{\theta}) + k \log(n) \end{equation} \]

\(k\) is the number of parameters as before, but we also multiply this by the log of the number of data points \(n\). Theoretically, BIC provides an approximation of the marginal likelihood under specific assumptions about the prior distribution over parameters \(P(\theta)\). However, this interpretation is not without controversy and the assumptions are hardly ever unpacked (see Ch.11 of Lewandowsky and Farrell 2010 for a discussion). Yet in practice, we can simply consider BIC to be a more strict approach to penalizing for complexity compared to AIC, such that an overfit model is less likely to come out on top.

Let’s now run the model fitting procedure for each of the 4 simulated agents, and compare AIC vs. BIC.

fitDF <- data.frame() #initialize data frame
for (targetAgent in c(1,2,3,4)){
  #Q learning
  MLE_q <- optim(par=c(1,1), fn = Qlikelihood,lower =  c(0,0), upper= c(1,Inf), method = 'L-BFGS-B', data = subset(simDF,agent == targetAgent)) #estimate parameters
  qDF <- data.frame(id=targetAgent, nLL = MLE_q$value, alpha = MLE_q$par[1], beta = MLE_q$par[2], eta = NA, model = 'Q-learning', k=2) #put results into a data frame
  #Value shaping
  MLE_VS <- optim(par= c(1,1,1), fn = VSlikelihood,lower = c(0,0, 0), upper= c(1,Inf, Inf), targetAgent =targetAgent, method = 'L-BFGS-B', data = simDF) #estimate parameters
  vsDF <- data.frame(id=targetAgent, nLL = MLE_VS$value, alpha = MLE_VS$par[1], beta = MLE_VS$par[2], eta = MLE_VS$par[3], model = 'Value Shaping', k=3) #put results into a dataframe
  fitDF <- rbind(fitDF, qDF, vsDF) #combine together
}

#compute AIC and BIC
fitDF$AIC <- (2*fitDF$nLL) + (2* fitDF$k)
fitDF$BIC <- (2*fitDF$nLL) + (fitDF$k * log(nrow(subset(simDF, agent == 1)))) 

fitDF <- fitDF %>% pivot_longer(cols = AIC:BIC, values_to = "fit", names_to = 'metric') #pivot to longer data frame for plotting

#Plot data
ggplot(fitDF, aes(x = model, y = fit, color = factor(id)))+
  theme_classic()+
  geom_line(aes(group=id), color = 'black', alpha = 0.2) + #lines connecting participants
  geom_point()+ #dot per participant
  facet_wrap(~metric) +
  labs(x = 'Model', y = 'Goodness of Fit (lower is better)')+
  scale_color_brewer(palette = 'Dark2',  name = 'agent')+
  theme(legend.position='right', strip.background = element_blank())

Model comparison with 4 simulated Q-learning agents. Each dot is single simulated agent, while the connecting line shows the relative difference in either AIC or BIC between models.

While there is substantial variability between simulated agents, we see that the true generating Q-learning model tends to produce better fits (lower values). This is much more the case when using BIC, which penalizes more heavily for complexity. A sample size of only 4 agents is too small to perform statistical analyses about whether the advantage for the Q-learning model is meaningful. Yet, these preliminary results provide early evidence that the correct Q-learning model can be recovered.

Next, let’s explore how cross-validation more directly avoids the perils of overfitting through out-of-sample predictions

Motivations

To begin, let’s address the main motivations behind using Bayesian model estimation.

First of all, Bayesian methods provide a probability distribution over parameters rather than single point estimate via MLE. This provides a more accurate picture of which parameters best describe human behavior and is more robust to issues such as the presence of outliers and bimodality in the data. Secondly, we are able to model hierarchical relationships allowing us to infer the characteristics of a general population from a sample of participants, each of whom may have individual differences in how they learn and make decisions. We can also hierarchically model within-subject manipulations, to describe how experimental manipulations change individual parameter estimates, even if there is a lot of variability in individual baselines. Lastly, the Bayesian approach also provides a natural solution to overfitting via Bayesian Occam’s Razor, since more complex models also have increased flexibility. This increased flexibility is penalized during Bayesian inference, since we evaluate the model across the entire range of parameter values.

Yet, for all these advantages, exact Bayesian inference is typically computational intractable. This motivates why we use Markov Chain Monte Carlo (MCMC) sampling to provide an approximation of an analytically intractable probability distribution using a collection of sampled data points.

MCMC sampling

MCMC combines two basic concepts:

1. A Markov Chain is a sequential process, where each random sample is used as a stepping stone to generate the next sample. By choosing the appropriate transition dynamics, the Markov Chain will gradually match the target distribution we are trying to approximate as it reaches equilibrium.
2. Monte Carlo sampling is a foundational concept in statistics named after a famous casino in Monaco. The basic idea is that a large enough number of randomly drawn samples will approximate the underlying distribution. For instance, flipping a coin enough times will provide an approximation of the underlying probability of heads vs. tails.

These two process are combined together to perform Bayesian inference about the entire posterior distribution over parameters:

\[ P(\theta|D,m) \propto P(D|\theta,m) P(\theta,m) \] Previously, we only used used a maximum likelihood point estimate of the parameters \(\hat{\theta}\). Here, we are interested in estimating a full distribution of parameters \(P(\theta|D,m)\). This is sometimes simply known as the posterior, since it is a function of the prior distribution \(P(\theta)\) (i.e., our initial guess about the parameters before we see the data) and the likelihood of our observed data \(P(D|\theta,m)\) (i.e., how well a given model and set of parameters predicts the data).

In this illustration below, we illustrate the Metropolis-Hastings algorithm for performing MCMC, which is one of the simplest implementations. We first sample an initial set of parameters from the prior distribution \(\theta^0\sim P(\theta)\). Then, we iteratively sample some new parameters from a proposal distribution \(\theta^i \sim P(\theta^i|\theta^{i-1})\). This step can be as simple as adding some Gaussian noise to the previous sample \(\theta^i \sim \mathcal{N}(\theta^{i-1}, \sigma)\). Each new sample is considered as a candidate, which we accept with a probability proportional to whether it provides a better account of the data (i.e., whether the likelihood of the new sample \(P(D|\theta^{i},m)\) is higher than the old sample \(P(D|\theta^{i-1},m)\)). Thus, we are more likely to accept new samples that have a higher likelihood (filled in circles) and less likely to accept samples with lower likelihood (empty circles). We omit some details for simplicity, but in the end, we have a collection of samples (dots) that approximates the posterior distribution (histogram).

An overview of MCMC sampling, adapted from Lee, Sung, and Choi (2015)

Introduction to STAN

To give you a concrete example of how to build and fit an MCMC model, we use STAN to describe a Q-learning model, we use cmdstanr as an interface to R to fit the model.

The code block below describes our Q learning model, where we are fitting group-level estimates about the \(\alpha\) and \(\beta\) parameters. For more detailed tutorials on STAN, please consider the following resources: https://mc-stan.org/users/documentation/tutorials.

// Q_learning model
data {
  int<lower = 0> All;  // number of observations
  int<lower = 0> Nsub;  // number of subjects
  int<lower = 0> Ncue;  // number of options
  int<lower = 0> Ntrial;  // number of trials per subject
  int<lower = 0> Nround;  // number of rounds per subject
  array[All] int<lower = 0> sub; // subject index
  array[All] int<lower = 0> Y; // index of chosen option: 0 => missing ()
  array[All] int<lower = 0> trial; // trial number
  array[All] int<lower = 0> thisRound; // round number
  array[All] real payoff; // payoff
}

parameters {
  real mu_alpha;
  real mu_beta;
  real<lower=0> s_alpha;
  real<lower=0> s_beta;
  
  // Varying effects clustered on individual, defined as deviations from grand mean
  // We do the non-centered version with Cholesky factorization
  matrix[2,Nsub] z_ID;               //Matrix of uncorrelated z - values
  vector<lower=0>[2] sigma_ID;       //SD of parameters among individuals
  cholesky_factor_corr[2] L_Rho_ID;    // This is the Cholesky factor: if you multiply this matrix and it's transpose you get correlation matrix

}

transformed parameters {
  matrix[Nsub,2] v_ID; // Matrix of varying effects for each individual for alpha, beta
  array[Nsub, Nround] matrix[Ncue, Ntrial] Q;  // Q-values for each target
  array[Nsub, Nround] matrix[Ncue, Ntrial] q; // softmax choice (log-)probability transformed from Q values
  vector[Nsub] logit_alpha;  // learning rate -- logit scale
  vector[Nsub] log_beta;  // softmax parameter -- log scale
  vector<lower = 0, upper = 1>[Nsub] alpha;  // learning rate
  vector<lower = 0>[Nsub] beta;  // inverse temperature (softmax)
  
  v_ID = ( diag_pre_multiply( sigma_ID , L_Rho_ID ) * z_ID )'; // the ' on the end of this line is a transposing function
  for(i in 1:Nsub) {
    logit_alpha[i] = mu_alpha + s_alpha * v_ID[i, 1];
    log_beta[i] = mu_beta + s_beta * v_ID[i, 2];
    beta[i] = exp(log_beta[i]);
    alpha[i] = 1/(1+exp(-logit_alpha[i]));
  }

  for(idx in 1:All) {
    // SETTING INITIAL Q-VALUE
    if(trial[idx] == 1)
      Q[sub[idx], thisRound[idx]][1:Ncue, trial[idx]] = rep_vector(0, Ncue);
    // CHOICE PROBABILITY
    q[sub[idx], thisRound[idx]][1:Ncue, trial[idx]] = 
      log_softmax( 
        Q[sub[idx], thisRound[idx]][, trial[idx]] * beta[sub[idx]] 
        );

    // Q-VALUE UPDATE
    if(trial[idx] < Ntrial) {
      Q[sub[idx], thisRound[idx]][, (trial[idx]+1)] = Q[sub[idx], thisRound[idx]][, trial[idx]]; // copy Q(t) to Q(t+1)
      // update Q on the chosen option
      if(Y[idx]>0) { // if Y == 0, it means no choise was made (missed trials)
        // Updating chosen option's Q-value
        Q[sub[idx], thisRound[idx]][Y[idx], (trial[idx]+1)] =
              (1-alpha[sub[idx]])*Q[sub[idx], thisRound[idx]][Y[idx], trial[idx]] + alpha[sub[idx]]*payoff[idx];
      }
    }
   
  }
}

model {
  // position and scales of the learning parameters
  mu_alpha ~ std_normal(); // prior
  mu_beta ~ normal(0, 5); // prior
  s_alpha ~ exponential(1); //std_normal(); // prior
  s_beta ~ exponential(1); //std_normal(); // prior
  
  // Varying effects
  to_vector(z_ID) ~ std_normal();
  sigma_ID ~ exponential(1);
  L_Rho_ID ~ lkj_corr_cholesky(2); // A weakly informative prior that accounts for the possibility that alpha and beta are correlated

  for(idx in 1:All) {
    if( Y[idx] > 0 ) {
      target += categorical_lpmf( Y[idx] | exp(q[sub[idx], thisRound[idx]][,trial[idx]]) );
    }
  }
}

generated quantities {
  vector[Nsub] log_lik;
  vector[All] reward_pred_err;
  matrix[2,2] Rho_ID; // the correlation matrix
  
  Rho_ID = multiply_lower_tri_self_transpose(L_Rho_ID);
  
  log_lik = rep_vector(0, Nsub); // initial values for log_lik
  reward_pred_err = rep_vector(-100, All); // initial values for log_lik
  for(idx in 1:All) {
    if( Y[idx] > 0 ) {
      reward_pred_err[idx] = payoff[idx] - Q[sub[idx], thisRound[idx]][Y[idx], trial[idx]];
      log_lik[sub[idx]] = log_lik[sub[idx]] + q[sub[idx], thisRound[idx]][Y[idx],trial[idx]];
    }
  }
}

Now let’s run the MCMC estimation on the MCMC estimation on the data.

simulated_data <- readRDS('data/simChoicesQlearning.Rds') # it should be replaced 
simulated_data_for_fit <- simulated_data %>% dplyr::filter(chosen == 1)

# Basic Q-learning model
stanmodel_Q_learning <- cmdstan_model('codeSnippets/Q_learning.stan') #compile model

#MCMC parameters, Note that we run a small number of samples since this is just an illustrative example
chains <- 4
parallel_chains <- 4
thin <- 1
iter_warmup <- 500 
iter_sampling <- 1000  

# preparing data fed to Stan
sim_dt <- list(
  All = nrow(simulated_data_for_fit)
  , Nsub = length(table(simulated_data_for_fit$agent)) # number of agents
  , Nround = length(table(simulated_data_for_fit$round)) # number of rounds
  , Ncue = 10 # number of options
  , Ntrial = 25 # number of trials
  , Ncon = 1 # number of global distributions 
  , sub = simulated_data_for_fit$agent
  , thisRound = simulated_data_for_fit$round
  , Y = simulated_data_for_fit$action # Y = 1, 2, 3, ,,, Ncue
  , trial = simulated_data_for_fit$trial
  , payoff = simulated_data_for_fit$reward # reward
)

# -- MCMC sampling --
fit <- stanmodel_Q_learning$sample(data = sim_dt, seed = 777, chains = chains, parallel_chains = parallel_chains, thin = thin, iter_warmup = iter_warmup, iter_sampling = iter_sampling)
fit$save_object(file = 'data/stanFit.stanmodel') #save to data 

Now let’s take a look at the results. Note that since both \(\alpha\) and \(\beta\) are bounded variables, the STAN code first transformed them into an unbounded range. We first need to transform them back using the inverse transformation (logit transform for \(\alpha\) and exponentiation for \(\beta\)) in order to examine them.

fit <- readRDS('data/stanFit.stanmodel') #read object from data, Note that this was removed from git because the model file is too large
#Extract samples
draws_df <- fit$draws(variables = c('mu_alpha', 'mu_beta'), format = "df") #extract the MCMC samples of the group-level alpha and beta estimates, turning it into a dataframe
#Transform parameters back to their original values
draws_df$alpha <- logisticFunc(draws_df$mu_alpha) #use an logit transform
draws_df$beta <- exp(draws_df$mu_beta) #use an exponentiation

#Pivot to a longer dataframe for plotting
plottingDF <- draws_df %>% pivot_longer(cols = c('alpha', 'beta'), names_to = 'parameter', values_to='estimate')

trueParams <- data.frame(parameter = c('alpha', 'beta'), estimate = c(.9,1))

ggplot(plottingDF, aes(x = estimate,  fill = parameter))+
  geom_histogram(aes(y = ..density..), color = 'black', bins = 30)+
  geom_vline(data = trueParams, aes(xintercept = estimate), color = 'black', linetype = 'dashed')+
  facet_grid(~parameter)+
  theme_classic()+
  xlab('Posterior Estimate')+
  ylab('Density')+
  scale_fill_brewer(palette = 'Dark2')+
  theme(legend.position = 'none', strip.background = element_blank())

The dashed lines indicate the true parameter values that generated the data. So it seems like our MCMC model did a good job of recovering the true parameter values.

Bayesian model comparison

For actually performing Bayesian model comparison, modern methods may seem quite daunting and quickly launch you off of a complexity curve. But in general, the intuitions are quite simple. We simply use a form of cross-validation and out-of-sample predictions to test goodness of fit. But rather than using the MLE as before, we use the full posterior distribution over parameters.

We include the code block below to give you an idea about how to access goodness of fit measures for a Bayesian STAN model, which are already computed when fitting the model. One of the easiest to interpret outputs is looic, which like AIC or BIC describes goodness of fit with lower values being better. However, we recommend referring to Vehtari, Gelman, and Gabry (2017) and McElreath (2020) for more details about how to interpret these quantities.

fit$loo()

## 
## Computed from 4000 by 4 log-likelihood matrix
## 
##          Estimate   SE
## elpd_loo   -385.0 31.4
## p_loo         1.5  0.3
## looic       770.0 62.8
## ------
## Monte Carlo SE of elpd_loo is NA.
## 
## Pareto k diagnostic values:
##                          Count Pct.    Min. n_eff
## (-Inf, 0.5]   (good)     0      0.0%   <NA>      
##  (0.5, 0.7]   (ok)       3     75.0%   481       
##    (0.7, 1]   (bad)      1     25.0%   120       
##    (1, Inf)   (very bad) 0      0.0%   <NA>      
## See help('pareto-k-diagnostic') for details.

While the Bayesian framework is the most robust, it is also the most complex approach to computational modeling. As we’ve seen above, the most complex method isn’t always the best one. Rather, we recommend that you use whichever method is right for the job. Ultimately, the correct choice of methods is whatever you can justify to yourself and the scientific community.