[논문리뷰] Style-Specific Neurons for Steering LLMs in Text Style Transfer

• LLM(Large Language Model)을 이용한 TST(Text Style Transfer)에서 스타일 특정 뉴런을 활용하는 방법 sNeuron-TST 제안
• 출력 문장의 스타일 다양성과 유창성(fluency)을 높이기 위해 source 및 target 스타일과 관련된 뉴런을 식별
• source 스타일 전용 뉴런을 비활성화 → target 스타일 단어의 확률을 높이는 전략 사용
• 이로 인한 fluency 저하 문제를 해결하기 위해 개선된 contrative decoding 방법 제안

---

Introduction

• TST의 목표는 문장의 의미는 유지하면서 원본 스타일에서 목표 스타일로 변환하여 문장을 생성하는 것
• LLM이 TST를 포함하여 다양한 NLP 태스크에서 예외적으로 높은 성능을 보임
• 그러나 LLM은 스타일을 변환하는 것보다 기존 문장의 의미 보존에 더 집중하는 경향이 있음
• 저자의 분석에서, LLaMA-3를 대상으로 생성된 문장을 확인했을 때, 34%가 input과 동일하게 생성(task: polite → impolite)
• decoding 과정에서 대상 스타일과 일치하는 단어 생성을 향상시키는 건 TST에서 중요한 과제로 남음
• 최근 LLM은 TST에서 크게 두 가지 방법으로 적용됨
  1. full-parameter or prameter-efficient fine-tuning(single-style or parallel-style text data)
  2. the robust in-context learning을 활용하여 zero-shot 또는 few-shot learning을 위한 전문적인 프롬프트 생성
• (1)은 상당한 데이터와 resources가 필요하고, (2)는 프롬프트의 민감성으로 사소한 변경도 결과에 영향을 줄 수 있다는 단점
• 따라서 저자는 LLM을 사용하되, Neuron analysis를 적용
• Neuron analysis 측면에서 다음 두 가지 research questions:
  1. Do LLMs possess neurons that specialize in processing style-specific text?
  2. If such neurons exist, how can we optimize their utilization during the decoding process to steer LLMs in generating text that faithfully adheres to the target style?
• LLM이 스타일과 관련된 뉴런을 가지고 있는지, 가지고 있다면 그 뉴런을 decoding 과정에서 어떻게 활용할 수 있을지가 중요
• 이 두 가지 질문을 통해 저자는 sNeuron-TST를 제안

Contributions

1. 스타일 특화 뉴런(style-specific neuron)을 이용하여 LLMs의 TST 성능을 올리는 첫 연구
2. source style과 target style로 동시에(overlap) 활성화되는 뉴런을 제거하는 것의 중요성 강조
3. Dola에서 영감을 받아 향상된 대조적 디코딩 방법을 도입하여 목표 스타일과 관련된 단어 생성을 높이고, 생성된 텍스트의 유창성 개선

---

Method

• Goal: zero-shot 설정에서 유창한 텍스트 생성을 유지하면서 목표 스타일에만 맞춘 어휘 생성을 위해 스타일별 뉴런 식별
• 이러한 목표를 달성하기 위해 activation value를 이용하여 스타일별 뉴런을 식별
• 그리고 source와 target 스타일 뉴런의 overlap을 피하는 것(to avoid)이 왜 중요한 지 설명
• 또 오직 source 스타일과 관련된 뉴런을 deactivated 했을 때 target 스타일과 관련된 단어의 생성 확률이 증가함을 확인
• 마지막으로 유창성을 향상시키기 위해 최근 연구인 contrastive decoding, Dola를 TST에 적용

 

1. Idenifying Style-Specific Neurons: 스타일 특화 뉴런 식별

• 뉴런 Neurons은 일반적으로 사람이 해석할 수 있는(interpretable) 개념에 매핑되는 feature extractors로 이해
• 뉴런은 단일 feature가 아니라 여러 개의 feature를 enconde하여 해석 가능성을 복잡하게 만들 수 있음
  • "However, neurons can exhibit polysemy, where a single neuron my encode multiple features, (...)"
