Model

A small, calibrated PyTorch text classifier with categorical variables, trained reproducibly and evaluated at every level of the classification.
Authors
Affiliation

Innovation Lab, Insee

Innovation Lab, Insee

1 Why a small model

Large pre-trained language models such as CamemBERT were tested early in the project. On APE coding, they gave only a small gain over a fastText-style model: about 0.76 against 0.75 accuracy on the same test set (Faria et al. 2024). That gain did not justify their inference cost under the production constraints (unit calls, CPU only, a few hundred milliseconds).

Faria, Thomas, Nathan Randriamanana, and Tom Seimandi. 2024. Retraining Strategies for an Economic Activity Classification Model. European Conference on Quality in Official Statistics (Q2024), Estoril. https://inseefrlab.github.io/codif-ape-prez/slides/estoril-paper/index.pdf.
Joulin, Armand, Edouard Grave, Piotr Bojanowski, and Tomas Mikolov. 2016. Bag of Tricks for Efficient Text Classification. https://arxiv.org/abs/1607.01759.

The production model is therefore a shallow neural network of about one million parameters, in the fastText family (Joulin et al. 2016). It is now implemented in PyTorch with torchTextClassifiers.

2 Architecture

flowchart LR
  T["Activity description<br/>(+ complementary text)"] --> TOK["Tokenizer<br/>(sub-words / n-grams)"]
  TOK --> EMB["Token embeddings<br/>(dim 80)"]
  EMB --> AGG["Mean pooling"]
  C["Categorical variables<br/>CJ, NAT, TYP, SRF, CRT"] --> CEMB["One embedding<br/>per variable (dim 16)"]
  AGG --> CAT["Concatenation"]
  CEMB --> CAT
  CAT --> HEAD["Linear head<br/>~745 classes"]
  HEAD --> SM["Softmax<br/>top-5 + confidence index"]
Figure 1: Model architecture.

The model has three parts:

  • Text. The description is split into sub-word tokens. The tokenizer, either a WordPiece vocabulary or fastText-style character n-grams, is trained on the training data and versioned with it. Sub-words make the model robust to typos, inflections and new words. The token embeddings are averaged into a single vector for the text.
  • Categorical variables. In the fastText era, auxiliary variables were appended to the text as artificial tokens (NAT_21 TYP_X …), which was a workaround. They now each have their own embedding, learned jointly with the text representation.
  • Head. A linear layer maps the combined representation to the sub-classes of the target classification.

The same library also supports self-attention layers, label attention and Hugging Face tokenizers. These options are evaluated on the validation set. A richer architecture is adopted only if it fits within the latency budget.

3 torchTextClassifiers

torchTextClassifiers (documentation, PyPI) is “a unified, extensible framework for text classification with categorical variables, built on PyTorch and PyTorch Lightning” (Couralet and Tayebjee 2025). It was born from the needs of APE coding and is developed as a general-purpose, open-source tool. It sits on the critical path: every prediction goes through it.

Couralet, Cédric, and Meilame Tayebjee. 2025. torchTextClassifiers: A Unified Framework for Text Classification in PyTorch. PyData Global. https://inseefrlab.github.io/codif-ape-prez/slides/pydata-global/.

3.1 Why a dedicated library

When Meta archived the fastText repository in March 2024, the production model was left depending on a compiled C++ library with no maintenance. fastText had other drawbacks:

  • auxiliary variables could only be handled by appending them to the text;
  • it offered no explainability;
  • its scores were poorly calibrated;
  • it could not train on GPU.

A first PyTorch reimplementation, torchFastText, reproduced the fastText architecture. It then grew into a general framework and was renamed torchTextClassifiers in 2025.

There is a reason why the model is a package rather than code inside the training repository. The same model class has to exist in the training repository, in the API, in production and in batch jobs. A versioned package on PyPI, mirrored in INSEE’s internal repository manager, is the single source of truth shared by all of them.

3.2 Design

The library splits a text classifier into interchangeable components:

Table 1: torchTextClassifiers building blocks.
Component Options
Tokenizer fastText-style character n-grams, WordPiece, or any Hugging Face tokenizer
Token embedder Embedding bag, optionally with self-attention layers (“make your own small BERT”)
Sentence embedder Mean, first or last token, or label attention (one representation per class)
Categorical variables One embedding per variable, combined with the text by sum, average-and-concatenate, or concatenation
Head Linear classifier, multi-class or multi-label; a multi-level variant exploits the hierarchy

A high-level wrapper takes care of encoding the raw categorical values and the labels, training (with optional sample weights), prediction with top-k outputs, and saving and loading. Two forms of explainability are built in:

  • token-level attributions with integrated gradients (Captum);
  • label attention, which shows which words drive each candidate class.

The intended users are teams that need accurate text classification but cannot deploy large models: small CPU budgets, unit latency constraints, closed environments. Beyond APE, the library is used for other classifications at INSEE, such as COICOP for product descriptions on receipts. It aims to become a common tool for automatic coding in official statistics.

