TokenAlign: frozen Qwen features for label-description classification

TokenAlign scores candidate descriptions through token-level matching. A frozen Qwen3-1.7B encoder supplies features. A 2,236,929-parameter trainable head adds two shared bidirectional blocks, then averages each label token's strongest cosine match in the document. There is no final scalar MLP scorer or class-specific output matrix. This is an independent research prototype inspired by ColBERT-style late interaction. It is not affiliated with JEV and does not reproduce an undisclosed JEV architecture or training recipe.

This repository contains only the selected alignment head, in Safetensors. It also needs the separately downloaded Qwen3-1.7B backbone (about 3.4 GB of FP16 weights). It is not a standalone 2.24M-parameter language model.

Installation

Use Python 3.12 for the tested configuration. Clone the model repository:

git clone https://huggingface.co/BechirTrabelsi1/tokenalign-qwen3-1.7b-banking77
cd tokenalign-qwen3-1.7b-banking77
python3.12 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
pip install -e . --no-deps

The release uses its own Python API. It does not currently support transformers.pipeline, AutoModel.from_pretrained, or the hosted Hugging Face Inference API. No trust_remote_code=True is required.

Python inference

from tokenalign import TokenAlignClassifier

model = TokenAlignClassifier.from_pretrained(
    'BechirTrabelsi1/tokenalign-qwen3-1.7b-banking77', device='auto'
)
result = model(
    'I was charged twice for the same card purchase.',
    ['duplicate card payment', 'cash withdrawal declined', 'card not received'],
    explain=True,
)
print(result['label'])
print(result['scores'])
print(result['matches'])

Pass '.' instead of the model ID to load a local checkout. Install this package's code first. The API downloads only config.json and model.safetensors from that repo, then loads the pinned backbone from Qwen. local_files_only=True prevents downloads when both are cached. cache_dir= sets a custom Hugging Face cache directory.

Load the classifier once and reuse it. It caches the projected/contextualized label representations for the eight most recent candidate lists. Calling model.clear_label_cache() releases them. A list of input strings returns a list of results. Scores use a softmax over the supplied labels and sum to one; these are not independent multi-label probabilities. A changed candidate set changes them.

CLI and persistent demo

python -m tokenalign --model . --device mps --local-only \
  --text 'I was charged twice for the same card purchase.' \
  --labels 'duplicate card payment' 'cash withdrawal declined' 'card not received' \
  --explain

Replace --text ... with --interactive to keep the model resident and enter new messages repeatedly. Empty input exits. Use --device cpu without Metal; CUDA is also supported by the code but was not tested for this release. --local-only requires cached backbone files, so omit it on a fresh installation.

The API reports truncation. The 128-token limit includes the input prefix. Longer inputs can lose decisive evidence. Supply at least two distinct, nonempty candidate descriptions. The model always chooses one of them; it has no trained abstention or out-of-domain detector.

Architecture

1. Encode each document and each label description independently with frozen Qwen.
2. Retain the last-layer, 2,048-dimensional token states. Exclude prompt boilerplate, special tokens, and pure punctuation from direct scoring using tokenizer offsets.
3. Apply a shared trainable input LayerNorm and a 2,048-to-256 projection.
4. Apply two shared pre-norm bidirectional self-attention blocks, separately within each document/description. Each block uses four heads, FFN width 1,024, GELU, residual connections, and dropout 0.1 during training.
5. Apply separate Q/K projections and L2-normalize token vectors.
6. Form label/document cosine similarities. Average the maximum over document tokens for each real label token. Multiply by a learned positive scale capped at 100, then divide logits by the saved calibration temperature before softmax.

For candidate c: logit[c] = scale * mean_i(max_j(cosine(q[c,i], k[j]))). The score uses raw similarities, not the sum of normalized attention weights. The contextual blocks do not change Qwen's causal mask or weights. No candidate attends to another candidate. Token matches are scoring contributions, not proof of a causal or linguistic explanation.

The implementation separates input normalization from its learned affine transform because a local PyTorch 2.6 Metal reproduction produced incorrect affine gradients for native LayerNorm on frozen inputs. CPU/Metal gradient checks cover the workaround.

Training data and protocol

Training uses Banking77, 10,003 source training messages and 3,080 source test messages across 77 English banking intents. Seed 42 selects disjoint groups of 50 training, 12 validation, and 15 unseen test labels. Exact text deduplication normalizes case/whitespace.

The two validation partitions share categories but contain disjoint messages. Descriptions use category names in the fixed template The customer's banking issue concerns {category with spaces}. No test-category examples train the head. Label descriptions can enter inference and frozen caching. Zero-shot means unseen by head training; backbone pretraining exposure is unknown.

Backbone revision: 70d244cc86ccca08cf5af4e1e306ecf908b1ad5e. Dataset revision: 57ec275d8078af65b7731c2a98be812d844a6d6b. The source hashes and per-partition example hashes are in evaluation/. Raw dataset messages are not included in this release. Download them from the attributed source.

Training uses AdamW, learning rate 0.0003, weight decay 0.01, batch size 32, gradient norm clipping at 1.0, and a maximum of 50 epochs. Patience is eight epochs without a lower validation NLL. The selected checkpoint is epoch 2; training stopped after epoch 10. Qwen runs in FP16 on Metal and the head in FP32. Features are cached on CPU in FP16. The post-training calibration temperature is 0.9440608621.