• 의도하지 않은 변경 없이 LLM의 특정 기능을 선택적으로 수정하려면 모호하지 않은 뉴런을 식별하고 제거하는 것이 필수적

1.1. Neruons in LLMs: LLM에서 뉴런

• LLM을 이루는 FFN(Feed-Forward Network)에서 특정 뉴런을 activation or deactivation 하는 건 모델 출력에 중요한 영향
• 여기에 영감을 받아, 저자는 LLM의 FFN modules에서 스타일을 가리키는  식별을 목표로 함
• the activation values: $a^{(j)}=\text{act_fn}(W^{(j)}a^{(j-1)}+b^{(j)})$
  • $W$: weight, $b$: biases, $j$: layer, $a$: the activation value of previous layer, $\text{act_fn}$: activation function(GLU in llama)
  • 뉴런은 activation value가 0 이상일 때 active 되었다고 함

1.2. Neuron Selection: 뉴런 선택

• Tang et al. (2024)에서 언어 특화 뉴런을 식별하는 방법을 소개하며, 대략 25%의 뉴런이 언어별로 overlap됨을 확인
• 그러나 해당 연구에서는 overlap이 되는 들이 암시하는 성능에 관해서는 평가하지 않음
• 저자는 이전 연구의 방법을 적용하여 스타일 코퍼스를 통해 직접적으로 스타일 특화 뉴런들의 overlap을 조사

• Figure 2를 보면 상반되는 스타일 간 overlap 되는 뉴런의 비율이 높음(색칠된 그래프)을 확인 가능
• 특히 Politics에서는 "democratic"과 "republican" 스타일의 overlap 뉴런 비율이 95%
• 이런 높은 overlap이 성능에서 부정적인 영향을 미침을 확인 가능(Table 3)
  • Table 3에서 "with" 열이 overlap되는 뉴런을 삭제했을 때 style transfer accuracy 성능

• 서로 다른 스타일의 뉴런 간 overlap을 없애기 위해, 스타일 특화 뉴런과 그들 간의 교차점 식별
• $A$, $B$로 표시되는 두 개의 스타일이 있다고 가정했을 때, 두 스타일의 코퍼스를 LLM에 별도로 입력
• 그런 다음 activation value가 0을 초과하는 뉴런을 선택하여 $S_A$, $S_B$로 표시되는 뉴런 집합 구성
• 그 후, $S_A$, $S_B$ 내의 activation value를 내림차순으로 정렬하고 top-$k$($k=500n,n\in\{1,2,3,...,20\}$을 validation set으로) → $S'_A$, $S'_B$
• 마지막으로 이 집합에서 스타일 중 오직 하나의 스타일과 관련된 뉴런으로 다시 분리(disjoint)
  
  • $N_A=S'_A\setminus S'_B$, $N_B=S'_B\setminus S'_A$
