BERT 뉴스 분류 #5 최종 모델 구조, 학습 설정, 실험 기록, ONNX 변환 #

#2026-03-01

#1 모델 구조

입력 텍스트
    ↓ 토크나이저 (max_length=192, Dynamic Padding)
[input_ids, attention_mask]  shape: (batch, ≤192)
    ↓ BERT Embedding Layer  ← lr = 0.99e-5 (가장 작게)
토큰 벡터  shape: (batch, ≤192, 768)
    ↓ Transformer × 12 (레이어 0~11) ← 레이어별 차등 lr (1.10e-5 ~ 3.50e-5)
문맥화된 벡터  shape: (batch, ≤192, 768)
    ↓ [CLS] 토큰만 추출
문장 벡터  shape: (batch, 768)
    ↓ [EnhancedClassifier]
    ↓ Dropout(0.3)
    ↓ Linear(768 → 256) + GELU
    ↓ Dropout(0.3)
    ↓ Linear(256 → 4)
logits  shape: (batch, 4)
    ↓ softmax
확률  shape: (batch, 4)   예: [0.02, 0.05, 0.91, 0.02] → Business

뉴스 기사 텍스트가 들어온다. 토크나이저가 이걸 토큰 번호 배열로 바꾸는데, max_length=192로 설정되어 있다. 512가 아니라 192다. 실제 데이터를 분석해서 95%의 기사가 180 토큰 이내라는 걸 확인한 뒤 내린 결정이다.

DataCollatorWithPadding이 배치를 구성할 때 그 배치 안에서 가장 긴 시퀀스에 맞춰 동적으로 패딩하므로, 192보다 짧은 배치에서는 더 적은 패딩이 들어간다. 이 두 최적화만으로 Self-Attention의 계산량이 원래의 7분의 1 이하로 줄어든다.

토큰 번호 배열이 BERT의 임베딩 레이어에 들어간다. 여기서 각 토큰 번호가 768차원 벡터로 변환된다. 이 레이어의 학습률은 0.99e-5로 전체에서 가장 작다. 토큰의 기본 표현은 뉴스 분류를 하든 다른 과제를 하든 거의 같아야 하니까 최소한으로만 조정한다.768차원 벡터들이 12개의 Transformer 레이어를 순서대로 통과한다. 각 레이어에는 차등 학습률이 적용되어 있다. 레이어 0은 1.10e-5, 레이어 11은 3.50e-5다. 한 레이어 올라갈 때마다 0.9의 역수만큼 학습률이 커진다. 하위 레이어의 범용 언어 지식은 살짝만, 상위 레이어의 과제 특화 표현은 적극적으로 조정하는 전략이다.

12개 레이어를 통과하면 (batch, ≤192, 768) 형태의 텐서가 나온다. 여기서 [CLS] 토큰 위치의 768차원 벡터만 뽑아낸다. 이 벡터가 문장 전체의 의미를 담고 있다. Self-Attention을 12번 거치면서 문장의 모든 단어가 서로의 문맥을 흡수했고, 그 결과가 [CLS] 위치에 응축된 것이다.

이 768차원 벡터가 EnhancedClassifier로 들어간다. 기본 BERT의 단순한 Linear(768→4) 대신, 768→256→4 구조로 중간 단계를 둔 분류 헤드다. 첫 번째 Dropout(0.3)이 과적합을 방지하고, Linear(768→256)이 차원을 축소하면서 GELU 활성화가 비선형성을 도입하고, 두 번째 Dropout(0.3)을 거쳐 최종 Linear(256→4)가 4개 클래스의 logit을 만든다. Softmax가 이 logit을 확률 분포로 변환하면 뉴스 카테고리 예측이 완성된다.

Before: input_length=512, 헤드=Linear(768→4), 전 레이어 동일 lr
After:  input_length=192, 헤드=768→256→4,    레이어별 차등 lr (decay=0.9)

기본 BERT에서 바뀐 것을 정리하면 세 가지다. 입력 길이가 512에서 192로 줄었고, 분류 헤드가 단일 선형 레이어에서 2층 비선형 네트워크로 확장되었고, 전체 모델에 동일하게 적용되던 학습률이 레이어별로 차등 적용된다. 각각은 독립적인 최적화이지만 합쳐지면 시너지를 낸다.

