DEVOCEAN

들어가며

안녕하세요. SK텔레콤 AI Comm개발본부 소속 고유경입니다. 저희 본부에서는 에이닷전화 서비스의 지속적인 발전을 위해 다양한 기술을 개발하고 있습니다.

최근 생성형 AI와 LLM의 발전이 주목받고 있지만, 모든 NLP 작업에서 LLM이 항상 최선의 선택이 되는 것은 아닙니다. 특히, 사용자 프라이버시 보호가 중요한 서비스나 인터넷 연결이 원활하지 않은 환경에서는 서버를 거치지 않고 단말기에서 직접 추론할 수 있는 온디바이스 모델이 현실적인 대안이 됩니다.  또한, 텍스트 분류와 같이 비교적 명확한 태스크를 수행하는 경우, LLM을 단말에서 실행하는 것은 과도한 자원 소모를 초래할 수 있습니다. 이에 반해 BERT 계열의 경량화된 모델은 상대적으로 적은 연산 비용으로도 충분한 성능을 발휘할 수 있어 온디바이스 AI 서비스에 적합한 솔루션이 될 수 있습니다.

이번 포스팅에서는 한국어 모델 중 KoBERT를 경량화한 DistilKoBERT를 이용하여 스팸 문자를 분류하는 모델을 만들고, 이를 TensorFlow Lite로 변환한 뒤 안드로이드 단말에서 실행하기까지의 과정을 살펴보겠습니다.

1. 스팸 분류 모델 만들기

1.1. 데이터셋 구성

스팸 문자 분류 모델 학습 및 평가를 위해 저희 본부에서 자체 확보/생성한 정상 문자 3,718건과 한국인터넷진흥원의 휴대전화 스팸트랩 문자 데이터에서 샘플링한 스팸 문자 3,718건을 결합하여 총 7,436건의 데이터셋을 구축하였습니다.

표 1에서 확인할 수 있듯이, 정상 문자는 일반적인 대화나 서비스 안내 메시지로 구성되어 있으며, 스팸 문자는 광고·사기 문구 및 URL 링크가 포함된 경우가 많습니다.

표1. 정상/스팸 문자 예시

전체 데이터를 각각 4,758/1,190/1,488건으로 분절하여 학습, 검증, 평가용 데이터셋을 구성하겠습니다. 문자 메시지를 토크나이징하고 HuggingFace Datasets 라이브러리를 활용해 모델 학습에 적합한 형태로 전처리합니다. 토크나이저는 이후 설명할 DistilKoBERT 모델의 토크나이저를 사용합니다.

import re
import pandas as pd
from transformers import AutoTokenizer
from datasets import Dataset, DatasetDict

train_df = pd.read_csv("train.csv")
valid_df = pd.read_csv("valid.csv")
test_df = pd.read_csv("test.csv")

def clean_text(text):
    return re.sub(r"[^ㄱ-ㅎ가-힣a-zA-Z0-9\s]", "", text) 

def preprocess_data(df, tokenizer):
    msg_list = [clean_text(msg) for msg in df["message"].tolist()]
    tokenizer = AutoTokenizer.from_pretrained("monologg/distilkobert", trust_remote_code=True)
    tokenized_data = tokenizer(msg_list, truncation=True, padding='max_length', max_length=256)

    dataset_dict = {
        "text": df["message"].tolist(),
        "label": df["label"].tolist(),
        "input_ids": tokenized_data["input_ids"],
        "attention_mask": tokenized_data["attention_mask"],
    }
    return Dataset.from_dict(dataset_dict)

train_dataset = preprocess_data(train_df, tokenizer)
valid_dataset = preprocess_data(valid_df, tokenizer)
test_dataset = preprocess_data(test_df, tokenizer)

datasets = DatasetDict({'train': train_dataset, 'validation': valid_dataset, 'test': test_dataset})

1.2. 모델 학습 및 평가

스팸 분류 모델로는 SKT의 KoBERT를 경량화한 DistilKoBERT를 활용하고자 합니다. DistilKoBERT는 지식 증류(Knowledge Distillation) 기법을 적용해 KoBERT 모델 대비 약 2배 작고 빠른 경량 모델로, 성능 저하를 최소화하면서도 추론 속도를 향상시킨 장점이 있습니다. (참고로 원조 DistilBERT는 BERT 대비 모델 크기를 40% 줄이고 속도를 60% 향상시키면서도 성능을 97% 이상 유지한 것으로 보고된 바 있습니다.)

