One of the most important features of today’s generative models is their ability to take unstructured, partially unstructured, or poorly structured inputs and convert them into structured objects that conform to specific schemas — relational-database fixed schemas, document store flexible schemas, function signatures, API specifications, and so on.
Large language models (LLMs) can perform this task if prompted with all the schema specifications and instructions on how to process the input. In addition, most of today’s LLMs include a dedicated JSON mode or structured-outputs mode that abstracts away part of this prompt engineering from users.
However, this approach has a few limitations. First is the cost of using LLMs when they need to scale to databases with millions or billions of records or requests; second is the potential complexity of the prompt engineering; and third is the limited complexity of the schemas that the built-in JSON and structured-outputs modes can support.
In two recent papers we published at the Conference on Empirical Methods in Natural Language Processing (EMNLP) and on ArXiv, we presented a native approach to this problem in the form of a specialized lightweight structured-object language model (SoLM). Unlike general-purpose LLMs, SoLM is trained to generate objects only in a specific schema. SoLM’s contributions include a special training method called self-supervised denoising and a special decoding method for use at inference time called confidence-aware substructure beam search (CABS), which mitigates hallucination risks.
In experiments, we found that SoLM’s output accuracy matched or exceeded that of state-of-the-art LLMs, while its cost efficiency was an order of magnitude better. We also found that, on the problem of product attribute generation, the CABS decoding approach improved recall by 16.7% relative to conventional beam search decoding, when precision was fixed at 90%.
Applications
In our paper, we unify multiple seemingly unrelated AI/ML problems under the structured-output formulation. For instance, a challenge arises when the structured object has multiple facets, or redundant pieces of information that depend on each other. One facet of the object might be a long descriptive text in natural language; other facets might be short type-constrained structured facts.
These types of multifaceted objects occur commonly in listing scenarios (products, houses, jobs, etc.), where the object contains a descriptive section and a section that lists key attributes. SoLM allows us to generate an object with these various types of facets while ensuring both relative consistency within the object and absolute consistency with respect to world knowledge.
The typical use of a structured-output model involves feeding it a blurb of unstructured data and letting it generate the corresponding structured object. In our paper, we also propose using SoLM as what we call a self-regeneration machine. In this case, we simply feed the model an object that is already structured according to the schema, and we let the model regenerate it end-to-end.
Here, the task is no longer to structure the input but to clean, normalize, correct, and/or complete it and make it self-consistent. Of course, the input can include any combination of an already structured record and additional unstructured content, or it can include a record that structured according to a different schema. SoLM is agnostic as to input and will always generate a clean record according to the target schema.
The self-regeneration machine can solve multiple problems at once: completion of missing facts, correction of incorrect facts, normalization of unnormalized facts, completion of missing descriptions, and correction of inaccurate information in descriptions. All these tasks are interdependent and introduce dependency loops when treated independently (e.g., should one extract facts from descriptions or write descriptions based on facts?). Self-regeneration solves for these dependencies in the most natural way.
Innovations
To train the SoLM model, we use self-supervised denoising. The idea is to use any sample of objects from an existing database, introduce artificial noise into these objects, and train the model to recover their original forms. The model thus learns to enhance the quality of any object we feed into it. By making the noise more aggressive — for instance, by completely removing the structure of the object or randomly shuffling tokens — we teach the model not only to enhance the quality of an existing object but also to operate on completely unstructured input.
Although LLMs are trained simply to generate the most likely next token in a sequence of token, at inference time, they typically use different decoding strategies in order to select outputs. One of the most popular is beam search decoding, in which the model considers several candidate sequences in parallel, ultimately selecting the sequence that has the highest cumulative probability. (Simply selecting the highest-probability token at each turn — greedy decoding — doesn’t guarantee the highest-probability sequence of tokens over some fixed number of turns.) The number of sequences that the model considers at once is called the width of the beam.
SoLM’s output consists of a sequence of key-value pairs, where the key is some data type from the schema — say, “brand” in the schema for product listings — and the value is the value for that type — say, the brand of a particular produce. We also use special tokens (“<SEP>” and “<END>”) to separate keys and values.
In confidence-aware substructure beam search, the key-value pair, rather than the token, is the atomic component of the beam search. The probability of the key-value pair can be inferred from the LLM’s confidence in its output. But we also experimented with a separately trained confidence score model, which takes as input the intermediate representation produced by one of the LLM’s inner layers. In practice, this approach worked better than relying directly on the model’s confidence scores.
In our papers, we show that a SoLM model with seven billion parameters matches or outperforms various prompt-engineering techniques on much larger foundational models, across metrics such as the completeness of facts, the correctness of facts, and the quality and factuality of the descriptive content. With CABS decoding, we further increase the correctness of the facts by removing facts that were hallucinated during decoding.