[Transformer] C로 Transformer 구현하기

Project Goal

Transformer를 C로 구현하기
- 검증 방법 : basic이 되는 python code와의 output 비교를 통해 정확도 평가

What is the basis of this project?

• Architecture : Transformer
• Paper : Attention is All You Need 
• Transformer Pytorch Code : harvaldnlp

   (+) My github code : https://github.com/je0nsye0n/Transformer_C

---

Pytorch Code 수정

기본적인 틀은 Pytorch코드와 C코드의 output이 같은지 확인하는 것이다.

이 조건이 성립되기 위해서는 동일한 아키텍처를 설계해야 하며, Input과 Linear 연산에서 발생하는 Weight와 bias 값을 동일하게 먹여야 한다. 때문에 모든 Linear연산에 대해 발생하는 값들을 저장해주는 과정이 필요하다.

 

Linear 연산마다 위와 같은 방식으로 데이터를 저장해주면, 아래와 같이 필요한 데이터들을 사용할 수 있는 것을 확인할 수 있다. (예시로 Decoder의 Data를 가져왔다)

 

C 코드 작성

본 포스팅은 transformer pytorch에 대한 코드를 이해했다는 전제하에 진행된다.

 

Model

 

Model 구조를 통해 확인할 수 있듯이 우리는 크게 Embedding, Multi-Head Attention, Norm, Feed Forward, Linear, Softmax를 구현해야 한다. Transformer에서 주요 역할을 하는 Embedding, MHA module, FFN module에 대해서 상세한 설명을 하겠다.

 

Positional Encoding

void compute_positional_encoding(float **pos_emb) {
    for (int t = 0; t < max_len; t++) {
        for (int d = 0; d < d_model; d += 2) { // 짝수 인덱스
            float div_term = exp(-log(10000.0) * (float)d / d_model);
            float angle = t * div_term;
            pos_emb[t][d] = sin(angle);
            if (d + 1 < d_model) { // 홀수 인덱스
                pos_emb[t][d + 1] = cos(angle);
            }
        }
    }
}

void TransformerEmbedding() {
   
    compute_positional_encoding(pos_emb);

    // 입력 시퀀스를 토큰 임베딩으로 변환
    for (int b = 0; b < batch_size; b++) {
        for (int t = 0; t < seq_len; t++) {
            int token_id = (int)input[b][t]; // 정수형 토큰 ID
            for (int d = 0; d < d_model; d++) {
                emb_output[b][t][d] = tok_emb[token_id][d] + pos_emb[t][d];
            }
        }
    }
}

 

 

Multi-Head Attention

• Multi-Head Attention

이 부분은 크게 4가지다.

input으로 들어온 값을 각각 Q,K,V 별로 Linear 연산

헤드 크기에 맞게 linear 값을 split

Scaled-dot Attention으로 값 전달

원래의 size로 concat

 

void MultiHeadAttention(float ***input, float ***enc_src, float ***output,
                        Linear *linear_q, Linear *linear_k, Linear *linear_v, Linear *linear_concat){
    Results results;
    allocate_results(&results, batch_size, header, seq_len, d_k);
    
    float ****output_tmp, ***output_tmp2, ***tmp;
    float ****head1, ****head2;
    data_allocate_4d(&head1,batch_size,seq_len,header,d_k);
    data_allocate_4d(&head2,batch_size,header,seq_len,d_k);
    data_allocate_4d(&output_tmp,batch_size,header,seq_len,d_k);
    data_allocate_3d(&tmp, batch_size, seq_len, d_model);
    data_allocate_3d(&output_tmp2, batch_size, seq_len, d_model);

    for(int a=0; a<3; a++){
        if(enc_src==NULL){
            if(a==0) LinearMapping(linear_q,input,tmp);
            if(a==1) LinearMapping(linear_k,input,tmp);
            if(a==2) LinearMapping(linear_v,input,tmp); 
        }
        else{
            if(a==0) LinearMapping(linear_q,input,tmp);
            if(a==1) LinearMapping(linear_k,enc_src,tmp);
            if(a==2) LinearMapping(linear_v,enc_src,tmp); 
        }

        // split
        for(int i=0; i<batch_size; i++){
            for(int j=0; j<seq_len; j++){
              for(int k=0; k<header; k++){
                for(int l = 0; l<d_k; l++){
                    head1[i][j][k][l] = tmp[i][j][k*d_k+l];
                    }
                }
            }
        }
        for (int i = 0; i < batch_size; i++) {
            for (int j = 0; j < seq_len; j++) {
                for (int k = 0; k < header; k++) {
                    for (int l = 0; l < d_k; l++) {
                        head2[i][k][j][l] = head1[i][j][k][l];
                    }
                }
            }
        }
       
        for (int i = 0; i < batch_size; i++) {
            for (int j = 0; j < header; j++) {
                for (int k = 0; k < seq_len; k++) {
                    if(a==0) memcpy(results.Q[i][j][k], head2[i][j][k], d_k * sizeof(float));
                    if(a==1) memcpy(results.K[i][j][k], head2[i][j][k], d_k * sizeof(float));
                    if(a==2) memcpy(results.V[i][j][k], head2[i][j][k], d_k * sizeof(float));
                }
            }
        }
    }      
    
    /*attention*/
    for(int i=0; i<batch_size; i++){
        Attention_h(results.Q[i],results.K[i],results.V[i],output_tmp[i]);
    }

    /*concat*/
    for (int i = 0; i < batch_size; i++) {
        for (int j = 0; j < seq_len; j++) {
            for (int h = 0; h < header; h++) { // header 별 d_k 연결
                for (int l = 0; l < d_k; l++) {
                    output_tmp2[i][j][h * d_k + l] = output_tmp[i][h][j][l]; 
                }
            }
        }
    }    

    //data_print_3D(output_tmp2,batch_size,seq_len,d_model);
    LinearMapping(linear_concat,output_tmp2,output);
}

 

