Stop Jostling: Adaptive Negative Sampling Reduces the Marginalization of Low-Resource Language Tokens by Cross-Entropy Loss

AI summary

Overview

Research area: Natural Language Processing — multilingual and low-resource language modeling, specifically the training dynamics of cross-entropy loss and negative sampling on token embeddings.

Technical level: Intermediate. The paper assumes familiarity with language model training, logits, softmax, weight tying, cross-entropy loss, temperature scaling, and sampling methods such as nucleus (top-p) and min-p sampling.

Scope: The paper identifies a specific noise source it calls "marginalization" — the pushing-away of non-target embeddings by cross-entropy loss — and proposes a logit thresholding technique to reduce its disproportionate harm to rare and low-resource-language tokens, validated with a character-level GPT-2 model on a simulated multilingual dataset.

What This Paper Is About

Neural language models learn poorly on low-resource languages because tokens from those languages appear rarely in training data. The paper argues that cross-entropy loss forces every non-target token embedding slightly away from the model's hidden state at each training step, and this "marginalization" hurts rare tokens far more than frequent ones because rare tokens receive few other, meaningful updates. The goal is a simple modification — thresholding logits so that very improbable non-target tokens are ignored entirely — that lets rare tokens learn cleaner representations.

Key Contributions

1. Identification of marginalization as a distinct problem. The paper formalizes the three ways embedding matrices are optimized during training — context alignment, target alignment, and non-target marginalization — and isolates the third as the specific source of noise degrading rare token representations.

2. A thresholding technique to mitigate it. The method sets logits below a threshold defined as logits_t[x_t] - margin to negative infinity before computing cross-entropy, making the marginalization gradient exactly zero for those tokens.

3. Guidance for choosing the new margin hyperparameter. The paper derives that choosing margin = T × ln(((N−1) × top_p) / (1 − top_p)) guarantees that thresholded tokens fall outside the top-p nucleus, and also gives a min-p analogue (margin = T × ln(p_base)). Worked example: for T = 0.9, top_p = 0.99, N = 100,000, the estimated margin is approximately 14.50.

4. Separated Embeddings (SE) as a complementary fix. Because AdamW updates momenta across all rows of a single embedding matrix, the paper stores each token's embedding independently so AdamW can skip updates for tokens with no gradients, isolating the effect.

Main Findings

• Marginalization is measurable and real in the low-resource setting. For the low-resource language tokens, the 14 least frequent tokens have a "Ratio of Average" below 1, and all tokens have an "Average of Ratio" below 1, indicating marginalization gradients dominate the useful alignment gradients.

• Thresholding increases useful gradient influence. The proposed method increases the "Average of Ratio" by 45% and the "Ratio of Average" by 12% on average.

• Reducing the margin monotonically improves low-resource metrics up to a point. With a conservative margin = 8 results are barely changed; performance improves as the margin drops through 4, 2, 1 and reaches its best result at margin = 0.6 with separated embeddings.

• Low-resource accuracy and perplexity close much of the gap with the high-resource language. The conclusion reports low-resource accuracy rising from 0.31 (baseline) to 0.49, close to the high-resource accuracy of 0.53, and low-resource PPL_best falling from 10.63 to 6.11, approaching the high-resource PPL_best of 4.97. (The results table lists the strongest low-resource PPL_best as 6.17 for the Proposed+SE margin = 0.6 configuration.)

• Separated Embeddings multiply the gains. SE yields ×0.72 PPL_best, ×1.22 Accuracy, ×1.09 Recall@5, and ×1.14 MRR relative to the corresponding non-SE setting. The paper notes SE only affects optimization of unselected tokens, suggesting the improvement comes from reduced marginalization rather than contrastive learning.

• Raw perplexity looks worse without temperature correction. Because thresholding leaves an "unreliable tail" of un-marginalized low-probability tokens, raw PPL rises with aggressive margins (e.g. low-resource PPL 19.94 at margin = 0.6), but after applying the best temperature T_best the proposed models beat the baseline PPL_best (6.17 versus 10.63).

• The model learns cross-language character relationships without a parallel corpus. In a top-3 cosine-similarity analysis, the thresholded + SE model places capitalizations and "translations" of a character among the top-3 neighbors in 10 out of 12 cases, versus 0 out of 12 for the baseline.

• Thresholded models produce more isotropic embeddings. Measured by I(W) following Mu and Viswanath (2018), thresholded models show better values than the baseline.

