* Enhancing Instruction-Following Capabilities in Seq2Seq Models (ArXiv 2025)

https://www.dropbox.com/scl/fi/oqbatf064r82nh0k30uz6/arxiv25_Steering_Seq2Seq_Beyond_Memorization.pdf?rlkey=4g8t4f86x12nsbrzhlqdsbpe0&dl=0

이 논문의 핵심은 **“decoder-only LLM에서 효과적이었던 DoLa를 FLAN-T5 같은 encoder–decoder(Seq2Seq) 모델에 적용하면 왜 잘 작동하지 않는가?”**를 분석하고, 그 분석을 바탕으로 gradient 기반 activation steering을 제안하는 것입니다.

논문의 가장 중요한 결과는 FLAN-T5-Large의 MemoTrap 정확도를 52.0% → 99.7%로 향상시켰다는 것입니다.


1. 연구 질문과 핵심 아이디어

FLAN-T5는 instruction tuning을 했음에도 다음과 같은 상황에서 instruction을 무시할 수 있습니다.

Instruction과 pretraining에서 학습한 memorized continuation이 충돌하는 경우

예를 들어 유명한 속담의 끝을 일부러 일반적인 끝말이 아닌 다른 단어로 완성하라고 하면, 모델이 instruction보다 암기된 속담을 출력하는 MemoTrap 문제가 발생합니다.

논문의 연구 흐름은 두 단계입니다.

Step 1. DoLa를 FLAN-T5에 적용
        ↓
왜 잘 되거나 실패하는지 layer-wise 분석
        ↓
T5 intermediate logits가 매우 불안정하다는 사실 발견

Step 2. logit-level contrast 대신
        representation-level intervention

        Gradient-Based Activation Steering

즉, DoLa 자체를 개선하는 논문이라기보다 DoLa를 diagnostic probe로 사용하여 Seq2Seq decoder의 내부 동역학을 분석한 뒤 activation steering을 제안하는 논문으로 보는 것이 정확합니다. 


2. 왜 T5에서 DoLa가 문제가 되는가?

먼저 이 부분을 이해해야 논문의 방법론이 명확해집니다.

2.1 원래 DoLa

DoLa(Decoding by Contrasting Layers)의 기본 가정은:

Transformer의 깊이가 증가하면서 next-token representation이 점진적으로 성숙한다.

예를 들어 decoder-only 모델이라면

h(4)h(8)h(12)h(16)h(N)h^{(4)}\rightarrow h^{(8)}\rightarrow h^{(12)} \rightarrow h^{(16)} \rightarrow h^{(N)}

으로 갈수록 prediction이 refine된다고 볼 수 있습니다.

각 layer hidden state를 LM head에 넣어서

qj(x)=softmax(WLMhj)q_j(x)=\operatorname{softmax}(W_{\mathrm{LM}}h_j)

라는 중간 prediction을 얻습니다.

최종 layer distribution을 qNq_N 이라고 하면, 후보 intermediate layer 가운데 final distribution과 가장 다른 layer를 찾습니다.

M=argmaxjJJSD(qN,qj)M=\arg\max_{j\in J}\operatorname{JSD}(q_N,q_j)

여기서 JSD는 Jensen-Shannon divergence입니다.

그리고

~=N+λ(NM)\tilde{\ell}=\ell_N+\lambda(\ell_N-\ell_M)

로 decoding합니다.

이를 직관적으로 쓰면

Nfinal prediction+λ(NM)mature information\underbrace{\ell_N}_{\text{final prediction}}+\lambda\underbrace{(\ell_N-\ell_M)}_{\text{mature information}}

입니다.

즉 premature layer에도 이미 강했던 prediction은 상대적으로 억제하고, late layer에서 새롭게 강해진 prediction을 강조합니다.


3. T5에 DoLa를 어떻게 적용했는가?

여기서 decoder-only와 T5의 중요한 차이가 있습니다.

T5는 EncoderHE\text{Encoder}\rightarrow H_E 를 먼저 만든 후, 모든 decoder layer가 encoder representation에 cross-attention합니다.

즉,

