{"slug": "constraining-output-space-for-slm-narrow-automation-optimization", "title": "Constraining Output Space for SLM Narrow Automation Optimization", "summary": "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.", "body_md": "# Constraining Output Space for SLM Narrow Automation Optimization\n\nThis 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.\n\nSo 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.\n\nThe 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.\n\nThis 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.\n\nTo set a level playing field, all benchmarks below use ** Qwen2.5-0.5B-Instruct** in float16 through\n\n**, running on an M2 Macbook Air with 24GB RAM and a 16-core Neural Engine.**\n\n[Hugging Face Transformers](https://huggingface.co/docs/transformers/index)First, setup a Python environment and install your requirements:\n\n```\npip install torch transformers accelerate\n```\n\n## # Why Constraining the Output Space?\n\nA classification task has a fixed answer set. If you are routing tickets into `billing`\n\n, `technical`\n\n, or `account`\n\n, 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.\n\nThis fails in two ways simultaneously. First, it is slow: `generate()`\n\nruns 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.\n\nThe 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.\n\n## # Parsing Free Text\n\nHere is the naive version, generating free text and parsing it after the fact. Save it to file and run it from the command line.\n\n``` python\nimport os\nimport time\nimport torch\nfrom transformers import AutoTokenizer, AutoModelForCausalLM\n\nMODEL_ID = \"Qwen/Qwen2.5-0.5B-Instruct\"\n\ntorch.set_num_threads(os.cpu_count() or 1)\n\ntokenizer = AutoTokenizer.from_pretrained(MODEL_ID)\nif tokenizer.pad_token is None:\n    tokenizer.pad_token = tokenizer.eos_token\n\nmodel = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32)\nmodel.eval()\n\n# our toy data to classify (600 records)\nLABELS = [\"billing\", \"technical\", \"account\"]\ntickets = [\n    \"My card was charged twice for the same invoice.\",\n    \"The mobile app crashes whenever I open the settings page.\",\n    \"I need to change the email address on my profile.\",\n] * 200\n\ndef build_prompt(ticket):\n    messages = [\n        {\n            \"role\": \"system\",\n            \"content\": \"You classify support tickets. Answer with exactly one of: billing, technical, account.\",\n        },\n        {\"role\": \"user\", \"content\": f\"Ticket: {ticket}\\nCategory:\"},\n    ]\n    return tokenizer.apply_chat_template(\n        messages, tokenize=False, add_generation_prompt=True\n    )\n\n# nothing constrains the output here, so we let the model write a short answer and search it for a label (tokens++)\n# each new token costs its own forward pass, and one ticket per call means no batching to amortize that (time++)\nprompts = [build_prompt(t) for t in tickets]\npredictions = []\n\n# time inference\nstart = time.time()\n\nfor n, prompt in enumerate(prompts, start=1):\n\n    # this loop runs for minutes on CPU, so report progress rather than sitting silent\n    if n % 50 == 0:\n        rate = (time.time() - start) / n\n        print(f\"  {n}/{len(prompts)} tickets ({rate:.2f}s each)\", flush=True)\n    inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n    with torch.inference_mode():\n        output = model.generate(\n            **inputs,\n            max_new_tokens=8,\n            do_sample=False,\n            pad_token_id=tokenizer.eos_token_id,\n        )\n\n    # generate() returns prompt + continuation, so slice the prompt off before decoding\n    generated = output[0, inputs[\"input_ids\"].shape[1] :]\n    text = tokenizer.decode(generated, skip_special_tokens=True).strip().lower()\n\n    # substring match against the label list\n    predictions.append(next((label for label in LABELS if label in text), \"UNPARSED\"))\n\nduration = time.time() - start\n\n# output task metrics\nprint(f\"Free-form generation took: {duration:.2f} seconds\")\nprint(f\"Unparseable outputs: {predictions.count('UNPARSED')} / {len(predictions)}\")\n\n# sample of inference output\nfor ticket, label in zip(tickets[-3:], predictions[-3:], strict=True):\n    print(f\"{ticket} -> {label}\")\n```\n\nOutput:\n\n```\nLoading weights: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 290/290 [00:01<00:00, 263.89it/s]\n  50/600 tickets (0.19s each)\n  100/600 tickets (0.19s each)\n  150/600 tickets (0.19s each)\n  200/600 tickets (0.19s each)\n  250/600 tickets (0.19s each)\n  300/600 tickets (0.19s each)\n  350/600 tickets (0.19s each)\n  400/600 tickets (0.19s each)\n  450/600 tickets (0.19s each)\n  500/600 tickets (0.19s each)\n  550/600 tickets (0.19s each)\n  600/600 tickets (0.19s each)\nFree-form generation took: 134.01 seconds\nUnparseable outputs: 0 / 600\nMy card was charged twice for the same invoice. -> billing\nThe mobile app crashes whenever I open the settings page. -> technical\nI need to change the email address on my profile. -> technical\n```\n\nWhile the the entirety of the batch came back in a shape the parser could handle, we will note the 134 second execution time.\n\n## # Constraining the Output Space\n\nNow let's try a constrained version, which scores the label set directly from a single forward pass:\n\n``` python\nimport os\nimport time\nimport torch\nfrom transformers import AutoTokenizer, AutoModelForCausalLM\n\nMODEL_ID = \"Qwen/Qwen2.5-0.5B-Instruct\"\n\ntorch.set_num_threads(os.cpu_count() or 1)\n\ntokenizer = AutoTokenizer.from_pretrained(MODEL_ID)\n\nmodel = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32)\nmodel.eval()\n\n# our toy data to classify (600 records)\nLABELS = [\"billing\", \"technical\", \"account\"]\ntickets = [\n    \"My card was charged twice for the same invoice.\",\n    \"The mobile app crashes whenever I open the settings page.\",\n    \"I need to change the email address on my profile.\",\n] * 200\n\ndef build_prompt(ticket):\n    messages = [\n        {\n            \"role\": \"system\",\n            \"content\": \"You classify support tickets. Answer with exactly one of: billing, technical, account.\",\n        },\n        {\"role\": \"user\", \"content\": f\"Ticket: {ticket}\\nCategory:\"},\n    ]\n    return tokenizer.apply_chat_template(\n        messages, tokenize=False, add_generation_prompt=True\n    )\n\n# prompt ends with \"<|im_start|>assistant\\n\", so the model's next token starts the label\n# comparing the logits of each label's FIRST token is enough to pick a winner, provided those first tokens are distinct\nlabel_first_ids = [tokenizer.encode(label, add_special_tokens=False)[0] for label in LABELS]\nassert len(set(label_first_ids)) == len(LABELS), (\n    \"Labels share a first token; score full label sequences instead (see notes).\"\n)\nlabel_first_ids = torch.tensor(label_first_ids, device=model.device)\n\n# one forward pass per ticket, no generation loop: the decision contained entirely in the next-token logits\nprompts = [build_prompt(t) for t in tickets]\npredictions = []\nconfidences = []\n\n# time inference\nstart = time.time()\nfor prompt in prompts:\n    inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n    with torch.inference_mode():\n        logits = model(**inputs).logits[0, -1, :]\n    # softmax over just the label logits, so the probabilities sum to 1 across the candidates\n    probs = torch.softmax(logits[label_first_ids].float(), dim=-1)\n    best = int(probs.argmax())\n    predictions.append(LABELS[best])\n    confidences.append(float(probs[best]))\nduration = time.time() - start\n\n# output task metrics\nprint(f\"Constrained scoring took: {duration:.2f} seconds\")\nprint(f\"Unparseable outputs: {predictions.count('UNPARSED')} / {len(predictions)}\")\n\n# sample of inference output\nfor ticket, label, confidence in zip(\n    tickets[-3:], predictions[-3:], confidences[-3:], strict=True\n):\n    print(f\"{ticket} -> {label} (confidence {confidence:.3f})\")\n```\n\nOutput:\n\n```\nLoading weights: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 290/290 [00:01<00:00, 285.60it/s]\nConstrained scoring took: 94.51 seconds\nUnparseable outputs: 0 / 600\nMy card was charged twice for the same invoice. -> billing (confidence 0.793)\nThe mobile app crashes whenever I open the settings page. -> technical (confidence 0.798)\nI need to change the email address on my profile. -> technical (confidence 0.673)\n```\n\nThis 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.\n\nSome further explanation of the code above:\n\n- Reading\n`logits[0, -1, :]`\n\ngives the model's unnormalized distribution over the next token. Everything`generate()`\n\nwould do afterward is unnecessary when the answer is one of three known strings. - Indexing that vector at\n`label_first_ids`\n\nand taking`argmax`\n\nmakes 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`\n\nby 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.\n- Mind the tokenization. Most byte-level BPE tokenizers treat\n`\" billing\"`\n\nand`\"billing\"`\n\nas distinct tokens, so encode the variant the model would actually emit after your prompt. The chat template ends with`\"<|im_start|>assistant\\n\"`\n\n, so the next token follows a newline and carries no leading space, hence`encode(label)`\n\nrather than`encode(\" \" + label)`\n\n. 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 (\n`\"refund_request\"`\n\nand`\"refund_status\"`\n\n, for instance), the assertion fires. Either rename the labels to single distinct tokens (such as`A`\n\n,`B`\n\n,`C`\n\n) with a legend in the prompt, or score the full label sequences instead of the first token.\n\n## # Wrapping Up\n\nThis 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.\n\nA 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.\n\n(\n\n[Matthew Mayo](https://www.kdnuggets.com/wp-content/uploads/./profile-pic.jpg)\n\n[) holds a master's degree in computer science and a graduate diploma in data mining. As managing editor of](https://twitter.com/mattmayo13)\n\n**@mattmayo13**[KDnuggets](https://www.kdnuggets.com/)&\n\n[Statology](https://www.statology.org/), and contributing editor at\n\n[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.", "url": "https://wpnews.pro/news/constraining-output-space-for-slm-narrow-automation-optimization", "canonical_source": "https://www.kdnuggets.com/constraining-output-space-small-language-model-narrow-automation-optimization", "published_at": "2026-08-13 12:00:17+00:00", "updated_at": "2026-08-13 12:35:16.311313+00:00", "lang": "en", "topics": ["machine-learning", "large-language-models", "natural-language-processing", "ai-tools"], "entities": ["Qwen2.5-0.5B-Instruct", "Hugging Face Transformers", "M2 MacBook Air", "Python"], "alternates": {"html": "https://wpnews.pro/news/constraining-output-space-for-slm-narrow-automation-optimization", "markdown": "https://wpnews.pro/news/constraining-output-space-for-slm-narrow-automation-optimization.md", "text": "https://wpnews.pro/news/constraining-output-space-for-slm-narrow-automation-optimization.txt", "jsonld": "https://wpnews.pro/news/constraining-output-space-for-slm-narrow-automation-optimization.jsonld"}}