• Scaled-dot Attention

이 부분은 이제 header별로 잘려들어온 input 값을 Q,K,V 값과 연산을 해주면 된다. 아래의 연산 방식을 C로 짜주었다.

yukyunglee 깃허브&nbsp

void Attention_h(float ***query, float ***key, float ***value, float ***output) {
    
    float score, scale = 1.0f / sqrt((float)d_k);
    float ***tmp;
    data_allocate_3d(&tmp, header, seq_len, seq_len);

    // attention score 구하기
    for (int i = 0; i < header; i++) {
        for (int j = 0; j < seq_len; j++) {
            for (int k = 0; k < seq_len; k++) {
                score = 0.0f;
                for (int l = 0; l < d_k; l++) {
                    score += query[i][j][l] * key[i][k][l];
                }
                tmp[i][j][k] = score * scale;  // 결과를 output에 저장
            }
        }
        Softmax(tmp[i]);
    }

    // attention value 구하기
    for(int i=0; i<header; i++){
        for(int j=0; j<seq_len; j++){
            for(int k=0; k<d_k; k++){
                score = 0.0f;
                for(int l=0; l<seq_len; l++){
                    score += tmp[i][j][l] * value[i][l][k];
                }
                output[i][j][k] = score;
            }
        }
    }
}

 

Feed Forward

이 부분은 input을 Linear 연산(완전 연결)을 통해 1차적으로 확장하고 ReLU를 적용 후 다시 차원을 축소해주면 된다. 

void feedforward(Linear *linear1, Linear *linear2, 
    float ***input, float ***output) {

    float hidden_tmp[batch_size][seq_len][hidden];

    // 1단계: 확장 (Feedforward1) 및 ReLU 적용
    for (int i = 0; i < batch_size; i++) {
        for (int j = 0; j < seq_len; j++) {
            for (int k = 0; k < hidden; k++) {
                float O = linear1->bias[k];

                for (int l = 0; l < d_model; l++) {
                    O += input[i][j][l] * linear1->weight[k][l];
                }

                hidden_tmp[i][j][k] = fmax(O, 0.0);  // ReLU 적용
            }
        }
    }

    // 2단계: 축소 (Feedforward2)
    for (int i = 0; i < batch_size; i++) {
        for (int j = 0; j < seq_len; j++) {
            for (int k = 0; k < d_model; k++) {
                float O = linear2->bias[k];

                for (int l = 0; l < hidden; l++) {
                    O += hidden_tmp[i][j][l] * linear2->weight[k][l];
                }

                output[i][j][k] = O;
            }
        }
    }
}

 

Results

아래와 같은 식으로 레이어 별 Python과 C의 출력 결과를 비교하여 동일한지 확인하면서 진행하였다.

 

---

C로 Transformer의 동작 과정을 순차적으로 작성하는 것이 구조를 이해하는데 있어 큰 도움이 되었다. 

다음으로는 실제 Pre-traing 모델(Vision Transformer)을 데이터를 먹여서 동일 동작을 하도록 구현해보고자 한다.

기본 구조는 유사하기 때문에 어렵지 않게 해낼 수 있을 것이라 생각한다.