Target-Aware Weighted Training (TAWT)

• TAWT is a framework that adaptively weights instances, tasks, and tokens by incorporating target task properties into the training objective.
• It utilizes representation-based distance metrics and mirror descent updates to minimize target risk and improve transfer learning efficiency.
• Empirical results show TAWT enhances performance in sequence tagging and NMT, yielding significant improvements in BLEU scores and F1 metrics in low-resource settings.

Target-Aware Weighted Training (TAWT) comprises a suite of principled weighting strategies for supervised learning that explicitly incorporate target task or target token properties into the training objective. Its central aim is to improve transfer or generalization performance on a target domain by adaptively weighting instances, tasks, or tokens according to their relevance or alignment with the target distribution. TAWT is instantiated in both cross-task learning and neural sequence modeling, with each setting providing rigorous methodology and empirical benefits by minimizing a representation- or token-aware distance to the target (Chen et al., 2021, Gu et al., 2020).

1. Formal Definitions and Problem Setting

TAWT in cross-task learning considers $T$ source tasks, indexed by $t \in \{1,\ldots,T\}$, and a designated target task ($t=0$), each with its own data distribution $\mathcal{D}_t$ over input–output pairs $(x, y) \in \mathcal{X} \times \mathcal{Y}$. Supervised data from the target is limited: $n_0 \ll n_t$ for source $t$. The goal is to learn a rich encoder $\phi: \mathcal{X} \to \mathbb{R}^r$ (e.g., BERT up to the penultimate layer) and a simple task-specific head $f: \mathbb{R}^r \to \mathcal{Y}$ (e.g., linear-softmax), minimizing population risk $$\mathcal{J}_t(\phi, f) = \mathbb{E}_{(x, y) \sim \mathcal{D}_t}[\ell(f \circ \phi(x), y)]$$ over a bounded loss $t \in \{1,\ldots,T\}$0.

In neural machine translation (NMT), TAWT is instantiated at the token level: let $t \in \{1,\ldots,T\}$1 be parallel bitext, with $t \in \{1,\ldots,T\}$2 a token sequence. For each target vocabulary token $t \in \{1,\ldots,T\}$3, a weight $t \in \{1,\ldots,T\}$4 is assigned and used in a weighted cross-entropy objective

$t \in \{1,\ldots,T\}$5

Weights $t \in \{1,\ldots,T\}$6 are set adaptively according to token frequency statistics in $t \in \{1,\ldots,T\}$7 (Gu et al., 2020).

2. Representation-Based Task Distance and Weighted Objective

TAWT for cross-task transfer formalizes the intuitive notion of “task closeness” using a representation-based task distance. For $t \in \{1,\ldots,T\}$8 (the simplex of task weights), define

$t \in \{1,\ldots,T\}$9

and for a candidate $t=0$0, its one-step risk on target is $t=0$1. The $t=0$2-representation-based task distance is defined as

$t=0$3

where $t=0$4 is the target-optimal encoder. TAWT adaptively selects $t=0$5 so as to minimize this distance, thereby bridging source representations with the specific geometry of the target (Chen et al., 2021).

The training objective itself is formulated as a bi-level optimization:

• Outer problem: minimize target risk $t=0$6,
• Inner problem: for fixed $t=0$7, minimize the weighted source risk to obtain $t=0$8.

In the token-level NMT case, $t=0$9 is chosen so that more semantically important but low-frequency tokens receive higher emphasis during training, thereby addressing imbalances in loss attribution inherent to standard cross-entropy.

3. Algorithmic Implementation

In cross-task TAWT, Algorithm 1 alternates three main steps for $\mathcal{D}_t$0:

1. Representation Update: With current $\mathcal{D}_t$1, perform $\mathcal{D}_t$2 steps of SGD on $\mathcal{D}_t$3.
2. Target Head Optimization: With new $\mathcal{D}_t$4 fixed, update $\mathcal{D}_t$5 on $\mathcal{D}_t$6.
3. Weight Update via Mirror Descent: Calculate for each source $\mathcal{D}_t$7 a dissimilarity gradient $\mathcal{D}_t$8 (approximated via cosine similarity of representation gradients), then update $\mathcal{D}_t$9.

This approach can be integrated into any deep learning framework, requiring only per-task tracking of encoder gradients, periodic simplex projections, and cosine-scaled updates to $(x, y) \in \mathcal{X} \times \mathcal{Y}$0—yielding $(x, y) \in \mathcal{X} \times \mathcal{Y}$1 computational cost per epoch (Chen et al., 2021).

In NMT, the algorithm comprises a one-time frequency analysis over the training set, weight computation, and a standard seq2seq loop with weighted loss aggregation. The pseudocode presented in (Gu et al., 2020) illustrates that a single array lookup for each token suffices to apply TAWT.

