Self-Distillation for Continual Learning in Masked Language Models

Motivation

• Known: masked language modeling으로 사전학습된 backbone(BERT류, ESM2 등)은 supervised fine-tuning을 통해 downstream task에 효과적으로 적응할 수 있으며, autoregressive LLM에서는 demonstration-conditioned teacher를 이용한 Self-Distillation Fine-Tuning(SDFT)이 제안된 바 있다.
• Gap: 기존 SDFT는 in-context learning 능력을 가진 autoregressive LLM과 demonstration-conditioned prompting에 의존하므로, scalar property prediction을 수행하고 프롬프트 기반 조건화가 자연스럽지 않은 masked language model 기반 continual fine-tuning에는 직접 적용할 수 없으며, 이 설정에서 forgetting을 완화하는 간단한 메커니즘이 부재하다.
• Why: 새로운 실험 assay가 계속 등장하는 단백질 특성 예측과 같이 순차적으로 새로운 supervised task에 적응해야 하는 실용적 상황에서, 저장 예제 없이 이전 task 성능을 유지하면서도 새 task를 학습할 수 있는 경량 continual learning 기법은 실제 배포 환경에서 중요한 실용적 가치를 갖는다.
• Approach: student 모델의 supervised MSE loss에, EMA teacher와 student 간 token-level 출력 분포의 KL divergence 페널티를 결합한 SDFT-MLM 목적함수를 제안하고, 이를 4개 protein property task에 걸친 3가지 순서의 순차적 fine-tuning 시나리오에서 표준 SFT와 비교 평가한다.

Achievement

Figure 1. Spearman trajectories for AAV →Stability →GFP →Meltome.

1. forgetting 완화: 세 가지 task 순서 모두에서 SDFT-MLM이 표준 SFT 대비 앞선 단계에서 학습된 task들의 최종 test Spearman correlation을 개선하여 catastrophic forgetting을 완화함을 보였다.
2. 최신 task 성능 유지: 동시에 가장 최근에 학습된 task에 대한 성능은 표준 SFT와 비교해 유사한 수준으로 유지되어, retention 개선이 plasticity 희생 없이 이루어짐을 보였다.
3. 순서 강건성: 세 가지 서로 다른 task 순서(AAV→Stability→GFP→Meltome 등)에서 일관된 개선 경향을 확인하여 방법의 task-order에 대한 강건성을 시사했다.

How

• 35M-parameter ESM2를 encoder로 사용하고, 각 task마다 2-layer MLP 기반 regression head를 [CLS] 임베딩 위에 부착
• AAV packaging, protein stability, GFP fluorescence, Meltome thermostability 4개 task를 3가지 순서로 배열해 4단계 순차 학습 진행 (각 단계 최대 1,500 gradient steps)
• SFT baseline: L_MSE만 최적화
• SDFT-MLM: 각 stage 시작 시 student를 복사해 teacher 초기화, teacher는 gradient 업데이트 없이 EMA(β=0.99)로 student를 추적
• teacher와 student의 token-level logits 분포 간 non-padding position 평균 KL divergence(L_KL)를 계산하고, L_SDFT-MLM = L_MSE + λ_KL·L_KL (λ_KL=1)로 결합
• AdamW, lr 5e-5, weight decay 0.01, gradient clipping 1.0, cosine schedule(10% warmup)로 학습, 50 step마다 Spearman correlation으로 평가

Originality

• autoregressive LLM 전용으로 제안된 SDFT의 핵심 아이디어를 masked language model 설정으로 확장, demonstration-conditioned teacher 대신 EMA teacher와 token-level KL distillation을 사용하는 방식으로 재구성
• rehearsal buffer나 이전 task 데이터 저장 없이, teacher-student self-distillation만으로 순차 적응 시 forgetting을 완화하는 간단한 정규화 메커니즘 제안
• protein language model의 continual property regression이라는 구체적이고 실용적인 생물학적 사례 연구를 통해 방법을 검증, task 순서에 따른 강건성까지 체계적으로 평가

Limitation & Further Study

• 4개 protein task가 완전히 독립적이지 않고 단백질 fitness/stability와 broader하게 연관되어 있어, 순수한 catastrophic forgetting을 완전히 통제된 형태로 검증한 것은 아님(저자도 명시)
• 35M-parameter의 비교적 작은 ESM2 모델 하나에 대해서만 검증되어, 더 큰 모델이나 다른 도메인(자연어 등)의 masked language model에 일반화되는지는 불분명
• λ_KL=1, β=0.99 등 하이퍼파라미터를 고정값으로만 사용했고 다른 가중치 조합에 대한 ablation이 부족함
• rehearsal 기반 continual learning 방법이나 EWC 등 기존 continual learning baseline과의 정량적 비교가 제시되지 않아 상대적 우수성 판단이 제한적
• 후속 연구로 더 많은 task 수, 더 큰 backbone, 다른 도메인(NLP downstream task 등)으로의 확장 및 하이퍼파라미터 민감도 분석이 필요

Evaluation

총평: 기존 autoregressive LLM용 self-distillation 기법을 masked language model 및 protein 도메인으로 자연스럽게 확장한 간결하고 실용적인 아이디어로, workshop paper 수준에서 명확한 실험적 증거를 제시하지만 baseline 비교와 스케일 검증이 제한적이어서 추가 검증이 필요하다.