3.3 In the APE pipeline

  • codif-ape-train builds and trains the model with TTC, and wraps it in the MLflow pyfunc.
  • codif-ape-api installs TTC so that it can load that pyfunc, but does not call it directly.
  • Compared with fastText, TTC gives the same automation rate, much better calibration, explainability and GPU training, for an inference time that remains within budget (see Table 2).

4 Training

Training is run by codif-ape-train, with these settings:

  • Configuration: Hydra, which handles the choice of classification (NAF rev. 2 or NAF 2025), the model, the tokenizer and the optimisation.
  • Framework: PyTorch Lightning.
    • Optimiser: Adam.
    • Learning rate: reduced when the validation loss plateaus.
    • Early stopping on the validation loss.
    • The best checkpoint is kept.
  • Sample weights: rows created by business rules or sampled from particular sources can be weighted up or down in the loss (see Retraining protocol).
  • Execution: an Argo workflow runs on a GPU node and trains the NAF rev. 2 and NAF 2025 models in parallel. Each run starts from a clean clone of the repository and a locked environment.
  • Tracking: each run logs the following in MLflow:
    • the full configuration;
    • its input datasets;
    • the training curves;
    • the evaluation metrics and diagnostic figures;
    • the packaged model.

5 Confidence index and calibration

For each declaration, the model returns its five most likely codes and their probabilities. The decision to code automatically relies on a confidence index (CI): the gap between the two highest probabilities, \(\text{CI} = p_{(1)} - p_{(2)}\).

  • If the CI is at or above a threshold \(\tau\), the first code is assigned automatically.
  • Otherwise, the declaration goes to a human coder, who sees the five suggestions.

The threshold sets the trade-off between automation rate and accuracy of automatic coding. With fastText, a threshold of about 0.75 sent roughly 20 % of declarations to manual coding (Randriamanana and Faria 2025). Since the score distribution changes from one model to the next, the threshold must be recalibrated for every new model on a representative sample. Reusing the previous threshold would shift the manual coding workload in an uncontrolled way.

This is why calibration matters. A probability of 0.9 should be right about 90 % of the time. fastText’s one-versus-all scores were poorly calibrated. The PyTorch model, trained with a softmax cross-entropy loss, is well calibrated out of the box, as measured on a test sample of 300,000 descriptions (Tayebjee et al. 2025). Calibration is measured at every run with the expected calibration error (ECE) and the Brier score. Going beyond a single threshold, conformal prediction is explored under Future work.

6 Evaluation

Every training run logs accuracy at each level of the classification, calibration metrics and the accuracy–automation curve. What these metrics mean, which data they are computed on, and how a model is judged are covered in the Evaluation chapter.

7 Results

7.1 fastText vs. PyTorch

Table 2: Comparison of the two model generations (Tayebjee et al. 2025).
Tayebjee, Meilame, Cédric Couralet, Nathan Randriamanana, and Julien Pramil. 2025. fastText, Tout Juste En Production Et Déjà Obsolète ? Présentation Du Package torchTextClassifiers Pour La Modernisation de La Codification Automatique via l’exemple de l’APE. Journées de méthodologie statistique (JMS). https://inseefrlab.github.io/codif-ape-prez/papers/jms-torchft/index.pdf.
fastText torchTextClassifiers
Automation rate ~80 % ~80 %
Mean inference time (CPU, unit call) < 100 ms < 300 ms
Calibration Limited Excellent
Categorical variables Appended to the text Dedicated embeddings
Explainability None Integrated gradients, label attention
GPU training No Yes
Maintenance Archived upstream (2024) Maintained in-house, open source

The PyTorch model trades some latency for maintainability, calibration and the ability to evolve. The latency remains within the production budget.

7.2 NAF 2025

Accuracy on the test set of the NAF 2025 retraining, by level (Randriamanana and Faria 2025):

Randriamanana, Nathan, and Thomas Faria. 2025. Changement de Nomenclature Statistique à l’ère Du Machine Learning : Réentraînement Basé Sur Les LLM Pour La Révision de La NAF. Journées de méthodologie statistique (JMS). https://inseefrlab.github.io/codif-ape-prez/papers/jms2025-llm-relabeling/index.pdf.
Table 3: Test accuracy of the models retrained on NAF 2025 data, before data editing.
Model Section Division Group Class Sub-class
fastText 0.855 0.833 0.797 0.784 0.776
torchTextClassifiers 0.92 0.89 0.86 0.84 0.83

Applying business rules to the training data (data editing) raised sub-class accuracy to about 0.88–0.90 for both models. The gain is large but must be read with care: rules also change the test labels, so part of the gain reflects closer agreement with the rules rather than better generalisation. The retraining protocol is designed precisely to separate the two effects.

8 Future work

The production model is deliberately simple. Two lines of work look at how to go further without breaking the production constraints: a richer architecture, and a better use of the model’s uncertainty.