#

#2 학습 설정

training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="steps",
    eval_steps=200,
    save_strategy="steps",
    save_steps=200,
    per_device_train_batch_size=16,     # max_length 단축 → 배치 크기 2배 가능
    per_device_eval_batch_size=32,
    num_train_epochs=3,                 # Gradual Unfreezing 없으므로 일반 학습
    learning_rate=3.5e-5,               # 베이스 lr (레이어별로 차등 적용)
    weight_decay=0.01,
    warmup_ratio=0.1,
    lr_scheduler_type="cosine",
    label_smoothing_factor=0.1,
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    greater_is_better=True,
    seed=SEED,
    report_to="mlflow",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=small_train_dataset,
    eval_dataset=small_test_dataset,
    compute_metrics=compute_metrics,
    data_collator=data_collator,           # Dynamic Padding
    optimizers=(optimizer, None),           # 차등 학습률 옵티마이저
    callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
)

per_device_train_batch_size=16은 이전의 8에서 두 배로 늘어났다. 이게 가능한 이유가 max_length를 192로 줄인 덕분이다. 시퀀스가 짧아지면 Self-Attention 행렬이 작아지고 GPU 메모리를 덜 먹으니까, 같은 메모리로 더 많은 샘플을 한 배치에 넣을 수 있다. 배치가 커지면 그래디언트 추정이 더 안정적이 되어 학습이 부드러워진다. 입력 길이 최적화가 단순히 속도만 빠르게 하는 게 아니라 배치 크기라는 하이퍼파라미터의 선택지까지 넓혀주는 연쇄 효과가 있는 것이다.

num_train_epochs=3은 전체 데이터를 세 바퀴 도는 것이다. 이전에 1 에포크만 돌았던 것보다 늘었지만, Early Stopping이 과적합을 감시하고 있으니 3 에포크를 다 채우지 않고 중간에 멈출 수도 있다. 넉넉하게 잡아놓고 Early Stopping이 최적 시점을 찾게 하는 전략이다.

learning_rate=3.5e-5는 베이스 학습률이다. 이 값이 레이어 11에 적용되고, 아래로 갈수록 0.9씩 곱해져서 줄어든다. weight_decay=0.01은 L2 정규화로 가중치가 불필요하게 커지는 것을 방지한다. warmup_ratio=0.1은 전체 학습 스텝의 10%를 워밍업에 쓴다. 학습 초반에 랜덤 초기화된 분류 헤드가 엉뚱한 그래디언트를 역전파할 때 사전학습 가중치를 보호하기 위해서다. lr_scheduler_type=“cosine"은 워밍업 후 학습률이 코사인 곡선을 따라 감쇠하게 만든다. 후반부까지 미세 조정이 지속되도록 하는 선택이다.

label_smoothing_factor=0.1은 정답 레이블을 [0, 0, 1, 0]에서 [0.025, 0.025, 0.925, 0.025]로 부드럽게 만들어서 모델의 과도한 확신을 방지한다. e valuation_strategy=“steps"와 eval_steps=200은 200 스텝마다 검증 성능을 체크한다. load_best_model_at_end=True와 metric_for_best_model=“f1"은 F1 기준으로 가장 좋았던 체크포인트를 최종 모델로 선택한다. EarlyStoppingCallback(early_stopping_patience=3)은 F1이 3번 연속 개선되지 않으면 학습을 중단한다.

optimizers=(optimizer, None)에서 optimizer는 레이어별 차등 학습률이 적용된 AdamW다. None은 학습률 스케줄러를 Trainer의 기본 설정(위에서 지정한 코사인 감쇠)에 맡긴다는 뜻이다.

이 설정들이 합쳐지면 학습의 전체 생애가 이렇게 그려진다. 시작할 때 학습률이 0에서 출발해서 워밍업 기간 동안 올라간다. 목표 학습률에 도달하면 본격적인 학습이 시작되는데, 각 레이어가 자기 깊이에 맞는 다른 학습률로 조정된다. 학습이 진행되면서 코사인 곡선을 따라 전체 학습률이 서서히 줄어든다. 200 스텝마다 검증 F1을 체크하면서 최고 성능 모델을 저장한다. 3번 연속 개선이 없으면 자동으로 멈추고 최고 시점의 모델을 복원한다.