hlSelfAttnCrossAttn(HE)FFNh_l\rightarrow\text{SelfAttn}\rightarrow\text{CrossAttn}(H_E)\rightarrow\text{FFN}

구조가 layer마다 반복됩니다.

따라서 decoder-only 모델과 달리

h1h2hNh_1\rightarrow h_2\rightarrow\cdots\rightarrow h_N

이 단순한 progressive refinement라고 보기 어렵습니다. 매 layer에서 encoder 정보가 다시 들어오기 때문입니다. 논문도 이것이 intermediate representation을 non-monotonic하게 변화시킬 수 있다고 지적합니다.

그래도 분석을 위해 각 decoder layer residual stream에 shared LM head를 적용합니다.

qj=softmax(WLMhj)q_j=\operatorname{softmax}(W_{\mathrm{LM}}h_j)

그리고 decoder layer만 premature candidate로 사용합니다.

후보 layer는 단순하게 even-indexed layers로 정합니다.

ModelDecoder layersDoLa candidate
FLAN-T5-Small80,2,4,6
Base120,2,4,6,8,10
Large / XL240,2,4,…,22

4. DoLa 실험 결과

IFEval에서 먼저 실험합니다.

전체 결과

ModelBaselineDoLa
FLAN-T5-Small20.3822.90
FLAN-T5-Base26.3823.02
FLAN-T5-Large26.8626.74
FLAN-T5-XL28.7825.30

DoLa는 T5에서 일관되게 좋아지지 않습니다.

Small에서는 +2.52%p지만 Base와 XL에서는 오히려 성능이 떨어집니다.

그런데 category별로 분석하면 흥미로운 현상이 나타납니다.


5. 왜 어떤 instruction에서는 DoLa가 성공하는가?

논문은 크게 두 가지 사례를 분석합니다.

Case A. Keyword constraint — DoLa 성공

예를 들어 반드시

“lacking”

이라는 단어를 포함해야 하는 instruction입니다.

논문의 Figure 5에서 이 token의 layer-wise rank를 추적합니다.

대략적으로

intermediate layers

180
 ↓
108
 ↓
164
 ↓
1156
 ↓
3350
 ↓
3992
 ↓
4832
 ↓
5585
 ...

final layer
5

처럼 엄청나게 흔들립니다. 

핵심은 instruction token "lacking"late layer에서 갑자기 강해진다는 것입니다.

따라서

N(lacking)M(lacking)0\ell_N(\text{lacking})-\ell_M(\text{lacking})\gg 0

가 됩니다.

DoLa는 이 차이를 증폭하므로

~(lacking)=N(lacking)+λ[N(lacking)M(lacking)]\tilde{\ell}(\text{lacking})=\ell_N(\text{lacking})+\lambda[\ell_N(\text{lacking})-\ell_M(\text{lacking})]

결과적으로 "lacking"rank 5 → rank 1로 올라갑니다.

즉,

instruction signal이 late-emerging signal일 때 DoLa가 잘 작동합니다.


6. 반대로 DoLa가 instruction을 망가뜨리는 경우

Start-End constraint에서는 반대입니다.

예:

“… End the blog post with ‘Naomi thanks you for reading.’ No other words should follow…”

Baseline은

Naomi thanks you for reading.

만 출력합니다.

IFEval metric상 instruction을 정확히 만족합니다.

하지만 DoLa는 훨씬 자연스러운 blog post를 생성하면서 마지막에 요구된 문구를 정확히 지키지 않습니다.

왜 그럴까요?

이 경우 instruction signal은 이미 intermediate layer에도 존재합니다.

즉, M(instruction)N(instruction)\ell_M(\text{instruction})\approx\ell_N(\text{instruction}) 입니다.

따라서, NM\ell_N-\ell_M 에서 instruction component는 상쇄됩니다.

반면 late layer에서 새롭게 생긴 natural language generation signal

N(natural)M(natural)0\ell_N(\text{natural})-\ell_M(\text{natural})\gg0

입니다.

결과적으로 DoLa가 오히려

instruction compliance가 아니라 natural continuation을 증폭

시킵니다.

이 분석이 논문의 activation steering으로 이어지는 가장 중요한 동기입니다.


