翻訳待ち：Constraining Output Space for SLM Narrow Automation Optimization

AI サービスが一時的に利用できないため、復旧後に翻訳を補完します。

--> Constraining Output Space for SLM Narrow Automation Optimization - KDnuggets --> Join Newsletter 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 Hugging Face Transformers, running on an M2 Macbook Air with 24GB RAM and a 16-core Neural Engine. 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. 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 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: 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 "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 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 "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 (@mattmayo13) holds a master's degree in computer science and [truncated for AI cost control]