Research map / Large language models / Instruction and fine-tuning
TAIA: Large Language Models are Out-of-Distribution Data Learners
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes TAIA, which fine-tunes all parameters but at inference keeps only the fine-tuned self-attention weights and reverts FFN to pretrained, improving LLM generalization when fine-tuning data is out-of-distribution.
Problem
Fine-tuning LLMs on domain-mismatched or out-of-distribution (OOD) data degrades downstream performance and causes catastrophic forgetting, which is a problem in data-scarce specialized domains.
Approach
Train all parameters (attention and FFN) during fine-tuning, but at inference discard the FFN updates and keep only the fine-tuned attention weights, reusing pretrained FFN. Attention updates carry instruction-following gains while FFN updates disrupt pretrained knowledge when data is OOD; training FFN still provides the non-linearity needed for learning.
Key results
- On Alpaca-GPT4 with LoRA, TAIA on Qwen1.5-7B averages 51.27 vs 46.50 for vanilla LoRA across seven test sets.
- With CoT-Collection, TAIA on Qwen1.5-1.8B reaches 35.54 average vs 27.68 for vanilla LoRA, and Llama3-8B TAIA improves over vanilla in MATH and SVAMP.
- On Advbench red-teaming for LLaMA2-7B-chat, TAIA lowers explicit-harm attack success from 84.59 (vanilla) to 8.27 while AlpacaEval rises from 7.46 to 9.94.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train all parameters, then infer with updated attention only instead of vanilla fine-tuning inference | Qwen1.5-1.8B LoRA · Alpaca-GPT4; seven-task average | average score: 32.98 -> 35.21 | Helped |
| Train all parameters, then infer with updated attention only instead of vanilla fine-tuning inference | Qwen1.5-7B LoRA · CoT-Collection; seven-task average | average score: 33.91 -> 48.57 | Helped |
| Train all parameters, then infer with updated attention only instead of vanilla fine-tuning inference | LLaMA3-8B LoRA · Alpaca-GPT4; seven-task average | average score: 51.99 -> 54.54 | Helped |
| Train all parameters, then infer with updated attention only instead of vanilla fine-tuning inference | LLaMA3-8B MoLoRA · Alpaca-GPT4; seven-task average | average score: 52.71 -> 54.32 | Helped |
| Train all parameters, then infer with updated attention only instead of vanilla fine-tuning inference | LLaMA3-8B LoRA · CoT-Collection; seven-task average | average score: 47.00 -> 51.74 | Helped |
| Train all parameters, then infer with updated attention only instead of vanilla fine-tuning inference | LLaMA3-8B MoLoRA · CoT-Collection; seven-task average | average score: 48.25 -> 53.24 | Helped |
| Retain only attention updates at inference rather than vanilla inference | Qwen1.5-1.8B LoRA · Medical Collection; six-task average | average score: 33.44 -> 44.37 | Helped |
| Train only attention parameters instead of training all parameters and keeping only attention updates at inference (ablation) | TAIA · Medical Collection; six-task average | average score: 44.37 -> 35.12 | Hurt |
| Train only FFN parameters instead of training all parameters and keeping only attention updates at inference (ablation) | TAIA · Medical Collection; six-task average | average score: 44.37 -> 33.86 | Hurt |
| Replace attention-focused training/inference with FFN-focused training/inference (ablation) | TAIA · Medical Collection; six-task average | average score: 44.37 -> 39.10 | Hurt |
| Train only attention parameters (ablation) | Qwen1.5-1.8B LoRA · OpenMath; six-task average | average score: 40.55 -> 38.88 | Hurt |
| Train only FFN parameters (ablation) | Qwen1.5-1.8B LoRA · OpenMath; six-task average | average score: 40.55 -> 38.88 | Hurt |
Surprising findings
- Fine-tuning only attention (TOA) performs worse than vanilla fine-tuning on OpenMath, indicating FFN training is needed for non-linear exploration.
- Under full fine-tuning, TAIA still beats vanilla but falls below the base model on knowledge-intensive tasks such as MMLU and MMedBench.
Limitations
- Underperforms vanilla fine-tuning when the downstream task requires learning new knowledge (e.g., SQuAD v2.0/XSum on Qwen1.5-1.8B).
- Under full fine-tuning, shows significant catastrophic forgetting on knowledge-intensive tasks, still below base model.
- The attention/FFN split is coarse; the authors note finer-grained parameter selection is needed.
Experimental setup
Datasets: Alpaca-GPT4 Alpaca-GPT4-zh CoT-Collection OpenMath MMedBench CMExam MedMCQA XSum SQuAD v2.0 ShareGPT-52K Anthropic red team Advbench
Benchmarks: MATH BBH CommonsenseQA LogiQA SVAMP MMedBench MMLU CMMLU C-Eval SQuAD v2.0 XSum AlpacaEval
Metrics: accuracy exact match average score attack success rate ROUGE win rate
Baselines: Vanilla fine-tuning L2 regularization EWC Self-Distill LoRACL TOA TOF TAIF Base model
Builds on: LoRA Mixture-of-LoRA RoPE Elastic Weight Consolidation
Closest papers by meaning
- CorDA: Context-Oriented Decomposition Adaptation of Large Language Models for Task-Aware Parameter-Efficient Fine-tuning (NeurIPS 2024 · similarity 0.91)
- Supervised Knowledge Makes Large Language Models Better In-context Learners (ICLR 2024 · similarity 0.90)
- Data Efficient Adaptation in Large Language Models via Continuous Low-Rank Fine-Tuning (NeurIPS 2025 · similarity 0.90)
- When Scaling Meets LLM Finetuning: The Effect of Data, Model and Finetuning Method (ICLR 2024 · similarity 0.89)
- Tuning LayerNorm in Attention: Towards Efficient Multi-Modal LLM Finetuning (ICLR 2024 · similarity 0.89)
- LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models (ICLR 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.