# SURGELLM: Rethinking Multi-Task Evaluation through Task-Aware Feature Gating with Class-Balanced Normalization

arXiv: https://arxiv.org/abs/2606.24259 · HTML: https://mlsanity.com/paper/2606.24259/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 13 of its 48 references that are in the map.
- First posted on arXiv 23 Jun 2026.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)

## Builds on (most cited first, 13 of 13)

- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Scaling Instruction-Finetuned Language Models](https://mlsanity.com/paper/2210.11416/) (2022)
- [Multi-Task Learning as Multi-Objective Optimization](https://mlsanity.com/paper/1810.04650/) (2018)
- [DetectGPT: Zero-Shot Machine-Generated Text Detection using Probability Curvature](https://mlsanity.com/paper/2301.11305/) (2023)
- [A Watermark for Large Language Models](https://mlsanity.com/paper/2301.10226/) (2023)
- [Multi-Task Learning with Deep Neural Networks: A Survey](https://mlsanity.com/paper/2009.09796/) (2020)
- [Efficiently Identifying Task Groupings for Multi-Task Learning](https://mlsanity.com/paper/2109.04617/) (2021)
- [Training Very Deep Networks](https://mlsanity.com/paper/1507.06228/) (2015)
- [Understanding and Improving Information Transfer in Multi-Task Learning](https://mlsanity.com/paper/2005.00944/) (2020)
- [Auto-Lambda: Disentangling Dynamic Task Relationships](https://mlsanity.com/paper/2202.03091/) (2022)

## Abstract (arXiv)

Fine-tuned encoders deployed across heterogeneous NLP tasks face three compounding problems: mismatched inductive biases, class-imbalance corruption of feature statistics, and no mechanism to condition attention on external lexical knowledge. We introduce surgellm, a unified transformer framework that addresses each with a dedicated lightweight module: a surgical feature gate (learned per-dimension sigmoid over curated lexical indicators and texttt[CLS]; provably degenerates to identity when features are uninformative), task-conditioned prefix tokens (quantized feature values and task identity prepended to every input), and Instance-Weighted Normalization (IWN; removes class-prior bias from gate statistics). We prove an excess-risk bound linking gate benefit to surgical feature alignment. Across four tasks, SST-2, multi-hop retrieval, LLM-prompt attribution, and authorship detection, covering 17,830 examples and eleven model variants over three seeds, the IWN variant achieves macro-F1 0.940 (+0.036 over the strongest non-IWN baseline; +0.130 on authorship detection). A random-vocabulary control (-0.028 avg. F1) confirms gains are lexical, not parametric. Code, vocabularies, and a 99.5%-recovery auto-extraction recipe are released.