Hugging Face의 Transformers 라이브러리를 활용해 학습과 평가를 진행하겠습니다. monologg/distilkobert 모델을 불러와 출력층에 2개 클래스(정상/스팸)를 구분하는 classification head를 추가하여 분류 태스크에 맞게 파인튜닝합니다.

import numpy as np
from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score

tokenizer = AutoTokenizer.from_pretrained("monologg/distilkobert", trust_remote_code=True)
model = AutoModelForSequenceClassification.from_pretrained("monologg/distilkobert")

# TrainingArguments 설정
training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="epoch",
    save_strategy="epoch",
    logging_dir="./logs",
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    num_train_epochs=10,
    learning_rate=2e-5,
    weight_decay=0.01,
    save_total_limit=2,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    report_to="none"
)

# 평가 함수 정의
def compute_metrics(pred):
    preds = np.argmax(pred.predictions, axis=1)
    labels = pred.label_ids

    precision = precision_score(labels, preds, average="macro")
    recall = recall_score(labels, preds, average="macro")
    f1 = f1_score(labels, preds, average="macro")
    accuracy = (preds == labels).mean()

    return {
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall,
        "f1": f1
    }

# Trainer 초기화
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=datasets["train"],
    eval_dataset=datasets["validation"],
    compute_metrics=compute_metrics
)

# 모델 학습
trainer.train()

# 테스트 데이터 평가
results = trainer.evaluate(datasets["test"])
print(results)

# 모델, 토크나이저 저장
output_dir = "saved_model/spam_cls_distilbert"
trainer.save_model(output_dir)
tokenizer.save_vocabulary(output_dir)

학습 과정에서는 정확도(Accuracy)를 주요 지표로 모니터링하며 총 10 epoch 학습을 진행합니다. 실제 학습 결과, epoch이 진행될수록 Valid(검증) 정확도가 향상되다가 일정 시점 이후 성능이 약 99%에 수렴하는 것을 확인하였습니다. 최종적으로 Valid 데이터셋과 Test 데이터셋 기준 약 99%의 정확도를 기록하였으며, 이로써 스팸 분류 모델 구축이 완료되었습니다.

2. 온디바이스용 모델 변환

2.1. TFLite 변환

이제 학습된 PyTorch 모델을 안드로이드 앱에서 사용할 수 있도록 TFLite 포맷으로 변환할 차례입니다.

온디바이스 추론을 위한 런타임에는 PyTorch Mobile이나 ONNX Runtime과 같은 여러 옵션이 있지만, 본 포스팅에서는 모델 파일을 TFLite(TensorFlow Lite)로 변환 후 LiteRT 런타임에서 실행하고자 합니다. LiteRT는 모바일·엣지 환경에 최적화된 경량 런타임으로 다음과 같은 장점이 있습니다.

• 모델 용량이 작고 배포가 용이함

• Android Hardware Accelerator 지원 (NNAPI, GPU, EdgeTPU 등)

• 다양한 양자화 옵션 제공

기존에는 PyTorch 모델을 TFLite로 변환하려면 PyTorch → ONNX → TensorFlow Graph → TFLite의 복잡한 과정을 거쳐야 했습니다. 그러나 Google의 AI Edge Torch 라이브러리를 활용하면 이러한 중간 단계를 생략하고 PyTorch 모델을 바로 TFLite 포맷으로 변환할 수 있습니다. 특히, ONNX 변환 과정에서 발생할 수 있는 채널 순서 차이로 인한 불필요한 Transpose 연산 삽입이나 변환 오류 문제를 방지할 수 있다는 점에서 AI Edge Torch는 효율적인 솔루션입니다.

모델 변환을 위해서는 저장된 PyTorch 모델과 샘플 입력이 필요합니다. ai_edge_torch.convert()함수에 평가(eval) 모드의 모델 객체와 샘플 입력 텐서를 전달합니다. 변환이 완료되면 모델은 "spam_cls_distilbert.tflite" 파일로 저장되며 약 108MB의 용량을 갖습니다.

import ai_edge_torch
from transformers import AutoModelForSequenceClassification