8.1 Architecture: Julien Bernardi’s internship

Julien Bernardi’s internship (Innovation Lab, 2026) benchmarked lightweight architectures built with torchTextClassifiers. The report is online at julber95.github.io/interpretable-text-classification. This is exploratory work.

Validating torchTextClassifiers. Before NAF, the architectures were trained on two public benchmarks with known reference results: Amazon reviews (5 classes) and CLINC150 (150 intents). On Amazon reviews, the deeper models reach about 62–63 % accuracy, close to the 63.3 % of multilingual BERT, against about 59 % for the fastText-style model. These results validated the library’s implementation.

Architecture innovations. The internship compared three extensions of the production fastText-style model:

  • FATE (FastText Attentive Transformer Encoder): self-attention layers on top of the token embeddings, before pooling. It lowers the loss on Amazon reviews, but brings nothing on CLINC150;
  • label attention: one learned query per class replaces mean pooling, so the model looks at the words that matter for each candidate code;
  • multi-level classification: one head per level of the classification, trained jointly.

On NAF 2025, the main findings are:

  • Label attention needs data. With 160,000 training descriptions it is clearly behind the fastText-style model. On the full training set (1.2 million descriptions) it catches up, or slightly overtakes it, but takes 3 to 6 times longer to train.
  • Its attention weights explain the prediction well. Deleting the words it attends to lowers the predicted probability faster than deleting the words chosen by integrated gradients.
  • Multi-level heads help when the sub-class is wrong. At coarse levels, they beat simply truncating the predicted sub-class. Among wrong sub-class predictions, the share that is wrong at every level falls from about 47 % to 29–34 %. At the sub-class level, both approaches tie.
  • Data matters more than architecture for rare codes. The accuracy of a sub-class depends mainly on how many training examples it has.

As for the options already in the library (above), an innovation is adopted only if it fits within the CPU latency budget, which the internship did not measure, and passes the quality gates of the retraining protocol. Multi-level heads suggest an untested use: when the sub-class is uncertain, show the manual coder a reliable division or class.

8.2 Conformal prediction

The confidence index answers one question: can this declaration be coded automatically? It says little about the cases sent to manual coding. The coder always sees the same five suggestions, whether the model hesitates between two codes or has no idea.

Conformal prediction (Angelopoulos and Bates 2023) replaces this fixed list with a statistical guarantee. A threshold is computed on a calibration set held out from training. For a new description, the model then returns a set of codes that contains the right one with probability at least \(1 - \alpha\), for example 95 %. The guarantee holds for any model, provided the calibration data and the production data are exchangeable. The size of the set carries the uncertainty: one or two codes when the model is sure, many more when it is not.

Angelopoulos, Anastasios N., and Stephen Bates. 2023. “Conformal Prediction: A Gentle Introduction.” Foundations and Trends in Machine Learning 16 (4): 494–591. https://doi.org/10.1561/2200000101.

Exploratory work in codif-ape-train (src/conformal_pred) applies split conformal prediction to the trained model:

  • Calibration on manual cases only. Only validation declarations whose CI is below the threshold, those that would go to manual coding, are used. This selection depends only on the model’s output, so the guarantee still holds within the subset. It therefore applies exactly to the declarations that human coders see.
  • Three scores are compared:
    • \(1 - \hat{p}(y \mid x)\), which tends to give the smallest sets;
    • the cumulative probability of the codes ranked up to the true one (adaptive prediction sets, Romano et al. (2020)), which adapts the size of the set to the difficulty of each case;
    • the rank of the true code, which amounts to choosing the length \(k\) of a top-\(k\) list with a guarantee.
  • Evaluation on the test set, by empirical coverage and mean set size.
Romano, Yaniv, Matteo Sesia, and Emmanuel J. Candès. 2020. “Classification with Valid and Adaptive Coverage.” Advances in Neural Information Processing Systems (NeurIPS). https://arxiv.org/abs/2006.02544.

Several questions remain open before any production use:

  • Conditional coverage. The guarantee is marginal: 95 % on average over all declarations, possibly much less for some sections. Calibrating separately by section (Mondrian conformal prediction) would fix this, at the cost of more calibration data.
  • Drift. The way declarants describe their activity changes over time (Monitoring), which weakens exchangeability. Like \(\tau\), the conformal threshold must be recomputed for every new model on recent data, and its coverage must be monitored.
  • Interface. A set of variable size changes the API contract, which currently returns a fixed top 5 (Serving), and the manual coding interface.
  • Choosing \(\tau\). The same tools could set the threshold \(\tau\) itself, with a guarantee on the error rate of automatic coding (conformal risk control, Angelopoulos et al. (2024)).
Angelopoulos, Anastasios N., Stephen Bates, Adam Fisch, Lihua Lei, and Tal Schuster. 2024. “Conformal Risk Control.” International Conference on Learning Representations (ICLR). https://arxiv.org/abs/2208.02814.
Back to top