I’m working on a multi‑class text classification problem where the input consists of very short descriptions (often only a few words) and the goal is to predict the correct category. I’m currently using an XGBoost classifier for the final prediction layer. for embeddings I am using e5 large model..
The main challenges I’m facing are:
- The descriptions are extremely short and many classes share similar vocabulary. Because of this, the model sometimes produces high‑confidence but incorrect predictions when common words appear across multiple classes.
-
The dataset is imbalanced: some classes have significantly more training samples than others, but the distribution of the prediction data is different. This causes the model to over‑predict certain classes.
-
I experimented with TF‑IDF features, embeddings from a generative AI model and a combination of TF‑IDF + embeddings. However, combining both actually reduced accuracy. I also tried downsampling majority class but it not really make any difference. Current model accuracy is around 30%.
I’m looking for advice on:
-
How to improve classification when different classes share highly overlapping vocabulary.
-
How to reduce high‑confidence wrong predictions.
-
What techniques work well for class imbalance and distribution shift between training and prediction data.
Why use xgboost instead of neural network for a text classification problem?
Also, have you seen this?
did you try using data augmentation ?
also how imbalanced is your data?
as I can see it is category classification, I wonder what is your algorithm behind and hyperparameters to check the variation in training.
Another important query is to check how have you divided the datasets. between validation, test and training or it is only validation and training only
Hi, I chose xgboost because the labeled dataset is small (only around 700 records). Since I’m using pretrained e5-large embeddings, xgb provides a strong baseline and is faster to retrain which fits my human in the loop retraining workflow. Also, many classes share overlapping terminology and the current dataset does not have enough examples showing those terms across different commodities. Before changing the model, I am focusing on improving the training data through the retraining loop and adding more labeled examples.
Hi, I have not tried augmentation yet. For manufacturing text, I think generic techniques like random synonym replacement can change the technical meaning. are there any domain specific augmentation techniques you would recommend?
For class distribution the largest classes have around 400 and 250 samples while smaller classes have around 70 and 20 samples. However, I think the bigger challenge is description diversity rather than class imbalance. The training data contains many similar and repetitive descriptions, while the inference set (10,000 records) has much more varied wording.
Because the model sees limited patterns during training, it learns strong associations between certain terms and a class. If a term only appears under one class in training and never appears under another class, the model will make a high confidence prediction based on that learned pattern. It seems more like a data coverage issue than a model issue but I am looking for ways to improve this.
pipeline uses e5-large embeddings from the description feature followed by an xgb multiclass classifier. current XGBoost parameters
n_estimators=300, max_depth=4, learning_rate=0.05, subsample=0.85, colsample_bytree=0.75, reg_alpha=2, reg_lambda=3, min_child_weight=3
and yeah split I use an 80-20 stratified split on the labeled data for training and validation. I also run inference on around 10,000 unseen records but the model gets confused by overlapping terminology across classes and makes high confidence incorrect predictions. as more labeled data becomes available through the human in theloop process, I plan to maintain a separate holdout test set.
hello @Shahina_Rajamani
thank you for the information.
how is the class distribution of 700 labelled dataset??
looks like xboost isn’t good practice for text classification with highly imbalanced dataset, instead use a pre-trained language model combined with semi-supervised learning or loss adjustment.
kindly provide information on class distribution of 700 labelled dataset, so that I can further suggest you.
Also provide information on what kind of model architecture you have created. if its a privacy concern, then you can personal DM the necessary information by clicking on my name and then message.
Regards
Dr. Deepti
I have updated my previous message with the all the details you requested for.
Please use a transformer based approach if applicable. It considers word position (via attention). Subword tokenizers might also be of help.
If it’s too much, an RNN based approach will still do good.
Thanks for the suggestion. My current pipeline already uses a pretrained transformer e5 large model to generate contextual embeddings so it benefits from selfattention, positional encoding and also the subword tokenization. as I am using XGBoost as the downstream classifier rather than fine tuning the transformer end to end, are you saying that I should be replacing XGBoost as well with a finetuned transformer classification head to improve performance?
Freeze all weights except the final classification layer. Is it safe to assume that you are familiar with deep learning specialization? The reason for asking is that my suggestion touches upon a concept called transfer learning. There are techniques mentioned in the specialization on how to prioritize fine tuning a model.
hi @Shahina_Rajamani
looking at your codes helps more to understand the problem in the approach.
When working with deep contextual embeddings like e5-large paired with tree-based models like XGBoost, accuracy roadblocks typically occur because trees are fundamentally poorly suited for continuous, highly correlated, dense vector manifolds.
For better accuracy, consider the following pipeline adjustments:
1. Swap the downstream model architecture
-
Issue -tree-based models partition feature spaces along orthogonal axes (per-dimension thresholds). Dense embeddings don’t store information in individual indices, the meaning is stored in the relative angles and geometric relationships across dimensions.
-
probable approach replace XGBClassifier with a Logistic Regression (with L2 regularization) or a simple Multi-Layer Perceptron (MLPClassifier). Linear models and shallow neural networks almost always outperform gradient-boosted trees on top of Transformer embeddings.
from sklearn.linear_model import LogisticRegression
# Typically provides a 3-8% absolute accuracy boost over trees on dense vectors
commodity_model = LogisticRegression(
max_iter=1000,
class_weight='balanced',
C=1.0,
random_state=42
)
- Fine-tune text normalization for E5
-
another issue: your code strips out all non-alphanumeric characters: F.regexp_replace(F.col("clean_text"), "[^a-zA-Z0-9 ]", "").
-
proable apprach - transformer architectures (like BERT/E5) rely heavily on punctuation, slashes, and structural symbols to build structural semantic meaning—especially in engineering or manufacturing part descriptions (line "BOLT, M6x12" vs "BOLT M612").
-
*what to do " -try bypassing the regex character stripping entirely. Let the e5-large tokenizer handle raw strings natively, as its vocabulary is trained on uncleaned internet text.
3. using joint hierarchical context
-
issue - your data explicitly contains a broader grouping called part_family along with the target commodity. Right now the model tries to predict commodity purely based on text, ignoring this helpful metadata.
-
what you could try - inject the hierarchical context into the transformer. Concatenate your metadata columns straight into your text string before generating embeddings.
# Enriching the context text before embedding generation
embedding_input = [
f"passage: Family: {row['part_family']} | Description: {row['clean_text']}"
for _, row in train_pdf.iterrows()
]
- tune the embedding extraction profile
-
important issue - right now configuration utilizes intfloat/e5-large. E5 models are asymmetric and require strict prefix prompts to trigger correct vector projection fields.
-
what you could do - ensure your training text blocks use the exact same prefix as your validation/testing blocks. If your text is short (like part descriptions), experiment with the query: prefix instead of passage: —the query field is optimized for short search phrases and fragments.
- address Out-of-Vocabulary (OOV) codes
-
metadata issue handling -industrial part descriptions often contain highly specific serial codes or part numbers (like "VND-771-A"). Dense semantic models often map these alphanumeric strings to generic or muddy embedding regions.
-
you can try - to build a hybrid model. Combine your dense embeddings with a classic sparse text matrix (TF-IDF with character n-grams) using a FeatureUnion. This ensures that specific textual keywords or sub-strings force exact structural classification matches even if the semantic model gets confused.
I am sending you the hybrid text formatting code by DM for your privacy concern.
regards
Dr. Deepti
about handling hyperparameter much efficiently in an imbalanced dataset especially keeping in mind about overfitting and underfitting, instead of using standard LogisticRegression useLogisticRegressionCV. this will automatically test multiple C values using stratified cross-validation.
from sklearn.linear_model import LogisticRegressionCV
print("Optimizing model via cross-validation...")
# Automatically sweeps through multiple regularization values
commodity_model = LogisticRegressionCV(
Cs=[0.01, 0.1, 1.0, 10.0],
cv=3, # 3-fold stratified cross-validation
max_iter=1000,
class_weight='balanced',
scoring='f1_macro', # Optimizes for imbalanced classes
n_jobs=-1, # Uses all worker CPU cores
random_state=42
)
commodity_model.fit(X_train, y_train)
# Display the top performing C value
print(f"Optimal C regularization parameter: {commodity_model.C_[0]}")
- ** handling large imbalanced dataset
Using .toPandas() crashes your driver node when processing massive datasets. Moving embedding generation into a distributed Spark Pandas UDF splits the text across worker nodes.
import pandas as pd
import pyspark.sql.functions as F
from pyspark.sql.types import ArrayType, FloatType
# Define distributed execution worker function
@F.pandas_udf(ArrayType(FloatType()))
def generate_embeddings_udf(text_series: pd.Series) -> pd.Series:
from sentence_transformers import SentenceTransformer
# Load model locally inside the worker memory node
model = SentenceTransformer("intfloat/e5-large")
embeddings = model.encode(
text_series.tolist(),
batch_size=64, # High batch size for faster GPU/CPU processing
show_progress_bar=False,
normalize_embeddings=True
)
return pd.Series(list(embeddings))
# 1. Build hybrid string column natively inside Spark
spark_df_with_context = train_df.withColumn(
"hybrid_text",
F.concat(
F.lit("passage: Family: "), F.col("part_family"),
F.lit(" | Description: "), F.col("clean_text")
)
)
# 2. Scale embedding generation across the entire cluster
distributed_embeddings_df = spark_df_with_context.withColumn(
"embedding",
generate_embeddings_udf(F.col("hybrid_text"))
)
# Now you can safely save directly to a delta table
# distributed_embeddings_df.write.format("delta").save(OUTPUT_PATH)
try these approaches, hope it helps, Good luck @Shahina_Rajamani 
Regards
Dr. Deepti
@Deepti_Prasad Just wanted to share a quick update, switching to LogisticRegressionCV with TF‑IDF/char‑ngrams has significantly improved Stage1 accuracy even without using embeddings. The automatic hyperparameter sweep over multiple C values (with stratified CV and macroF1 scoring) made the model much more stable on our imbalanced classes and reduced both overfitting and underfitting.
The TF-IDF + LR CV baseline is already outperforming the embedding only model for several classes especially where descriptions are short or noisy and lexical features carry more signal. I’m continuing to work on improving accuracy further 
Thanks much for the guidance