Sohaib Arif — AI Safety Research

Notes on AI safety for local & embedded models

Experimental Setup and results for tool use verification project

This post describes how the experiments were set up and the rationale behind the design decisions. The goal of the project is to check if local models, specifically the gemma4 family follows instructions and uses provided arithmetic tools properly and what kind of success and failure modes are in this case. Just to briefly re-iterate why the LLM needs to use the tools for the provided task.

Consistency: if the LLM uses a human or AI written + human verified function it will get the same results every time. Models are probabilistic, you can set a seed value to induce some amount of stability but that may not be enough.

Reduced Token Use: using a pre-built tool and using it correctly means the LLM will use less tokens trying to do the math in its head. This can especially be a problem with reasoning or thinking turned on.

Better Debugging: if the LLM is restricted to tools for calculations, it means that you can go back to the internals and check what is wrong more easily instead of theorizing what could have gong wrong or applying more difficult time consuming techniques like linear probing to debug.

Additional Calculation abilities: Since this is being done on a smaller GPU and I can go down to even embedded level devices, it opens up the possibility of adding more complex matrix calculations or using scipy to create and then run formulas with verification that the correct ones are being used.

Now that that is briefly covered, lets turn to the experimental setup:

Models used:

The models used were gemma4:e4b, gemma4:26b and gemma4:12b. Why Gemma, why these models and why these sizes? The original idea for these experiments came out of my previous project which was smart punching bag, that project will be described in more detail in a later blog post and I will edit this one to add a link when it is ready. That project was specifically about Kaggle/Google's Gemma4Good Hackathon which required the use of Gemma 4 series of models.

Explaining some specifics of the models:

model_name: this is the name of the model. As of this writing it is gemma4:e4b, gemma4:12b and gemma4:26b. e4b means it has 4 billion effective parameters, 12b means 12 billion parameters and 26b means 26 billion parameters. Larger models take up more memory and take longer to respond but in theory are more accurate. The results show this is not necessarily always the case. Details on the model can be found on the model card

temperature: Normally this is taken as a measure of how deterministic or random the model is. Most books and documentation I had read so far mention using a lower temperature for more deterministic and a higher temperature for more creative outputs.

From the book Deep Learning With Python Third Edition by François Chollet section 16.2.4

You’ll notice we added a parameter called temperature. We can use this to sharpen or widen our probability distribution so our sampling explores our distribution less or more.

If we pass a low temperature, we will make all logits larger before the softmax function, which makes our most likely output even more likely. If we pass a high temperature, our logits will be smaller before the softmax, and our probability distribution will be more spread out.

At a high temperature, our outputs no longer resemble English, settling on seemingly random tokens. At a low temperature, our model behavior starts to resemble greedy search, repeating certain patterns of text over and over.

This is why the experiments were originally done at temperatures 0.0, 0.1, 0.2 and 0.3. At 0 I was expecting very little randomness and likely a lot of repetition, and I was expecting more reasonable results at 0.2 and 0.3. But I later read the model card in more depth and found that Gemma 4 recommends a temperature of 1.0 so I added that. The model card also recommends setting two other parameters top_k and top_p. In the initial experiments I left those as defaults but for the 1.0 ones, I changed top_k to 64 and top_p to 0.95. Changint the top_k sets the number of next words being considered by the model. So in the case of temperature of 1.0, and top_k of 64 we consider the whole distribution as set of possible outputs then we reduce that to only the top 64 most likely words. Then top_p reduces that to the 95% most likely words among that 64 possible words. This interactive graph generated by claude shows the impact of these three parameters.

Two additional parameters were added later. These are context size and number of tokens to predict. context size number of previous tokens used to generate the next one, listed as num_ctx in ChatOllama class. num predict number of tokens to generate at max, listed as num_predict in ChatOllama class. The defaults for both are set by the original model and are not the ones mentioned in the ChatOllama documentation. Context size is very important since it includes the thought tokens and without a large enough "compute budget" for it, it can miss parts of thought and repeat the same statement again and again. Aside from these numerical parameters, there is a parameter for thinking or reasoning. Gemma4 supports thinking in some models so for the evaluation I ran the tools with thinking on and off for each model. Note that this does increase the time required to generate answers and I have seen that misconfiguring it by applying the model card's method as opposed to setting reasoning to True or False. The reason is that it is very easy to mistakenly put both the think tag from the model card and turn on the reaasoning without knowing which one is actually providing the reasoning. Double think tags can send the models into a loop where it thinks on top a previous think call and then starts mistaking its thoughts with the original input.

Note that I only changed the num_ctx and num_predict after seeing bad results from gemma4:26b several times especially with thinking because it could not reason properly and kept repeating the input list on one of the more complex questions.

How the model was called:

