SToRI: Semantic Token Reweighting Explained

• Semantic Token Reweighting (SToRI) is a framework that assigns explicit weights to tokens, enhancing interpretability, sample efficiency, and attribute control in neural models.
• It integrates seamlessly into existing pipelines by modulating loss functions, attention scores, and token distributions with minimal architectural changes.
• Empirical results show SToRI achieves significant gains in applications like medical vision-language modeling and controlled generation, with improvements up to 10× in low-data regimes.

Semantic Token Reweighting (SToRI) encompasses a collection of methods for modulating the influence of individual tokens or token types in deep neural architectures, particularly in context-sensitive natural language and multimodal processing. SToRI frameworks have been developed to introduce explicit, semantically motivated importance weights at the granularity of subwords or phrases; this reweighting can be data-driven, rule-based, or controller-driven (e.g., guided by external classifiers or user preferences). The principal objective is to boost sample efficiency, enforce attribute-level constraints, or provide interpretability and controllability with minimal architectural change. SToRI has been applied in domains ranging from medical vision-language modeling and interpretable representation learning to reinforcement learning with verifiable rewards and attribute-controlled generation.

1. Mathematical Foundations and Core Variants

The unifying principle behind SToRI is the explicit assignment of per-token weights (or, more generally, mask tensors) that modulate loss contributions, generation probabilities, or attention weights. Let $x = (x_1, \ldots, x_T)$ be the target token sequence.

• Loss Reweighting: Instead of minimizing the uniform cross-entropy

$$L_{ce}(x) = -\sum_{t=1}^T \log p_\theta(x_t | x_{<t}, I),$$

SToRI introduces token weights $\lambda_t$:

$$L_{SToRI}(x;\theta) = -\frac{1}{\Lambda} \sum_{t=1}^T \lambda_t \log p_\theta(x_t | x_{<t}, I),\quad \Lambda = \sum_{t=1}^T \lambda_t,$$

with $\lambda_t$ chosen by matching to semantically salient tokens (e.g., clinical keywords) and $\gamma>1$ an upweight factor for critical terms (Weers et al., 22 Apr 2026).

• Self-Attention Reweighting: In text encoders such as CLIP, SToRI multiplies each key’s attention scores by a learned or user-specified $w_n \geq 0$ within the self-attention softmax for token $n$:

$$\hat{A} = \mathrm{softmax}(Q K^\top + \mathrm{diag}(\log w))$$

This enables explicit control over the influence of each semantic element on the final token or [EOS] embedding (Kim et al., 2024).

• Probabilistic Token Reweighting in Controlled Generation: For constraint-driven generation, SToRI forms a modified token distribution

$$q(w|x_{<t}) \propto p_{LM}(w|x_{<t}) \cdot \exp\left(\lambda\, \nabla_e V(e_{<t})^\top e_w \right)$$

where $$L_{ce}(x) = -\sum_{t=1}^T \log p_\theta(x_t | x_{<t}, I),$$0 is a differentiable attribute verifier and $$L_{ce}(x) = -\sum_{t=1}^T \log p_\theta(x_t | x_{<t}, I),$$1 tunes the strength of semantic control (Ahmed et al., 4 May 2025).

2. Identification and Source of Token Weights

The assignment of token reweighting is governed by domain knowledge, external controllers, or end-to-end data-driven optimization.