output_dir = "saved_model/spam/spam_cls_distilbert"
model = AutoModelForSequenceClassification.from_pretrained(output_dir).eval()

edge_model = ai_edge_torch.convert(model, sample_input)
edge_model.export(f"{output_dir}/spam_cls_distilbert.tflite")

2.2. 양자화(Quantization) 적용

양자화(Quantization)는 모델의 숫자 표현을 더 적은 비트(bit)로 변환하여 모델 크기를 줄이고 연산 속도를 높이는 기술입니다. 일반적으로 딥러닝 모델은 Float32(32비트 부동소수점)으로 학습되지만, 온디바이스 환경에서는 메모리 사용량을 줄이고 CPU/GPU 연산을 최적화하기 위해 양자화를 적용하는 것이 효과적입니다.

그림1. INT8 Quantization (출처: Datature 블로그)

양자화에는 여러 가지 방식이 있지만, 본 포스팅에서는 학습 후 적용할 수 있는 Post-Training Dynamic Range Quantization(동적 범위 양자화, 이하 DRQ)을 사용합니다. DRQ는 모델을 변환하는 과정에서 가중치를 INT8(8비트 정수)로 변환하고 활성화 함수의 출력값의 경우 실행 시점에 동적으로 INT8로 변환하는 방식입니다. 이를 통해 모델 용량을 최대 4배까지 줄이고 CPU 연산에서 양자화 연산 커널을 사용해 추론 속도를 높일 수 있습니다​.

AI Edge Torch에서는 TFLite 변환 과정에 비교적 손쉽게 DRQ를 적용할 수 있습니다. tf.lite.Optimize.DEFAULT 옵션을 설정하여 Dynamic Range Quantization을 활성화한 설정 값을 변환 함수의 파라미터로 제공하기만 하면 됩니다.

변환이 완료된 모델은 "spam_cls_distilbert_int8.tflite" 파일로 저장되며 양자화 적용 전 파일 용량보다 약 4배 작은 28MB의 용량을 갖습니다.

import ai_edge_torch
import tensorflow as tf

tfl_converter_flags = {'optimizations': [tf.lite.Optimize.DEFAULT]}

tfl_drq_model = ai_edge_torch.convert(
    model, sample_input, _ai_edge_converter_flags=tfl_converter_flags
)
tfl_drq_model.export(f"{output_dir}/spam_cls_distilbert_int8.tflite")

DRQ에서는 가중치만 INT8로 변환되므로, 모델의 정확도가 원본과 거의 동일하게 유지됩니다. 실제로 양자화된 모델이 원본 모델과 성능 차이가 없는지 확인하기 위해 TFLite로 변환한 모델을 실행하고 정확도를 비교한 결과, TFLite로 변환한 모델은 원본과 동일한 성능을 유지했으며, 양자화 옵션을 적용한 모델도 정확도 차이가 0.07%p에 불과하여 성능 저하가 미미한 것으로 확인되었습니다. 즉, 양자화를 통해 원본 모델과 거의 동일한 성능을 유지하면서도 모델 크기를 줄이고 추론 속도를 향상시키는 최적화 효과를 얻을 수 있습니다.

표2. 원본, 양자화 모델 성능 비교

이렇게 Float32 모델과 INT8 양자화 모델 두 가지 버전의 스팸 문자 분류 모델이 준비되었습니다.

이제 두 모델을 안드로이드 앱에 탑재하여 실제 동작을 검증하고 추론 속도, 메모리 사용량 등의 리소스를 측정해보겠습니다.

3. 온디바이스 실행

3.1. LiteRT Text Classification 예제 앱

TFLite 모델을 안드로이드 앱에서 실행하기 위해, Google AI Edge 깃헙에서 제공하는 LiteRT API Samples의 Text Classification 예제 앱을 활용하고자 합니다. 이 예제 앱에서는 TensorFlow Lite의 Task Library를 활용하여, 메타데이터(metadata)가 포함된 TFLite 모델을 간편하게 불러와 추론하는 기능을 제공합니다. Task Library를 사용하면 BertNLClassifier와 같은 API를 통해 Java/Kotlin 코드에서 간편하게 텍스트 분류 기능을 구현할 수 있습니다. 하지만 본 포스팅에서 학습한 모델을 적용하기에는 메타데이터 작성 도구가 DistilKoBERT 모델과 호환되지 않는다는 한계점이 있습니다.

