online softmax - beyondnlp/nlp GitHub Wiki

https://github.com/karpathy/llama2.c/blob/b3c4b6c3c4bbff42e5211293280307019368ccb5/run.c#L197

위 코드는 karpathy가 만든 llama2.c 코드 중 softmax를 구현부입니다.

void softmax(float* x, int size) {
    // find max value (for numerical stability)
    float max_val = x[0];
    for (int i = 1; i < size; i++) {
        if (x[i] > max_val) {
            max_val = x[i];
        }
    }
    // exp and sum
    float sum = 0.0f;
    for (int i = 0; i < size; i++) {
        x[i] = expf(x[i] - max_val);
        sum += x[i];
    }
    // normalize
    for (int i = 0; i < size; i++) {
        x[i] /= sum;
    }
}

softmax는 모든 숫자를 확률 벡터로 변환하기 위해 사용됩니다. 입력 벡터를 모두 지수로 계산을 합니다. 그런 다음 전체합으로 각 벡터를 나누는 것으로 값을 계산할수 있습니다.

def EXP( b ):
    alpha=2.7182818284590452353602874713526624977 # =e^1
    return alpha ** b

def softmax_normal(x):
    exp_x=[]
    for nx in x:
        e1=EXP( nx )
        exp_x.append( e1 )
    sum_exp_x = 1e-10
    for i in exp_x: sum_exp_x = sum_exp_x + i
    y=[]
    for es in exp_x: y.append(es/sum_exp_x)
    return y

softmax_normal()함수는 가장 기본이 되는 구현 방법입니다.

아래는 지수함수의 특성으로 인해 overflow나 underflow를 막기 위해

분모가 되는 sum_exp_x에 1e-10와 같은 작은 값을 더하고 overflow를 막기 위해 벡터중 가장큰 값을 나누게 됩니다. https://upload.wikimedia.org/wikipedia/commons/thumb/c/c6/Exp.svg/2560px-Exp.svg.png

def softmax_max(x):

    max_one = -math.inf
    for nx in x:
        if nx > max_one : max_one = nx

    exp_x=[]
    for nx in x:
        new_nx = nx-max_one
        e1=EXP( new_nx )
        exp_x.append( e1 )
    sum_exp_x = 1e-10
    for i in exp_x: sum_exp_x = sum_exp_x + i
    print('[max]\tsum_exp_x', sum_exp_x)
    y=[]
    for es in exp_x: y.append(es/sum_exp_x)
    return y

위에 설명한 softmax_max의 경우 max value를 찾기 위해 벡터를 한번 순회하게 되는데 계산 속도를 빠르게 하기 위해 아래와 같이 for loop를 하나 제거한 버전이 flash attention과 같은 속도 최적화에 사용이 됩니다.

def softmax_online(x):
    N = len(x)

    # Initialize variables
    m = [-1] * N
    d = [0] * N
    y = [0] * N

    old_max = -math.inf
    new_max = -math.inf
    exp_sum_x = 0;
    for i in range(N):
        new_max = get_max( old_max, x[i] )
        a=EXP( old_max - new_max )
        b=EXP( x[i] - new_max )
        exp_sum_x = exp_sum_x * a + b
        old_max = new_max

    for i in range(N): y[i] = EXP( x[i] - old_max ) / exp_sum_x;

    return y

• old_max - new_max는 0이상의 값을 가진다.( new_max는 old_max보다 항상 크기 때문에 )
• 즉 a는 exp(0이상의값)=1 이상의 값이 된다.
• a가 1이면 이전 값에 b의 값만 exp_sum_x에 더해지는 상황
• exp(3) * exp(2) = exp(5)

[3, 5, 1]이 입력벡터일때 최대값은 5가 된다. 모든 벡터는 최대값을 뺀 형태로 지수함수를 거치게 된다. [-2, 0, -4]의 형태로 각 벡터가 지수 함수를 거친값을 모두 더해 분모를 완성한다.

softmax_max에서 첫번째 아이템 -2를 기준으로 생각하면 exp(-2)=0.1353352832366127이 된다.

그러면 online_softmax()를 기준으로 보면 [3, 5, 1]에 '3'을 계산할때는 max값이 3이라서 exp(0)이 된다. softmax_max와의 차이는 -2가 생기는데 '5'를 계산시점에 sum=sum*a + b로 계산, a=old_max - new_max로 계산 위에서 언급한 차이(-2) 여기서 계산이 된다고 보면 된다.

• a=old_max-new_max는 항상 0이상이 된다
• 즉 exp(a)는 항상 1 이상을 갖게된다
• 만약 old_max와 new_max가 같으면(업데이트가 없으면) exp(a)는 1이 된다
• 현재값만 new_max로 계산하려 더한다.
• 업데이트가 되면, new_max가 커지면 old_max와의 차이만큼 sum에 곱해진 값이 sum에 다시 추가되는 로직이다.