• The monolingual control is worse than the multilingual baseline. Training only on the low-resource 2% of steps gives low-resource PPL 12.24 and accuracy 0.2851, worse than the baseline's 10.65 and 0.3147, and far below the Proposed+SE results.

Methodology in Plain English

The researchers used a small character-level language model built on Karpathy's nanoGPT, trained on Shakespeare text. Shakespeare has 65 unique characters with a very imbalanced frequency distribution, which makes it a convenient toy stand-in for a multilingual corpus. To simulate a second, low-resource language, they modified the tokenizer so that in 2% of randomly selected training sentences, each character ID was increased by 65, creating a second language with token IDs from 65 to 129. The 2% ratio was chosen because French is roughly 2% of words in GPT-3's pre-training data.

The models compared were: a baseline GPT-2 architecture with 800k parameters and weight tying; a monolingual model trained only on the low-resource 2% of steps; the proposed model with thresholding at margins between 0 and 8; and the proposed model with separated embeddings.

The modification itself is simple. After the classification head produces logits for a position, the algorithm sets a threshold equal to the correct token's logit minus the margin, and replaces any logit below that threshold with negative infinity. Those tokens contribute zero gradient, so they are not pushed away from the current context. Tokens above the threshold — like "Bob" in the paper's song-title example — still get marginalized normally. The paper also derives a formula that ties the margin to temperature and top-p (or min-p) so that tokens inside the nucleus used at generation time are never thresholded.

Why This Matters

Impact on research: The paper reframes negative sampling from a tool for training efficiency or contrastive learning into a tool for improving the representations of the negative samples themselves, claiming to be the first to show this. It connects to earlier work on the representation degeneration problem, anisotropy, and adaptive gradient gating (AGG), which the authors describe as similar in spirit but more complex because it requires counting token frequencies during training.

Real-world applications:

• Improving language model quality for low-resource and underrepresented languages, where large corpora are unavailable.
• Better cross-lingual token representations in multilingual models, potentially helping transfer between languages with shared scripts.
• More balanced token embeddings in any model with a highly skewed vocabulary distribution, such as character-level or byte-level models.
• Reduced representation degeneration in production embedding tables, since the paper reports more isotropic I(W) values.

Industry relevance: Thresholding is a few lines of code applied to logits and adds no parameters; the paper notes a PyTorch implementation appears in the appendix. The separated embeddings modification is more invasive because it changes how weight tying is stored, but it produced the largest gains in the experiments.

Future Directions

• Scaling up. The authors note all experiments used a small model and dataset, and that larger models may behave differently; they argue intuitively that a larger embedding dimension d_model could increase the benefit of thresholding, and provide margin estimates for GPT-2 family dimensions (768, 1024, 1280, 1600) in the appendix.
• Real languages and real tokenizers. Testing only used simulated multilingual data with a 1:1 character correspondence between languages; the paper explicitly asks for evaluation on natural languages and with subword tokenization such as BPE, noting that tokenizer entropy correlates with model performance.
• Downstream task evaluation. The paper states that validation-set improvements do not guarantee downstream gains and that further task-level testing is needed.
• Principled margin selection. The paper does not cover choosing the optimal margin and instead offers only a way to bound the search range; finding the best value remains open.
• Broader architectures. Because the method modifies logits, the authors suggest it could apply beyond decoder transformers, and the appendix reports a model without weight tying.

Target Audience

Researchers and engineers working on multilingual or low-resource language modeling, particularly those training their own models or fine-tuning embedding tables; practitioners interested in the mechanics of cross-entropy loss and negative sampling; and anyone studying token embedding geometry, anisotropy, and representation degeneration. Readers should be comfortable with softmax, logits, gradient flow through embedding matrices, and sampling methods like top-p and min-p.

Authors’ abstract

Neural language models often struggle with low-resource languages due to the limited availability of training data, making tokens from these languages rare in the training set. This paper addresses a specific challenge during training: rare tokens are disproportionately affected by marginalization, which prevents them from learning effectively. We propose a thresholding technique that reduces the impact of this marginalization, allowing rare tokens to benefit from more meaningful alignment. Through experiments with a character-level language model, we demonstrate that this method significantly improves performance on low-resource language validation data. This work is the first to show how negative sampling can be applied to improve the representation of rare tokens by limiting the harmful influence of excessive marginalization, offering a new approach to enhancing language model performance for underrepresented languages.