7. 논문의 핵심 발견: T5의 intermediate logits는 volatile하다

저자들의 결론은 다음과 같습니다.

Decoder-only에서는 대략

immatureless immaturemature\text{immature}\rightarrow\text{less immature}\rightarrow\text{mature}

라는 progression을 기대할 수 있습니다.

그러나 T5 decoder에서는

hlcross-attentionhl+1cross-attentionhl+2h_l\overset{\text{cross-attention}}{\longrightarrow}h_{l+1}\overset{\text{cross-attention}}{\longrightarrow}h_{l+2}

과정에서 encoder representation이 계속 유입됩니다.

그래서 instruction token probability가

p2,p4,p6,p8,p10p_2 \uparrow,\quad p_4\downarrow,\quad p_6\uparrow,\quad p_8\downarrow\downarrow,\quad p_{10}\uparrow

처럼 non-monotonic하게 변합니다.

실제로 late layer로 갈수록 JSD도 크게 증가합니다.

따라서 T5에서는

intermediate LM-head projection = “모델이 그 시점에서 믿고 있는 next token”

이라는 해석 자체가 decoder-only보다 훨씬 약합니다.

이것이 이 논문의 mechanistic observation 중 가장 중요한 부분입니다.


8. Gradient-Based Activation Steering

이제 논문의 핵심 제안 방법입니다.

DoLa의 문제는

logitNlogitM\text{logit}_N-\text{logit}_M

이라는 output-space contrast에 의존한다는 것입니다.

저자들은 아예 hidden representation을 직접 움직입니다.


9. Contrastive target 구성

MemoTrap example마다 두 개의 token을 정의합니다.

x+=instruction-following tokenx^+=\text{instruction-following token}

x=memorized/default tokenx^-=\text{memorized/default token}

즉 모델 내부에는

instruction
     ↓
   x+

vs.

memorized prior
     ↓
   x-

라는 경쟁이 존재한다고 봅니다.


10. Contrastive loss

논문은 약간 흥미로운 방향으로 loss를 정의합니다.

L=logP(x)logP(x+)L=\log P(x^-)-\log P(x^+)

즉 loss가 커지는 방향은

x,x+x^- \uparrow,\qquad x^+\downarrow

입니다.

따라서 decoder layer l의 hidden state에 대한 gradient gl=hlLg_l=\nabla_{h_l}L

memorized answer 쪽으로 representation을 움직이는 방향

입니다.

그렇다면 instruction-following 방향은 반대입니다.

gl-g_l 입니다.

이 부분이 방법론의 핵심입니다.


11. Steering vector mining

하나의 example에서 얻은 gradient를 그대로 사용하지 않습니다.

100개 MemoTrap example을 사용합니다.

각 example i에 대해

gl(i)=hl[logP(xi)logP(xi+)]g_l^{(i)}=\nabla_{h_l}[\log P(x_i^-)-\log P(x_i^+)]

를 계산합니다.

그리고 평균을 냅니다.

vtask=1Ni=1NhlLi\boxed{v_{\text{task}}=-\frac1N \sum_{i=1}^{N}\nabla_{h_l}L_i}

여기서 N=100입니다.

이 평균이 중요한 이유가 있습니다.

각 example은 서로 다른 proverb이고 서로 다른 target token을 갖습니다.

따라서 개별 token semantics는 averaging 과정에서 상당 부분 상쇄되고,

vtaskshared instruction-compliance directionv_{\text{task}}\approx\text{shared instruction-compliance direction}

이 남을 것이라는 것이 저자들의 가설입니다.

즉, 이 vector를 특정 단어의 steering vector가 아니라

“memorization보다 instruction을 따르는 방향”

으로 해석합니다.


12. Inference-time intervention

얻은 vector를 normalize합니다.

v^=vtaskvtask\hat v=\frac{v_{\text{task}}}{\|v_{\text{task}}\|}

그리고 inference에서 decoder hidden state를

hl=hl+αv^\boxed{h_l’=h_l+\alpha\hat v}

로 변경합니다

