RでTransformerのAttention機構:Softmax

softmax <- function(x) {
  exp(x) / sum(exp(x))
}

scores <- c(2.0, 1.0, 0.5, -1.0)
probs <- softmax(scores)

cat("=== 入力 (スコア) ===\n")
print(scores)
cat("\n=== Softmax適用後 (確率) ===\n")
print(round(probs, 4))
cat(sprintf("\n合計 (必ず1になる): %.6f\n", sum(probs)))

# 単純な正規化: 各値を合計で割るだけ
plain_normalize <- function(x) {
  x / sum(x)
}

cat("=== 単純な正規化 vs Softmax の比較 ===\n")

# 負の値を含む例
scores_with_negative <- c(3.0, -2.0, 1.0)

cat("入力:", scores_with_negative, "\n")
cat(
  "単純な正規化:", round(plain_normalize(scores_with_negative), 4),
  " <- 負の「確率」が出てしまう\n"
)
cat(
  "Softmax:    ", round(softmax(scores_with_negative), 4),
  " <- 全て正の値で合計1\n"
)

# 値の差の強調具合を比較する例
#
# 以下の2つを意図的に対比させています:
#   scores_close: 差が小さく(0.1刻み)、値自体も小さい
#     -> e^x はxが小さい範囲ではテイラー展開 e^x ≈ 1+x により
#        ほぼ線形になるため、単純な正規化とSoftmaxの結果は
#        ほとんど一致する (指数関数の非線形性がまだ効かない)
#   scores_far: 差が大きい(4刻み)
#     -> e^x の非線形性がはっきり効き、Softmaxは単純な正規化
#        よりも明確に「勝者(最大値)」に確率を強く寄せる
#
# つまりSoftmaxが単純な正規化と違う結果を生むのは、あくまで
# 「値の差が大きい」場合に限られます

scores_close <- c(1.0, 1.1, 1.2) # わずかな差
scores_far <- c(1.0, 5.0, 9.0) # 大きな差

cat("\n差が小さい入力:", scores_close, "\n")
cat("  単純な正規化:", round(plain_normalize(scores_close), 4), "\n")
cat("  Softmax:    ", round(softmax(scores_close), 4), "\n")

cat("\n差が大きい入力:", scores_far, "\n")
cat("  単純な正規化:", round(plain_normalize(scores_far), 4), "\n")
cat("  Softmax:    ", round(softmax(scores_far), 4), "\n")

x <- c(2.0, 1.0, 0.5, -1.0)
x_shifted <- x + 100 # 全体に定数100を加える

cat("=== シフト不変性の確認 ===\n")
cat("元の入力の softmax:        ", round(softmax(x), 6), "\n")
cat("+100 した入力の softmax:   ", round(softmax(x_shifted), 6), "\n")

x_large <- c(1000, 999, 998)

cat("=== オーバーフロー問題の実演 ===\n")

naive_softmax <- function(x) {
  exp(x) / sum(exp(x)) # 対策なし
}

cat("対策なしの softmax:\n")
result_naive <- naive_softmax(x_large)
print(result_naive) # NaN になる (Inf / Inf = NaN)

# 数値的に安定な softmax: 計算前に最大値を引く
stable_softmax <- function(x) {
  x_shifted <- x - max(x) # 最大値を0にする (シフト不変性より結果は不変)
  exp(x_shifted) / sum(exp(x_shifted))
}

cat("\n数値的に安定な softmax (最大値を引いてから計算):\n")
result_stable <- stable_softmax(x_large)
print(round(result_stable, 4))

softmax_with_temperature <- function(x, T = 1) {
  softmax(x / T)
}

base_scores <- c(3.0, 1.0, 0.2, -1.0)
temperatures <- c(0.3, 1.0, 3.0, 10.0)

temp_results <- lapply(temperatures, function(T) {
  data.frame(
    category = factor(seq_along(base_scores)),
    probability = softmax_with_temperature(base_scores, T),
    temperature = paste0("T = ", T)
  )
})
temp_df <- do.call(rbind, temp_results)
temp_df$temperature <- factor(temp_df$temperature, levels = paste0("T = ", temperatures))

cat("=== 温度による分布の変化 ===\n")
for (T in temperatures) {
  cat(sprintf(
    "T = %-5s : %s\n", T,
    paste(round(softmax_with_temperature(base_scores, T), 3), collapse = "  ")
  ))
}