IRT Linking and Gradient Asymmetry: Diagnostic Guide

library(mixedsubjectsirt)
library(ggplot2)

# Apply (A, B) and return a standardised data frame of item parameters.
# Guards against degenerate linking constants (Inf, NaN, non-positive A) that
# can arise from unusual mirt fits on different platforms.
apply_link <- function(source, A, B, slope_lower = 1e-4) {
  if (!is.finite(A) || A <= 0 || !is.finite(B)) {
    # Degenerate constants: fall back to identity transform
    A <- 1
    B <- 0
  }
  pars <- data.frame(
    item = source$item,
    a    = pmax(slope_lower, source$a / A),
    b    = A * source$b + B,
    stringsAsFactors = FALSE
  )
  # Guard b and d against any residual non-finite values
  pars$b <- ifelse(is.finite(pars$b), pars$b, 0)
  pars$d <- -pars$a * pars$b
  list(pars = pars, A = A, B = B)
}

link_mean_mean <- function(source, target) {
  A <- mean(source$a) / mean(target$a)
  B <- mean(target$b) - A * mean(source$b)
  apply_link(source, A, B)
}

link_mean_sigma <- function(source, target) {
  sd_src <- sd(source$b)
  # If source difficulties have no variance (degenerate mirt fit), fall back
  # to mean-mean linking which does not depend on sd(b).
  if (!is.finite(sd_src) || sd_src < 1e-6) {
    return(link_mean_mean(source, target))
  }
  A <- sd(target$b) / sd_src
  B <- mean(target$b) - A * mean(source$b)
  apply_link(source, A, B)
}

# Stocking-Lord with bounded optimization to prevent item-level overcorrection.
# A is constrained to [0.4, 2.5]; large-variance items can otherwise receive
# inflated linked discriminations that destabilize the downstream M-step.
link_stocking_lord <- function(source, target,
                                theta_grid = seq(-4, 4, by = 0.05),
                                A_bounds   = c(0.4, 2.5),
                                B_bounds   = c(-4,  4)) {
  w <- dnorm(theta_grid) / sum(dnorm(theta_grid))

  tcc <- function(pars, theta) {
    eta <- outer(theta, pars$a, `*`) +
      matrix(pars$d, nrow = length(theta), ncol = nrow(pars), byrow = TRUE)
    rowSums(plogis(eta))
  }

  tcc_target <- tcc(target, theta_grid)

  criterion <- function(params) {
    A <- exp(params[1])
    B <- params[2]
    sum(w * (tcc_target - tcc(apply_link(source, A, B)$pars, theta_grid))^2)
  }

  mm   <- link_mean_mean(source, target)
  init <- c(log(mm$A), mm$B)

  # L-BFGS-B on log(A) keeps A > 0 and enforces the stated bounds
  opt <- optim(
    par     = init,
    fn      = criterion,
    method  = "L-BFGS-B",
    lower   = c(log(A_bounds[1]), B_bounds[1]),
    upper   = c(log(A_bounds[2]), B_bounds[2]),
    control = list(factr = 1e-14, maxit = 1000)
  )

  apply_link(source, exp(opt$par[1]), opt$par[2])
}

rmse <- function(x, y) sqrt(mean((x - y)^2))

n_human     <- 400
n_generated <- 1200
n_items     <- 8

true_pars <- data.frame(
  item = paste0("Item", seq_len(n_items)),
  a    = seq(0.8, 1.6, length.out = n_items),
  d    = seq(-1.1, 1.1, length.out = n_items)
)
true_pars$b <- -true_pars$d / true_pars$a

# Human responses
theta_human <- rnorm(n_human)
observed    <- simulate_2pl(theta_human, true_pars)

# LLM: ~10% attenuated discrimination per item, +0.25 logit intercept shift
llm_pars_true <- true_pars
llm_pars_true$a <- pmax(0.4, 0.9 * true_pars$a + rnorm(n_items, 0, 0.05))
llm_pars_true$d <- true_pars$d + 0.25 + rnorm(n_items, 0, 0.15)
llm_pars_true$b <- -llm_pars_true$d / llm_pars_true$a

# Paired predictions and two generated datasets
predicted         <- simulate_2pl(theta_human, llm_pars_true)
generated_matched <- simulate_2pl(rnorm(n_generated),                    llm_pars_true)
generated_shifted <- simulate_2pl(rnorm(n_generated, mean = 0.2, sd = 0.9), llm_pars_true)

human_pars       <- fit_2pl(observed,         technical = list(NCYCLES = 500))$pars
llm_raw_matched  <- fit_2pl(generated_matched, technical = list(NCYCLES = 500))$pars
llm_raw_shifted  <- fit_2pl(generated_shifted, technical = list(NCYCLES = 500))$pars

methods <- c("mean_mean", "mean_sigma", "stocking_lord")