논문에서는 여러 α\alpha를 실험하고, α=1000\alpha=1000 이 안정적으로 좋은 결과를 보였다고 보고합니다.

Figure 6에서는 주로 α{500,1000,1500}\alpha\in\{500,1000,1500\}를 비교합니다.


13. 이 방법과 일반적인 CAA steering의 차이

일반적인 Contrastive Activation Addition은 보통

v=E[h(x+)]E[h(x)]v=E[h(x^+)]-E[h(x^-)]

처럼 activation difference를 사용합니다.

반면 이 논문은

v=E[h(logP(x)logP(x+))]v=-E[\nabla_h(\log P(x^-)-\log P(x^+))]

입니다.

즉,

CAA

representation difference\text{representation difference} 를 찾는다면,

이 논문

representation을 어느 방향으로 움직이면 output preference가 바뀌는가?\text{representation을 어느 방향으로 움직이면 output preference가 바뀌는가?}

를 gradient로 찾습니다.

따라서 후자는 output behavior에 대해 좀 더 직접적인 local causal/control direction이라고 해석할 수 있습니다.


14. MemoTrap 실험 설정

MemoTrap의 Proverb Ending subset을 사용합니다.

  • 100 examples → steering vector mining
  • 별도의 300 examples → free-generation evaluation

즉 steering vector를 만드는 데이터와 test data를 분리했습니다.

모델은

  • FLAN-T5-Small
  • FLAN-T5-Base
  • FLAN-T5-Large

입니다.


15. 핵심 실험 결과

결과가 상당히 큽니다.

ModelBaselineBest SteeringImprovement
Small60.3%92.3%+32.0%p
Base65.7%98.7%+33.0%p
Large52.0%99.7%+47.7%p

특히 Large:

52.0%99.7%52.0\%\rightarrow 99.7\%

즉 거의 benchmark를 해결합니다.

흥미롭게도 model size가 커졌다고 baseline MemoTrap 성능이 좋아지는 것도 아닙니다.

Base=65.7%>Large=52.0%\text{Base}=65.7\%>\text{Large}=52.0\%

인데, 이는 memorization prior가 강한 larger model에서 instruction conflict가 더 심해질 가능성과 연결해서 볼 수 있습니다. 다만 이 인과 해석 자체를 논문이 직접 입증한 것은 아닙니다.


16. 가장 중요한 실험: 어느 layer에 steering해야 하는가?

이 논문의 mechanistic finding에서 상당히 흥미로운 결과입니다.

최적 layer:

ModelDecoder depthBest steering layer
Small82
Base125
Large2410

즉, 대략

전체 decoder depth의 40% 전후\boxed{\text{전체 decoder depth의 40\% 전후}}

에서 최고 성능을 냅니다.


17. 왜 mid-layer인가?

논문의 설명은 representation의 informativeness와 malleability의 trade-off입니다.

Early layer

representation이 아직 충분히 semantic하지 않습니다.

low information+high flexibility\text{low information}+\text{high flexibility}

따라서 vector를 넣어도 downstream layer에서 cross-attention/FFN 등에 의해 희석될 수 있습니다.

Mid layer

high enough information+high enough flexibility\boxed{\text{high enough information}+\text{high enough flexibility}}

instruction과 memorization의 경쟁 구조가 형성되어 있지만 아직 output trajectory가 고정되지 않았습니다.

따라서 intervention 효과가 가장 큽니다.

Late layer

흥미롭게도 token logits는 여전히 크게 변하고 있습니다.

하지만 representation 자체는 이미 특정 generation trajectory로 committed된 상태라고 저자들은 해석합니다.

high information+low flexibility\text{high information}+\text{low flexibility}

그래서 steering 효과가 다시 감소합니다.

이것을 논문은 사실상

volatile predictions ≠ flexible representations

라고 구별합니다.

이 관찰이 꽤 중요합니다.


18. 전체 방법을 수식으로 정리하면

이 논문의 gradient steering은 사실 상당히 간단합니다.

Training/mining 단계

각 contrastive example:

(xi,xi+,xi)(x_i,x_i^+,x_i^-)

에 대해

