OASIS Introduces Attention Abstention to Boost Long-Context and Low-Bit Robustness

Enhancing model expressiveness allows Transformers to read and integrate information more flexibly. Attention Residuals (AttnResidual) extend this direction by expanding attention from inter-token information exchange to network depth, enabling each sub-layer to learn which historical layer outputs to select.

However, this new selection capability introduces new challenges. In the AttnResidual settings evaluated by the research team, attention concentration on initial tokens became more pronounced, and activation outliers increased significantly. These changes further manifested as performance degradation in low-bit quantization and long-context evaluations.

A research team from Northwestern University, Illinois Institute of Technology, Rutgers University, and other institutions proposed OASIS. This approach provides an explicit null option for such dual-level information routing and adjusts depth-level branch selection based on token-level null probabilities. The study has been accepted by NeurIPS 2026.

Paper: https://arxiv.org/abs/2605.17887

Code: https://github.com/robinzixuan/OASIS

Compared to Vanilla AttnResidual, OASIS achieves an average reduction of 95.90% in activation kurtosis and 81.75% in maximum activation magnitude across LLaMA-3.2-1B, Qwen3-0.6B, and Phi-4. It also reduces W8A8 perplexity by 82.00% and improves W4A4 GSM8K Pass@1 by an average relative gain of 42.11%. Note that these improvements are relative ratios; absolute accuracy changes are detailed in the tables below.

From Token Selection to Cross-Layer Selection

Standard Transformer self-attention determines which tokens a given token should read. AttnResidual adds a selection mechanism along the depth dimension: for the same token at the current sub-layer, it decides which representations generated by previous sub-layers should be used.

In standard residual connections, the initial embedding and outputs from prior sub-layers are accumulated with fixed weights. AttnResidual replaces this accumulation with a learnable weighted mixture, feeding the mixed result into the current attention or MLP sub-layer. The depth query is a trainable vector for the target sub-layer, keys come from normalized historical outputs, and values are the historical outputs themselves.

Thus, the model performs two types of selection simultaneously: token-level selection chooses context positions, while depth-level selection chooses historical computation results. Both selections are normalized via softmax, where standard softmax requires the sum of candidate weights to equal 1.

If a single read only requires a small contribution, the model must still distribute all probability mass among real tokens. Similarly, if some historical branches provide little useful information, the depth layer must still allocate all mass among real branches. Both normalizations lack an explicit "select none" option.

Sinks and Anomalous Activations Under Dual Normalization

Attention sinks refer to certain positions that consistently absorb large amounts of attention, even when those positions carry limited information. Attention maps in the paper show that in the studied AttnResidual models, vertical bands of concentration on sequence-starting tokens are stronger than in original Transformers and persist across network depths.

Figure 1: Attention diagnostics from the paper. Top row: Original Transformer; Bottom row: AttnResidual. Layers 0, 9, and 14 are shown. Vertical bands on starting tokens reflect attention concentration. Source: Paper Figure 2.

Anomalous activations expand the numerical range, making it harder for limited quantization scales to accommodate most ordinary values. For example, in LLaMA-3.2-1B, Vanilla AttnResidual has an average activation kurtosis of 91.71 and a maximum activation magnitude of 84.43; OASIS reduces these to 3.02 and 12.25, respectively.

The paper also provides conditional mechanistic analysis: under specified separation conditions for value vectors, standard softmax may need to concentrate large probability mass on low-norm candidates to produce sufficiently small attention outputs. A similar conditional concentration constraint exists for depth-wise mixing. This motivates handling both token and depth normalizations simultaneously.

Two Null Routing Mechanisms Plus One Cross-Layer Coordination

OASIS first adopts Softmax₁, adding a zero-value null option for real candidates. For real candidate i, its weight is calculated as follows:

The 1 in the denominator corresponds to a null candidate with a fixed logit of 0. The weights of real candidates can sum to less than 1, with the remaining mass assigned to the null option, thereby reducing the total contribution of real content. The version using Softmax₁ at both token and depth levels is referred to as AoS in the paper.

However, having independent null options at both layers does not resolve how the two layers coordinate. If attention heads in one branch have already selected the null option frequently, the depth routing might still assign high weight to that branch. Therefore, OASIS introduces token-to-depth null coupling.

Specifically, at each token position, OASIS averages the null probabilities across attention heads within a source branch to obtain a branch statistic $\psi_i$. This is compared against the average of candidate transformation branches, and the difference is used to adjust the depth logits:

