Building an AI Agent for Android, Part 23: On-Device Inference — Model Size, Latency, and What Actually Runs

This is part of an ongoing series documenting the architecture of StareBrain, an on-device AI agent for Android. Previous parts covered the overlay UI, drag-to-reposition, auto-collapse, voice input, session handoff, natural language planning, and plan validation. This part covers the inference layer itself: how we run a model locally, and the tradeoffs that decision forces.

Why On-Device

The short answer: privacy and latency. An agent that controls your device can see your screen, read your notifications, and act on your behalf. Sending that data to a remote API is a trust boundary most users won't cross, and rightly so. On-device inference keeps everything local — no data leaves the phone.

The longer answer: network calls inside an agentic loop compound badly. If each step in a 5-step task has a 400ms round-trip, that's two seconds of network overhead before you count inference time. On a mid-range device with a good local model, inference is slower per token — but it's consistent, offline-capable, and doesn't stall when signal drops.

The Model Size Problem

Android devices in 2024 run the gamut from 3GB RAM to 16GB. A model that fits comfortably on a Pixel 8 Pro will OOM on a mid-range Galaxy. This forces a real decision early: which device tier are you targeting, and what does that mean for model quality?

StareBrain currently targets devices with 6GB+ RAM. That rules out some budget devices but covers the majority of active Android users. Within that constraint, we tested several model sizes:

We settled on 3B for the default, with a flag to enable 7B on high-RAM devices. The 1B model is fast but doesn't reliably parse ambiguous instructions into valid step sequences — it collapses multi-step plans into single actions too often.

Loading the Model

We use MediaPipe's LLM Inference API, which handles tokenization, KV-cache management, and GPU delegation. The model is bundled as a .task file and loaded at service start:

class OnDeviceModelClient(private val context: Context) {

    private var llmInference: LlmInference? = null

    fun initialize() {
        val options = LlmInference.LlmInferenceOptions.builder()
            .setModelPath(modelPath())
            .setMaxTokens(1024)
            .setTopK(40)
            .setTemperature(0.7f)
            .setRandomSeed(42)
            .build()

        llmInference = LlmInference.createFromOptions(context, options)
    }

    private fun modelPath(): String {
        return context.filesDir.absolutePath + "/models/starebrain_3b.task"
    }

    fun generateResponse(prompt: String): String {
        return llmInference?.generateResponse(prompt)
            ?: throw IllegalStateException("Model not initialized")
    }

    fun release() {
        llmInference?.close()
        llmInference = null
    }
}

Initialization happens once when AgentService starts, not per-request. Cold loading a 3B model takes 2–4 seconds on most devices — acceptable at service start, unacceptable mid-task.

Latency Budget

An agentic loop has a latency budget. Each step costs:

1. Context assembly — building the prompt from AgentContext (~5ms)
2. Inference — first token latency + generation time (~340ms–1200ms depending on output length)
3. Response parsing — extracting structured output from model response (~10ms)
4. Action execution — the actual device action (varies)

For a 5-step task, inference alone is 1.7–6 seconds. This is the number users feel. The mitigation isn't making inference faster (you can't, within a fixed model size) — it's making the wait feel intentional.

The overlay's step progress display (from Part 16) handles this: showing "Step 2 of 5: Opening Settings" gives the user a frame for the latency. An indeterminate spinner for 4 seconds feels broken. A progress indicator for 4 seconds feels like work happening.

Keeping the Context Small

Inference time scales with prompt length. A bloated context — full conversation history, verbose system prompt, uncompressed task state — adds tokens that slow every step.

AgentContext (introduced earlier in the series) is designed to stay lean:

data class AgentContext(
    val deviceState: DeviceState,         // current app, screen title
    val recentActions: List<ActionRecord>, // last 3 actions only
    val userIntent: String,               // original instruction
    val confidenceLevel: ContextConfidence,
    val capturedAt: Long
)

We cap recentActions at 3. The model doesn't need the full history to plan the next step — it needs enough to avoid repeating itself and to detect if a previous step failed. More than 3 actions adds tokens without meaningfully improving planning quality, based on our testing.

Failure Modes Specific to On-Device

Remote APIs fail predictably: timeouts, rate limits, network errors. On-device inference fails in ways that are harder to handle:

OOM during inference. If another app spikes memory mid-task, inference can be killed by the OS. We catch this and surface it as a task failure with a "device resources unavailable" message rather than a crash.

Thermal throttling. On sustained tasks, the device heats up and the GPU throttles. Inference that took 340ms at step 1 can take 900ms by step 8. We don't solve this — we just track per-step latency and surface a warning if average latency doubles mid-task.

Model output degradation. At high temperature or with certain prompt structures, the 3B model produces malformed JSON that parsePlanResponse (from Part 21) can't parse. PlanResult.Failure handles this, but the retry path re-runs inference, which doubles the cost. The fix is tighter output constraints in the system prompt — explicitly enumerating the output schema reduces malformed responses significantly.

// In buildPlanningPrompt():
"""
Respond ONLY with valid JSON matching this exact schema:
{
  "steps": [
    {
      "action": "string",
      "targetPackage": "string or null",
      "parameters": {},
      "requiresConfirmation": boolean
    }
  ]
}
Do not include explanation, commentary, or any text outside the JSON object.
"""

This single addition dropped malformed response rate from ~12% to ~2% in our internal testing.

What's Next

Part 24 will cover accessibility service integration — how StareBrain reads screen content and dispatches actions using AccessibilityService, and the permission flow that precedes it.

The model runs. The plan is validated. The next piece is the hands: actually touching the UI on behalf of the user.

StareBrain is an on-device AI agent for Android. If you're building in this space or have questions about the inference layer, I'm in the comments.