Li=logP(xi|xi)logP(xi+|xi)L_i=\log P(x_i^-|x_i)-\log P(x_i^+|x_i)

decoder layer l:

gi(l)=hlLig_i^{(l)}=\nabla_{h_l}L_i

instruction direction:

vl=1Nigi(l)v_l=-\frac1N\sum_i g_i^{(l)}

normalize:

v^l=vlvl\hat v_l=\frac{v_l}{\|v_l\|}

Inference

hl=hl+αv^l\boxed{h_l’=h_l+\alpha\hat v_l}

입니다.

따라서 모델 parameter를 업데이트하지 않는 inference-time intervention입니다.


19. 논문의 전체적인 해석

논문의 결과를 하나의 그림으로 요약하면 다음과 같습니다.

              FLAN-T5 decoder depth
──────────────────────────────────────────▶

Early                 Mid                  Late
│                      │                     │
│ representation       │ instruction vs      │ output trajectory
│ forming              │ memorization        │ largely committed
│                      │ distinction         │
│                      │ emerges             │
│
▼                      ▼                     ▼

Low semantic          Informative +          Informative but
information           malleable              less malleable

      weak steering    ★ BEST STEERING ★     weak steering

그리고 DoLa 관점에서는

decoder-only LLM

premature ───────────────▶ mature
             비교적 안정적 progression


FLAN-T5

layer 2 ↑
layer 4 ↓↓
layer 6 ↑
layer 8 ↓
layer 10 ↑↑
layer 12 ↓
...
final ↑↑↑

cross-attention 때문에
intermediate token distribution이 volatile

하므로

final logitspremature logits\text{final logits}-\text{premature logits}

에 의존하는 DoLa가 안정적이지 않습니다.


20. 논문의 장점과 한계

논문의 가장 큰 장점은 DoLa의 실패를 단순한 negative result로 끝내지 않고 activation dynamics 분석 → steering 위치 선정 → intervention으로 연결했다는 점입니다. 특히 “middle layers are informative but still malleable”이라는 결과는 activation steering 및 mechanistic interpretability 관점에서 흥미롭습니다. 저자들도 최종적으로 mid-depth activations를 중요한 locus of controllability로 해석합니다.

반면 실험적 한계도 상당히 명확합니다. MemoTrap은 instruction-following 전체를 대표하기보다 memorized continuation과 explicit instruction이 단일 token 수준에서 충돌하는 매우 깨끗한 설정입니다. steering direction 역시 task-specific이고 다른 instruction type으로의 generalization은 검증되지 않았으며, multi-token/sequence-level constraint도 다루지 않습니다. 저자들도 이 점들을 limitation으로 명시합니다.

특히 연구 관점에서는 다음 문제가 중요합니다.

vl=E[hlL]v_l=-E[\nabla_{h_l}L]

이 정말 general instruction-following feature인지, 아니면 단순히

proverb completion suppression\text{proverb completion suppression}

혹은

xx+x^- \rightarrow x^+

에 특화된 평균 gradient인지는 현재 실험만으로 완전히 분리되지 않습니다.

따라서 52→99.7%라는 숫자는 매우 인상적이지만 “general instruction-following이 거의 해결되었다”라고 해석해서는 안 됩니다.


이 논문의 핵심 contribution을 3개로 압축하면

  1. DoLa를 FLAN-T5 Seq2Seq 구조에 적용하고, decoder-only와 달리 intermediate LM-head projections이 매우 volatile하다는 것을 관찰했습니다.
  2. 이를 통해 Seq2Seq에서는 early-vs-final logit contrast가 instruction signal을 안정적으로 분리하지 못한다는 것을 보였습니다.
  3. 대신

vl=E[hl(logP(x)logP(x+))]\boxed{v_l=-E[\nabla_{h_l}(\log P(x^-)-\log P(x^+))]}

라는 gradient-based instruction-compliance steering direction을 만들고 mid-decoder에 주입하여 MemoTrap에서 최대 **52.0% → 99.7%**를 달성했습니다.


게시됨

카테고리

작성자

댓글

답글 남기기

이메일 주소는 공개되지 않습니다. 필수 필드는 *로 표시됩니다