If the null probability is higher than the branch average, the depth logit is decreased; if lower, it is increased accordingly. The initial embedding branch lacks token attention null statistics, so its logit remains unchanged. The adjusted logits then pass through depth Softmax₁ to yield the final mixing weights.

Figure 2: The two-level routing structure of OASIS. Both Softmax₁ instances provide a null option, and the coupling on the right injects token-level null statistics into depth logits. Source: Paper Figure 1.

OASIS retains the original training objectives, modifying only the routing mechanism. The token-level null corresponds to a zero-value contribution, while the depth-level null attenuates the mixed contribution from actual source branches; the depth null itself does not imply directly preserving the previous layer's hidden state. Experiments adapt these changes through continued training, followed by quantization evaluation.

Component ablation directly validates the contribution of coupling

The team validated the method on LLaMA-3.2-1B, Qwen3-0.6B, and Phi-4, using BookCorpus and Wiki40B/en for continued training, and compared it against standard AttnResidual, OutEffHop, pruning, and gating schemes. The experiments employed three independent random seeds.

First, consider the W4A4 GSM8K ablation for Qwen3-0.6B. Adding only the token null, adding only the depth null, adding both independent nulls, and adding coupling yielded the following results respectively:

Table 1: Component ablation for Qwen3-0.6B. Higher GSM8K Pass@1 is better. Data from Paper Table 8.

The two independent nulls already improved performance, and coupling further increased it by 4.67 percentage points on top of that. Full OASIS improved by 13.21 percentage points over Vanilla AttnResidual. These results directly support the role of cross-layer coordination.

Preserving more model capability under low-bit quantization

In language modeling and mathematical reasoning, OASIS also improved post-quantization performance. W8A8 indicates that both weights and activations use 8-bit precision, while W4A4 uses 4-bit for both; the W4A4 results in the paper employ OmniQuant. The table below excerpts the averages from Paper Table 2.

Table 2: Quantization results for three models. Lower PPL is better; higher GSM8K Pass@1 is better.

Taking Qwen3-0.6B as an example, OASIS achieves 62.68% on GSM8K in FP16, retaining 58.88% after quantizing to W4A4; Vanilla AttnResidual drops from 62.58% in FP16 to 45.67%. While their full-precision performances are similar, the retention of capability under low-bit settings differs significantly.

These improvements also appear in other tasks. For Qwen3-0.6B, ThePile W8A8 PPL decreased from 44.52 to 11.58; W4A4 accuracy on MMLU, TriviaQA, and MATH500 improved from 33.21%, 12.09%, and 15.41% to 41.00%, 16.14%, and 21.51%, respectively.

Recovering degraded performance in 12K long contexts

Long-context experiments compared models trained with a maximum sequence length of 12K. On RULER and LongBench, Vanilla AttnResidual performed significantly lower than the original backbone. OASIS mitigated this degradation and clearly outperformed AoS, which features only independent null routing.

Table 3: Long-context results for 12K tokens. Values are raw scores from the paper; higher is better. Data from Paper Table 3.

Relative to the gap between Vanilla AttnResidual and the original backbone, OASIS recovered 77.3% and 80.3% of the RULER performance loss, and 89.1% and 91.8% of the LongBench performance loss for the two models, respectively.

Global attention statistics also supported the mitigation of attention sinks. These statistics cover all layers, attention heads, query positions, and evaluation samples: the average mass assigned to BOS by Qwen3-0.6B dropped from 0.4990 to 0.0013, and for LLaMA-3.2-1B it dropped from 0.5630 to 0.0099, indicating that the change is not limited to a few visualization examples.

Stabilizing the newly added routing capabilities

Starting from the dual normalization of AttnResidual, OASIS connects expressiveness, information routing, and low-bit robustness. After the model gains a more flexible selection interface, it still needs to express how much candidate content should contribute and how different selection levels coordinate.

Explicit null routing provides an interface to attenuate real content contributions, while token-to-depth coupling allows cross-layer selections to utilize token-level null signals.

Component ablations, quantization evaluations, and long-context experiments collectively demonstrate that this coordination helps the studied models stabilize the utilization of their newly added routing capabilities.

References:

https://arxiv.org/abs/2605.17887

Chasing ASI Instantly

⭐ Like, Share, and "Watching" – One Click Combo ⭐

Star us to lock in instant updates from New Intelligence Era!