all_links <- list(
  matched = list(
    mean_mean     = link_mean_mean    (llm_raw_matched, human_pars),
    mean_sigma    = link_mean_sigma   (llm_raw_matched, human_pars),
    stocking_lord = link_stocking_lord(llm_raw_matched, human_pars)
  ),
  shifted = list(
    mean_mean     = link_mean_mean    (llm_raw_shifted, human_pars),
    mean_sigma    = link_mean_sigma   (llm_raw_shifted, human_pars),
    stocking_lord = link_stocking_lord(llm_raw_shifted, human_pars)
  )
)

# Linking constants
const_tab <- do.call(rbind, lapply(c("matched", "shifted"), function(case) {
  do.call(rbind, lapply(methods, function(m) {
    data.frame(case = case, method = m,
               A = round(all_links[[case]][[m]]$A, 4),
               B = round(all_links[[case]][[m]]$B, 4))
  }))
}))
knitr::kable(const_tab, row.names = FALSE,
             caption = "Linking constants (A, B)")

n_quad <- 11
quad   <- make_quadrature(n_quad)

q_obs  <- mixedsubjectsirt:::build_quadrature_summary(observed,  human_pars, quad)
q_pred <- mixedsubjectsirt:::build_quadrature_summary(predicted, human_pars, quad,
                                                       weights = q_obs$weights)

gradient_analysis <- function(gen_data, linked_pars, label) {
  q_gen <- mixedsubjectsirt:::build_quadrature_summary(gen_data, linked_pars, quad)
  g_obs  <- mixedsubjectsirt:::gradient_expected_counts(q_obs$counts,  human_pars)
  g_gen  <- mixedsubjectsirt:::gradient_expected_counts(q_gen$counts,  human_pars)
  g_pred <- mixedsubjectsirt:::gradient_expected_counts(q_pred$counts, human_pars)

  data.frame(
    config    = label,
    item      = human_pars$item,
    grad_a_combined_0.5 = round(g_obs + 0.5 * (g_gen - g_pred), 4)[seq_len(n_items)],
    g_obs  = round(g_obs[seq_len(n_items)], 4),
    g_gen  = round(g_gen[seq_len(n_items)], 4),
    g_pred = round(g_pred[seq_len(n_items)], 4)
  )
}

grad_rows <- rbind(
  gradient_analysis(generated_matched, human_pars,                  "unlinked"),
  gradient_analysis(generated_matched, all_links$matched$mean_mean$pars,     "mean_mean"),
  gradient_analysis(generated_matched, all_links$matched$mean_sigma$pars,    "mean_sigma"),
  gradient_analysis(generated_matched, all_links$matched$stocking_lord$pars, "stocking_lord")
)

# Show combined gradient for Items 5 and 8 — the problem items
grad_items <- grad_rows[grad_rows$item %in% c("Item5", "Item8"),
                         c("config","item","g_obs","g_gen","g_pred",
                           "grad_a_combined_0.5")]
knitr::kable(grad_items, row.names = FALSE,
  col.names = c("Config","Item","∇L_obs","∇L_gen","∇L_pred","Combined (λ=0.5)"),
  caption   = paste("Gradient of discrimination a for the two problematic items at",
                    "starting parameters. Negative combined gradient pushes a upward."))

lambda_grid <- c(0, 0.02, 0.05, 0.10, 0.15, 0.20, 0.30, 0.50)

sweep_lambda <- function(gen_data, linked_pars) {
  # Validate linked_pars: if any parameter is non-finite (can happen when
  # sd(b) ~ 0 on some platforms and apply_link fallback wasn't triggered),
  # substitute human_pars so counts are always valid.
  if (!all(is.finite(c(linked_pars$a, linked_pars$d)))) {
    linked_pars <- human_pars
  }
  q_gen <- mixedsubjectsirt:::build_quadrature_summary(gen_data, linked_pars, quad)
  lapply(lambda_grid, function(lam) {
    fit <- tryCatch(
      mixedsubjectsirt:::fit_from_counts(
        q_obs$counts, q_pred$counts, q_gen$counts,
        initial_pars = human_pars, lambda = lam, control = list(maxit = 500)),
      error = function(e) list(
        item_pars = data.frame(item = human_pars$item,
                               a = rep(NA_real_, nrow(human_pars)),
                               b = rep(NA_real_, nrow(human_pars)),
                               d = rep(NA_real_, nrow(human_pars))),
        value = NA_real_, convergence = 99L)
    )
    data.frame(
      lambda  = lam,
      rmse_a  = if (anyNA(fit$item_pars$a)) NA_real_ else rmse(fit$item_pars$a, true_pars$a),
      rmse_d  = if (anyNA(fit$item_pars$d)) NA_real_ else rmse(fit$item_pars$d, true_pars$d),
      max_a   = if (anyNA(fit$item_pars$a)) NA_real_ else max(fit$item_pars$a),
      conv    = fit$convergence
    )
  })
}

sweep_results <- do.call(rbind, lapply(c("matched","shifted"), function(case) {
  gen_data <- if (case == "matched") generated_matched else generated_shifted
  do.call(rbind, lapply(c("unlinked", methods), function(m) {
    lp <- if (m == "unlinked") human_pars else all_links[[case]][[m]]$pars
    rows <- do.call(rbind, sweep_lambda(gen_data, lp))
    rows$method <- m
    rows$case   <- case
    rows
  }))
}))