• 일반적인 BERT 기반 분류 모델은 input_ids, attention_mask, token_type_ids 세 개의 입력을 사용합니다.

• 반면, DistilKoBERT는 token_type_ids 없이 input_ids와 attention_mask만 입력으로 사용합니다.

이처럼 Task Library와 호환되는 TFLite 모델을 만들려면 모델 구조에 맞는 메타데이터를 추가해야 하지만, metadata_writers의 커스터마이징이 어렵다는 한계가 있습니다. DistilKoBERT처럼 일반적인 BERT 모델과 입력 형식이 다른 경우, 기존 메타데이터 작성 도구로는 사용자가 원하는 입력 텐서 형식을 직접 정의하거나 수정하기 매우 어렵습니다. 이러한 한계를 극복하기 위해 Task Library를 활용하는 대신 TFLite Interpreter를 직접 호출하여 추론을 수행하는 방식을 활용해보고자 합니다.

3.2. LiteRT Interpreter를 활용한 온디바이스 추론

LiteRT 예제 앱에서는 Task Library의 기본 토크나이저를 사용합니다. 이 토크나이저는 띄어쓰기 기준으로 단어를 분절한 후, 사전에 존재하는 단어만을 매칭합니다.

이러한 방식은 BERT 계열 모델이 요구하는 서브워드 기반 토크나이징을 지원하지 않기 때문에, 모델의 기대 결과와 다른 출력이 나올 수 있습니다.

또한, 단어사전에 정확히 매칭되는 표현이 없는 경우 UNK 처리되어 성능에 영향을 줄 수 있습니다.

따라서, DistilKoBERT의 토크나이저와 유사한 방식의 토크나이징 로직을 직접 구현하여 입력 배열을 생성하였습니다.

이제 토큰 ID 배열을 모델에 입력하고, LiteRT Interpreter로 TFLite 모델을 실행하여 예측 결과를 얻을 차례입니다.

LiteRT Interpreter는 Google AI Edge LiteRT를 기반으로 TFLite 모델을 실행하는 엔진입니다.

기존의 TensorFlow Lite Interpreter와 유사하지만, 더 가볍고 빠르며, 온디바이스 AI 추론에 최적화된 구조를 가지고 있습니다.

또한, 다중 입출력을 지원하므로 input_ids와 attention_mask를 동시에 모델에 전달할 수 있어 BERT 기반 모델 실행에 적합합니다.

val inputs = arrayOf(input_ids, attention_mask)
interpreter!!.runForMultipleInputsOutputs(inputs, outputs) // 추론 실행

3.3. 실행 결과

이제 LiteRT 앱을 빌드하고 실행해보겠습니다. 텍스트 필드에 정상 및 스팸 예시 문자를 입력하고 'Classify' 버튼을 탭하면, 표3의 화면과 같이 분류 결과가 표시됩니다.

• 모델 Output: Softmax 적용 전 Logit 값으로, 점수가 더 높은 라벨이 최종 예측 결과로 선택됩니다.

• 추론 속도: 모델에 데이터를 입력한 후 결과를 반환하기까지 소요되는 시간입니다.

  DistilBERT와 INT8 양자화된 DistilBERT 모델은 모두 동일한 예측 결과를 반환하였으며, INT8 양자화 모델은 원본 DistilBERT 모델 대비 약 2.5배 빠른 추론 속도를 보였습니다.

  모델 크기가 줄어든 만큼 추론 시간이 단축되었으며, 이를 통해 온디바이스 환경에서 양자화의 최적화 효과를 확인할 수 있었습니다.

표3. 정상/스팸 예시 문자 분류 결과

4. 단말 리소스 측정

4.1. TensorFlow Lite Benchmark Tool이란

모바일 환경에서는 메모리 사용량을 최소화하고, 추론 속도를 높이며, 배터리 소모를 줄이는 것이 필수적입니다.

그렇다면 안드로이드 단말에서 DistilKoBERT 모델을 실행할 경우, CPU/메모리 등의 리소스를 얼마나 소모할까요? 양자화 전/후 모델의 리소스 점유 비율 차이는 어느정도일까요?