• 실제 코드 상 구현

    def llama_forward(self, x):
        gate_up, _ = self.gate_up_proj(x)  # b, l, 2i
        i = gate_up.size(-1)
        gate_up[:, :, : i // 2] = torch.nn.SiLU()(gate_up[:, :, : i // 2])
        activation = gate_up[:, :, : i // 2].float() # b, l, i
        ...
        over_zero[idx, :] += (activation > 0).sum(dim=(0,1))
        x = gate_up[:, :, : i // 2] * gate_up[:, :, i // 2 :]
        x, _ = self.down_proj(x)
        return x

• 위의 코드를 보면, forward에서 활성화 함수(SiLU)를 통과한 후 각 레이어 마다(over_zero[idx, :]) 개별 뉴런이 activation value가 0 이상으로 활성화 되는 빈도수를 저장하는 것을 확인 가능

## 1. positive neurons
        positive_negative_difference = positive_activation['over_zero'] - negative_activation['over_zero']
        positive_negative_difference_sorted = torch.topk(positive_negative_difference.view(-1), args.threshold)
        positive_neurons = positive_negative_difference_sorted.indices.tolist()
        
        final_positive_neurons_list = [[] for _ in range(num_layers)]
        for pn in positive_neurons:
            row_index = pn // intermediate_size
            col_index = pn % intermediate_size
            final_positive_neurons_list[row_index].append(col_index)
        
        final_positive_neurons = []
        for fpn in final_positive_neurons_list:
            if fpn:
                tmp_list = sorted(fpn)
                final_positive_neurons.append(torch.tensor(tmp_list).long())
            else:
                final_positive_neurons.append(torch.tensor(fpn).long())

• 이렇게 저장한 정보는 지정한 threshold 만큼의 개수로 topk를 뽑음(많이 활성화된 뉴런만)
• 활성화된 뉴런은 $N_A=S'_A\setminus S'_B$, $N_B=S'_B\setminus S'_A$와 같이 서로 겹치는 뉴런 없이 처리하여 activation_mask를 생성하고, 이를 generation할 때 적용

for _, (pos, neg) in enumerate(zip(positive_neurons_list, negative_neurons_list)):
            tmp_inter = list(set(pos).intersection(set(neg)))
            # intersection_list.append(tmp_inter)
            final_inter_tensor.append(torch.tensor(tmp_inter).long())
            
            tmp_pos_non_inter = list(set(pos).difference(set(neg)))
            # positive_non_inter.append(tmp_pos_non_inter)
            pos_non_inter_tensor.append(torch.tensor(tmp_pos_non_inter).long())
            
            tmp_neg_non_inter = list(set(neg).difference(set(pos)))
            # negative_non_inter.append(tmp_neg_non_inter)
            neg_non_inter_tensor.append(torch.tensor(tmp_neg_non_inter).long())

2. Deactivating Source Style Neurons: 입력 스타일 뉴런 비활성화

• 특정 스타일과 관련된 뉴런을 식별한 후 일반적인 방법은 모델의 forward pass 중 activation value를 0으로 설정하여 deactivate
• 그러나 뉴런은 민감한 구성 요소이므로 특정 기능과 관련된 뉴런을 비활성화하면 성능이 크게 저하할 수도 있음
• TST task에서 source 및 target 스타일 뉴런을 비활성화하는 효과를 조사하기 위해 formality, politeness 대상 실험 수행

• Table 1을 통해 다음 2가지의 사실을 알 수 있음
  • target 스타일 뉴론을 유지하고 source 스타일 뉴런을 deactivating하는 건 target 스타일을 생성하는 정확도 향상
    • 반대로 target 스타일 뉴론을 deactivating 하면 target 스타일을 생성하는 정확도가 감소
    • 이는 target 스타일 뉴론을 deactivating 하면 decoding 중에 LLM이 target 스타일 단어를 생성하는 능력 감소
  • 다만 Fluency가 감소
    • 이는 주로 deactivating된 뉴론이 decoding 중 단어 분포에 미치는 영향이 상당하기 때문
    • 모델은 더 높은 확률로 deactivating되지 않은 스타일의 단어를 생성하는 경향이 있어 deactivating 되지 않은 스타일의 단어가 단순히 연결된 텍스트가 생성되어 유창성이 손상
    • Figure 1과 같이 source 스타일 뉴런을 deactivating한 후 생성된 텍스트에는 "neither"과 "quality"가 모두 포함
    • 이는 문장의 유창성을 유지하지 못하는 두 개의 target 스타일 단어

 

3. Contrastive Decoding for TST: TST에 Contrastive Decoding 적용

• LLM의 출력을 더 약하고 작은 모델과 비교하여 다음 단어를 예측할 확률을 조정하는 Contrastive Decoding(CD)은 유창성과 일관성(coherence)을 향상시키는 데 효과적
• 최근 Chuang et al. (2024)는 final layer와 early layer 간의 출력을 비교하여 우수한 결과를 얻는 Dola 제안
• 뉴런 deactivating 중 관찰되는 fluency 문제를 완화하기 위해 Dola를 TST에 적용

Figure 1

• 위의 예시에서도, source style(positive) neurons을 deactivating 함으로 target style(negative) words(neither, poor)가 생성이 되어 오히려 바른 문장이 생성이 되지 않음(input 문장의 negative 변환이 되려면 output 문장에서 poor → quality로 생성이 되어야 함)
• 이와 같은 정확성과 유창성 저하를 막기 위해 본 논문에서는 CD 적용

3.1. Dola

• background
  • token sequence $\{x_1, x_2, ..., x_{t-1}\}$와 LLM의 총 레이어 수 $N$이 주어지면 $j$번째 transformer layer의 다음 토큰 $x_t$ 확률을 미리 계산 가능(=early exit, Schuster et al., 2022)(*early exit: LLM의 추론 과정을 최적화하기 위한 기술)
    • $p^j(x_t | x_{<t}) = \text{softmax}(\phi (h_t^{(j)}))_{x_t}$
    • $h_t$는 hidden states로 embedding layer에서 얻을 수 있음
    • $\phi$는 vocabulary head로 token의 확률을 예측하는 데 사용
• goal: final layer와 set of early layers의 정보를 contrast하는 것
• why: LLM의 hallucinate(사실과 일치하지 않는 내용 생성) 문제 해결
• motivation: 초기 레이어의 "저수준 정보"와 후기 레이어의 "의미론적 정보 이용"
  • 참고 논문: What do you learn from context? Probing for sentence structure in contextualized word representations(2019)
• method
  • Next token prediction
    • 기존: 마지막 레이어의 출력에 vocabulary head 적용
    • DoLA: 여러 후보 레이어 각각에 대해 vocab head 적용
  • Premature layer 선택(마지막 레이어와 대조할 레이어 선택)
    • 각 후보 레이어의 출력 분포와 마지막 레이어의 출력 분포 간 JSD(Jensen-Shannon Divergence) 계산
    • JSD가 가장 큰 레이어를 선택하여 대조 디코딩

 

3.2. Our Adaptation to TST

• 이 논문에서는 위와 같은 대조 디코딩 방식 DoLA를 다음과 같이 변형하여 적용
• Candidate layer selection
  • DoLA는 일반적으로 초기 레이어를 사용
  • sNeuron-TST는 스타일 관련 정보가 후기 레이어에 집중되어 있음을 확인(Figure 3)

• Next token prediction
  • source style neuron 비활성화 후, 생성된 target style 토큰 확률 변화 분석를 분석하여 토큰을 두 가지로 분류
  • 1. style-independent words(function word): 초기 레이어부터 일정한 확률 가진 토큰
  • 2. target-style words: 후기 레이어에서 확률이 급격히 높아진 토큰
    → 이 경우, 후보 레이어 중 JSD 거리가 가장 큰 레이어를 선행 레이어로 선택하여 CD

 

---

Experiments

Datasets

• 6가지의 TST tasks에 대한 제안 방법 평가 수행
• 각 벤치마크는 두 가지의 스타일을 포함하며, 총 12가지의 TST 가능
• Formality, Politeness, Toxicity, Authorship, Politics, Sentiment
• 모든 데이터는 공개적으로 사용 가능
• 전처리: 120자를 초과하는 문장 제거 / 중복 문장 제거 / 많은 특수 기호를 포함하는 문장 제거

 

results

 

 

1. Overall Performance

• APE, AVF 및 PNMA는 다국어 기능을 향상하지만 content preservation metric을 제외하고는 TST에서는 성능 저하
• contrastive decoding을 하지 않은 baseline은 fluency 저하
• 전반적으로 overlap neurons을 제거하고, contrastive decoding을 사용하는 것이 정확성과 유창성에서 baseline을 능가

2. Content Preservation 

• LLaMA-3 및 baseline이 높은 content preservation을 보임
• 이는 copy problem → 생성된 텍스트는 원래 의미를 유지하려는 경향으로 스타일을 정하는 문체를 무시하는 경향

3.Different Directions

• 동일 스타일에서 서로 다른 방향으로 transfer 시, 상당한 성능 불일치
• 이는 LLM의 학습데이터가 주로 긍정적인 코퍼로 구성된 이유