diff --git a/R/milestone.R b/R/milestone.R
index 3cc3d91e..d79dfa9a 100644
--- a/R/milestone.R
+++ b/R/milestone.R
@@ -33,6 +33,8 @@
 #' - `surv_diff` - Survival difference between the experimental and control arm.
 #' - `std_err0` - Standard error of the control arm.
 #' - `std_err1` - Standard error of the experimental arm.
+#' @references Klein, John P., et al. "Analyzing survival curves at a fixed point in time."
+#' Statistics in medicine 26.24 (2007): 4505-4519.
 #'
 #' @export
 #'
@@ -46,27 +48,32 @@ milestone <- function(data, ms_time) {
   fit_res <- summary(fit, time = ms_time, extend = TRUE)
 
   # Survival difference
-  diff_survival <- fit_res$surv[2] - fit_res$surv[1]
+  surv_col <- fit_res$surv[1]
+  surv_exp <- fit_res$surv[2]
+  diff_survival <- surv_exp - surv_col
 
   # Indicator whether the std is NA or not
-  na_col <- is.na(fit_res$std.err[1])
-  na_exp <- is.na(fit_res$std.err[2])
+  var_col <- fit_res$std.err[1]^2
+  var_exp <- fit_res$std.err[2]^2
+
+  na_col <- is.na(var_col)
+  na_exp <- is.na(var_exp)
+
+  sigma2_col <- var_col / (surv_col^2)
+  sigma2_exp <- var_exp / (surv_exp^2)
 
   # Calculate the test statistics
   if (na_col + na_exp == 2) {
     z <- -Inf
   } else {
-    term1 <- if (na_col) 0 else fit_res$std.err[1]
-    term2 <- if (na_exp) 0 else fit_res$std.err[2]
-    var_survival <- term1^2 + term2^2
-    z <- diff_survival / sqrt(var_survival)
+    z <- (log(-log(surv_exp)) - log(-log(surv_col)))^2 / (sigma2_exp / (log(surv_exp))^2 + sigma2_col / (log(surv_col))^2)
   }
 
   ans <- data.frame(
     method = "milestone", z = z, ms_time = ms_time,
-    surv0 = fit_res$surv[1], surv1 = fit_res$surv[2],
+    surv_col = surv_col, surv_exp = surv_exp,
     surv_diff = diff_survival,
-    std_err0 = fit_res$std.err[1], std_err1 = fit_res$std.err[2]
+    std_err_col = fit_res$std.err[1], std_err_exp = fit_res$std.err[2]
   )
   return(ans)
 }
diff --git a/man/milestone.Rd b/man/milestone.Rd
index 0f039c30..ee7d3ee5 100644
--- a/man/milestone.Rd
+++ b/man/milestone.Rd
@@ -37,3 +37,7 @@ sim_pw_surv(n = 200) |>
   cut_data_by_event(150) |>
   milestone(10)
 }
+\references{
+Klein, John P., et al. "Analyzing survival curves at a fixed point in time."
+Statistics in medicine 26.24 (2007): 4505-4519.
+}
diff --git a/man/sim_gs_n.Rd b/man/sim_gs_n.Rd
index ee612317..ac7d80c1 100644
--- a/man/sim_gs_n.Rd
+++ b/man/sim_gs_n.Rd
@@ -83,9 +83,9 @@ fail_rate <- define_fail_rate(
 )
 
 # Other related parameters
-alpha <- 0.025 # type I error
-beta <- 0.1 # type II error
-ratio <- 1 # randomization ratio (exp:col)
+alpha <- 0.025 # Type I error
+beta <- 0.1 # Type II error
+ratio <- 1 # Randomization ratio (experimental:control)
 
 # Define cuttings of 2 IAs and 1 FA
 # IA1