이를 측정하기 위한 방안으로 Google에서는 TensorFlow Lite Benchmark Tool(이하 TFLite Benchmark Tool)을 제공하고 있습니다. TFLite Benchmark Tool은 온디바이스에서 실행 속도, 메모리 사용량 등을 평가할 수 있는 벤치마크 앱입니다. 이 앱을 활용하여 모델 실행 중 표4에 나열된 지표들을 측정할 수 있습니다.

표4. TFLite Benchmark Tool 측정 지표

4.2. 측정 결과

갤럭시 S25 (SM-S931N) 단말을 활용하여 벤치마크를 실행한 결과는 표5와 같습니다.

표 5. 양자화 전/후 단말 리소스 측정 결과

양자화 전/후 모델을 비교해보면 추론 속도와 메모리 사용량에서 양자화의 큰 이점을 확인할 수 있습니다.

양자화된 모델은 평균 10ms 수준의 빠른 추론 속도를 기록하며, 이는 기존 모델 대비 약 4배 향상된 성능입니다.  또한, 메모리 사용량이 65MB 이하로 줄어 원본 모델 대비 약 3배 감소하여, 모바일 환경에서도 보다 안정적으로 실행될 수 있음을 입증하고 있습니다.

반면, 초기화 속도*는 FP32 모델보다 다소 증가하는 경향이 있습니다. 이는 FP32 모델이 메모리에 로드된 즉시 실행 가능한 것과 달리, INT8 모델은 초기 로드 시 연산 최적화를 위한 추가 변환 과정이 필요하기 때문입니다. 그럼에도 이러한 초기화 속도의 차이는 실제 추론 성능 향상과 메모리 절감 효과를 고려했을 때 충분히 감수할 수 있는 수준으로 판단됩니다. 

마지막으로 CPU 사용량을 비교했을 때, 양자화 이후에도 CPU 사용량 감소는 크지 않았으나 연산 최적화로 인해 일정 부분 CPU 부담이 줄어든 것을 확인할 수 있습니다.

마치며

지금까지 한국어 DistilKoBERT를 활용한 스팸 문자 분류 모델을 구축하고, 이를 온디바이스 환경에 최적화하여 안드로이드 앱에 통합하는 전 과정을 살펴보았습니다. Pytorch로 학습한 모델을 TFLite로 변환 및 양자화하여 모바일에서도 경량으로 실행할 수 있음을 확인하였습니다.

온디바이스에서 BERT 계열 모델을 활용하면 서버/클라우드에 의존하지 않고도 실시간 NLP 기능을 제공할 수 있어 프라이버시 보호, 지연 감소 등의 이점을 얻을 수 있습니다. 비록 최신 LLM에 비하면 규모는 작지만 분류나 QA와 같은 특화된 태스크에서는 충분한 성능을 발휘할 수 있습니다. 다만, 온디바이스 AI를 실제 서비스에 적용하기 위해서는 단순히 모델 성능뿐만 아니라 아래 항목들을 고려한 의사 결정이 필요합니다.

• 서비스 관점에서의 효용성

  • 앱 내부에 AI 모델을 탑재했을 때, 기존 로직과의 통합 과정에서 성능 저하나 과부하가 발생하지 않는지

  • 모델이 요구하는 연산량과 실제 단말의 처리 성능이 적절하게 균형을 이루는지

  • 양자화를 적용했을 때 성능 저하 없이 실시간 처리가 가능한지

• 단말에 따른 성능 차이

  • 지원 가능한 단말의 범위는 어디까지인지

  • 고성능 단말과 보급형 단말의 성능 차이가 용납 가능한 수준인지

  • 메모리 제한이 있는 환경에서 다른 앱 프로세스와 동시에 실행되면서 발생할 수 있는 문제점은 없는지

앞으로 TensorFlow Lite 등의 모바일 ML 툴킷의 발전으로 인해 더 큰 모델도 단말에서 원활하게 돌아가거나 메타데이터 표준이 개선되어 모델 커스터마이징이 용이하게 된다면, 온디바이스 AI의 활용 범위는 더욱 넓어질 것입니다. 이번 포스팅이 온디바이스 NLP를 고려하는 분들에게 유용한 참고자료가 되었기를 바라며 글을 마치겠습니다.