Typed Decisions (Laya)

A decision model answers a bounded question about a state with a typed result instead of generating text: one option out of a list (choice), a level on an ordinal scale (score), or the probability that a statement is true (boolean probability). The result carries a confidence and takes milliseconds of encoder time, not seconds of decoding. The three modules below put that on the JVM path with the same contract for every model.

Coordinates

dependencies {
    implementation(platform("sk.ainet.transformers:skainet-transformers-bom:0.63.3"))
    implementation("sk.ainet.transformers:skainet-transformers-decisions-api")      // the contract, no engine dependency
    implementation("sk.ainet.transformers:skainet-transformers-decisions-runtime")  // validation + acceptance policy
    implementation("sk.ainet.transformers:skainet-transformers-inference-laya")     // Laya as a DecisionAdapter
}

The contract (sk.ainet.decisions)

A DecisionRequest names a model, a state and the questions:

val request = DecisionRequest.builder("ticket-4711")
    .model("laya-multilingual")
    .state("Customer writes: my invoice shows the old address although I changed it last week.")
    .question(ChoiceQuestion("department", "Which department handles this request?", listOf(
        ChoiceOption("billing", "billing and invoices"),
        ChoiceOption("account", "account data and profile"),
        ChoiceOption("shipping", "delivery and logistics"),
    )))
    .question(ScoreQuestion("urgency", "How urgent is this request?", listOf(
        ScoreLevel("can wait"), ScoreLevel("this week"), ScoreLevel("today"), ScoreLevel("now"),
    )))
    .question(BooleanProbabilityQuestion("escalate", "Should a human look at this request?",
        BooleanCriteria(whenFalse = "no", whenTrue = "yes")))
    .build()

The state is DecisionState.text(…​) or DecisionState.structured(jsonElement); the InferencePolicy decides whether an over-long state is rejected (REJECT, the default) or truncated (TRUNCATE_STATE). A request fails before any inference when an option set is invalid, a primitive is unsupported or the option count exceeds the adapter’s capabilities.

Every question ends in exactly one QuestionOutcome:

Accepted(decision)

Decision.Choice(optionId), Decision.Score(level) or Decision.Boolean(value), each with its Prediction (probabilities, margin, expected score, confidence).

Abstained(reason, candidate)

The policy did not accept the candidate: BELOW_CONFIDENCE, BELOW_MARGIN, BETWEEN_THRESHOLDS, NO_POLICY, … The candidate stays visible, so a caller can show "unclear" instead of dropping the item.

Failed(error)

ContextOverflow, NumericalFailure, UnsupportedCapability, … A failure never masquerades as a decision.

DecisionResult.provenance records the model id, backend and precision, the preprocessing version, the calibration status, the policy id and version, the duration and one QuestionTrace per question (sequence length, state tokens kept and dropped).

The runtime (sk.ainet.decisions.runtime)

val agent = LayaModelLoader.load(Path.of("/models/laya/multilingual"), ctx, exactGelu = true)
val runtime = DecisionRuntime(listOf(LayaDecisionAdapter(agent, modelId = "laya-multilingual")))
val policy = AcceptancePolicy(
    id = "triage", version = "1",
    choice = ChoiceAcceptance(minConfidence = 0.2, minMargin = 0.1),
    score = ScoreAcceptance.MinTopProbability(0.35),
    boolean = BooleanAcceptance(lower = 0.4, upper = 0.6),
)
runtime.openSession(ModelSelector.Installed("laya-multilingual")).use { session ->
    val result = session.decide(request, policy)
    when (val outcome = result.question("department")) {
        is QuestionOutcome.Accepted -> println((outcome.decision as Decision.Choice).optionId)
        is QuestionOutcome.Abstained -> println("unclear: ${outcome.reason}")
        is QuestionOutcome.Failed -> println("failed: ${outcome.error}")
    }
}

Thresholds are inclusive. A choice is accepted when its entropy confidence (1 − H(p) / log k) reaches minConfidence and the gap to the runner-up reaches minMargin; a score when the mass on its top level reaches MinTopProbability (or its normalized entropy stays under MaxNormalizedEntropy); a boolean when P(true) lies outside the open interval (lower, upper). Without a policy every question abstains with NO_POLICY. AcceptanceEngine is public, so a caller can re-apply a stricter policy to the recorded candidates.

DecisionRuntime holds one adapter per installed model id and validates every request against the adapter’s AdapterCapabilities (primitives, languages, context and option budgets, maximum questions per request) before the adapter sees it. The conformance tests in decisions-runtime run against a FakeAdapter and are the contract any adapter must satisfy.