sweep_results$method_f <- factor(sweep_results$method,
  levels = c("unlinked","mean_mean","mean_sigma","stocking_lord"),
  labels = c("Unlinked","Mean-mean","Mean-sigma","Stocking-Lord"))

# Recommended workflow: mean-sigma linking for the generated E-step,
# human parameters for observed and predicted E-steps, then power-tune lambda.
ms_linked_pars <- all_links$matched$mean_sigma$pars

# Build the three quadrature summaries with the correct parameterization for each.
# q_obs and q_pred use human parameters; q_gen uses the linked LLM parameters.
q_gen_linked <- mixedsubjectsirt:::build_quadrature_summary(
  generated_matched, ms_linked_pars, quad)

risk_tab <- do.call(rbind, lapply(c(0, 0.05, 0.10, 0.20, 0.30, 0.50), function(lam) {
  fit_counts <- tryCatch(
    mixedsubjectsirt:::fit_from_counts(
      q_obs$counts, q_pred$counts, q_gen_linked$counts,
      initial_pars = human_pars, lambda = lam,
      slope_upper = 4,                # prevents divergence at large lambda
      control = list(maxit = 500)),
    error = function(e) list(item_pars = data.frame(a = rep(NA_real_, n_items),
                                                     d = rep(NA_real_, n_items)))
  )

  # fit_mixed_subjects is used for vcov; ms_linked_pars for all three E-steps
  # is a proxy — the risk trend is the quantity of interest.
  fit_for_vcov <- tryCatch(
    fit_mixed_subjects(
      observed = observed, predicted = predicted, generated = generated_matched,
      lambda = lam, initial_pars = ms_linked_pars,
      n_quad = n_quad, slope_upper = 4, control = list(maxit = 200)),
    error = function(e) NULL
  )

  rmse_a <- if (anyNA(fit_counts$item_pars$a)) NA_real_ else
    round(rmse(fit_counts$item_pars$a, true_pars$a), 4)

  if (is.null(fit_for_vcov)) {
    return(data.frame(lambda = lam, rmse_a = rmse_a, mean_param_var = NA_real_))
  }

  tryCatch({
    Sigma <- vcov_mixed_subjects(fit_for_vcov)
    risk  <- ability_risk(observed, fit_for_vcov, vcov = Sigma)
    data.frame(lambda = lam, rmse_a = rmse_a,
               mean_param_var = round(risk$summary$mean_param_var, 6))
  }, error = function(e) {
    data.frame(lambda = lam, rmse_a = rmse_a, mean_param_var = NA_real_)
  })
}))

knitr::kable(risk_tab, row.names = FALSE,
  col.names = c("λ", "RMSE(a)", "Mean ability-score risk"),
  caption   = "Ability-score risk and parameter recovery — mean-sigma linking, matched case")

# Use human_pars (fitted human 2PL MLE) as evaluation point
n_generated_val <- n_generated  # 1200
upper_bound     <- n_generated / (n_human + n_generated)  # N/(n+N)

# Direct comparison: frozen EC with slope cap vs MML without cap
fit_frozen <- fit_mixed_subjects(
  observed = observed, predicted = predicted, generated = generated_matched,
  lambda = 0.2, initial_pars = human_pars,
  n_quad = n_quad, slope_upper = 4, control = list(maxit = 200))

fit_mml <- fit_mixed_subjects_mml(
  observed = observed, predicted = predicted, generated = generated_matched,
  lambda = 0.2, initial_pars = human_pars,
  n_quad = n_quad, control = list(maxit = 200))

comp <- data.frame(
  item   = human_pars$item,
  true_a = true_pars$a,
  frozen_a = round(fit_frozen$item_pars$a, 3),
  mml_a    = round(fit_mml$item_pars$a,    3)
)
knitr::kable(comp, row.names = FALSE,
  caption = "Item discrimination: true vs. frozen-EC (slope_upper=4) vs. MML at lambda=0.2")

# Lambda sweep: MML without slope cap
mml_sweep <- do.call(rbind, lapply(lambda_grid, function(lam) {
  fit <- tryCatch(
    fit_mixed_subjects_mml(
      observed = observed, predicted = predicted, generated = generated_matched,
      lambda = lam, initial_pars = human_pars,
      n_quad = n_quad, control = list(maxit = 300)),
    error = function(e) NULL
  )
  if (is.null(fit)) return(data.frame(lambda=lam, rmse_a=NA_real_, max_a=NA_real_))
  data.frame(
    lambda = lam,
    rmse_a = round(rmse(fit$item_pars$a, true_pars$a), 4),
    max_a  = round(max(fit$item_pars$a), 3)
  )
}))

knitr::kable(mml_sweep, row.names = FALSE,
  col.names = c("λ", "RMSE(a)", "max(a)"),
  caption = "MML parameter recovery across lambda — no slope_upper needed")