Ollama explanation: This is a Python implementation of lammaa.cpp. Originally made to use Laama 2 on desktop environments this now supports most LLMs. This is the option I used for the project since it was the one mentioned in some of the texts I had been reading, it has reasonable defaults and it is customizable with the parameters I mentioned above. What other options are available (vLLM, Transformers): Other options to run an LLM locally are to use vLLM or directly using Transformers. Both have their own advantages with Transformers being very customizable and having the ability to quantize a model or apply parameter efficient fine tuning, which I might add later. Langchain and Langgraph explanation LangChain and LangGraph: While the model can be run directly using Ollama, adding LangGraph allows adding tool use and creating a graph structure with steps that run from one to the next or loop back around until a stop condition is met. LangGraph is for defining such graphs and LangChain is a related library to define tasks in LCEL. Because these experiments required tool-use, LangGraph was more suitable for this task. You can learn more about both on LangChain Academy specifically the intro to langgraph course but be warned that while the courses are free, they do keep mentioning paid subscriptions or free tiers of subscriptions including for a sort of unrelated paid debugging tool with a free tier called LangSmith. I have been working on a fork of the foundations course for LangChain Academy using only local model and removing the subscriptions which I may add to the blog later.

Tools

The tools are the functions that the LLM can call, they start off in stage_00 with a few scalar operations and an operation to apply another op to a list or to reduce a list using an operation. Note that at this stage I deliberately did not give the list functions the ability to use a while loop to see how different local LLMs handle that. The tools folder also has other stages already listed which would add, remove or change tools to test out overspecified and underspecified tools.

Concretely this is the list of tools:

TOOLS: dict[str, Callable] = {
    #list functions 
    "apply_over_list": apply_over_list,
    "reduce_over_list": reduce_over_list,
    "remove_from_end_of_list": remove_from_end_of_list,
    "remove_from_list": remove_from_list,
    "add_lists": add_lists,
    "remove_last_n": remove_last_n,
    "remove_first_n": remove_first_n,
    #combination of string and list 
    "split_string": split_string,
    "convert_str_to_int": convert_str_to_int,
    #scalar only functions 
    "multiply": multiply,
    "divide": divide,
    "add": add,
    "sub": sub,
    "pow": pow,
    "greater_than": greater_than,
    "less_than": less_than,
    "convert_equals": equals
}

The graph

The graph here means the sequence of instructions performed by the LLM. In this case the basic structure is very simple. An "assistant" node which takes the instruction and calls relevant tools which are bound to the graph and then the tool can give the result as a string back to the assistant which can use it to form an answer or circle back and do another tool call until it has an answer. Importantly, the assistant has an instance of ChatOllama which has the parameters fedined in the parameters section. These are the params being swept over for the evaluation. I later composed this into a reusable class where the tool set and model can be defined and where the graph can be invoked directly, meaning it will run the full steps and return an answer or it.

Concretely these three functions define the graph in this file:

def assistant(self,state:MessagesState):
        sys_msg:list[AnyMessage] = [SystemMessage(content=self.sys_msg)]
        result = self.llm_with_tools.invoke(input=sys_msg+state['messages'])
        return {"messages":result}

def get_graph_builder(self):
        graph_builder = StateGraph(MessagesState)
        graph_builder.add_node("assistant", self.assistant)
        graph_builder.add_node("tools",ToolNode(self.tool_list))
        graph_builder.add_conditional_edges("assistant", tools_condition)
        graph_builder.add_edge(START, "assistant")
        graph_builder.add_edge("tools", "assistant")
        graph_builder.add_edge("assistant",END)
        return graph_builder
def build_graph(self):
        builder = self.get_graph_builder()
        self.graph = builder.compile(checkpointer=self.llm_memory)

The self is there because these are part of the GenerateGraph class which allows me to change the system message, the set of provided tools, and the model and its parameters easily.

Questions

Right now there are only 5 questions but they are designed to test out specific topics. Having a small list has the bonus of being able to run the entire grid of test parameters in a reasonable amount of time, especially since the larger models can take a long time to respond. Here is the question list so far:

questions = {
    "simple_question" :["testing basic tool call what is 1024 multiplied by 64."],
    "text_simple_question" = ["How many 'r' s are there in the word 'strawberrry'. Don't change the spelling and use the provided tools."],
    "multiple_step_question" ["Add 3 and 4. Multiply the output by 2. Divide the output by 5"],
    "complex_question": [
        "While the total is less than 10,000 multiply the values in the list [102, 100, 200, 150, 101, 90, 101,100,110,110,] by 10 and keep adding. What does the total come out to BEFORE reaching 10,000. Use the provided tools.",
        "From the original list remove the first 3 values and append [99,80,75] to the list. Then answer the same question. Remember to use the provided tools to answer not your memory"
    ]
}

In the snippet I am defining these as a list but in the code they are defined as a class for easy use with the IDE's text text completion.

Now for what each of these test:

question 1 The first question is a simple scalar question, it can be expanded with a list of scalar operations to check all functions. But it just checks if the tools are being called for simple cases.

question 2 This is a simple question that can be done using one tool that splits the strings or using two calls, one to split the string and the other to count a particular letter. It checks weather the model falls into the trap of just using the correct spelling of strawberry and reporting 3 r's, it can call split string and then count the number of 'r's in memory, or it can do the split and then try to call a reduce list tool which exists but it would fail to use any condition.

question 3 This is a multiple step question using only the scalar operation tools. It would need two tool calls where the output of the first call is a parameter to the second call, or it could a way to merge the two tools when given the additional tools to do so.