• Keyword/Span-Based Rule Sets: In domains such as medical report generation, pre-defined sets of salient keywords $$L_{ce}(x) = -\sum_{t=1}^T \log p_\theta(x_t | x_{<t}, I),$$2 (e.g., quantitative and diagnostic terms) are compiled. All subwords comprising these phrases are assigned upweight factor $$L_{ce}(x) = -\sum_{t=1}^T \log p_\theta(x_t | x_{<t}, I),$$3, producing a set $$L_{ce}(x) = -\sum_{t=1}^T \log p_\theta(x_t | x_{<t}, I),$$4. Non-salient tokens retain weight 1. This form is effective when high-priority semantic classes are well characterized (Weers et al., 22 Apr 2026).
• Gradient or Sensitivity-Based Selection: In reinforcement learning with verifiable rewards (RLVR) over multimodal LLMs, tokens are selected by computing (a) tokenwise entropy $$L_{ce}(x) = -\sum_{t=1}^T \log p_\theta(x_t | x_{<t}, I),$$5 for reasoning-related uncertainties, and (b) visual sensitivity $$L_{ce}(x) = -\sum_{t=1}^T \log p_\theta(x_t | x_{<t}, I),$$6 for perception tokens sensitive to visual input. Critical tokens are dynamically chosen via thresholding top $$L_{ce}(x) = -\sum_{t=1}^T \log p_\theta(x_t | x_{<t}, I),$$7 or $$L_{ce}(x) = -\sum_{t=1}^T \log p_\theta(x_t | x_{<t}, I),$$8 quantiles for each category (Lu et al., 26 Mar 2026).
• Learned or User-Driven Weights: In text encoders like CLIP, per-token weights $$L_{ce}(x) = -\sum_{t=1}^T \log p_\theta(x_t | x_{<t}, I),$$9 can be learned by fine-tuning (with parameters $\lambda_t$0, $\lambda_t$1) for data-driven emphasis, or specified by the user to align retrieval with subjective preferences (Kim et al., 2024).
• Verifier-Guided Tilting: For semantic control, the magnitude and direction of the reweighting for each token is derived from the gradient of the verifier function $\lambda_t$2 with respect to the sequence embedding space (Ahmed et al., 4 May 2025).

3. Algorithmic Integration and Implementation

SToRI methods are designed for minimal disruption of existing pipelines.

• Training Loop Integration: In medical VLM fine-tuning, the SToRI loss replaces CE without changing forward/backward passes or optimizer logic. Token matching is performed at each sample, and weighted loss is accumulated batch-wise (Weers et al., 22 Apr 2026).
• Plug-and-Play in RLVR Objectives: For policy-gradient style objectives (like PPO, GRPO, DAPO), token-weights $\lambda_t$3 are inserted multiplicatively in existing loss computations. Only token masking and reduction steps are changed; no extra reward terms or modules are added (Lu et al., 26 Mar 2026).
• Self-Attention Modification: In transformer-based text encoders, SToRI modifies the attention softmax computation in one or more layers, allowing for fine-grained or global control with negligible computational cost (Kim et al., 2024).
• Sampling-Time Control: For attribute-constrained LM sampling, SToRI modulates $\lambda_t$4 at every step using the verifier-tilted score, exploiting gradient-based signal and top-K truncation for efficiency (Ahmed et al., 4 May 2025).

4. Quantitative Effects and Empirical Results

SToRI consistently yields measurable gains in sample efficiency, attribute adherence, and interpretability. Notable quantitative findings include:

Domain	Model/Task	Main Effect or Metric	Best Reported Gain
Medical VLM	Llama3-3B + LoRA on AMD reports (Weers et al., 22 Apr 2026)	F1_macro for staging/biomarkers (low data)	SToRI (γ>1) at 10% data outperforms baseline at 100%
CLIP Few-Shot	CLIP ViT-L/14, 1–16 shot (Kim et al., 2024)	Few-shot classification accuracy	SToRI improves or matches TaskRes; 1-shot: 82.01% vs. 81.70%
RLVR/MM LLMs	Qwen2.5-VL, RLVR (Lu et al., 26 Mar 2026)	Multimodal reasoning benchmarks (MathVerse, etc)	ToR-GRPO +2.2 pt HalluBench, +1.5 pt WeMath over GRPO
Attribute LM	Toxicity/sentiment-controlled (Ahmed et al., 4 May 2025)	Attribute satisfaction in LM generation	>95% satisfaction with minimal quality loss

Absolute performance improvements are most pronounced in the low-data regime, e.g., sample efficiency gains up to $\lambda_t$5 for medical report VLMs with SToRI. In RLVR, optimizing only reasoning or only perception tokens is significantly inferior to joint token-type reweighting, with typical drops of 2–3 F1 points when ablated.