4. Non-Asymptotic Generalization Guarantees

TAWT in cross-task learning provides non-asymptotic learning-theoretic guarantees. Under standard Lipschitz and VC-type assumptions, and a “transferability” condition relating source and target representability, the sample-split version yields

$(x, y) \in \mathcal{X} \times \mathcal{Y}$2

where $(x, y) \in \mathcal{X} \times \mathcal{Y}$3 are complexity parameters, $(x, y) \in \mathcal{X} \times \mathcal{Y}$4 bounds the weight range, and $(x, y) \in \mathcal{X} \times \mathcal{Y}$5 is the residual task distance (Chen et al., 2021). If $(x, y) \in \mathcal{X} \times \mathcal{Y}$6 is chosen such that $(x, y) \in \mathcal{X} \times \mathcal{Y}$7, TAWT can approach the statistical efficiency of knowing the target-optimal representation a priori.

5. Empirical Results and Analysis

In (Chen et al., 2021), TAWT is evaluated on four sequence tagging tasks (PoS, chunking, predicate detection, NER) using English data from Ontonotes 5.0 and CoNLL-2000, with target sets of size 100–500 for a low-resource setting. Baselines include single-task, uniform pre-training, and joint training. TAWT, both as weighted pre-training and weighted joint training, achieves gains exceeding +3 percentage points F1 over uniform-weight baselines (Table 1), with further improvements in low source:target data-ratio regimes.

Ablation experiments reveal that adaptive, dynamic $(x, y) \in \mathcal{X} \times \mathcal{Y}$8 is superior or equivalent to fixed weighting. Performance gains are maximized for low-resource targets or when the ratio of source-to-target data is small, and diminish as the joint training baseline improves or target data becomes sufficient.

In NMT (Gu et al., 2020), TAWT—implemented as frequency-based token-weighted loss—achieves consistent gains in BLEU versus multiple baselines:

• +0.88 BLEU on ZH→EN, +0.49 BLEU on EN→RO, +0.39 BLEU on EN→DE compared to fine-tuned Transformer baselines, with statistical significance ($(x, y) \in \mathcal{X} \times \mathcal{Y}$9).
• Gains are especially pronounced on “rare-heavy” sentences, with BLEU increases up to +1.68 (ZH→EN), +1.02 (EN→RO), +0.52 (EN→DE).
• Lexical diversity metrics (TTR, HD-D, MTLD) are consistently boosted by TAWT.

Both exponential and chi-square weighting schemes are used, with the maximum weight set to $n_0 \ll n_t$0 and a single scalar $n_0 \ll n_t$1 tuned per language pair.

6. Implementation and Practical Guidelines

TAWT is computationally lightweight:

• In cross-task learning, the main overhead is the $n_0 \ll n_t$2 cost per epoch for cosine-similarity-based gradient calculations. Batch sizes for gradient estimation can be small (e.g., 64 tokens), and only one new scalar hyperparameter $n_0 \ll n_t$3 is introduced.
• In token-level NMT, the weighting involves only a preprocessing frequency count and an array lookup per token, introducing negligible overhead and no model parameter increase.

Recommended best practices include:

• Use TAWT when target supervision is scarce and related source tasks or sequence-level knowledge is available.
• Initialise $n_0 \ll n_t$4 source weights as uniform or proportional to $n_0 \ll n_t$5; update with mirror descent per epoch or every few epochs.
• For TAWT in NMT, set $n_0 \ll n_t$6 such that $n_0 \ll n_t$7, and tune only the scalar $n_0 \ll n_t$8 parameter.
• Monitoring $n_0 \ll n_t$9 dynamics can inform pruning of source tasks with vanishing influence, as justified by the representation-based distance formulation (cf. Prop. 3.4 of (Chen et al., 2021)).
• Fine-tuning shared representations on the target is critical when the residual distance is non-negligible.

7. Connections, Extensions, and Related Work

TAWT generalizes classical weighted training by integrating target-awareness either via cross-task representation geometry or through explicit token-level rebalancing. In contrast to ad hoc importance sampling or uniform task pooling, TAWT leverages theoretically motivated proxies—representation gradients or frequency-based signal—to guide the model toward target-aligned optima.

The TAWT formalism has influenced practices in transfer and multi-task learning where adaptive weighting is critical for low-resource robustness and calibration. Notably, token-level variants provide a rigorous alternative to sentence-level rare sampling, entropy regularization, or linear token scaling, achieving both overall performance and rare-token precision gains with minimal procedural complexity (Gu et al., 2020). TAWT further introduces representation-based task distance as a diagnostic metric for transferability—enabling empirical and theoretical assessment of source-target alignment (Chen et al., 2021).