Constraining Output Space for SLM Narrow Automation Optimization

Constraining the output space instead of parsing generated text can make small language models (SLMs) faster and more reliable for narrow automation tasks, according to a technical article benchmarking Qwen2.5-0.5B-Instruct on an M2 MacBook Air. The technique runs one forward pass and scores candidate labels directly, avoiding slow multi-token generation and error-prone regex parsing, and yields a calibrated confidence score as a byproduct.

Constraining Output Space for SLM Narrow Automation Optimization This article will kick off a series on narrow automation optimization for SLMs, and as the first entry will cover one of the more most useful techniques for doing so: constraining the output space instead of parsing generated text. So much of the attention in applied AI goes to frontier-scale reasoning; however, a large share of the actual production workload necessity in industry is far less glamorous: narrow automation. Tasks that fall into this category include properly routing a support ticket, extracting a field from a form, tagging a document, and flagging a record for human review. These tasks all have common characteristics: constrained input, a fixed output space, and enormous call volume. They are also exactly the types of tasks that are well-suited to small language models SLMs . A model that fits comfortably on one GPU, or is even CPU-bound, and is able to return an answer in milliseconds can often be the correct engineering choice over an API call to a large language model LLM that could cost a thousand times more per item. The trouble is that teams tend to carry frontier-model habits over to small models. They write long conversational prompts and let the model generate free-form text before hunting through it with regular expressions. They call the model once at a time from inside a Python loop. Against a local SLM, these types of inefficiencies are accentuated: when a single forward pass takes ten milliseconds, everything you wrap around that forward pass becomes the bottleneck; loose output handling turns directly into measurable error rates. This article will kick off a series on narrow automation optimization for SLMs, and as the first entry will cover one of the more most useful techniques for doing so: constraining the output space instead of parsing generated text. To set a level playing field, all benchmarks below use Qwen2.5-0.5B-Instruct in float16 through , running on an M2 Macbook Air with 24GB RAM and a 16-core Neural Engine. Hugging Face Transformers https://huggingface.co/docs/transformers/index First, setup a Python environment and install your requirements: pip install torch transformers accelerate Why Constraining the Output Space? A classification task has a fixed answer set. If you are routing tickets into billing , technical , or account , there are exactly three valid outputs and no others. Yet the standard pattern is to ask the model to write the answer, generate a handful of tokens, and then search the resulting string for something recognizable. This fails in two ways simultaneously. First, it is slow: generate runs one sequential forward pass per output token, so asking for eight tokens costs roughly eight times the compute of asking for the answer directly. Second, it is unreliable: a small model will happily reply with "Sure This looks like a billing issue.", or "Billing/Account", or a category you never defined. Every one of those responses requires either a fallback rule or a retry, and every fallback rule is a place for error accumulation. The fix is to stop generating and start scoring. Run one forward pass, read the model's next-token distribution, and restrict your decision to the token IDs of your candidate labels. The answer becomes impossible to get wrong structurally, and you get a calibrated confidence score as a byproduct. Parsing Free Text Here is the naive version, generating free text and parsing it after the fact. Save it to file and run it from the command line. python import os import time import torch from transformers import AutoTokenizer, AutoModelForCausalLM MODEL ID = "Qwen/Qwen2.5-0.5B-Instruct" torch.set num threads os.cpu count or 1 tokenizer = AutoTokenizer.from pretrained MODEL ID if tokenizer.pad token is None: tokenizer.pad token = tokenizer.eos token model = AutoModelForCausalLM.from pretrained MODEL ID, dtype=torch.float32 model.eval our toy data to classify 600 records LABELS = "billing", "technical", "account" tickets = "My card was charged twice for the same invoice.", "The mobile app crashes whenever I open the settings page.", "I need to change the email address on my profile.", 200 def build prompt ticket : messages = { "role": "system", "content": "You classify support tickets. Answer with exactly one of: billing, technical, account.", }, {"role": "user", "content": f"Ticket: {ticket}\nCategory:"}, return tokenizer.apply chat template messages, tokenize=False, add generation prompt=True nothing constrains the output here, so we let the model write a short answer and search it for a label tokens++ each new token costs its own forward pass, and one ticket per call means no batching to amortize that time++ prompts = build prompt t for t in tickets predictions = time inference start = time.time for n, prompt in enumerate prompts, start=1 : this loop runs for minutes on CPU, so report progress rather than sitting silent if n % 50 == 0: rate = time.time - start / n print f" {n}/{len prompts } tickets {rate:.2f}s each ", flush=True inputs = tokenizer prompt, return tensors="pt" .to model.device with torch.inference mode : output = model.generate inputs, max new tokens=8, do sample=False, pad token id=tokenizer.eos token id, generate returns prompt + continuation, so slice the prompt off before decoding generated = output 0, inputs "input ids" .shape 1 : text = tokenizer.decode generated, skip special tokens=True .strip .lower substring match against the label list predictions.append next label for label in LABELS if label in text , "UNPARSED" duration = time.time - start output task metrics print f"Free-form generation took: {duration:.2f} seconds" print f"Unparseable outputs: {predictions.count 'UNPARSED' } / {len predictions }" sample of inference output for ticket, label in zip tickets -3: , predictions -3: , strict=True : print f"{ticket} - {label}" Output: Loading weights: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 290/290 00:01<00:00, 263.89it/s 50/600 tickets 0.19s each 100/600 tickets 0.19s each 150/600 tickets 0.19s each 200/600 tickets 0.19s each 250/600 tickets 0.19s each 300/600 tickets 0.19s each 350/600 tickets 0.19s each 400/600 tickets 0.19s each 450/600 tickets 0.19s each 500/600 tickets 0.19s each 550/600 tickets 0.19s each 600/600 tickets 0.19s each Free-form generation took: 134.01 seconds Unparseable outputs: 0 / 600 My card was charged twice for the same invoice. - billing The mobile app crashes whenever I open the settings page. - technical I need to change the email address on my profile. - technical While the the entirety of the batch came back in a shape the parser could handle, we will note the 134 second execution time. Constraining the Output Space Now let's try a constrained version, which scores the label set directly from a single forward pass: python import os import time import torch from transformers import AutoTokenizer, AutoModelForCausalLM MODEL ID = "Qwen/Qwen2.5-0.5B-Instruct" torch.set num threads os.cpu count or 1 tokenizer = AutoTokenizer.from pretrained MODEL ID model = AutoModelForCausalLM.from pretrained MODEL ID, torch dtype=torch.float32 model.eval our toy data to classify 600 records LABELS = "billing", "technical", "account" tickets = "My card was charged twice for the same invoice.", "The mobile app crashes whenever I open the settings page.", "I need to change the email address on my profile.", 200 def build prompt ticket : messages = { "role": "system", "content": "You classify support tickets. Answer with exactly one of: billing, technical, account.", }, {"role": "user", "content": f"Ticket: {ticket}\nCategory:"}, return tokenizer.apply chat template messages, tokenize=False, add generation prompt=True prompt ends with "<|im start| assistant\n", so the model's next token starts the label comparing the logits of each label's FIRST token is enough to pick a winner, provided those first tokens are distinct label first ids = tokenizer.encode label, add special tokens=False 0 for label in LABELS assert len set label first ids == len LABELS , "Labels share a first token; score full label sequences instead see notes ." label first ids = torch.tensor label first ids, device=model.device one forward pass per ticket, no generation loop: the decision contained entirely in the next-token logits prompts = build prompt t for t in tickets predictions = confidences = time inference start = time.time for prompt in prompts: inputs = tokenizer prompt, return tensors="pt" .to model.device with torch.inference mode : logits = model inputs .logits 0, -1, : softmax over just the label logits, so the probabilities sum to 1 across the candidates probs = torch.softmax logits label first ids .float , dim=-1 best = int probs.argmax predictions.append LABELS best confidences.append float probs best duration = time.time - start output task metrics print f"Constrained scoring took: {duration:.2f} seconds" print f"Unparseable outputs: {predictions.count 'UNPARSED' } / {len predictions }" sample of inference output for ticket, label, confidence in zip tickets -3: , predictions -3: , confidences -3: , strict=True : print f"{ticket} - {label} confidence {confidence:.3f} " Output: Loading weights: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 290/290 00:01<00:00, 285.60it/s Constrained scoring took: 94.51 seconds Unparseable outputs: 0 / 600 My card was charged twice for the same invoice. - billing confidence 0.793 The mobile app crashes whenever I open the settings page. - technical confidence 0.798 I need to change the email address on my profile. - technical confidence 0.673 This took about 30% less time, with the potential of failure eliminated. Further tests show the time ratio holds at scale, and also that messing with the ticket text can expose the failures in the naive version that are caught with the second version. I'll leave testing this to the reader. Some further explanation of the code above: - Reading logits 0, -1, : gives the model's unnormalized distribution over the next token. Everything generate would do afterward is unnecessary when the answer is one of three known strings. - Indexing that vector at label first ids and taking argmax makes an out-of-vocabulary answer structurally impossible. The model is no longer allowed to be creative about formatting, which is why the unparseable count is 0 / 600 by construction rather than by luck. - The softmax over the restricted logits is a useful confidence check. Practically speaking, you could route anything below a threshold you choose — say, 0.6 as a starting point — to a human queue rather than allowing a low-confidence label to flow downstream in the workflow. - Mind the tokenization. Most byte-level BPE tokenizers treat " billing" and "billing" as distinct tokens, so encode the variant the model would actually emit after your prompt. The chat template ends with "<|im start| assistant\n" , so the next token follows a newline and carries no leading space, hence encode label rather than encode " " + label . Mix this up and the script runs fine; however, you end up scoring three tokens the model was never going to emit. - If two labels share a first token "refund request" and "refund status" , for instance , the assertion fires. Either rename the labels to single distinct tokens such as A , B , C with a legend in the prompt, or score the full label sequences instead of the first token. Wrapping Up This has been our first attempt at optimizing SLMs for narrow automation, and our target technique this time was constrained scoring . This technique replaces free-form generation and string parsing with a single forward pass restricted to the valid label set. By implementing it, we can make malformed output structurally impossible while handing you a confidence score for routing edge cases to humans. A small language model, such as the 0.5B parameter model we used today, becomes a practical production choice for narrow automation once the code around it stops treating it like a generic chatbot, and stops interacting with it like it would ChatGPT. With an enforced output contract, the small model stops being a compromise and starts being the obvious answer. Matthew Mayo https://www.kdnuggets.com/wp-content/uploads/./profile-pic.jpg holds a master's degree in computer science and a graduate diploma in data mining. As managing editor of https://twitter.com/mattmayo13 @mattmayo13 KDnuggets https://www.kdnuggets.com/ & Statology https://www.statology.org/ , and contributing editor at Machine Learning Mastery https://machinelearningmastery.com/ , Matthew aims to make complex data science concepts accessible. His professional interests include natural language processing, language models, machine learning algorithms, and exploring emerging AI. He is driven by a mission to democratize knowledge in the data science community. Matthew has been coding since he was 6 years old.