The DSL (sk.ainet.decisions.dsl)

Like the NN DSL and the data pipeline DSL, the decision contract has a Kotlin DSL that produces the same data classes as the constructors above; nothing in a block exists that the request model cannot express, so serialisation and provenance are unaffected. Three entry points:

import sk.ainet.decisions.*            // the contract and the typed readback
import sk.ainet.decisions.dsl.*        // decisionSchema, decisionRequest, acceptancePolicy
import sk.ainet.decisions.runtime.*    // openSession(schema), decide(schema, ...)

// Questions and policy declared once, states supplied later.
val triage = decisionSchema("commit-triage") {
    model("laya-multilingual")
    choice("kind", "What kind of change is this commit?", "feature", "bugfix", "refactor", "docs", "test", "chore")
    score("risk", "How risky is this change for the team?", "trivial", "low", "moderate", "high", "critical or breaking")
    probability("standup", "Is this commit worth mentioning in the team's daily stand-up?")
    probability("adr", "Does this commit contain or imply an architecture decision?") {
        whenFalse = "no decision"; whenTrue = "a decision is made or implied"
    }
    policy(version = "1") {                       // id defaults to the schema id
        choice { minConfidence = 0.2; minMargin = 0.1 }
        score { minTopProbability(0.35) }
        boolean { abstainBetween(0.4, 0.6) }
    }
}

runtime.openSession(triage).use { session ->
    for (commit in commits) {
        val result = session.decide(triage, "commit-${commit.id}", commit.message)
        println("${commit.id}: ${result.summary()}")   // kind=bugfix risk=1 standup=? adr=false
        if (result.boolean("standup") == false) continue
        val kind = result.choice("kind") ?: "unclear"
    }
}

// One request, inline.
val request = decisionRequest("ticket-4711") {
    model("laya-multilingual")
    state("Customer writes: my invoice shows the old address although I changed it last week.")
    choice("department", "Which department handles this request?") {
        option("billing", "billing and invoices")
        option("account", "account data and profile")
        option("shipping", "delivery and logistics")
    }
    score("urgency", "How urgent is this request?", "can wait", "this week", "today", "now")
    probability("escalate", "Should a human look at this request?")
    truncateState()
}

// A policy on its own.
val strict = acceptancePolicy("triage", "2") { choice { minConfidence = 0.5 } }

  • decisionSchema(id) { } fixes the questions, the policy, the model and the inference policy and stamps out one DecisionRequest per state with schema.request(requestId, state); DecisionRuntime.openSession(schema) and DecisionSession.decide(schema, requestId, state) in decisions-runtime use them directly, the schema policy applying unless a call overrides it.

  • decisionRequest(id) { } builds a single request; acceptancePolicy(id, version) { } a policy. A score { } block takes exactly one criterion and boolean { } needs both thresholds, otherwise InvalidPolicyException; a primitive without a block keeps abstaining with NO_POLICY.

  • choice/score accept a block (option(id, description), level(description)) or a vararg of ids whose text doubles as the description; probability takes the plain no/yes criteria unless its block sets whenFalse/whenTrue.

  • Typed readback on DecisionResult: choice(id), score(id), boolean(id) return the accepted value or null for an abstention or failure and throw for an unknown id or another primitive; probabilityTrue(id) reads the accepted or abstained candidate; outcome.candidate, accepted(id), isComplete and the one-line summary() complete the set.

The scopes carry a @DslMarker, so a nested block cannot call the functions of an enclosing scope by accident.

Laya (sk.ainet.models.laya)

skainet-transformers-inference-laya wraps the Laya checkpoints (Apache 2.0): the multilingual one (convaiinnovations/laya, an mmBERT-base encoder, 22 layers, hidden 768) and the English one (convaiinnovations/laya-typed-decisions, ModernBERT-large, 28 layers, hidden 1024) — a ModernBERT-family encoder (GeGLU, RoPE on every layer, a local attention window on two of every three layers) followed by a type embedding, a two-layer decision head and a scorer. Every size comes from the checkpoint’s own encoder/config.json. The port follows the reference build_sequence exactly, including CPython json.dumps float formatting for structured states (PythonJson), the 48-token option cap and the state-type-dependent truncation direction.

  • LayaCheckpointFiles.read(dir) reads rl_agent_config.json, encoder/config.json and the tokenizer pair; LayaModelLoader.load(dir, ctx) adds model.safetensors (F16, up-cast to FP32) and returns a LayaAgent.

  • LayaDecisionAdapter(agent, modelId) is the DecisionAdapter. Option ids are the model-facing keys; the descriptions are what the model reads. Calibration is UNVERIFIED unless the checkpoint ships temperatures.

  • Exact (erf) GELU is required for parity with the reference (exactGelu = true, the engine’s gelu(approximate = false) since SKaiNET 0.59); the tanh approximation moves the predictions by about 7e-4.