5. Interpretability, Controllability, and Analysis

SToRI confers interpretability by making token importances explicit and controllability by allowing manual or programmatic adjustment.

• Token Attribution: Learned or assigned weights in SToRI (e.g., for CLIP or RLVR models) allow direct inspection of which tokens influenced a classification, retrieval, or policy decision. Ablations confirm that only semantically meaningful tokens are upweighted—random tokens, when explicitly reweighted, yield no performance improvement (Kim et al., 2024).
• Control Interfaces: User-driven reweighting in retrieval settings precipitates smooth, monotonic changes in outputs (e.g., top-ranked images in CelebA or CUB reflect emphasis on “blonde hair” or “eyeglasses” as weights vary), enabling real-time preference steering (Kim et al., 2024).
• Proxy Effectiveness: In RLVR, diagnostic analysis showed that tokens with high entropy (reasoning) and high visual-sensitivity (perception) must be reweighted jointly. Pure strategies fail to achieve optimal accuracy, supporting the coupling hypothesis (Lu et al., 26 Mar 2026).
• Constraint Satisfaction: In controlled generation, an appropriate setting of the reweighting intensity $\lambda_t$6 achieves constraint adherence >95% while largely preserving linguistic fluency. Top-K truncation provides further control over diversity and computational cost (Ahmed et al., 4 May 2025).

6. Domain Applications and Limitations

SToRI has been deployed across multiple modalities and domains:

• Medical Vision-Language Modeling: Upweighting semantically critical diagnostic tokens in ophthalmological report generation yields pronounced sample efficiency, especially with limited labeled data (Weers et al., 22 Apr 2026).
• Multimodal Chain-of-Thought Reasoning: Reweighting reasoning and perception tokens in RLVR with ToR (Token Reweighting) advances the state of the art in explicit visual reasoning tasks by balancing factual grounding with symbolic inference processes (Lu et al., 26 Mar 2026).
• Controllable Representation Learning: In CLIP and related VLMs, SToRI delivers interpretable embeddings and enables user-guided semantic retrieval (Kim et al., 2024).
• Semantic Control of Language Generation: SToRI probabilistically tilts token-level distributions toward constraint satisfaction as quantified by attribute verifiers, applicable in toxicity, politeness, and topic adherence (Ahmed et al., 4 May 2025).

Limitations identified across studies include dependency on the presence of discriminative tokens (CLIP), lack of invention capability (cannot create new information), and, when operating in pretrained black-box spaces, inability to address upstream data/model biases. Layer-position ablation in CLIP SToRI shows little sensitivity, but token reweighting must persist across blocks for full effect (Kim et al., 2024). For RLVR, further extensions could consider more granular region- or span-based reweighting and dynamic, attribution-based schedules (Lu et al., 26 Mar 2026).

7. Extensions and Outlook

Future directions for SToRI include:

• Fine-Grained and Region-Based Reweighting: Extending beyond subword and word tokens to visual regions (e.g., SAM-masks), phrase-level spans, and interconnected multimodal semantic units.
• Dynamic Attribution and Online Scheduling: Leveraging gradient-based or attributional signals for online adaptation of weights rather than fixed or purely data-driven assignment.
• Sparsity and Regularization: Controlling the sparsity or spread of token weights via regularizers to further improve interpretability or tailor to user-comprehensible emphasis patterns.
• Integration with Black-Box Encoders: Adapting SToRI strategies to vision encoders, multi-modal transformers, and beyond, while addressing inherited biases and distributional artifacts.

A plausible implication is that SToRI will serve as a foundation for robust, general-purpose interpretability and control across language, vision, and multi-agent reasoning environments, given its lightweight integration and empirically demonstrated efficacy in high-stakes and data-limited scenarios (Weers et al., 22 Apr 2026, Lu et al., 26 Mar 2026, Ahmed et al., 4 May 2025, Kim et al., 2024).