question 4 Similar to the previous one, this is a two part question but it is divided into two questions and each of those have two parts. Both questions are also slightly ambegious. First question checks if it can apply an operation to a list and then, without the ability to do a while loop, it can continue doing an operation until a condition is met.

To answer the second question, it has to use the original list again without explicitly being provided said list. Then it must apply a list operation to another list and then it must infer from the last question what other tasks need to be done.

Failure modes seen:

While the main test is weather the LLM uses the tool and that it uses it correctly there are additional questions. Here are some failure modes I have seen:

  • It hallucinates the tool name or sub-tool name in cases where one tool is supposed to be used with another
  • It does the entire calculation in its head
  • It does the calculations in its head while still doing tool calls
  • It tries to do the calculation in its head but runs out of compute budget
  • It tries to call the wrong tool, gets the error for it and fails gracefully by making a different tool call from a list provided by the error string

Entry point for the code

The main input to the code is the experiments folder, the highest number indicates the the experiment that is currently being worked on, planned or complete. As of this writing this is stage_03_add_thinking.py.

Logging

Logging of results is done by sending the results to the results folder labeled with the stage and then the date/time of when the experiments were run. This was actually one of the more difficult parts of the experiments so part of the logging code was AI generated up till Stage 3 before refactoring the experiments in that stage to split the logging from the experiment runner manually. One problem with the logging code generated with Calude Opus 5 was that it was mixing up the formula mentioned in the gemma 4 model card and the parameters used in the ChatOllama class. Specifically it was adding the thinking tag manually and then trying to strip out the thinking results manually before the next call. It also tried to get the tool names via introspection instead of using a dictionary of the named tools. A lot of code was generated for the logs at the same time so this is what I was referring to in my readme when I said allowing too much AI generated code could loose control of the experiments. So I will be refactoring it manually and add the results this week and add analysis for stage 3 experiment results.

Results Stage 02

Stage 02 results show the case when thinking is not enabled and the only changed parameter is the temperature. The sweep runs across 2 models and 4 temperatures, running 5 questions each. All cases are with thinking off and are not using temp 1.0 which was added in stage 3.

question 1 result The stage_02 the last results summary shows that for both models, all of the tool calls were correct across both models and the 4 temperatures. This was the simple question to check if the model could call the multiply tool correctly on a scalar.

question 2 result The second question starts showing errors for the tool calls only at the gemma4:e26b. Only the run with temperature 0.2 was correct for this question and it required 6 tool calls because it did multiple incorrect ones. With the e4b variant across all temperatures it answered correctly and did between 1 and 3 tool calls. For the 0.2 specifically it only had to use 1 tool call which was the split string. Looking at the actual conversation it tried to call a non-existing tool "count_occurrences". When that did not work it called remove from list and then counted the number of r's it had to remove. This is actually quite clever since it followed the instruction without needing an obvious tool name or combination of tools.

With a temperature of 0 on e4b it tried to do 3 tool calls, the first one was correct but then it tried to use the equals call with the apply to list function without first turning the characters into their ordinal values which did result in an error. Then it tried to call the equals function on the list itself which is completely illogical. Despite that it had the already split up string in its memory and just used that to recover and answer correctly.

question 3 result Question 3 was a multi step question but both parts were simple scalar operations. Interestingly both models across all temperatures completed the task in the same number of tool calls and correctly called the same tools. While it was a multi-step question it was also a simple and well specified one so I think that lends to this result.

question 4 result For this one, at temperature 0, the tool call completely failed for gemma4:e4b variant since it returned no answer. At temperature 0.1 on the same variant it started to repeat a pattern that I saw in many other cases. It correctly applied the apply_over_list with the multiply operator and then it summed the result to check if it was less than 10,000. Then it skipped doing subsequent tool calls and just showed the calculation in memory. So basically it "maliciously" followed the instruction to use its tools, the tool use shows up in the log but then it did the calculation without using tools. There is an indicator that a model does in both e4b and 26b. If the final text answer shows the calculation it is usually an indicator in this stage of the experiments that it completed the task in memory after trying the tool.

question 5 result Question 5 follows on from 4 because it uses the same list. To do that it uses the same thread id for memory and thus the logged results reference both question 4 and 5. Across both models and all temperatures the remove_first_n tool and add lists tool is called correctly with the exception of gemma4:e4b at temperature 0. The 26b variant on temperature 0 has an interesting result here because it actually created its own while loop in a way by calling the add operation multiple times removing one number each time the result exceed 10,000. As with the previous question the indicator that it used tools here is that the answer in only this case is not showing the calculation.

Conclusion

So far I can see gemma:26b is better at following instructions, including being able to build its own while condition to call tools, it is not a consistent result. No variant in these results mentioned that question 4 might be underspecified. There are reasons for this inconsentency. For one, the 26b variant did not have enough memory and the context size was not set appropriately for the larger model. That said, the e4b variant under the same defaults and low temperature showed more abality to run the tasks without making up tool names in any of these results. When it did fail, especially on temperature 0, it just did not return any result which might be easier to debug than checking for all hallucinated names.

large-language-modellocal llm

← All posts