#

#3 평가 지표

def compute_metrics(eval_pred):
    predictions = np.argmax(logits, axis=-1)
    precision, recall, f1, _ = precision_recall_fscore_support(
        labels, predictions, average="weighted"
    )
    acc = accuracy_score(labels, predictions)
    return {"accuracy": acc, "precision": precision, "recall": recall, "f1": f1}

compute_metrics 함수가 매 평가마다 네 가지 지표를 계산한다.

Accuracy는 전체 예측 중 맞힌 비율이다. 가장 직관적이지만 가장 거친 지표다. AG News는 4개 클래스가 균등하니까 Accuracy만으로도 어느 정도 유용하지만, 클래스가 불균형한 경우에는 오도할 수 있다.

Precision은 모델이 특정 클래스라고 판정한 것 중 실제로 그 클래스인 비율이다. 모델이 “이건 경제 기사다"라고 100번 말했을 때 그중 95번이 정말 경제 기사면 Precision이 0.95다. 이 지표가 높으면 모델의 판정을 신뢰할 수 있다.

Recall은 실제로 특정 클래스인 것 중 모델이 잡아낸 비율이다. 경제 기사가 100건 있는데 모델이 그중 90건을 경제라고 맞췄으면 Recall이 0.90이다. 이 지표가 높으면 모델이 놓치는 게 적다.

F1은 Precision과 Recall의 조화평균이다. 둘 중 하나만 높고 다른 하나가 낮으면 F1도 낮아진다. 두 지표의 균형을 하나의 숫자로 요약하는 것이다. Early Stopping의 기준으로 F1을 쓰는 이유가 여기 있다. Precision만 높이려면 확실한 것만 판정하고 애매한 건 전부 다른 클래스로 밀면 되고, Recall만 높이려면 닥치는 대로 그 클래스라고 찍으면 된다. 둘 다 편향된 전략이다. F1은 이런 편향을 허용하지 않고 균형 잡힌 성능을 요구한다.

weighted 평균은 각 클래스의 샘플 수에 비례해서 가중 평균을 내는 것이다. AG News는 클래스가 균등하니까 macro 평균(단순 평균)과 거의 같은 값이 나온다. 하지만 코드에서 weighted를 쓰는 건, 나중에 불균형한 데이터를 다루게 될 때도 같은 코드를 쓸 수 있도록 범용성을 확보하는 좋은 습관이다.

#

#4 MLflow: 실험 기록

with mlflow.start_run(run_name="bert-enhanced-gradual-unfreeze") as run:
    mlflow.log_params({
        "model_name": MODEL_NAME,
        "max_length": 192,            # 최적화된 값 기록
        "frozen_layers": "0-5",       # 동결 전략 기록
        "lr_decay": 0.9,              # 차등 학습률 감쇠율
        "label_smoothing": 0.1,
        "classifier_hidden": 256,
        "classifier_dropout": 0.3,
    })
    trainer.train()
    mlflow.log_metrics({...})
    mlflow.pytorch.log_model(model, "model")

MLflow에 기록하는 파라미터를 보면, 이번에는 단순한 학습률이나 에포크 수 같은 표준 하이퍼파라미터뿐 아니라 구조적 결정까지 기록한다. max_length=192, frozen_layers=“0-5”, lr_decay=0.9, classifier_hidden=256, classifier_dropout=0.3, label_smoothing=0.1 같은 것들이다.

이 기록이 왜 중요한가? 실험을 여러 번 반복하다 보면 “아, 저번에 192로 줄이고 Dropout 0.3으로 했을 때 F1이 얼마였지?“라는 질문이 반드시 나온다. 기록이 없으면 기억에 의존해야 하고, 기억은 틀린다. MLflow에 모든 설계 결정과 그 결과를 자동으로 기록해두면, 나중에 웹 인터페이스에서 실험들을 나란히 놓고 “어떤 조합이 가장 좋았는가"를 정확히 비교할 수 있다. 모델 아티팩트 자체도 저장한다. 나중에 이 모델을 다시 불러올 수 있고, 다른 모델과 성능을 비교할 수 있다. 연구 재현성의 기본 인프라인 셈이다.

#

#5 ONNX 변환