Verification against the fixtures shipped with the checkpoint (copied under llm-inference/laya/src/jvmTest/resources/laya-fixtures/, Apache 2.0): 26 of 26 packed sequences, the 40-text tokenizer corpus (CJK, RTL, combining marks, emoji, literal <mask>), and all six predict.json cases within 5e-5. The model-gated tests run when LAYA_MULTILINGUAL_PATH points at the checkpoint directory.

Expert heads and routing

A Laya checkpoint is an encoder plus a small decision head (type embedding, two transformer blocks, scorer), and the head is what a publisher fine-tunes for a domain: the encoder stays frozen, the head is copied, trained on domain data and saved on its own as expert.safetensors (the type_emb., head., scorer. tensors, act_head. ignored) with an expert.json (name, domains, its own calibration temperatures). sk.ainet.models.laya runs any number of such heads on one encoder:

  • LayaExpertLoader.load(metaJson, weights, base, ctx) binds an expert head for a loaded base LayaAgent (weightFormFor = LayaWeightForms.q8_0For packs its projections like the base’s); LayaExpert holds the LayaDecisionHead and a LayaDecoder with the expert’s temperatures over the base agent’s budgets. LayaWeightPreparation.prepareHead / LayaDecisionModel.buildHead are the steps underneath; a head prepared on its own reproduces the base model bit for bit.

  • LayaDecisionModel.logitsFrom(encoded, sequence, type, head) runs a head on an encoder output: the encoder runs once per question (the sequence contains the question) and every head that should answer it reads the same hidden states. LayaAgent.encodeFor / logitsFrom do the same over the padded batch plan.

  • LayaRouterSpec(question, options, experts, threshold) is the prompted router: the base head answers a choice over fine-grained kinds of text (each with a description) about the state, every kind maps to an expert, and the expert answers when it is loaded and the top probability reaches the threshold — otherwise general, the base head. Routing is sequence-level, one head per request; this is not a sparse mixture of experts.

  • LayaMixtureAgent(base, experts, router).decide(state, questions, expert = null, compare = false) routes (one extra encoder pass for the router question) and answers with the chosen head; expert forces a head (the routing is still reported as overridden), compare adds the base head’s answers on the same encoder output. The result carries the routing (probabilities over the kinds, the pick, the threshold), the answers, and encoder / head milliseconds.

  • LayaMixtureDecisionAdapter(mixture, modelId, onRouted) is the DecisionAdapter of a mixture: every request is routed once and answered by the chosen head with its calibration; the contract has no field for the head that answered, so onRouted reports it.

The gated LayaMoeEvalTest replays the published evaluation of layaMOE (Apache-2.0; safety and customer_ops heads on convaiinnovations/laya-typed-decisions, 108 hand-labelled cases) and requires the same routing and the same answers as the PyTorch FP32 run; it needs LAYA_TYPED_DECISIONS_PATH and LAYA_MOE_EXPERTS_PATH.

Quantized lane

LayaModelLoader.load(files, weights, ctx, weightFormFor = LayaWeightForms.q8_0For) loads the encoder and head projections and the token table as Q8_0 at load (engine 0.61 quantize-at-load, SKEEP-007): about a quarter of the FP32 bytes, served by the packed kernels on every tier — on the JVM the FFM pack, on Android the JNI pack, in a browser the skainet-backend-wasm-simd sidecar after WasmSimd.install(). The contract of that lane is decision quality rather than FP32 parity: LayaQuantizedParityTest requires every predict.json choice to agree with the FP32 agent and no probability, score, noul value or confidence to shift by more than 0.05 (measured: 5/5 choices, max shift 0.035). On the eager JVM path this lane trades speed for memory — the packed kernels run one activation row at a time — so use it where the bytes matter: the browser, a phone, a small heap.

Current limits

  • The loader is common code (a RandomAccessSource plus the four side files as text); obtaining the bytes — a file, an app’s files directory, a download held in memory — is the caller’s part.

  • Eager FP32 on the CPU path by default; batching across the questions of one request since 0.61 saves about 10 % per question there. The compiled lane is an open item; the quantized lane is the section above.

  • The base checkpoint is a starting point the publisher expects to be specialised; keep option sets small (six or fewer) and treat abstentions as the normal case, not the exception.