Jisr-WordAlign-29M

The encoder half of the Jisr Arabic-English MT model, finetuned for word alignment. 29M parameters, 448-dim, 6 layers. Trained from scratch.

Cosine over hidden_states[6], max-pooled from subwords up to whitespace words, matched with SimAlign-style mutual argmax. No decoder, no search - the similarity matrix is the model.

Results

AER against gemma-4-31B two-pass gold, lower is better. Baseline is the same encoder before this training, i.e. the MT model straight out of the box:

The gain scales with how dialectal the eval is - training data was dialect - and MSA improves too, so it is not a trade. The sentence aligner (oddadmix/Jisr-Align-50M) is much worse at this task (0.259 / 0.443 / 0.508): contrastive training on a pooled vector costs token-level accuracy.

Gold: oddadmix/jisr-align-gold, config word.

Two things it does not do

One-to-many. Mutual argmax and itermax both blank a word's row and column once it is linked, so the output is always a matching. Arabic clitics (بيتها = "her house") genuinely need two links and cannot get them. That is the decoder, not the encoder - the similarity matrix usually has both cells hot - and it needs a different matching rule or morphological segmentation.

Guarantee its labels. The gold it was trained and measured on is LLM-adjudicated, and the two passes disagreed 21-27% of the time. Read AER here as a reliable ranking between models, not as ground truth.

Training

3 epochs of symmetric cross-entropy over the pooled word-similarity matrix, uniform target across each word's gold partners, 9999 dialect sentence pairs from oddadmix/dialectal-arabic-english-parallel adjudicated in two opposite- order passes. Batch 16, lr 1e-5, learnable temperature from 0.05, layer 6. Converged by ~step 1200 of 1779. Model selection ran on held-out GENERATED pairs; the three gold sets above were scored once from the final checkpoint, so none of them fed back into which checkpoint was kept.

Usage

Complete and self-contained - this is the whole aligner, not a sketch of it. It reproduces word_align.py from the training repo link-for-link.

import re

import numpy as np
import torch
from transformers import AutoTokenizer, MarianMTModel

M = "oddadmix/Jisr-WordAlign-29M"
tok = AutoTokenizer.from_pretrained(M)
enc = MarianMTModel.from_pretrained(M).model.encoder.eval()

TAG = {"ar": ">>ara<<", "en": ">>eng<<"}
STRIP = "،؛؟.,!?;:()[]{}\"'“”«»…-"
LAYER = 6

@torch.no_grad()
def words_and_states(text, lang):
    # Whitespace words, their subword indices, and L2-normalised states.
    tag = TAG[lang] + " "
    b = tok(tag + text, return_tensors="pt", truncation=True, max_length=256,
            return_offsets_mapping=True)
    off = b.pop("offset_mapping")[0].tolist()
    H = enc(**b, output_hidden_states=True).hidden_states[LAYER][0]
    H = H.float().numpy()

    # Keep a token if it ENDS past the tag, not if it STARTS past it: the
    # tokenizer folds the tag's trailing space into the first real word, so a
    # start-based test silently deletes word 0 of every sentence.
    keep = [t for t, (s, e) in enumerate(off) if e > s and e > len(tag)]
    span = [(max(0, off[t][0] - len(tag)), off[t][1] - len(tag)) for t in keep]
    H = H[keep]
    H = H / (np.linalg.norm(H, axis=1, keepdims=True) + 1e-9)

    words = [(m.group(0), m.start(), m.end()) for m in re.finditer(r"\S+", text)]
    idx = []
    for _, a, z in words:
        t = [k for k, (s, e) in enumerate(span) if s < z and e > a]
        # Drop punctuation-only subwords. A full stop is the SAME token in both
        # languages, so a max over subword pairs would hand it to every
        # sentence-final word. Kept only for words that are all punctuation.
        real = [k for k in t if text[span[k][0]:span[k][1]].strip(STRIP)]
        idx.append(real or t)
    return [w for w, _, _ in words], idx, H

def align(ar_text, en_text):
    wa, ia, A = words_and_states(ar_text, "ar")
    wb, ib, B = words_and_states(en_text, "en")
    S = A @ B.T
    W = np.full((len(wa), len(wb)), -1.0, dtype=np.float32)
    for i, ta in enumerate(ia):
        for j, tb in enumerate(ib):
            if ta and tb:
                W[i, j] = S[np.ix_(ta, tb)].max()   # MAX over subword pairs
    bi, bj = W.argmax(1), W.argmax(0)               # mutual argmax
    return [(wa[i], wb[bi[i]], float(W[i, bi[i]]))
            for i in range(len(wa)) if bj[bi[i]] == i and W[i, bi[i]] > -1]

for a, b, s in align(
        "تطفو الإبرة الفولاذيّة على الماء بسبب التوتّر السطحي.",
        "The steel needle floats on top of the water because of surface tension."):
    print(f"{a:<22} {b:<12} {s:.3f}")

تطفو                   floats       0.442
الإبرة                 needle       0.496
الفولاذيّة             steel        0.716
على                    on           0.682
الماء                  water        0.781
بسبب                   because      0.792
التوتّر                tension.     0.553
السطحي.                surface      0.609

The last two links cross, and that is correct: Arabic puts the adjective after the noun, so التوتّر السطحي is literally "the-tension the-surface".

Three details are load-bearing. Layer 6, not the last. Max over subword pairs, not mean - a content word splits unevenly across its subwords and averaging buries the piece that matched. And drop punctuation-only subwords before pooling: measured without that filter, الآن. linked to added. at 0.889 and outscored the correct الآن ~ now at 0.746.

© KAND CA 2026