학습이 끝나면 PyTorch 모델을 ONNX 형식으로 변환한다. PyTorch는 동적 그래프 방식이다. 코드를 실행할 때마다 연산 그래프가 새로 만들어진다. 이건 학습할 때는 편리하다. 조건문이나 반복문을 자유롭게 쓸 수 있고 디버깅이 쉽다. 하지만 추론할 때는 이 유연성이 오히려 오버헤드가 된다. 매번 그래프를 새로 만들 필요가 없는데도 Python 인터프리터가 개입하니까 느려지는 것이다.

ONNX 변환은 모델의 연산 그래프를 한 번 고정시키고, 그 위에 최적화를 적용한다. 연속된 작은 연산들을 하나의 큰 연산으로 합치고(연산 융합), 결과에 영향을 미치지 않는 연산을 제거하고, 메모리 접근 패턴을 최적화한다. Python 오버헤드가 완전히 사라지고 순수한 수치 연산만 남는다.

왜 ONNX가 빠른가?
- PyTorch: Python 오버헤드, 동적 그래프
- ONNX: 정적 그래프 최적화 (연산 융합, 불필요 연산 제거)

dummy_input = tokenizer(
    "This is a sample news article.",
    return_tensors="pt",
    padding="max_length",
    truncation=True,
    max_length=192,   # 학습 시 설정한 값과 동일하게
)

torch.onnx.export(
    model,
    (dummy_input["input_ids"], dummy_input["attention_mask"]),
    ONNX_PATH,
    input_names=["input_ids", "attention_mask"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids": {0: "batch_size", 1: "sequence"},
        "attention_mask": {0: "batch_size", 1: "sequence"},
        "logits": {0: "batch_size"},
    },
    opset_version=14,
)

dummy_input으로 샘플 입력을 만드는 이유는 ONNX가 연산 그래프를 추적하기 위해서다. 실제 입력을 한 번 흘려보내면서 어떤 연산이 어떤 순서로 일어나는지를 기록하고, 그 기록을 정적 그래프로 저장한다. max_length=192로 설정하는 건 학습 때와 동일한 입력 형태를 맞추기 위해서다.

dynamic_axes 설정이 중요한데, batch_size와 sequence 차원을 가변적으로 만든다. 실제 서비스에서는 한 번에 기사 1개만 올 수도 있고 10개가 묶여서 올 수도 있다. 기사 길이도 20 토큰일 수도 있고 190 토큰일 수도 있다. 이 차원들을 고정하면 항상 같은 크기의 입력만 받을 수 있어서 실용성이 떨어진다. dynamic_axes로 가변 크기를 지원하면 어떤 크기의 입력이든 유연하게 처리할 수 있다.

max_length를 192로 줄인 효과가 여기서도 나타난다. 192 토큰짜리 입력의 ONNX 추론은 512 토큰짜리보다 훨씬 빠르다. 입력 길이 최적화가 학습 속도뿐 아니라 서빙 속도에도 직접적으로 영향을 미치는 것이다. 학습 파이프라인에서 내린 하나의 결정이 서빙 파이프라인까지 이어지는 셈이다.

#

#6 정리

이 단계에서 만들어진 것은 하나의 완성된 분류 시스템이다. 데이터 분석에 기반한 입력 최적화(192 토큰, 동적 패딩), 문제에 맞게 설계된 분류 헤드(768→256→4), 레이어별 차등 학습률, Label Smoothing, 코사인 학습률 스케줄링, Early Stopping이 모두 합쳐져서 하나의 학습 설정을 이룬다. 학습 과정은 MLflow에 자동 기록되고, 최종 모델은 ONNX로 변환되어 빠른 서빙이 가능한 형태로 출력된다.

이전 챕터들에서 DNA 분석, 망막증 분류, 분자 생성을 할 때는 모델을 만들고 평가하는 것이 끝이었다. 이 프로젝트에서는 그 다음 단계까지 간다. 만든 모델을 실제 서비스에 투입하고, 성능을 감시하고, 필요하면 재학습하는 운영 시스템을 구축하는 것이다. 지금까지는 그 시스템의 첫 번째 부분인 “모델 학습"을 완성한 것이고, 다음 단계에서 서빙, 감시, 재학습으로 나아간다.