Evaluation

Macro F1 averages per-class F1 with equal class weights. Micro F1 aggregates true positives, false positives, and false negatives across all classes. In this exhaustive single-label multiclass task, micro F1 equals accuracy; the saved micro F1 values are derived from the original correct-example counts.

Brier is the mean sum of squared class-probability errors; ECE uses ten equal-width confidence bins. Lower is better for NLL, Brier, and ECE. Positive temperature scaling leaves predictions, accuracy, and both F1 scores unchanged.

This is one exploratory development split, not the standard 77-way Banking77 benchmark or a held-out architecture-selection study. The checkpoint was selected by validation NLL. Independent splits and additional domains are required for stronger generalization claims.

The original M4 Pro 24 GB run recorded batch-one inference p50 110.3 ms and p95 427.1 ms, across 30 messages with cached label features and an untimed warm-up. Those timings include the backbone and a CPU feature round-trip. They exclude model loading and are not measurements of this refactored persistent inference API.

Exploratory out-of-ontology label-description experiment

An additional small experiment was performed to test whether TokenAlign can use semantic descriptions for classes that are outside the Banking77 label ontology used during training.

This experiment is intentionally separate from the reported Banking77 evaluation above. It is a manually constructed exploratory test, not a standard benchmark.

Setup

Fifteen synthetic banking-related messages were created for concepts such as:

• mortgage applications
• joint accounts
• personal loans
• loan repayment schedules
• insurance claims
• tax documents
• cryptocurrency trading
• stock investment
• overdraft-limit changes

These candidate classes were not used as TokenAlign training classes.

For every test example, the input text and the set of candidate concepts were kept fixed. Only the textual representation of the candidate labels was changed.

Three formulations were evaluated:

1. Name — short category names, e.g. stock investment
2. Short description — e.g. buying company shares or stocks as an investment
3. Rich description — e.g. using the bank or investment service to purchase shares, stocks, or equities in publicly traded companies

No model weights were changed between the three runs.

Results

Providing semantic descriptions therefore improved top-1 accuracy from 66.67% to 80.00% on this small test set.

The average probability assigned to the expected class also increased from 0.551 with bare class names to 0.697 with richer descriptions.

Example behavior

Some errors made with bare category names were corrected when more semantic information was provided.

For example, an input about buying Apple and Microsoft shares was incorrectly classified as savings account when using bare class names. Replacing stock investment with a description referring explicitly to buying company shares or stocks changed the prediction to the intended class.

Likewise, an insurance-reimbursement request was initially classified as a generic complaint. Both short and rich semantic descriptions changed the prediction to insurance claim.

The experiment also showed that longer descriptions are not automatically better. TokenAlign averages MaxSim evidence across retained label tokens, so adding informative tokens may help while adding weak or redundant tokens can dilute otherwise strong matches.

Interpretation

These results are consistent with the intended behavior of TokenAlign: prediction depends on the semantic content of the supplied class descriptions rather than on a fixed learned output index.

However, this experiment should be interpreted only as an exploratory sanity check.

Important limitations include:

• only 15 manually written examples were evaluated;
• candidate descriptions were manually authored;
• the examples are still banking-domain inputs;
• description wording was not sampled across multiple independent paraphrases;
• probabilities were calibrated on the original development setting, not on these new labels;
• no claim of general out-of-domain or arbitrary-ontology generalization is made from this experiment.

A stronger evaluation would use multiple independently written descriptions per class, more examples per new class, multiple random candidate sets, and additional domains.

Limitations and intended use

Use for research and demos of English banking-intent classification with candidate label descriptions. General-domain, multilingual, multi-label, and open-set behavior remain unvalidated. Softmax confidence can be misleading for irrelevant labels, negation, ambiguous requests, and inputs beyond the training domain. Temperature calibration transfers imperfectly to new labels and candidate counts. Do not infer that a highlighted token is a faithful causal explanation.

Reproduction

python scripts/train.py --device mps --run-dir runs/reproduction-42
python -m unittest discover -s tests -v

Training requires a fresh output directory. It downloads pinned public data and backbone weights as needed. --local-only requires a prepared dataset and cached weights. The script trains only this contextual alignment architecture and saves only its best checkpoint. A fresh run may show floating-point/device variation.

Technical presentations

• Academic Beamer presentation (PDF)
• Editable standalone LaTeX source
• Technical PowerPoint

License and related work

Project code and the released head use Apache-2.0. Qwen3-1.7B has its own Apache-2.0 license. Banking77 remains under CC-BY-4.0. Dataset licensing does not change under this repository's code license. See NOTICE.

• Khattab and Zaharia, 2020. ColBERT.
• Wu, Su, and Juang, 2021. LABAN.
• Casanueva et al., 2020. Efficient Intent Detection with Dual Sentence Encoders.
• BehnamGhader et al., 2024. LLM2Vec.

The project is a ColBERT-inspired experimental combination. It does not claim to invent MaxSim, label-aware classification, or frozen-backbone adaptation.