Softmax에 대한 정리

소프트맥스는 로짓과 확률 사이의 관계를 정의합니다.

\[\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^n e^{z_j}}, \quad z_i = x_i - x_{\text{max}}\]

나이브한 소프트맥스를 파이토치 문법으로 구현하면 다음과 같습니다.

def naive_softmax(x):
	x_max = x.max(dim=1)[0]
  z = x - x_max[:, None]
  numerator = torch.exp(z)
  denominator = numerator.sum(dim=1)
  ret = numerator / denominator[:, None]
  return ret

나이브한 소프트맥스의 내부를 확인하면 3개의 루프로 구성이 됩니다.

연산 중 글로벌한 값(최대, 합) 이 있으며, 해당 값들 사이에 의존성이 존재하기 때문입니다.

\[m_i = \max(m_{i-1}, x_i), \quad \text{for } i = 1, \dots, N\] \[d_i = d_{i-1} + e^{x_i - m_N}, \quad \text{for } i = 1, \dots, N\] \[a_i = \frac{e^{x_i - m_N}}{d_N}, \quad \text{for } i = 1, \dots, N\]

def naive_softmax(x: torch.Tensor) -> torch.Tensor:
  rows, cols = x.shape

  row_max = torch.zeros(rows)
  for i in range(rows):
  	row_max[i] = float("-inf")
    for j in range(cols):
      if x[i][j] > row_max[i]:
        row_max[i] = x[i][j]

  numerator = torch.zeros_like(x)
  denominator = torch.zeros(rows)
  for i in range(rows):
    for j in range(cols):
      exp_val = torch.exp(x[i][j] - row_max[i]
      numerator[i][j] = exp_val
      denominator[i] += exp_val

  ret = torch.zeros_like(x)
  for i in range(rows):
    for j in range(cols):
      ret[i][j] = numerator[i][j] / denominator[i]

  return ret

루프를 통합할 수 있다면 연산 속도를 높일 수 있습니다.

Flash Attention 논문은 온라인 소프트맥스를 제안하였고, 최대값과 합의 루프를 하나로 구성 할 수 있음을 입증하였습니다.

\[m_i = \max(m_{i-1}, x_i), \quad d_i' = d_{i-1}' \cdot e^{m_{i-1} - m_i} + e^{x_i - m_i}, \quad \text{for } i = 1, \dots, N\] \[a_i = \frac{e^{x_i - m_N}}{d_N'}, \quad \text{for } i = 1, \dots, N\]

def online_softmax(x: torch.Tensor) -> torch.Tensor:
  rows, cols = x.shape

  row_max = torch.zeros(rows)
  denominator = torch.zeros(rows)
  for i in range(rows):
  	row_max[i] = float("-inf")
    for j in range(cols):
      if x[i][j] > row_max[i]:
        new_max = x[i][j]
      denominator[i] = denominator[i] * torch.exp(row_max[i] - new_max) + torch.exp(x[i][j] - new_max)
      row_max[i] = new_max
        
  ret = torch.zeros_like(x)
  for i in range(rows):
    for j in range(cols):
      ret[i][j] = torch.exp(x[i][j] - row_max[i]) / denominator[i]

  return ret