Farizeh/nllb-eng-arabic-wmt26

NLLB-200 Fine-Tuned: English → Arabic (WMT26)

Fine-tuned NLLB-200-3.3B on the WMT26 Low-Resource Arabic-Asian MT shared task (Task 1A: English → Arabic).

This repo also includes an NTK-Mirror controller (ntk_controller.pt) — a lightweight post-hoc channel controller (5 K scalars) that can be applied at inference time for a small additional BLEU gain without any retraining.

Quick start — plain NLLB (no NTK-Mirror)

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import torch

model_id = "Farizeh/nllb-eng-arabic-wmt26"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSeq2SeqLM.from_pretrained(model_id, torch_dtype=torch.float16)
model.eval().cuda()

tgt_lang_id = tokenizer.convert_tokens_to_ids("arb_Arab")

def translate(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True,
                       max_length=256).to("cuda")
    with torch.no_grad():
        out = model.generate(**inputs, forced_bos_token_id=tgt_lang_id,
                             num_beams=4, max_length=256)
    return tokenizer.decode(out[0], skip_special_tokens=True)

print(translate("Hello, how are you?"))

With NTK-Mirror controller

NTK-Mirror applies a learned per-channel scaling to the model's hidden states at inference time. No extra training is needed — just load the controller and attach the hook before generate(), then remove it after.

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from huggingface_hub import hf_hub_download
import torch

model_id = "Farizeh/nllb-eng-arabic-wmt26"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSeq2SeqLM.from_pretrained(model_id, torch_dtype=torch.float16)
model.eval().cuda()

# Load the NTK-Mirror controller
ctrl = torch.load(hf_hub_download(model_id, "ntk_controller.pt"), map_location="cpu")


def attach_ntk(model, ctrl):
    """Register NTK-Mirror forward hooks. Returns handles — call h.remove() when done."""
    layer_ch_map = {}
    for g, (l, c) in enumerate(zip(ctrl.layer_indices.tolist(),
                                    ctrl.channel_indices.tolist())):
        layer_ch_map.setdefault(l, []).append((g, c))

    raw, max_lg = ctrl.raw, ctrl.max_log_gate
    layers = [m for n, m in model.named_modules()
              if n.startswith(ctrl.layer_path + ".") and
                 n.count(".") == ctrl.layer_path.count(".") + 1]

    handles = []
    for i_layer, layer in enumerate(layers):
        entries = layer_ch_map.get(i_layer, [])
        if not entries:
            continue
        def hook(mod, inp, out, _e=entries):
            h    = out[0] if isinstance(out, tuple) else out
            rest = out[1:] if isinstance(out, tuple) else ()
            sv   = torch.ones(h.shape[-1], dtype=torch.float32, device=h.device)
            for g, c in _e:
                sv[c] = (max_lg * torch.tanh(raw[g].float())).exp()
            h_new = h * sv.to(h.dtype)
            return (h_new,) + rest if isinstance(out, tuple) else h_new
        handles.append(layer.register_forward_hook(hook))
    return handles


tgt_lang_id = tokenizer.convert_tokens_to_ids("arb_Arab")

def translate_ntk(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True,
                       max_length=256).to("cuda")
    handles = attach_ntk(model, ctrl)
    with torch.no_grad():
        out = model.generate(**inputs, forced_bos_token_id=tgt_lang_id,
                             num_beams=4, max_length=256)
    for h in handles:
        h.remove()
    return tokenizer.decode(out[0], skip_special_tokens=True)

print(translate_ntk("Hello, how are you?"))

How NTK-Mirror works

NTK-Mirror learns a small set of log-gate scalars (one per selected hidden channel) that rescale activations during the forward pass:

h'[layer, token, channel] = exp(max_log_gate × tanh(raw[g])) × h[layer, token, channel]

Only ~5 000 channels across all layers are selected (activation-magnitude scoring). The controller adds zero parameters to the model weights — it is stored separately and applied via a PyTorch forward hook at inference time.

Results (devtest, WMT26 Task 1A: EN → AR)

System	BLEU	chrF2	COMET
NLLB-FT base	18.52	52.72	87.47
NLLB-FT + NTK-Mirror	18.56	52.01	87.38

NTK-Mirror is approximately neutral on NLLB after joint multilingual fine-tuning. Larger gains (+2–10 BLEU) are observed when applying NTK-Mirror to causal LMs (Qwen2.5-7B).

Citation

@inproceedings{aldabbas-etal-2026-farabi,
    title     = "Lightweight Post-Hoc Channel Controllers for Low-Resource Arabic--Asian Machine Translation",
    author    = "Aldabbas, Farizeh and others",
    booktitle = "Proceedings of the Eleventh Conference on Machine Translation (WMT26)",
    year      = "2026",
}