SURGELLM: Rethinking Multi-Task Evaluation through Task-Aware Feature Gating with Class-Balanced Normalization | Noor Islam S. Mohammad, Uluğ Bayazıt | ResearchPod

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

Alex: Welcome to another episode of ResearchPod. Today, we're looking at why modern language models struggle with tasks that seem straightforward to humans, like identifying who wrote a specific text.

Sam: That's right. The paper introduces a framework called SURGELLM. It addresses a specific problem: why AI models that have been trained and fine-tuned on multiple tasks at once tend to miss surface-level clues in language — things like a writer's characteristic word choices or punctuation habits.

Alex: Is the core problem that these models are just too generic, or is there a deeper technical issue?

Sam: It's both, really. Imagine you're trying to learn to recognize different people's handwriting, but ninety percent of your practice examples are from the same person. You'd get very good at recognizing that one person, but you'd develop blind spots for everyone else. The paper describes this as "prior-contaminated" feature statistics — the model's internal sense of what's normal gets skewed by whichever type of example shows up most often, and it starts ignoring the subtle details that would help it distinguish the minority cases.

Alex: So it's like a camera lens automatically adjusting its exposure to the brightest part of a scene — and in doing so, losing all the detail in the shadows?

Sam: That's a good way to put it. To fix this, the authors introduce something called Instance-Weighted Normalization, or IWN. Instead of letting the most common category of example set the standard for what the model pays attention to, IWN recalculates that standard in a balanced way — giving equal weight to every category. It's a bit like a teacher who deliberately makes sure quieter students get called on just as often as the ones who always raise their hands.

Alex: And this is where the "surgical" part of the name comes in? How does the model actually use these more balanced features?

Sam: So once the model has a fairer internal picture, the authors add what they call a Surgical Feature Gate. Think of it as a small, learned filter sitting at the end of the model's process. On one side, you have the model's broad, general understanding of language. On the other, you have a compact set of very specific surface-level signals — things like how often a writer uses semicolons, or their tendency toward long sentences. The gate's job is to decide, for each of those specific signals independently, how much it should trust that signal versus its general training. It can say, "For punctuation habits, I'll lean heavily on the specific data. For overall topic, I'll trust my general training."

Alex: So it's not an all-or-nothing switch — it's more like a mixing board, where each channel has its own volume?

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

Noor Islam S. Mohammad, Uluğ Bayazıt

Jun 23, 2026·4:31·arxiv·Original Paper

SURGELLMSurgical Feature GateInstance-Weighted Normalization (IWN)Task-conditioned prefix tokensMulti-task LearningSurgical Feature Extraction

Listen

5 min

Summary

Research Question

Fine-tuned transformer encoders often struggle when deployed across heterogeneous NLP tasks due to mismatched inductive biases, class-imbalance corruption of feature statistics, and a lack of mechanisms to incorporate external lexical knowledge. This paper investigates whether a unified encoder can effectively handle diverse tasks—ranging from sentiment analysis to authorship detection—by explicitly injecting task-aware lexical features and correcting for class-prior bias.

Approach

The authors introduce SURGELLM, a framework that integrates three lightweight modules into any standard transformer encoder:

Surgical Feature Gate: A learned per-dimension sigmoid gate that fuses the [CLS] representation with a 16-dimensional vector of curated lexical indicators.
Task-Conditioned Prefix Tokens: Quantized feature values and task identifiers prepended to the input, allowing the model to condition its attention on task-specific lexical context.
Instance-Weighted Normalization (IWN): A technique that replaces standard global feature normalization with class-balanced statistics, preventing the gate from being biased by majority-class feature profiles.

The framework is evaluated across four tasks (SST-2, multi-hop retrieval, LLM-prompt attribution, and authorship detection) using eleven model variants, including BERT, RoBERTa, and T5.

Main Findings

SURGELLM-IWN-RoBERTa achieves an aggregate macro-F1 of 0.940, outperforming the strongest non-IWN baseline by 0.036. The most dramatic improvement occurs in authorship detection, where IWN corrects a significant class-imbalance bias, resulting in a 0.130 F1 gain. A random-vocabulary control confirms that these performance improvements are driven by the injection of meaningful lexical signals rather than additional parameters. The authors also provide a theoretical excess-risk bound linking the gate's performance to the alignment of these surgical features.

Why It Matters

This work demonstrates that even with powerful pre-trained encoders, explicit feature injection and statistical correction for class imbalance remain critical for multi-task performance. By providing a plug-and-play framework that is compatible with existing HuggingFace encoders, the authors offer a practical solution for practitioners dealing with heterogeneous, imbalanced datasets where standard fine-tuning fails to capture necessary surface-level linguistic patterns.

Key Terms and Definitions

Surgical Vocabulary — A set of curated lexical indicator groups used to compute 16-dimensional feature vectors that capture task-specific surface statistics.
Instance-Weighted Normalization (IWN) — A normalization technique that uses class-balanced mean and variance statistics to standardize input features, removing bias from majority-class priors.
Surgical Feature Gate — A learned mechanism that uses a sigmoid-based interpolation to fuse contextual encoder representations with external lexical features.
Task-Conditioned Prefix Tokens — A sequence of tokens prepended to the input that encodes task identity and feature counts, allowing the transformer to attend to these signals during processing.
Excess-Risk Bound — A theoretical result provided in the paper that links the benefit of the feature gate to the alignment between the surgical features and the task labels.

Abstract

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 \textbf{\surgellm}, a unified transformer framework that addresses each with a dedicated lightweight module: a \emph{surgical feature gate} (learned per-dimension sigmoid over curated lexical indicators and \texttt{[CLS]}; provably degenerates to identity when features are uninformative), \emph{task-conditioned prefix tokens} (quantized feature values and task identity prepended to every input), and \emph{Instance-Weighted Normalization} (IWN; removes class-prior bias from gate statistics). We prove an excess-risk bound linking gate benefit to \emph{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 \textbf{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.

SURGELLM: Rethinking Multi-Task Evaluation through Task-Aware Feature Gating with Class-Balanced Normalization: Research Paper Podcast